Any region handling (v2)

requirements/base.pip

@@ -17,7 +17,7 @@ neurom>=3,<4
 pandas
 placement_algorithm>=2.2.0
 PyYAML
-region_grower>=0.3.1
+region_grower>=0.4.0
 scipy
 seaborn
 tmd

src/synthesis_workflow/circuit.py

@@ -3,18 +3,20 @@
 
 import numpy as np
 import pandas as pd
+import yaml
 from atlas_analysis.planes.planes import create_planes as _create_planes
 from brainbuilder.app.cells import _place as place
+from region_grower.atlas_helper import AtlasHelper
 from tqdm import tqdm
 from voxcell import CellCollection
 from voxcell.nexus.voxelbrain import LocalAtlas
 
 L = logging.getLogger(__name__)
-LEFT = 0
-RIGHT = 1
+LEFT = "left"
+RIGHT = "right"
 
 
-def halve_atlas(annotated_volume, axis=0, side=LEFT):
+def halve_atlas(annotated_volume, axis=2, side=LEFT):
     """Return the half of the annotated volume along the x-axis.
 
     The identifiers of the voxels located on the left half or the right half
@@ -24,8 +26,8 @@ def halve_atlas(annotated_volume, axis=0, side=LEFT):
         annotated_volume: integer array of shape (W, L, D) holding the annotation
             of a brain region.
         axis: (Optional) axis along which to halve. Either 0, 1 or 2.
-            Defaults to 0.
-        side: (Optional) Either LEFT or RIGHT, depending on which half is requested.
+            Defaults to 2.
+        side: (Optional) Either 'left' or 'right', depending on which half is requested.
             Defaults to LEFT.
 
     Returns:
@@ -85,6 +87,22 @@ def create_atlas_thickness_mask(atlas_dir):
     return brain_regions.with_data(np.logical_and(~too_thick, isocortex_mask).astype(np.uint8))
 
 
+def get_regions_from_composition(cell_composition_path):
+    """Get list of region regex from cell_composition."""
+    with open(cell_composition_path, "r") as comp_p:
+        cell_composition = yaml.safe_load(comp_p)
+
+    region_regex = "@"
+    for reg in set(entry["region"] for entry in cell_composition["neurons"]):
+        if reg[-1] != "|":
+            reg += "|"
+        if reg[0] == "@":
+            region_regex += reg[1:]
+        else:
+            region_regex += reg
+    return region_regex[:-1]
+
+
 def build_circuit(
     cell_composition_path,
     mtype_taxonomy_path,
@@ -109,7 +127,7 @@ def build_circuit(
         atlas_url=atlas_path,
         mini_frequencies_path=None,
         atlas_cache=None,
-        region=region,
+        region=region or get_regions_from_composition(cell_composition_path),
         mask_dset=mask,
         density_factor=density_factor,
         soma_placement="basic",
@@ -153,8 +171,12 @@ def circuit_slicer(cells, n_cells, mtypes=None, planes=None, hemisphere=None):
     if mtypes is not None:
         cells = slice_per_mtype(cells, mtypes)
 
-    if hemisphere is not None and "hemisphere" in cells:
-        cells = cells[cells.hemisphere == hemisphere]
+    # TODO: rough way to split hemisphere, maybe there is a better way, to investigate if needed
+    if hemisphere is not None:
+        if hemisphere == "left":
+            cells = cells[cells.z < cells.z.mean()]
+        if hemisphere == "right":
+            cells = cells[cells.z >= cells.z.mean()]
 
     if planes is not None:
         # between each pair of planes, select n_cells
@@ -198,31 +220,25 @@ def _get_principal_direction(points):
     return v[:, w.argmax()]
 
 
-def get_centerline_bounds(atlas, layer, region=None, central_layer=5, hemisphere=None):
+def get_centerline_bounds(layer):
     """Find centerline bounds using PCA of the voxell position of a given layer in the region."""
-    if region is not None:
-        mask = atlas.get_region_mask(region)
-        mask.raw = np.array(mask.raw, dtype=int)
-        layer.raw = layer.raw * mask.raw
-
-    _ids = np.where(layer.raw == central_layer)
-    ids = np.column_stack(_ids)
+    _ls = np.unique(layer.raw[layer.raw > 0])
+    central_layer = _ls[int(len(_ls) / 2)]
+    ids = np.column_stack(np.where(layer.raw == central_layer))
     points = layer.indices_to_positions(ids)
-
-    ids = np.array(ids)
-    points = np.array(points)
-    if hemisphere is not None:
-        if hemisphere == "left":
-            point_mask = points[:, 2] < points[:, 2].mean()
-        if hemisphere == "right":
-            point_mask = points[:, 2] > points[:, 2].mean()
-        ids = ids[point_mask]
-        points = points[point_mask]
-    direction = _get_principal_direction(points)
-    _align = points.dot(direction)
+    _align = points.dot(_get_principal_direction(points))
     return ids[_align.argmin()], ids[_align.argmax()]
 
 
+def get_local_bbox(annotation):
+    """Compute bbox where annotation file is strictly positive."""
+    ids = np.where(annotation.raw > 0)
+    dim = annotation.voxel_dimensions
+    return annotation.offset + np.array(
+        [np.min(ids, axis=1) * dim, (np.max(ids, axis=1) + 1) * dim]
+    )
+
+
 def create_planes(
     layer_annotation,
     plane_type="aligned",
@@ -254,6 +270,8 @@ def create_planes(
         centerline_axis (str): (for plane_type = aligned) axis along which to create planes
     """
     if plane_type == "centerline":
