Implement older monocot pipeline

sleap_roots/bases.py

@@ -72,19 +72,6 @@ def get_base_tip_dist(
     return distances
 
 
-def get_lateral_count(pts: np.ndarray):
-    """Get number of lateral roots.
-
-    Args:
-        pts: lateral root landmarks as array of shape `(instance, node, 2)`.
-
-    Return:
-        Scalar of number of lateral roots.
-    """
-    lateral_count = pts.shape[0]
-    return lateral_count
-
-
 def get_base_xs(pts: np.ndarray, monocots: bool = False) -> np.ndarray:
     """Get x coordinates of the base of each lateral root.
 
@@ -247,9 +234,9 @@ def get_base_median_ratio(lateral_base_ys, primary_tip_pt_y, monocots: bool = Fa
     """Get ratio of median value in all base points to tip of primary root in y axis.
 
     Args:
-        lateral_base_ys: y-coordinates of the base points of lateral roots of shape
+        lateral_base_ys: Y-coordinates of the base points of lateral roots of shape
             `(instances,)`.
-        primary_tip_pt_y: y-coordinate of the tip point of the primary root of shape
+        primary_tip_pt_y: Y-coordinate of the tip point of the primary root of shape
             `(1)`.
         monocots: Boolean value, where false is dicot (default), true is rice.
 

sleap_roots/lengths.py

@@ -1,7 +1,7 @@
 """Get length-related traits."""
 import numpy as np
 from sleap_roots.bases import get_base_tip_dist
-from typing import Optional
+from typing import Union
 
 
 def get_max_length_pts(pts: np.ndarray) -> np.ndarray:
@@ -69,7 +69,8 @@ def get_root_lengths(pts: np.ndarray) -> np.ndarray:
     segment_lengths = np.linalg.norm(segment_diffs, axis=-1)
     # Add the segments together to get the total length using nansum
     total_lengths = np.nansum(segment_lengths, axis=-1)
-    # Find the NaN segment lengths and record NaN in place of 0 when finding the total length
+    # Find the NaN segment lengths and record NaN in place of 0 when finding the total
+    # length
     total_lengths[np.isnan(segment_lengths).all(axis=-1)] = np.nan
 
     # If there is 1 instance, return a scalar instead of an array of length 1
@@ -113,57 +114,44 @@ def get_root_lengths_max(pts: np.ndarray) -> np.ndarray:
 
 
 def get_grav_index(
-    primary_length: Optional[float] = None,
-    primary_base_tip_dist: Optional[float] = None,
-    pts: Optional[np.ndarray] = None,
-) -> float:
-    """Calculate the gravitropism index of a primary root.
+    lengths: Union[float, np.ndarray],
+    base_tip_dists: Union[float, np.ndarray],
+) -> Union[float, np.ndarray]:
+    """Calculate the gravitropism index of a root.
 
     The gravitropism index quantifies the curviness of the root's growth. A higher
     gravitropism index indicates a curvier root (less responsive to gravity), while a
     lower index indicates a straighter root (more responsive to gravity). The index is
-    computed as the difference between the maximum primary root length and straight-line
-    distance from the base to the tip of the primary root, normalized by the root length.
+    computed as the difference between the maximum root length and straight-line
+    distance from the base to the tip of the root, normalized by the root length.
 
     Args:
-        primary_length: Maximum length of the primary root. Used if `pts` is not
-            provided.
-        primary_base_tip_dist: The straight-line distance from the base to the tip of
-        the primary root. Used if `pts` is not provided.
-        pts: Landmarks of the primary root of shape `(instances, nodes, 2)`. If
-        provided, `primary_length` and `primary_base_tip_dist` are ignored.
+        lengths: Maximum length(s) of the root(s).
+        base_tip_dists: The straight-line distance(s) from the base to the tip of the
+        root(s).
 
     Returns:
-        float: Gravitropism index of the primary root, quantifying its curviness.
+        float or np.ndarray: Gravitropism index of the root(s), quantifying its (their)
+        curviness.
     """
-    # Use provided scalar values if available
-    if primary_length is not None and primary_base_tip_dist is not None:
-        max_primary_length = primary_length
-        max_base_tip_distance = primary_base_tip_dist
-
-    # Use provided pts array to compute required values if available
-    elif pts is not None:
-        if np.isnan(pts).all():
-            return np.nan
-        primary_length_max = get_root_lengths_max(pts=pts)
-        primary_base_tip_dist = get_base_tip_dist(pts=pts)
-        max_primary_length = np.nanmax(primary_length_max)
-        max_base_tip_distance = np.nanmax(primary_base_tip_dist)
-
-    else:
-        raise ValueError(
-            "Either both primary_length and primary_base_tip_dist, or pts"
-            "must be provided."
-        )
-
-    # Check for invalid values (NaN or zero lengths)
-    if (
-        np.isnan(max_primary_length)
-        or np.isnan(max_base_tip_distance)
-        or max_primary_length == 0
-    ):
-        return np.nan
-
-    # Calculate and return gravitropism index
-    grav_index = (max_primary_length - max_base_tip_distance) / max_primary_length
-    return grav_index
+    # Convert inputs to NumPy arrays for element-wise operations
+    lengths = np.asarray(lengths)
+    base_tip_dists = np.asarray(base_tip_dists)
+
+    # Initialize an array to store the gravitropism index, filled with NaN
+    grav_index = np.full(lengths.shape, np.nan)
+
+    # Identify valid and invalid values
+    valid_values = ~np.isnan(lengths) & ~np.isnan(base_tip_dists) & (lengths != 0)
+
+    # Calculate gravitropism index only for valid values
+    grav_index[valid_values] = (
+        lengths[valid_values] - base_tip_dists[valid_values]
+    ) / lengths[valid_values]
+
+    # If the input was a float, return a float; otherwise, return the NumPy array.
+    return (
+        grav_index
+        if grav_index.size > 1
+        else (grav_index.item() if valid_values else np.nan)
+    )

sleap_roots/networklength.py

@@ -3,7 +3,7 @@
 import numpy as np
 from shapely import LineString, Polygon
 from sleap_roots.lengths import get_root_lengths, get_max_length_pts
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Union, List
 
 
 def get_bbox(pts: np.ndarray) -> Tuple[float, float, float, float]:
@@ -60,44 +60,42 @@ def get_network_width_depth_ratio(
 
 
 def get_network_length(
-    primary_length: float,
-    lateral_lengths: Union[float, np.ndarray],
-    monocots: bool = False,
+    lengths0: Union[float, np.ndarray],
+    *args: Optional[Union[float, np.ndarray]],
 ) -> float:
     """Return the total root network length given primary and lateral root lengths.
 
     Args:
-        primary_length: Primary root length.
-        lateral_lengths: Either a float representing the length of a single lateral
-          root or an array of lateral root lengths with shape `(instances,)`.
-        monocots: A boolean value, where True is rice.
+        lengths0: Either a float representing the length of a single
+            root or an array of root lengths with shape `(instances,)`.
+        *args: Additional optional floats representing the lengths of single
+            roots or arrays of root lengths with shape `(instances,)`.
 
     Returns:
         Total length of root network.
     """
-    # Ensure primary_length is a scalar
-    if not isinstance(primary_length, (float, np.float64)):
-        raise ValueError("Input primary_length must be a scalar value.")
-
-    # Ensure lateral_lengths is either a scalar or has the correct shape
-    if not (
-        isinstance(lateral_lengths, (float, np.float64)) or lateral_lengths.ndim == 1
-    ):
-        raise ValueError(
-            "Input lateral_lengths must be a scalar or have shape (instances,)."
-        )
-
-    # Calculate the total lateral root length using np.nansum
-    total_lateral_length = np.nansum(lateral_lengths)
-
-    if monocots:
-        length = total_lateral_length
-    else:
-        # Calculate the total root network length using np.nansum so the total length
-        # will not be NaN if one of primary or lateral lengths are NaN
-        length = np.nansum([primary_length, total_lateral_length])
-
-    return length
+    # Initialize an empty list to store the lengths
+    all_lengths = []
+    # Loop over the input arrays
+    for length in [lengths0] + list(args):
+        if length is None:
+            continue  # Skip None values
+        # Ensure length is either a scalar or has the correct shape
+        if not (np.isscalar(length) or (hasattr(length, "ndim") and length.ndim == 1)):
+            raise ValueError(
+                "Input length must be a scalar or have shape (instances,)."
+            )
+        # Add the length to the list
+        if np.isscalar(length):
+            all_lengths.append(length)
+        else:
+            all_lengths.extend(list(length))
+
+    # Calculate the total root network length using np.nansum so the total length
+    # will not be NaN if one of primary or lateral lengths are NaN
+    total_network_length = np.nansum(all_lengths)
+
+    return total_network_length
 
 
 def get_network_solidity(
@@ -121,60 +119,34 @@ def get_network_solidity(
 
 
 def get_network_distribution(
-    primary_pts: np.ndarray,
-    lateral_pts: np.ndarray,
+    pts_list: List[np.ndarray],
     bounding_box: Tuple[float, float, float, float],
     fraction: float = 2 / 3,
-    monocots: bool = False,
 ) -> float:
     """Return the root length in the lower fraction of the plant.
 
     Args:
-        primary_pts: Array of primary root landmarks. Can have shape `(nodes, 2)` or
-            `(1, nodes, 2)`.
-        lateral_pts: Array of lateral root landmarks with shape `(instances, nodes, 2)`.
+        pts_list: A list of arrays, each having shape `(nodes, 2)`.
         bounding_box: Tuple in the form `(left_x, top_y, width, height)`.
         fraction: Lower fraction value. Defaults to 2/3.
-        monocots: A boolean value, where True indicates rice. Defaults to False.
 
     Returns:
         Root network length in the lower fraction of the plant.
     """
-    # Input validation
-    if primary_pts.ndim not in [2, 3]:
+    # Input validation for pts_list
+    if any(pts.ndim != 2 or pts.shape[-1] != 2 for pts in pts_list):
         raise ValueError(
-            "primary_pts should have a shape of `(nodes, 2)` or `(1, nodes, 2)`."
+            "Each pts array in pts_list should have a shape of `(nodes, 2)`."
         )
 
-    if primary_pts.ndim == 2 and primary_pts.shape[-1] != 2:
-        raise ValueError("primary_pts should have a shape of `(nodes, 2)`.")
-
-    if primary_pts.ndim == 3 and primary_pts.shape[-1] != 2:
-        raise ValueError("primary_pts should have a shape of `(1, nodes, 2)`.")
-
-    if lateral_pts.ndim != 3 or lateral_pts.shape[-1] != 2:
-        raise ValueError("lateral_pts should have a shape of `(instances, nodes, 2)`.")
-
+    # Input validation for bounding_box
     if len(bounding_box) != 4:
         raise ValueError(
-            "bounding_box should be in the form `(left_x, top_y, width, height)`."
+            "bounding_box must contain exactly 4 elements: `(left_x, top_y, width, height)`."
         )
 
-    # Make sure the longest primary root is used
-    if primary_pts.ndim == 3:
-        primary_pts = get_max_length_pts(primary_pts)  # shape is (nodes, 2)
-
-    # Make primary_pts and lateral_pts have the same dimension of 3
-    primary_pts = (
-        primary_pts[np.newaxis, :, :] if primary_pts.ndim == 2 else primary_pts
-    )
-
     # Filter out NaN values
-    primary_pts = [root[~np.isnan(root).any(axis=1)] for root in primary_pts]
-    lateral_pts = [root[~np.isnan(root).any(axis=1)] for root in lateral_pts]
-
-    # Collate root points.
-    all_roots = primary_pts + lateral_pts if not monocots else lateral_pts
+    pts_list = [pts[~np.isnan(pts).any(axis=-1)] for pts in pts_list]
 
     # Get the vertices of the bounding box
     left_x, top_y, width, height = bounding_box
@@ -185,7 +157,6 @@ def get_network_distribution(
         return np.nan
 
     # Convert lower bounding box to polygon
-    # Vertices are in counter-clockwise order
     lower_box = Polygon(
         [
             [left_x, top_y + (height - lower_height)],
@@ -197,7 +168,9 @@ def get_network_distribution(
 
     # Calculate length of roots within the lower bounding box
     network_length = 0
-    for root in all_roots:
+    for root in pts_list:
+        if root.shape[0] < 2:  # Skip if fewer than two points
+            continue
         root_poly = LineString(root)
         lower_intersection = root_poly.intersection(lower_box)
         root_length = lower_intersection.length
@@ -207,53 +180,27 @@ def get_network_distribution(
 
 
 def get_network_distribution_ratio(
-    primary_length: float,
-    lateral_lengths: Union[float, np.ndarray],
+    network_length: float,
     network_length_lower: float,
-    fraction: float = 2 / 3,
-    monocots: bool = False,
 ) -> float:
-    """Return ratio of the root length in the lower fraction over all root length.
+    """Return ratio of the root length in the lower fraction to total root length.
 
     Args:
-        primary_length: Primary root length.
-        lateral_lengths: Lateral root lengths. Can be a single float (for one root)
-            or an array of floats (for multiple roots).
         network_length_lower: The root length in the lower network.
-        fraction: The fraction of the network considered as 'lower'. Defaults to 2/3.
-        monocots: A boolean value, where True indicates rice. Defaults to False.
+        network_length: Total root length of network.
 
     Returns:
         Float of ratio of the root network length in the lower fraction of the plant
-        over all root length.
+            over the total root length.
     """
     # Ensure primary_length is a scalar
-    if not isinstance(primary_length, (float, np.float64)):
-        raise ValueError("Input primary_length must be a scalar value.")
-
-    # Ensure lateral_lengths is either a scalar or a 1-dimensional array
-    if not isinstance(lateral_lengths, (float, np.float64, np.ndarray)):
-        raise ValueError(
-            "Input lateral_lengths must be a scalar or a 1-dimensional array."
-        )
-
-    # If lateral_lengths is an ndarray, it must be one-dimensional
-    if isinstance(lateral_lengths, np.ndarray) and lateral_lengths.ndim != 1:
-        raise ValueError("Input lateral_lengths array must have shape (instances,).")
+    if not isinstance(network_length, (float, np.float64)):
+        raise ValueError("Input network_length must be a scalar value.")
 
     # Ensure network_length_lower is a scalar
     if not isinstance(network_length_lower, (float, np.float64)):
         raise ValueError("Input network_length_lower must be a scalar value.")
 
-    # Calculate the total lateral root length
-    total_lateral_length = np.nansum(lateral_lengths)
-
-    # Determine total root length based on monocots flag
-    if monocots:
-        total_root_length = total_lateral_length
-    else:
-        total_root_length = np.nansum([primary_length, total_lateral_length])
-
     # Calculate the ratio
-    ratio = network_length_lower / total_root_length
+    ratio = network_length_lower / network_length
     return ratio