Use NeuroCollage

requirements/base.pip

@@ -1,4 +1,4 @@
-atlas_analysis>=0.0.3
+atlas_analysis>=0.0.4
 bluepy>=2.3,<3
 bluepy-configfile
 bluepymm
@@ -15,6 +15,7 @@ matplotlib
 morph_tool>=2.9.0,<3
 morphio>=3,<4
 neuroc>=0.2.8,<1
+neurocollage>=0.2.0
 neurom>=3,<4
 pandas
 placement_algorithm>=2.2.0

src/synthesis_workflow/circuit.py

@@ -8,48 +8,14 @@
 from atlas_analysis.planes.planes import create_centerline
 from atlas_analysis.planes.planes import create_planes as _create_planes
 from brainbuilder.app.cells import _place as place
+from neurocollage.planes import get_cells_between_planes
+from neurocollage.planes import slice_n_cells
+from neurocollage.planes import slice_per_mtype
 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 = "left"
-RIGHT = "right"
-
-
-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
-    of the annotated volume are zeroed depending on which `side` is chosen.
-
-    Args:
-        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 2.
-        side: (Optional) Either 'left' or 'right', depending on which half is requested.
-            Defaults to LEFT.
-
-    Returns:
-        Halves `annotated_volume` where where voxels on the opposite `side` have been
-        zeroed (black).
-    """
-    assert axis in range(3)
-    assert side in [LEFT, RIGHT]
-
-    middle = annotated_volume.shape[axis] // 2
-    slices_ = [slice(0), slice(0), slice(0)]
-    for coord in range(3):
-        if axis == coord:
-            slices_[coord] = (
-                slice(0, middle) if side == RIGHT else slice(middle, annotated_volume.shape[axis])
-            )
-        else:
-            slices_[coord] = slice(0, annotated_volume.shape[coord])
-    annotated_volume[slices_[0], slices_[1], slices_[2]] = 0
-    return annotated_volume
 
 
 def create_atlas_thickness_mask(atlas_dir):
@@ -140,34 +106,6 @@ def build_circuit(
     )
 
 
-def slice_per_mtype(cells, mtypes):
-    """Selects cells of given mtype."""
-    return cells[cells["mtype"].isin(mtypes)]
-
-
-def slice_n_cells(cells, n_cells, random_state=0):
-    """Selects n_cells random cells per mtypes."""
-    if n_cells <= 0:
-        return cells
-    sampled_cells = pd.DataFrame()
-    for mtype in cells.mtype.unique():
-        samples = cells[cells.mtype == mtype].sample(
-            n=min(n_cells, len(cells[cells.mtype == mtype])), random_state=random_state
-        )
-        sampled_cells = sampled_cells.append(samples)
-    return sampled_cells
-
-
-def get_cells_between_planes(cells, plane_left, plane_right):
-    """Gets cells gids between two planes in equation representation."""
-    eq_left = plane_left.get_equation()
-    eq_right = plane_right.get_equation()
-    left = np.einsum("j,ij", eq_left[:3], cells[["x", "y", "z"]].values)
-    right = np.einsum("j,ij", eq_right[:3], cells[["x", "y", "z"]].values)
-    selected = (left > eq_left[3]) & (right < eq_right[3])
-    return cells.loc[selected]
-
-
 def circuit_slicer(cells, n_cells, mtypes=None, planes=None, hemisphere=None):
     """Selects n_cells mtype in mtypes."""
     if mtypes is not None:
@@ -184,28 +122,28 @@ def circuit_slicer(cells, n_cells, mtypes=None, planes=None, hemisphere=None):
         # between each pair of planes, select n_cells
         return pd.concat(
             [
-                slice_n_cells(get_cells_between_planes(cells, plane_left, plane_right), n_cells)
-                for plane_left, plane_right in tqdm(
-                    zip(planes[:-1:3], planes[2::3]), total=int(len(planes) / 3)
+                slice_n_cells(
+                    get_cells_between_planes(cells, plane["left"], plane["right"]), n_cells
                 )
+                for plane in planes
             ]
         )
     return slice_n_cells(cells, n_cells)
 
 
-def slice_circuit(input_mvd3, output_mvd3, slicer):
-    """Slices an mvd3 file using a slicing function.
+def slice_circuit(input_circuit, output_circuit, slicer):
+    """Slices a circuit file using a slicing function.
 
     Args:
-        input_mvd3 (str): path to input mvd3 file
-        output_mvd3 (str): path to ouput_mvd3 file
+        input_circuit (str): path to input file
+        output_circuit (str): path to output file
         slicer (function): function to slice the cells dataframe
     """
-    cells = CellCollection.load_mvd3(input_mvd3)
+    cells = CellCollection.load(input_circuit)
     sliced_cells = slicer(cells.as_dataframe())
     sliced_cells.reset_index(inplace=True, drop=True)
     sliced_cells.index += 1  # this is to match CellCollection index from 1
