From 924b6c60d64f7bac373ff55b5fb9a9d19ebf70be Mon Sep 17 00:00:00 2001
From: Eleftherios Zisis <eleftherios.zisis@epfl.ch>
Date: Mon, 15 Apr 2024 13:40:48 +0200
Subject: [PATCH] Speedup mtype density creation from probability maps using
 voxcell's ValueToIndexVoxels (#79)

---
 .../mtype_densities_from_map/create.py        | 28 +++++++++++--------
 setup.py                                      |  2 +-
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/atlas_densities/densities/mtype_densities_from_map/create.py b/atlas_densities/densities/mtype_densities_from_map/create.py
index 936005d..ff5eb02 100644
--- a/atlas_densities/densities/mtype_densities_from_map/create.py
+++ b/atlas_densities/densities/mtype_densities_from_map/create.py
@@ -18,6 +18,7 @@
 from atlas_commons.typing import FloatArray
 from joblib import Parallel, delayed
 from tqdm import tqdm
+from voxcell.voxel_data import ValueToIndexVoxels
 
 from atlas_densities.densities.mtype_densities_from_map.utils import (
     _check_probability_map_consistency,
@@ -95,34 +96,37 @@ def create_from_probability_map(  # pylint: disable=too-many-arguments
         .set_index("acronym")
         .loc[probability_map.index.get_level_values("region")]
         .reset_index()[["region", "id"]]
+        .drop_duplicates(subset="region")
+        .reset_index(drop=True)
     )
-    region_acronyms = set(region_info.region)
-
-    region_masks = {
-        region_acronym: annotation.raw == region_id
-        for _, region_acronym, region_id in region_info.itertuples()
-    }
+    annotation_index = ValueToIndexVoxels(annotation.raw)
 
     # ensure output directory exists
     Path(output_dirpath).mkdir(exist_ok=True, parents=True)
 
     def _create_densities_for_metype(metype: str) -> Optional[Tuple[str, str]]:
         coefficients: Dict[str, Dict[str, Any]] = {}
-        for region_acronym in region_acronyms:
+        for region_acronym in region_info.region:
             coefficients[region_acronym] = {
                 molecular_type: probability_map.at[(region_acronym, molecular_type), metype]
                 for molecular_type in list(molecular_type_densities.keys())
                 if (region_acronym, molecular_type) in probability_map.index
             }
 
-        metype_density = np.zeros(annotation.shape, dtype=float)
-        for region_acronym in region_acronyms:
-            region_mask = region_masks[region_acronym]
+        # perform the manipulation in the 1d flat array
+        metype_density = np.zeros(np.prod(annotation.shape), dtype=float)
+
+        for region_acronym, region_id in region_info.itertuples(index=False):
+
+            region_indices = annotation_index.value_to_1d_indices(region_id)
             for molecular_type, coefficient in coefficients[region_acronym].items():
                 if coefficient <= 0.0:
                     continue
-                density = molecular_type_densities[molecular_type]
-                metype_density[region_mask] += density[region_mask] * coefficient
+                density = annotation_index.ravel(molecular_type_densities[molecular_type])
+                metype_density[region_indices] += density[region_indices] * coefficient
+
+        # reshape the 1d metype_density array back to the annotation's shape
+        metype_density = annotation_index.unravel(metype_density)
 
         if np.any(metype_density):
             # save density file
diff --git a/setup.py b/setup.py
index 308d6ee..ee9c873 100644
--- a/setup.py
+++ b/setup.py
@@ -28,7 +28,7 @@
         # from the HiGHS library. We use the "highs" method in the densities module.
         "scipy>=1.6.0",
         "tqdm>=4.44.1",
-        "voxcell>=3.1.7",
+        "voxcell>=3.1.8",  # ValueToIndexVoxels ravel/unravel
     ],
     extras_require={
         "tests": [