Merge pull request #673 from malariagen/GH656_avoid_anjl_import

Replicate anjl params to avoid global import causing delays

Author: leehart

malariagen_data/anoph/distance.py

@@ -6,7 +6,6 @@
 import numba  # type: ignore
 import numpy as np
 from numpydoc_decorator import doc  # type: ignore
-import anjl.params  # type: ignore
 
 # Internal imports.
 from .snp_data import AnophelesSnpData
@@ -410,10 +409,10 @@ def plot_njt(
         metric: distance_params.distance_metric = distance_params.default_distance_metric,
         distance_sort: Optional[tree_params.distance_sort] = None,
         count_sort: Optional[tree_params.count_sort] = None,
-        center_x: anjl.params.center_x = 0,
-        center_y: anjl.params.center_y = 0,
-        arc_start: anjl.params.arc_start = 0,
-        arc_stop: anjl.params.arc_stop = 2 * math.pi,
+        center_x: distance_params.center_x = 0,
+        center_y: distance_params.center_y = 0,
+        arc_start: distance_params.arc_start = 0,
+        arc_stop: distance_params.arc_stop = 2 * math.pi,
         width: plotly_params.fig_width = 800,
         height: plotly_params.fig_height = 600,
         show: plotly_params.show = True,
@@ -426,8 +425,8 @@ def plot_njt(
         color_discrete_sequence: plotly_params.color_discrete_sequence = None,
         color_discrete_map: plotly_params.color_discrete_map = None,
         category_orders: plotly_params.category_order = None,
-        edge_legend: anjl.params.edge_legend = False,
-        leaf_legend: anjl.params.leaf_legend = True,
+        edge_legend: distance_params.edge_legend = False,
+        leaf_legend: distance_params.leaf_legend = True,
         legend_sizing: plotly_params.legend_sizing = "constant",
         thin_offset: base_params.thin_offset = 0,
         sample_sets: Optional[base_params.sample_sets] = None,
@@ -449,6 +448,10 @@ def plot_njt(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
     ) -> plotly_params.figure:
+        # Only import anjl if needed, as it requires a couple of seconds to compile
+        # functions.
+        import anjl  # type: ignore
+
         # Normalise params.
         if count_sort is None and distance_sort is None:
             count_sort = True

malariagen_data/anoph/distance_params.py

@@ -19,3 +19,20 @@
 ]
 
 default_nj_algorithm: nj_algorithm = "dynamic"
+
+center_x: TypeAlias = Annotated[int | float, "X coordinate where plotting is centered."]
+
+center_y: TypeAlias = Annotated[int | float, "Y coordinate where plotting is centered."]
+
+arc_start: TypeAlias = Annotated[int | float, "Angle where tree layout begins."]
+
+arc_stop: TypeAlias = Annotated[int | float, "Angle where tree layout ends."]
+
+edge_legend: TypeAlias = Annotated[
+    bool, "Show legend entries for the different edge (line) colors."
+]
+
+leaf_legend: TypeAlias = Annotated[
+    bool,
+    "Show legend entries for the different leaf node (scatter) colors and symbols.",
+]