-    CellCollection.from_dataframe(sliced_cells).save_mvd3(output_mvd3)
+    CellCollection.from_dataframe(sliced_cells).save(output_circuit)
     return sliced_cells
 
 

src/synthesis_workflow/extra/make_map.py

@@ -5,10 +5,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+from neurocollage.collage import get_annotation_info
 from pyquaternion import Quaternion
 
-from synthesis_workflow.validation import get_annotation_info
-
 matplotlib.use("Agg")
 
 

src/synthesis_workflow/synthesis.py

@@ -316,7 +316,7 @@ def create_axon_morphologies_tsv(
     else:
         L.info("Do not use placement algorithm for axons (use random choice instead)")
 
-        cells_df = CellCollection.load_mvd3(circuit_path).properties
+        cells_df = CellCollection.load(circuit_path).properties
         axon_morphs = pd.DataFrame(index=cells_df.index, columns=["morphology"])
 
         morphs_df = pd.read_csv(morphs_df_path)
@@ -520,6 +520,7 @@ def _process_scaling_rule(
 
     # Build the file list for each mtype
     morphs_df = pd.read_csv(morphs_df_path)
+
     file_lists = [
         (mtype, morphs_df.loc[morphs_df.mtype == mtype, morphology_path].to_list())
         for mtype in scaling_rules.keys()

src/synthesis_workflow/tasks/circuit.py

@@ -1,24 +1,25 @@
 """Luigi tasks for circuit and atlas processings."""
+import pickle
+from copy import deepcopy
 from functools import partial
+from pathlib import Path
 
 import luigi
+import pandas as pd
 import yaml
-from atlas_analysis.planes.planes import load_planes_centerline
-from atlas_analysis.planes.planes import save_planes_centerline
 from luigi.parameter import OptionalPathParameter
 from luigi.parameter import PathParameter
 from luigi_tools.parameter import RatioParameter
 from luigi_tools.task import ParamRef
 from luigi_tools.task import WorkflowTask
 from luigi_tools.task import copy_params
+from neurocollage import create_planes
+from neurocollage.planes import get_layer_annotation
 from voxcell import VoxelData
 
 from synthesis_workflow.circuit import build_circuit
 from synthesis_workflow.circuit import circuit_slicer
 from synthesis_workflow.circuit import create_atlas_thickness_mask
-from synthesis_workflow.circuit import create_planes
-from synthesis_workflow.circuit import get_layer_tags
-from synthesis_workflow.circuit import halve_atlas
 from synthesis_workflow.circuit import slice_circuit
 from synthesis_workflow.tasks.config import AtlasLocalTarget
 from synthesis_workflow.tasks.config import CircuitConfig
@@ -41,20 +42,21 @@ def requires(self):
 
     def run(self):
         """ """
-        annotation, layer_mapping = get_layer_tags(
-            CircuitConfig().atlas_path,
-            self.input().pathlib_path / CircuitConfig().region_structure_path,
+
+        annotation = get_layer_annotation(
+            {
+                "atlas": CircuitConfig().atlas_path,
+                "structure": self.input().pathlib_path / CircuitConfig().region_structure_path,
+            },
             CircuitConfig().region,
+            CircuitConfig().hemisphere,
         )
-        if CircuitConfig().hemisphere is not None:
-            annotation.raw = halve_atlas(annotation.raw, side=CircuitConfig().hemisphere)
-
         annotation_path = self.output()["annotations"].path
-        annotation.save_nrrd(annotation_path)
+        annotation["annotation"].save_nrrd(annotation_path)
 
         layer_mapping_path = self.output()["layer_mapping"].path
         with open(layer_mapping_path, "w", encoding="utf-8") as f:
-            yaml.dump(layer_mapping, f)
+            yaml.dump(annotation["mapping"], f)
 
     def output(self):
         """ """
@@ -120,8 +122,10 @@ def run(self):
             raise Exception("Number of planes should be larger than one")
 
         layer_annotation = VoxelData.load_nrrd(self.input()["annotations"].path)
+        layer_mappings = yaml.safe_load(self.input()["layer_mapping"].open())
+
         planes, centerline = create_planes(
-            layer_annotation,
+            {"annotation": layer_annotation, "mapping": layer_mappings},
             self.plane_type,
             self.plane_count,
             self.slice_thickness,
@@ -130,7 +134,8 @@ def run(self):
             self.centerline_axis,
         )
 
-        save_planes_centerline(self.output().path, planes, centerline)
+        with open(self.output().path, "wb") as f_planes:
+            pickle.dump({"planes": planes, "centerline": centerline}, f_planes)
 
     def output(self):
         """ """
@@ -173,6 +178,30 @@ def run(self):
             thickness_mask.save_nrrd(self.mask_path)
             thickness_mask_path = self.mask_path.stem
 
+        # create uniform densities if they don't exist
+        mtypes = SynthesisConfig().mtypes
+        if not mtypes:
+            mtypes = pd.read_csv(mtype_taxonomy_path, sep="\t")["mtype"].to_list()
+
+        for mtype in mtypes:
+            nrrd_path = Path(CircuitConfig().atlas_path) / f"[cell_density]{mtype.upper()}.nrrd"
+            if not nrrd_path.exists():
+                annotation = get_layer_annotation(
+                    {
+                        "atlas": CircuitConfig().atlas_path,
+                        "structure": self.input()["synthesis"].pathlib_path
+                        / CircuitConfig().region_structure_path,
+                    },
+                    CircuitConfig().region,
+                    CircuitConfig().hemisphere,
+                )
+                layer = mtype[1]
+                keys = [k + 1 for k, d in annotation["mapping"].items() if d.endswith(layer)]
+                density_annotation = deepcopy(annotation["annotation"])
+                density_annotation.raw[annotation["annotation"].raw == keys[0]] = 1000
+                density_annotation.raw[annotation["annotation"].raw != keys[0]] = 0
+                density_annotation.save_nrrd(str(nrrd_path))
+
         cells = build_circuit(
             self.input()["composition"].path,
             mtype_taxonomy_path,
@@ -193,15 +222,15 @@ def output(self):
     mtypes=ParamRef(SynthesisConfig),
 )
 class SliceCircuit(WorkflowTask):
-    """Create a smaller circuit .mvd3 file for subsampling.
+    """Create a smaller circuit file for subsampling.
 
     Attributes:
         mtypes (list): List of mtypes to consider.
     """
 
     sliced_circuit_path = PathParameter(
-        default="sliced_circuit_somata.mvd3",
-        description=":str: Path to save sliced circuit somata mvd3.",
+        default="sliced_circuit_somata.h5",
+        description=":str: Path to save sliced circuit.",
     )
     n_cells = luigi.IntParameter(
         default=10, description=":int: Number of cells per mtype to consider."
@@ -216,7 +245,8 @@ def requires(self):
 
     def run(self):
         """ """
-        planes = load_planes_centerline(self.input()["atlas_planes"].path)["planes"]
+        with open(self.input()["atlas_planes"].path, "rb") as f_planes:
+            planes = pickle.load(f_planes)["planes"]
 
         _slicer = partial(
             circuit_slicer,

src/synthesis_workflow/tasks/config.py

@@ -9,7 +9,6 @@
 import yaml
 from git import Repo
 from luigi.parameter import OptionalChoiceParameter
-from luigi.parameter import OptionalIntParameter
 from luigi_tools.parameter import ExtParameter
 from luigi_tools.target import OutputLocalTarget
 from luigi_tools.task import WorkflowTask
@@ -295,12 +294,6 @@ class PathConfig(luigi.Config):
     )
 
 
-class ValidationConfig(luigi.Config):
-    """Validation configuration."""
-
-    sample = OptionalIntParameter(default=None)
-
-
 class AtlasLocalTarget(OutputLocalTarget):
     """Specific target for atlas targets."""
 

src/synthesis_workflow/tasks/synthesis.py

@@ -458,8 +458,8 @@ class Synthesize(WorkflowTask):
     """
 
     out_circuit_path = PathParameter(
-        default="sliced_circuit_morphologies.mvd3",
-        description=":str: Path to circuit mvd3 with morphology data.",
+        default="circuit.h5",
+        description=":str: Path to circuit with morphology data.",
     )
     axon_morphs_base_dir = luigi.OptionalParameter(
         default=None,
@@ -526,7 +526,7 @@ def requires(self):
 
     def run(self):
         """ """
-        out_mvd3 = self.output()["out_mvd3"]
+        circuit = self.output()["circuit"]
         out_morphologies = self.output()["out_morphologies"]
         out_apical_points = self.output()["apical_points"]
         debug_scales = self.output().get("debug_scales")
@@ -553,7 +553,7 @@ def run(self):
             "tmd_parameters": self.input()["tmd_parameters"].path,
             "tmd_distributions": self.input()["tmd_distributions"].path,
             "atlas": CircuitConfig().atlas_path,
-            "out_cells": out_mvd3.path,
+            "out_cells": circuit.path,
             "out_apical": out_apical_points.path,
             "out_morph_ext": [str(self.ext)],
             "out_morph_dir": out_morphologies.path,
@@ -580,7 +580,7 @@ def run(self):
     def output(self):
         """ """
         outputs = {
-            "out_mvd3": SynthesisLocalTarget(self.out_circuit_path),
+            "circuit": SynthesisLocalTarget(self.out_circuit_path),
             "out_morphologies": SynthesisLocalTarget(PathConfig().synth_output_path),
             "apical_points": SynthesisLocalTarget(self.apical_points_path),
         }