Allow specifying the "root" name

atlas_densities/app/cell_densities.py

@@ -140,8 +140,14 @@ def app(verbose):
     help="Path where to write the output cell density nrrd file."
     "A voxel value is a number of cells per mm^3",
 )
+@click.option(
+    "--root-region-name",
+    type=str,
+    default="root",
+    help="Name of the region in the hierarchy",
+)
 @log_args(L)
-def cell_density(annotation_path, hierarchy_path, nissl_path, output_path):
+def cell_density(annotation_path, hierarchy_path, nissl_path, output_path, root_region_name):
     """Compute and save the overall mouse brain cell density.
 
     The input Nissl stain volume of AIBS is turned into an actual density field complying with
@@ -174,6 +180,7 @@ def cell_density(annotation_path, hierarchy_path, nissl_path, output_path):
         annotation.raw,
         _get_voxel_volume_in_mm3(annotation),
         nissl.raw,
+        root_region_name=root_region_name,
     )
     nissl.with_data(overall_cell_density).save_nrrd(output_path)
 
@@ -380,6 +387,12 @@ def glia_cell_densities(
     help="Path to the directory where to write the output cell density nrrd files."
     " It will be created if it doesn't exist already.",
 )
+@click.option(
+    "--root-region-name",
+    type=str,
+    default="root",
+    help="Name of the region in the hierarchy",
+)
 @log_args(L)
 def inhibitory_and_excitatory_neuron_densities(
     annotation_path,
@@ -389,6 +402,7 @@ def inhibitory_and_excitatory_neuron_densities(
     neuron_density_path,
     inhibitory_neuron_counts_path,
     output_dir,
+    root_region_name,
 ):  # pylint: disable=too-many-arguments
     """Compute and save the inhibitory and excitatory neuron densities.
 
@@ -442,6 +456,7 @@ def inhibitory_and_excitatory_neuron_densities(
         VoxelData.load_nrrd(nrn1_path).raw,
         neuron_density.raw,
         inhibitory_data=inhibitory_data(inhibitory_df),
+        root_region_name=root_region_name,
     )
 
     if not Path(output_dir).exists():

atlas_densities/densities/cell_density.py

@@ -1,25 +1,19 @@
 """Functions to compute the overall mouse brain cell density.
 """
-
-from typing import Dict
+from __future__ import annotations
 
 import numpy as np
 from atlas_commons.typing import AnnotationT, BoolArray, FloatArray
 from voxcell import RegionMap  # type: ignore
 
+from atlas_densities.densities import utils
 from atlas_densities.densities.cell_counts import cell_counts
-from atlas_densities.densities.utils import (
-    compensate_cell_overlap,
-    get_group_ids,
-    get_region_masks,
-    normalize_intensity,
-)
 
 
 def fix_purkinje_layer_intensity(
-    region_map: "RegionMap",
+    group_ids: dict[str, set[int]],
     annotation: AnnotationT,
-    region_masks: Dict[str, BoolArray],
+    region_masks: dict[str, BoolArray],
     cell_intensity: FloatArray,
 ) -> None:
     """
@@ -29,7 +23,8 @@ def fix_purkinje_layer_intensity(
     The array `cell_intensity` is modified in place.
 
     Args:
-        region_map: object to navigate the mouse brain regions hierarchy.
+        group_ids: a dictionary whose keys are group names and whose values are
+            sets of AIBS structure identifiers.
         annotation: integer array of shape (W, H, D) enclosing the AIBS annotation of
             the whole mouse brain.
         region_masks: A dictionary whose keys are region group names and whose values are
@@ -40,8 +35,8 @@ def fix_purkinje_layer_intensity(
             way that the Purkinje layer has a constant intensity value.
     """
 
-    group_ids = get_group_ids(region_map)
     purkinje_layer_mask = np.isin(annotation, list(group_ids["Purkinje layer"]))
+
     # Force Purkinje Layer regions of the Cerebellum group to have a constant intensity
     # equal to the average intensity of the complement.
     # pylint: disable=fixme
