speed-up find_soma:

- work with mask to avoid subsetting node table
- return empty array if no criteria given

navis/morpho/analyze.py

@@ -57,36 +57,61 @@ def find_soma(x: 'core.TreeNeuron') -> Sequence[int]:
         raise TypeError(f'Input must be TreeNeuron, not "{type(x)}"')
 
     soma_radius = getattr(x, 'soma_detection_radius', None)
-    soma_label = getattr(x, 'soma_detection_label', 1)
+    soma_label = getattr(x, 'soma_detection_label', None)
 
-    soma_nodes = x.nodes
+    check_labels = not isinstance(soma_label, type(None)) and 'label' in x.nodes.columns
+    check_radius = not isinstance(soma_radius, type(None))
+
+    # If no label or radius is given, return empty array
+    if not check_labels and not check_radius:
+        return np.array([], dtype=x.nodes.node_id.values.dtype)
+
+    # Note to self: I've optimised the s**t out of this function
+    # The reason reason why we're using a mask and this somewhat
+    # convoluted logic is to avoid having to subset the node table
+    # because that's really slow.
+
+    # Start with a mask that includes all nodes
+    mask = np.ones(len(x.nodes), dtype=bool)
+
+    if check_radius:
+        # When checking for radii, we use an empty mask and fill it
+        # with nodes that have a large enough radius
+        mask[:] = False
 
-    if not isinstance(soma_radius, type(None)):
         # Drop nodes that don't have a radius
-        soma_nodes = soma_nodes.loc[~soma_nodes.radius.isnull()]
+        radii = x.nodes.radius.values
+        has_radius = ~np.isnan(radii)
 
         # Filter further to nodes that have a large enough radius
-        if not soma_nodes.empty:
+        if has_radius.any():
             if isinstance(soma_radius, pint.Quantity):
                 if isinstance(x.units, (pint.Quantity, pint.Unit)) and \
                    not x.units.dimensionless and \
-                   not isinstance(x.units._magnitude, np.ndarray):
+                   not isinstance(x.units._magnitude, np.ndarray) \
+                   and x.units != soma_radius:  # only convert if units are different
                     # Do NOT remove the .values here -> otherwise conversion to units won't work
-                    is_large = soma_nodes.radius.values * x.units >= soma_radius
+                    is_large = radii * x.units >= soma_radius
                 else:
                     # If neurons has no units or if units are non-isotropic,
                     # assume they are the same as the soma radius
-                    is_large = soma_nodes.radius.values * soma_radius.units >= soma_radius
+                    is_large = radii >= soma_radius._magnitude
             else:
-                is_large = soma_nodes.radius >= soma_radius
+                is_large = radii >= soma_radius
 
-            soma_nodes = soma_nodes[is_large]
+            # Mark nodes that have a large enough radius
+            mask[is_large] = True
 
-    if not isinstance(soma_label, type(None)) and 'label' in soma_nodes.columns:
+    # See if we (also) need to check for a specific label
+    if check_labels:
         # Important: we need to use np.asarray here because the `label` column
         # can be categorical in which case a `soma_nodes.label.astype(str)` might
         # throw annoying runtime warnings
-        labels = np.asarray(soma_nodes.label).astype(str)
-        soma_nodes = soma_nodes[labels == str(soma_label)]
+        soma_node_ids = x.nodes.node_id.values[mask]
+        soma_node_labels = np.asarray(x.nodes.label.values[mask]).astype(str)
+
+        return soma_node_ids[soma_node_labels == str(soma_label)]
+    # If no labels to check we can return the mask directly
+    else:
+        return x.nodes.node_id.values[mask]
 
-    return soma_nodes.node_id.values