Improve neighbour-joining tree performance (#641)

* begin njt refactoring

* refactor plot_njt

* fix docs

* tweak

* add tests

* fix nan comparisons

* poetry update

* consistent color for NA

* fix karyotype categorical

* fix warning

* fix haplotype networks

Author: alimanfoo

docs/source/Af1.rst

@@ -134,6 +134,7 @@ Genetic distance and neighbour-joining trees (NJT)
     :toctree: generated/
 
     plot_njt
+    njt
     biallelic_diplotype_pairwise_distances
 
 Heterozygosity analysis

docs/source/Ag3.rst

@@ -144,6 +144,7 @@ Genetic distance and neighbour-joining trees (NJT)
     :toctree: generated/
 
     plot_njt
+    njt
     biallelic_diplotype_pairwise_distances
 
 Heterozygosity analysis

malariagen_data/ag3.py

@@ -392,17 +392,20 @@ def karyotype(
         ds_snps = self.snp_calls(
             region=region, sample_sets=sample_sets, sample_query=sample_query
         )
-        geno = allel.GenotypeDaskArray(ds_snps["call_genotype"].data)
-        pos = allel.SortedIndex(ds_snps["variant_position"].values)
-        samples = ds_snps["sample_id"].values
-        alts = ds_snps["variant_allele"].values.astype(str)
-
-        # subset to position of inversion tags
-        mask = pos.locate_intersection(inversion_pos)[0]
-        alts = alts[mask]
-        geno = geno.compress(mask, axis=0).compute()
 
         with self._spinner("Inferring karyotype from tag SNPs"):
+            # access variables we need
+            geno = allel.GenotypeDaskArray(ds_snps["call_genotype"].data)
+            pos = allel.SortedIndex(ds_snps["variant_position"].values)
+            samples = ds_snps["sample_id"].values
+            alts = ds_snps["variant_allele"].values.astype(str)
+
+            # subset to position of inversion tags
+            mask = pos.locate_intersection(inversion_pos)[0]
+            alts = alts[mask]
+            geno = geno.compress(mask, axis=0).compute()
+
+            # infer karyotype
             gn_alt = _karyotype_tags_n_alt(
                 gt=geno, alts=alts, inversion_alts=inversion_alts
             )
@@ -422,7 +425,8 @@ def karyotype(
                     "total_tag_snps": total_sites,
                 },
             )
-            kt_dtype = CategoricalDtype(categories=[0, 1, 2], ordered=True)
+            # Allow filling missing values with "<NA>" visible placeholder.
+            kt_dtype = CategoricalDtype(categories=[0, 1, 2, "<NA>"], ordered=True)
             df[f"karyotype_{inversion}"] = df[f"karyotype_{inversion}"].astype(kt_dtype)
 
         return df

malariagen_data/anoph/dipclust_params.py

@@ -1,6 +1,6 @@
 """Parameters for diplotype clustering functions."""
 
-from .diplotype_distance_params import distance_metric
+from .distance_params import distance_metric
 from .clustering_params import linkage_method