Feat: Extent scaling for basals

.pylintrc

@@ -19,7 +19,7 @@ max-args=10
 # with leading underscore
 ignored-argument-names=_.*
 # Maximum number of locals for function / method body
-max-locals=25
+max-locals=27
 # Maximum number of return / yield for function / method body
 max-returns=6
 # Maximum number of branch for function / method body

requirements/base.pip

@@ -23,7 +23,7 @@ numpy>=1.24.1
 pandas>=1.5.3
 placement_algorithm>=2.3.1
 PyYAML>=6
-region_grower>=1.0.0,<2
+region_grower>=1.2.0,<2
 scipy>=1.10
 seaborn>=0.12.2
 tmd>=2.3

src/synthesis_workflow/fit_utils.py

@@ -8,17 +8,21 @@
 from tmd.Population.Population import Population
 
 
-def _get_tmd_feature(input_population: Population, feature: str) -> np.array:
+def _get_tmd_feature(
+    input_population: Population, feature: str, neurite_type: str = "apical_dendrite"
+) -> np.array:
     """Returns a list of features using tmd."""
     f = [
-        tmd.methods.get_persistence_diagram(n.apical_dendrite[0], feature=feature)
+        tmd.methods.get_persistence_diagram(getattr(n, neurite_type)[0], feature=feature)
         for n in input_population.neurons
     ]
 
     return np.array([np.max(p) for p in f])
 
 
-def get_path_distances(input_population: Population) -> np.array:
+def get_path_distances(
+    input_population: Population, neurite_type: str = "apical_dendrite"
+) -> np.array:
     """Returns path distances using tmd.
 
     Args:
@@ -27,10 +31,12 @@ def get_path_distances(input_population: Population) -> np.array:
     Returns:
         list of path distances
     """
-    return _get_tmd_feature(input_population, "path_distances")
+    return _get_tmd_feature(input_population, "path_distances", neurite_type)
 
 
-def get_projections(input_population: Population) -> np.array:
+def get_projections(
+    input_population: Population, neurite_type: str = "apical_dendrite"
+) -> np.array:
     """Returns projections using tmd.
 
     Args:
@@ -39,7 +45,7 @@ def get_projections(input_population: Population) -> np.array:
     Returns:
         list of projections
     """
-    return _get_tmd_feature(input_population, "projection")
+    return _get_tmd_feature(input_population, "projection", neurite_type)
 
 
 def fit_function(x: float, slope: float) -> float:
