Nsetm 2196 expand output

README.rst

@@ -350,7 +350,7 @@ The molecular density of approx_lamp5 was calculated from the other molecular de
 
 which approximates the molecular density of lamp5.
 
-The command outputs the density files in the output-dir and a legend json file:
+The command outputs the density files in the output-dir and a metadata json file:
 
 .. code-block:: json
 

atlas_densities/app/mtype_densities.py

@@ -49,7 +49,10 @@
 from atlas_densities.densities.mtype_densities_from_map.create import (
     create_from_probability_map as create_from_map,
 )
-from atlas_densities.densities.mtype_densities_from_map.utils import check_probability_map_sanity
+from atlas_densities.densities.mtype_densities_from_map.utils import (
+    SYNAPSE_CLASSES,
+    check_probability_map_sanity,
+)
 from atlas_densities.densities.mtype_densities_from_profiles import DensityProfileCollection
 from atlas_densities.exceptions import AtlasDensitiesError
 
@@ -213,6 +216,12 @@ def _check_config_sanity(config: dict) -> None:
     required=True,
     help="Name and path to marker: ex: --marker pv path/pv.nrrd",
 )
+@click.option(
+    "--synapse-class",
+    type=click.Choice(list(SYNAPSE_CLASSES), case_sensitive=False),
+    required=True,
+    help="Target synapse class, the other will be skipped.",
+)
 @click.option(
     "--output-dir",
     required=True,
@@ -226,11 +235,12 @@ def _check_config_sanity(config: dict) -> None:
     help="Number of jobs to run in parallel.",
 )
 @log_args(L)
-def create_from_probability_map(
+def create_from_probability_map(  # pylint: disable=too-many-arguments
     annotation_path,
     hierarchy_path,
     probability_map,
     marker,
+    synapse_class,
     output_dir,
     n_jobs,
 ):  # pylint: disable=too-many-locals
@@ -251,7 +261,9 @@ def create_from_probability_map(
     for probability_map_path in probability_map:
         L.info("Loading probability map %s", probability_map_path)
         loaded_probability_map = pd.read_csv(probability_map_path)
-        loaded_probability_map.set_index(["region", "molecular_type"], inplace=True)
+        loaded_probability_map.set_index(
+            ["region", "molecular_type", "synapse_class"], inplace=True
+        )
         check_probability_map_sanity(loaded_probability_map)
         probability_maps.append(loaded_probability_map)
 
@@ -279,6 +291,7 @@ def create_from_probability_map(
             for (molecular_type, density) in molecular_type_densities.items()
         },
         probability_maps,
+        synapse_class,
         output_dir,
         n_jobs,
     )

atlas_densities/densities/mtype_densities_from_map/create.py

@@ -23,6 +23,7 @@
     _check_probability_map_consistency,
     _merge_probability_maps,
 )
+from atlas_densities.exceptions import AtlasDensitiesError
 
 if TYPE_CHECKING:  # pragma: no cover
     import pandas as pd
@@ -33,11 +34,12 @@
 L = logging.getLogger(__name__)
 
 