@@ -64,6 +59,7 @@ def compute_cell_density(
     annotation: AnnotationT,
     voxel_volume: float,
     nissl: FloatArray,
+    root_region_name: str,
 ) -> FloatArray:
     """
     Compute the overall cell density based on Nissl staining and cell counts from literature.
@@ -101,12 +97,12 @@ def compute_cell_density(
     """
 
     nissl = np.asarray(nissl, dtype=np.float64)
-    nissl = normalize_intensity(nissl, annotation, threshold_scale_factor=1.0, copy=False)
-    nissl = compensate_cell_overlap(nissl, annotation, gaussian_filter_stdv=-1.0, copy=False)
+    nissl = utils.normalize_intensity(nissl, annotation, threshold_scale_factor=1.0, copy=False)
+    nissl = utils.compensate_cell_overlap(nissl, annotation, gaussian_filter_stdv=-1.0, copy=False)
 
-    group_ids = get_group_ids(region_map)
-    region_masks = get_region_masks(group_ids, annotation)
-    fix_purkinje_layer_intensity(region_map, annotation, region_masks, nissl)
+    group_ids = utils.get_group_ids(region_map, root_region_name=root_region_name)
+    region_masks = utils.get_region_masks(group_ids, annotation)
+    fix_purkinje_layer_intensity(group_ids, annotation, region_masks, nissl)
     non_zero_nissl = nissl > 0
     for group, mask in region_masks.items():
         mask = np.logical_and(non_zero_nissl, mask)

atlas_densities/densities/fitting.py