@@ -74,7 +80,9 @@ def clean_outliers(
 
 
 def fit_path_distance_to_extent(
-    input_population: Population, outlier_percentage: int = 90
+    input_population: Population,
+    outlier_percentage: int = 90,
+    neurite_type: str = "apical_dendrite",
 ) -> Tuple[float, float]:
     """Returns slope and intercept of a linear fit.
 
@@ -85,12 +93,13 @@ def fit_path_distance_to_extent(
     Args:
         input_population: the population of neurons
         outlier_percentage: the percentage used to find and remove outliers
+        neurite_type: neurite_type to make the fit
 
     Returns: slope and intercept of the fit
     """
     # Compute path distances, projections using tmd
-    x = get_projections(input_population)
-    y = get_path_distances(input_population)
+    x = get_projections(input_population, neurite_type=neurite_type)
+    y = get_path_distances(input_population, neurite_type=neurite_type)
 
     # Clean data
     x_clean, y_clean = clean_outliers(x, y, outlier_percentage)

src/synthesis_workflow/synthesis.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import re
+from collections import defaultdict
 from functools import partial
 from pathlib import Path
 
@@ -410,17 +411,20 @@ def rescale_morphologies(
     return morphs_df
 
 
-def _fit_population(mtype, file_names):
+def _fit_population(mtype, neurite_type, file_names):
     # Load neurons
     if len(file_names) > 0:
         input_population = load_population(file_names, use_morphio=True)
     else:
-        return (mtype, None, None)
+        return (mtype, neurite_type, None, None)
 
     # Compute slope and intercept
-    slope, intercept = fit_path_distance_to_extent(input_population)
+    try:
+        slope, intercept = fit_path_distance_to_extent(input_population, neurite_type=neurite_type)
+    except IndexError:
+        return (mtype, neurite_type, None, None)
 
-    return mtype, slope, intercept
+    return mtype, neurite_type, slope, intercept
 
 
 def add_scaling_rules_to_parameters(
@@ -469,13 +473,17 @@ def _process_scaling_rule(
 
     # Add scaling rules to TMD parameters
     default_rules = scaling_rules.get("default") or {}
-
+    neurite_types_map = defaultdict(list)
     for mtype in tmd_parameters.keys():
         mtype_rules = scaling_rules.get(mtype) or {}
         neurite_types = set(list(default_rules.keys()) + list(mtype_rules.keys()))
         for neurite_type in neurite_types:
             default_limits = default_rules.get(neurite_type) or {}
             limits = mtype_rules.get(neurite_type) or {}
+
+            if "extent_to_target" in limits:
+                neurite_types_map[mtype].append(neurite_type)
+
             _process_scaling_rule(
                 tmd_parameters,
                 mtype,
@@ -508,32 +516,36 @@ def _process_scaling_rule(
     morphs_df = pd.read_csv(morphs_df_path)
 
     file_lists = [
-        (mtype, morphs_df.loc[morphs_df.mtype == mtype, morphology_path].to_list())
+        (mtype, neurite_type, morphs_df.loc[morphs_df.mtype == mtype, morphology_path].to_list())
         for mtype in scaling_rules.keys()
+        for neurite_type in neurite_types_map[mtype]
         if mtype != "default"
     ]
 
-    L.debug("Number of files: %s", [(t, len(f)) for t, f in file_lists])
+    L.debug("Number of files: %s", [(t, len(f)) for t, _, f in file_lists])
 
     # Fit data and update TMD parameters
-    for mtype, slope, intercept in Parallel(nb_jobs)(
-        delayed(_fit_population)(mtype, file_names) for mtype, file_names in file_lists
+    for mtype, neurite_type, slope, intercept in Parallel(nb_jobs)(
+        delayed(_fit_population)(mtype, neurite_type, file_names)
+        for mtype, neurite_type, file_names in file_lists
     ):
         if slope is None or intercept is None:
             L.debug(
-                "Fitting parameters not found for %s (slope=%s ; intercept=%s)",
+                "Fitting parameters not found for neurite_type %s and %s (slope=%s ; intercept=%s)",
+                neurite_type,
                 mtype,
                 slope,
                 intercept,
             )
             continue
         context = tmd_parameters[mtype].get("context_constraints", {})
-        neurite_type_params = context.get("apical_dendrite", {}).get("extent_to_target", {})
+        neurite_type_params = context.get(neurite_type, {}).get("extent_to_target", {})
         neurite_type_params.update({"slope": slope, "intercept": intercept})
-        context["apical_dendrite"]["extent_to_target"] = neurite_type_params
+        context[neurite_type]["extent_to_target"] = neurite_type_params
         tmd_parameters[mtype]["context_constraints"] = context
         L.debug(
-            "Fitting parameters for %s: slope=%s ; intercept=%s",
+            "Fitting parameters for neurite_type %s and %s: slope=%s ; intercept=%s",
+            neurite_type,
             mtype,
             slope,
             intercept,

src/synthesis_workflow/validation.py

@@ -598,7 +598,14 @@ def _generate_synthetic_random_population(
     return files, y_synth
 
 
-def _get_fit_population(mtype, files, outlier_percentage, tmd_parameters, tmd_distributions):
+def _get_fit_population(
+    mtype,
+    files,
+    outlier_percentage,
+    tmd_parameters,
+    tmd_distributions,
+    neurite_type="apical_dendrite",
+):
     """Get projections and path lengths of a given and a synthetic population."""
     # Load biological neurons
     return_error = (mtype, None, None, None, None, None, None)
@@ -615,8 +622,8 @@ def _get_fit_population(mtype, files, outlier_percentage, tmd_parameters, tmd_di
         return return_error + (f"No fit for mtype='{mtype}'",)
 
     # Get X and Y from biological population
-    x = get_path_distances(input_population)
-    y = get_projections(input_population)
+    x = get_path_distances(input_population, neurite_type)
+    y = get_projections(input_population, neurite_type)
     x_clean, y_clean = clean_outliers(x, y, outlier_percentage)
 
     # Create synthetic neuron population
@@ -645,6 +652,7 @@ def plot_path_distance_fits(
     region=None,
     outlier_percentage=90,
     nb_jobs=-1,
+    neurite_type="apical_dendrite",
 ):
     """Plot path-distance fits."""
     # Read TMD parameters
@@ -688,6 +696,7 @@ def plot_path_distance_fits(
                 outlier_percentage,
                 tmd_parameters[region][mtype],
                 tmd_distributions[region][mtype],
+                neurite_type=neurite_type,
             )
             for mtype, files in file_lists
         ):