-def create_from_probability_map(
+def create_from_probability_map(  # pylint: disable=too-many-arguments
     annotation: "VoxelData",
     region_map: "RegionMap",
     molecular_type_densities: Dict[str, FloatArray],
     probability_maps: List["pd.DataFrame"],
+    synapse_class: str,
     output_dirpath: str,
     n_jobs: int,
 ) -> None:
@@ -56,12 +58,13 @@ def create_from_probability_map(
             and whose values are 3D float arrays holding the density fields of the cells of the
             corresponding types (i.e., those cells reacting to the corresponding gene markers).
             Example: {"pv": "pv.nrrd", "sst": "sst.nrd", "vip": "vip.nrrd", "gad67": "gad67.nrrd"}
-        probability_map:
-            data frame whose rows are labeled by regions and molecular types and whose columns are
-            labeled by metypes ('|' separated).
+        probability_maps: list of data frames whose rows are labeled by regions and molecular
+            types and whose columns are labeled by metypes ('|' separated).
+        synapse_class: synapse class to use for density calculation
         output_dirpath: path of the directory where to save the volumetric density nrrd files.
             It will be created if it doesn't exist already. It will contain a volumetric density
             file of each metype appearing as column label of `probability_map`.
+        n_jobs: number of jobs to run in parallel
 
     Raises:
         AtlasBuildingTools error if
@@ -72,6 +75,20 @@ def create_from_probability_map(
     # pylint: disable=too-many-locals
     probability_map = _merge_probability_maps(probability_maps)
 
+    _check_probability_map_consistency(probability_map, set(molecular_type_densities.keys()))
+
+    # filter by synapse class
+    probability_map = probability_map[
+        probability_map.index.get_level_values("synapse_class") == synapse_class
+    ]
+    if probability_map.empty:
+        raise AtlasDensitiesError(
+            f"Filtering probability map by requested synapse_class {synapse_class} "
+            "resulted in empty probability map."
+        )
+    probability_map.index = probability_map.index.droplevel("synapse_class")
+
+    # get info on regions
     region_info = (
         region_map.as_dataframe()
         .reset_index()
@@ -81,13 +98,12 @@ def create_from_probability_map(
     )
     region_acronyms = set(region_info.region)
 
-    _check_probability_map_consistency(probability_map, set(molecular_type_densities.keys()))
-
     region_masks = {
         region_acronym: annotation.raw == region_id
         for _, region_acronym, region_id in region_info.itertuples()
     }
 
+    # ensure output directory exists
     Path(output_dirpath).mkdir(exist_ok=True, parents=True)
 
     def _create_densities_for_metype(metype: str) -> Optional[Tuple[str, str]]:
@@ -110,7 +126,7 @@ def _create_densities_for_metype(metype: str) -> Optional[Tuple[str, str]]:
 
         if np.any(metype_density):
             # save density file
-            metype_filename = f"{metype}_densities.nrrd"
+            metype_filename = f"{metype}_{synapse_class}_densities.nrrd"
             filepath = str(Path(output_dirpath) / metype_filename)
             annotation.with_data(metype_density).save_nrrd(filepath)
 
@@ -121,15 +137,27 @@ def _create_densities_for_metype(metype: str) -> Optional[Tuple[str, str]]:
     returns = Parallel(n_jobs=n_jobs, return_as="generator")(
         delayed(_create_densities_for_metype)(metype) for metype in probability_map.columns
     )
+
+    # construct metadata
     output_legend: Dict[str, Dict[str, str]] = defaultdict(dict)
     for return_value in tqdm(returns, total=len(probability_map.columns)):
         if return_value is not None:
             metype, filepath = return_value
             mtype, etype = metype.split(SEPARATOR)
             output_legend[mtype][etype] = filepath
 
-    L.info("Saving output legend.")
-    output_legend_filename = "output_legend.json"
-    filepath = str(Path(output_dirpath) / output_legend_filename)
+    metadata = {
+        "synapse_class": synapse_class,
+        "metype2file": output_legend,
+        "file2metype": {
+            filename: (mtype, etype)
+            for mtype, etypes in output_legend.items()
+            for etype, filename in etypes.items()
+        },
+    }
+
+    L.info("Saving metadata.")
+    metadata_filename = "metadata.json"
+    filepath = str(Path(output_dirpath) / metadata_filename)
     with open(filepath, "w", encoding="utf8") as file:
-        json.dump(output_legend, file, indent=4)
+        json.dump(metadata, file, indent=4)

atlas_densities/densities/mtype_densities_from_map/utils.py

@@ -9,6 +9,10 @@
 
 from atlas_densities.exceptions import AtlasDensitiesError
 
+SYNAPSE_CLASS_INH = "INH"
+SYNAPSE_CLASS_EXC = "EXC"
+SYNAPSE_CLASSES = {SYNAPSE_CLASS_INH, SYNAPSE_CLASS_EXC}
+
 L = logging.getLogger(__name__)
 
 
@@ -23,6 +27,7 @@ def check_probability_map_sanity(probability_map: "pd.DataFrame") -> None:
     Raises:
         AtlasDensitiesError if the sum of each row is not (close to) 1.0 or if `probability_map`
         holds a negative value.
+        AtlasDensitiesError if the synapse_class is not just "INH" or "EXC".
     """
     if not np.all(probability_map >= 0.0):
         raise AtlasDensitiesError("The probability map has negative values.")
@@ -32,6 +37,11 @@ def check_probability_map_sanity(probability_map: "pd.DataFrame") -> None:
             "The sum of each row is not 1.0. Consider renormalizing your input data frame with:\n"
             "df = df.div(np.sum(df, axis=1), axis=0)"
         )
+    if not {synapse_class for _, _, synapse_class in probability_map.index} <= SYNAPSE_CLASSES:
+        raise AtlasDensitiesError(
+            "The probability map has invalid value for synapse class."
+            f"Only {SYNAPSE_CLASSES} are allowed."
+        )
 
 
 def _check_probability_map_consistency(
@@ -49,7 +59,7 @@ def _check_probability_map_consistency(
     """
     check_probability_map_sanity(probability_map)
 
-    _, df_molecular_types = zip(*probability_map.index)
+    _, df_molecular_types, _ = zip(*probability_map.index)
     molecular_types_diff = molecular_types - set(df_molecular_types)
     if molecular_types_diff:
         L.info(

tests/app/test_mtype_densities.py

@@ -3,6 +3,7 @@
 import os
 from copy import deepcopy
 from pathlib import Path
+from unittest.mock import patch
 
 import numpy as np
 import numpy.testing as npt
@@ -117,6 +118,8 @@ def get_result_from_probablity_map_(runner):
             "--marker",
             "approx_lamp5",
             "approx_lamp5.nrrd",
+            "--synapse-class",
+            "EXC",
             "--output-dir",
             "output_dir",
         ],
@@ -148,14 +151,16 @@ def test_output(self):
             result = get_result_from_probablity_map_(runner)
             assert result.exit_code == 0
 
-            BPbAC = VoxelData.load_nrrd(str(Path("output_dir") / "BP|bAC_densities.nrrd"))
+            BPbAC = VoxelData.load_nrrd(str(Path("output_dir") / "BP|bAC_EXC_densities.nrrd"))
             assert BPbAC.raw.dtype == float
             npt.assert_array_equal(BPbAC.voxel_dimensions, self.data["annotation"].voxel_dimensions)
 
-            with open(str(Path("output_dir") / "output_legend.json"), "r") as file:
-                output_legend = json.load(file)
-            assert "BP" in output_legend
-            assert "bAC" in output_legend["BP"]
+            with open(str(Path("output_dir") / "metadata.json"), "r") as file:
+                metadata = json.load(file)
+            assert "BP" in metadata["metype2file"]
+            assert "bAC" in metadata["metype2file"]["BP"]
+            assert metadata["file2metype"][metadata["metype2file"]["BP"]["bAC"]] == ["BP", "bAC"]
+            assert "EXC" == metadata["synapse_class"]
 
     class Test_mtype_densities_from_composition:
         @pytest.fixture(scope="session")