+        if centerline_first_bound is None and centerline_last_bound is None:
+            centerline_first_bound, centerline_last_bound = get_centerline_bounds(layer_annotation)
         centerline = np.array(
             [
                 layer_annotation.indices_to_positions(centerline_first_bound),
@@ -262,13 +280,10 @@ def create_planes(
         )
 
     elif plane_type == "aligned":
-        _n_points = 10
-        centerline = np.zeros([_n_points, 3])
-        bbox = layer_annotation.bbox
+        centerline = np.zeros([2, 3])
+        bbox = get_local_bbox(layer_annotation)
         centerline[:, centerline_axis] = np.linspace(
-            bbox[0, centerline_axis],
-            bbox[1, centerline_axis],
-            _n_points,
+            bbox[0, centerline_axis], bbox[1, centerline_axis], 2
         )
     else:
         raise Exception(f"Please set plane_type to 'aligned' or 'centerline', not {plane_type}.")
@@ -285,3 +300,28 @@ def create_planes(
     planes_select_ids += list(planes_all_ids[int(id_shift / 2) - 1 :: id_shift])
     planes_select_ids += list(planes_all_ids[int(id_shift / 2) + 1 :: id_shift])
     return [planes[i] for i in sorted(planes_select_ids)], centerline
+
+
+def get_layer_tags(atlas_dir, region_structure_path, region=None):
+    """Create a VoxelData with layer tags."""
+    atlas_helper = AtlasHelper(LocalAtlas(atlas_dir), region_structure_path=region_structure_path)
+
+    brain_regions = atlas_helper.brain_regions
+    layers_data = np.zeros_like(brain_regions.raw, dtype="uint8")
+
+    region_mask = None
+    if region is not None:
+        region_mask = atlas_helper.atlas.get_region_mask(region).raw
+
+    layer_mapping = {}
+    for layer_id, layer in enumerate(atlas_helper.layers):
+        layer_mapping[layer_id] = atlas_helper.region_structure["names"].get(layer, str(layer))
+        region_query = atlas_helper.region_structure["region_queries"].get(layer, None)
+        mask = atlas_helper.atlas.get_region_mask(region_query).raw
+        if region_mask is not None:
+            mask *= region_mask
+        layers_data[mask] = layer_id + 1
+        if not len(layers_data[mask]):
+            L.warning("No voxel found for layer %s.", layer)
+    brain_regions.raw = layers_data
+    return brain_regions, layer_mapping

src/synthesis_workflow/extra/insitu_validation.py

@@ -17,7 +17,7 @@
 from voxcell import CellCollection
 from voxcell.nexus.voxelbrain import Atlas
 
-from synthesis_workflow.tools import get_layer_tags
+from synthesis_workflow.circuit import get_layer_tags
 
 matplotlib.use("Agg")
 
@@ -72,6 +72,7 @@ def plot_extents(results_df, pdf_name="extents.pdf", bins=200):
             pdf.savefig()
 
 
+# pylint: disable=too-many-locals
 def get_layer_morpho_counts(
     sonata_path,
     morphology_path,
@@ -81,6 +82,7 @@ def get_layer_morpho_counts(
     region=None,
     hemisphere="right",
     ext=".h5",
+    region_structure_path="region_structure.yaml",
 ):
     """For each layer, compute the fraction of crossing morphologies.
 
@@ -93,6 +95,7 @@ def get_layer_morpho_counts(
         region (str): is not None, must be the name of a region to filter morphologies
         hemisphere (str): left/right hemisphere
         ext (str): extension of morphology files
+        region_structure_path (str): path to region_structure.yaml file
 
     Returns:
         pandas.DataFrame: dafaframe with the counts, mtype and layers information
@@ -101,7 +104,7 @@ def get_layer_morpho_counts(
     cells = CellCollection.load(sonata_path).as_dataframe()
     morph_path = Path(morphology_path)
 
-    layer_annotation, _ = get_layer_tags(atlas_path)
+    layer_annotation, _ = get_layer_tags(atlas_path, region_structure_path)
 
     dfs = []
     for mtype in tqdm(cells.mtype.unique()):
@@ -205,6 +208,7 @@ def plot_layer_collage(
     shift=200,
     dpi=200,
     ext=".h5",
+    region_structure_path="region_structure.yaml",
 ):
     """Plot morphologies next to each other with points colored by layer.
 
@@ -221,15 +225,16 @@ def plot_layer_collage(
         shift (float): x-shift between morphologies
         dpi (int): dpi to save rasterized pdf
         ext (str): extension of morphology files
+        region_structure_path (str): path to region_structure.yaml file
     """
     if section_types is None:
         section_types = ["basal_dendrite", "apical_dendrite"]
     section_types = [getattr(SectionType, section_type) for section_type in section_types]
 
     cells = CellCollection.load(sonata_path).as_dataframe()
     morph_path = Path(morphology_path)
-    layer_annotation, _ = get_layer_tags(atlas_path)
-    atlas = AtlasHelper(Atlas().open(atlas_path))
+    layer_annotation, _ = get_layer_tags(atlas_path, region_structure_path)
+    atlas = AtlasHelper(Atlas().open(atlas_path), region_structure_path=region_structure_path)
 
     cmap = matplotlib.colors.ListedColormap(["C0", "C1", "C2", "C3", "C4", "C5", "C6"])
     with PdfPages(pdf_name) as pdf:

src/synthesis_workflow/synthesis.py

@@ -8,11 +8,14 @@
 import matplotlib
 import numpy as np
 import pandas as pd
+import yaml
 from joblib import Parallel
 from joblib import delayed
 from morphio.mut import Morphology
 from neuroc.scale import ScaleParameters
 from neuroc.scale import scale_section
+from neurom import load_morphology
+from neurom.check.morphology_checks import has_apical_dendrite
 from neurom.core.dataformat import COLS
 from neurots import extract_input
 from placement_algorithm.app import utils
@@ -29,14 +32,25 @@
 matplotlib.use("Agg")
 
 
-def get_neurite_types(morphs_df, mtypes):
-    """Get the neurite types to consider for PC or IN cells."""
-    return {
-        mtype: ["basal", "axon"]
-        if morphs_df.loc[morphs_df.mtype == mtype, "morph_class"].tolist()[0] == "IN"
-        else ["basal", "apical", "axon"]
-        for mtype in mtypes
-    }
+def get_neurite_types(morphs_df):
+    """Get the neurite types to consider for PC or IN cells by checking if apical exists."""
+
+    def _get_morph_class(path):
+        return "PC" if has_apical_dendrite(load_morphology(path)) else "IN"
+
+    morphs_df["morph_class"] = morphs_df["path"].apply(_get_morph_class)
+    neurite_types = {}
+    for mtype, _df in morphs_df.groupby("mtype"):
+        morph_class = list(set(_df["morph_class"]))
+
+        if len(morph_class) > 1:
+            raise Exception("f{mtype} has not consistent morph_class, we stop here")
+
+        if morph_class[0] == "IN":
+            neurite_types[mtype] = ["basal"]
+        if morph_class[0] == "PC":
+            neurite_types[mtype] = ["basal", "apical"]
+    return neurite_types
 
 
 def apply_substitutions(original_morphs_df, substitution_rules=None):
@@ -72,16 +86,27 @@ def _build_distributions_single_mtype(
     trunk_method="simple",
 ):
     """Internal function for multiprocessing of tmd_distribution building."""
-    morphology_paths = morphs_df.loc[morphs_df.mtype == mtype, morphology_path].to_list()
-    kwargs = {
-        "neurite_types": neurite_types[mtype],
-        "diameter_input_morph": morphology_paths,
-        "diameter_model": diameter_model_function,
-    }
-    if trunk_method != "simple":
-        kwargs["trunk_method"] = trunk_method
+    data = {}
+    for neurite_type in neurite_types[mtype]:
+        if "use_axon" in morphs_df.columns:
+            morphology_paths = morphs_df.loc[
+                (morphs_df.mtype == mtype) & morphs_df.use_axon, morphology_path
+            ].to_list()
+        else:
+            morphology_paths = morphs_df.loc[morphs_df.mtype == mtype, morphology_path].to_list()
+        kwargs = {
+            "neurite_types": neurite_types[mtype],
+            "diameter_input_morph": morphology_paths,
+            "diameter_model": diameter_model_function,
+        }
+        if trunk_method != "simple":
+            kwargs["trunk_method"] = trunk_method
 
-    return mtype, extract_input.distributions(morphology_paths, **kwargs)
+        _data = extract_input.distributions(morphology_paths, **kwargs)
+        data[neurite_type] = _data[neurite_type]
+        data["diameter"] = _data["diameter"]
+        data["soma"] = _data["soma"]
+    return mtype, data
 
 
 def build_distributions(
@@ -90,7 +115,7 @@ def build_distributions(
     neurite_types,
     diameter_model_function,
     morphology_path,
-    cortical_thickness,
+    region_structure_path,
     nb_jobs=-1,
     joblib_verbose=10,
     trunk_method="simple",
@@ -102,7 +127,10 @@ def build_distributions(
         morphs_df (DataFrame): morphology dataframe with reconstruction to use
         diameter_model_function (function): diametrizer function (from diameter-synthesis)
         morphology_path (str): name of the column in morphs_df to use for paths to morphologies
-        cortical_thickness (list): list of cortical thicknesses for reference scaling
+        region_structure_path (str): path to region_structure yaml file with thicknesses
+        nb_jobs (int): number of jobs to run in parallal with joblib
+        joblib_verbose (int): verbose level of joblib
+        trunk_method (str): method to set trunk on soma for synthesis (simple|3d_angle)
 
     Returns:
         dict: dict to save to tmd_distribution.json
@@ -115,11 +143,15 @@ def build_distributions(
         morphology_path=morphology_path,
         trunk_method=trunk_method,
     )
+    thicknesses = []
+    if Path(region_structure_path).exists():
+        with open(region_structure_path, "r") as r_p:
+            region_structure = yaml.safe_load(r_p)
+        thicknesses = [
+            region_structure["thicknesses"][layer] for layer in region_structure["layers"]
+        ]
 
-    tmd_distributions = {
-        "mtypes": {},
-        "metadata": {"cortical_thickness": cortical_thickness},
-    }
+    tmd_distributions = {"mtypes": {}, "metadata": {"cortical_thickness": thicknesses}}
     for mtype, distribution in Parallel(nb_jobs, verbose=joblib_verbose)(
         delayed(build_distributions_single_mtype)(mtype) for mtype in mtypes
     ):