@@ -620,7 +620,7 @@ def linear_fitting(  # pylint: disable=too-many-arguments
 
     L.info("Getting group names ...")
     # We want group region names to be stable under taking descendants
-    groups = get_group_names(region_map, cleanup_rest=True)
+    groups = get_group_names(region_map, cleanup_rest=True, root_region_name=region_name)
 
     L.info("Computing fitting coefficients ...")
     fitting_coefficients = compute_fitting_coefficients(groups, average_intensities, densities)

atlas_densities/densities/glia_densities.py

@@ -1,24 +1,20 @@
 """Functions to compute glia cell densities."""
+from __future__ import annotations
+
 import logging
-from typing import Dict, Set
 
 import numpy as np
 import voxcell
 from atlas_commons.typing import AnnotationT, FloatArray
 
-from atlas_densities.densities.utils import (
-    compensate_cell_overlap,
-    constrain_cell_counts_per_voxel,
-    get_group_ids,
-    normalize_intensity,
-)
+from atlas_densities.densities import utils
 
 L = logging.getLogger(__name__)
 
 
 def compute_glia_cell_counts_per_voxel(  # pylint: disable=too-many-arguments
     glia_cell_count: int,
-    group_ids: Dict[str, Set[int]],
+    region_map: voxcell.RegionMap,
     annotation: AnnotationT,
     glia_intensity: FloatArray,
     cell_counts_per_voxel: FloatArray,
@@ -34,9 +30,7 @@ def compute_glia_cell_counts_per_voxel(  # pylint: disable=too-many-arguments
             cell counts.
     Args:
         glia_cell_count: overall glia cell count found in the scientific literature.
-        group_ids: dictionary whose keys are group names and whose values are
-            sets of AIBS structure ids. These groups are used to identify the voxels
-            with maximum density (fiber tracts) and those of zero density (Purkinje layer).
+        region_map: object to navigate the mouse brain regions hierarchy
         annotation: integer array of shape (W, H, D) enclosing the AIBS annotation of the whole
             mouse brain.
         glia_intensity: float array of shape (W, H, D) with non-negative entries. The input
@@ -53,13 +47,13 @@ def compute_glia_cell_counts_per_voxel(  # pylint: disable=too-many-arguments
         layer).
     """
 
-    fiber_tracts_mask = np.isin(annotation, list(group_ids["Fiber tracts group"]))
+    fiber_tracts_mask = np.isin(annotation, list(utils.get_fiber_tract_ids(region_map)))
     fiber_tracts_free_mask = np.isin(
         annotation,
-        list(group_ids["Purkinje layer"]),
+        list(utils.get_purkinje_layer_ids(region_map)),
     )
 
-    return constrain_cell_counts_per_voxel(
+    return utils.constrain_cell_counts_per_voxel(
         glia_cell_count,
         glia_intensity,
         cell_counts_per_voxel,
@@ -74,11 +68,11 @@ def compute_glia_densities(  # pylint: disable=too-many-arguments
     annotation: AnnotationT,
     voxel_volume: float,
     glia_cell_count: int,
-    glia_intensities: Dict[str, FloatArray],
+    glia_intensities: dict[str, FloatArray],
     cell_density: FloatArray,
-    glia_proportions: Dict[str, str],
+    glia_proportions: dict[str, str],
     copy: bool = False,
-) -> Dict[str, FloatArray]:
+) -> dict[str, FloatArray]:
     """
     Compute the overall glia cell density as well as astrocyte, olgidendrocyte and microglia
     densities.
@@ -131,7 +125,7 @@ def compute_glia_densities(  # pylint: disable=too-many-arguments
         glia_densities[glia_type] = np.asarray(glia_densities[glia_type], dtype=np.float64)
     cell_density = np.asarray(cell_density, dtype=np.float64)
 
-    glia_densities["glia"] = compensate_cell_overlap(
+    glia_densities["glia"] = utils.compensate_cell_overlap(
         np.asarray(glia_densities["glia"], dtype=np.float64),
         annotation,
         gaussian_filter_stdv=-1.0,
@@ -144,19 +138,19 @@ def compute_glia_densities(  # pylint: disable=too-many-arguments
 
     glia_densities["glia"] = compute_glia_cell_counts_per_voxel(
         glia_cell_count,
-        get_group_ids(region_map),
+        region_map,
         annotation,
         glia_densities["glia"],
         cell_density * voxel_volume,
     )
     placed_cells = np.zeros_like(glia_densities["glia"])
     for glia_type in ["astrocyte", "oligodendrocyte"]:
-        glia_densities[glia_type] = normalize_intensity(
+        glia_densities[glia_type] = utils.normalize_intensity(
             np.asarray(glia_densities[glia_type], dtype=np.float64),
             annotation,
             copy=copy,
         )
-        glia_densities[glia_type] = compensate_cell_overlap(
+        glia_densities[glia_type] = utils.compensate_cell_overlap(
             glia_densities[glia_type],
             annotation,
             gaussian_filter_stdv=2.0,
@@ -168,7 +162,7 @@ def compute_glia_densities(  # pylint: disable=too-many-arguments
             glia_type,
             cell_count,
         )
-        glia_densities[glia_type] = constrain_cell_counts_per_voxel(
+        glia_densities[glia_type] = utils.constrain_cell_counts_per_voxel(
             cell_count,
             glia_densities[glia_type],
             glia_densities["glia"] - placed_cells,

atlas_densities/densities/inhibitory_neuron_density.py

@@ -89,6 +89,7 @@ def compute_inhibitory_neuron_density(  # pylint: disable=too-many-arguments
     neuron_density: FloatArray,
     inhibitory_proportion: Optional[float] = None,
     inhibitory_data: Optional[InhibitoryData] = None,
+    root_region_name: Optional[str] = None,
 ) -> FloatArray:
     """
     Compute the inhibitory neuron density using a prescribed neuron count and the overall neuron
@@ -115,6 +116,7 @@ def compute_inhibitory_neuron_density(  # pylint: disable=too-many-arguments
                 assigning the proportion of ihnibitory neurons in each group named by a key string.
             'neuron_count': the total number of inhibitory neurons (float).
             Used only if `inhibitory_proportion` is None.
+        root_region_name(str): Name of the root region in the hierarchy
 
     Returns:
         float64 array of shape (W, H, D) with non-negative entries.
@@ -138,7 +140,7 @@ def compute_inhibitory_neuron_density(  # pylint: disable=too-many-arguments
                 "Either inhibitory_proportion or inhibitory_data should be provided"
                 ". Both are None."
             )
-        group_ids = get_group_ids(region_map)
+        group_ids = get_group_ids(region_map, root_region_name=root_region_name)
         inhibitory_data["region_masks"] = get_region_masks(group_ids, annotation)
     else:
         inhibitory_data = {