From 260a25ae6a4b4f01bb7d0b2aca32be09dea02a81 Mon Sep 17 00:00:00 2001
From: MikeG <michael.gevaert@epfl.ch>
Date: Tue, 26 Mar 2024 12:16:43 +0100
Subject: [PATCH] Faster _compute_region_cell_counts (#72)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Faster _compute_region_cell_counts

Some timings:
    old: 391s
    new: 7.77 s ± 26.5 ms per loop
    So about 55 times faster
    maxium difference in cell_count: 9.313225746154785e-10

Code used for testing:
```
from atlas_densities.densities.inhibitory_neuron_densities_optimization import _compute_region_cell_counts
import voxcell
from atlas_densities.densities import utils

annotation = voxcell.VoxelData.load_nrrd('input-data/annotation_ccfv2_l23split_barrelsplit_validated.nrrd')
density = voxcell.VoxelData.load_nrrd('input-data/neuron_density.nrrd')
voxel_volume = annotation.voxel_volume / 1e9
rm =voxcell.RegionMap.load_json('input-data/hierarchy_ccfv2_l23split_barrelsplit.json')
hierarchy_info = utils.get_hierarchy_info(rm, root='root')

res = _compute_region_cell_counts(annotation.raw, density.raw, voxel_volume, hierarchy_info,)
```
---
 ...nhibitory_neuron_densities_optimization.py | 22 ++++++++++---------
 setup.py                                      |  2 +-
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/atlas_densities/densities/inhibitory_neuron_densities_optimization.py b/atlas_densities/densities/inhibitory_neuron_densities_optimization.py
index 675a657..0747918 100644
--- a/atlas_densities/densities/inhibitory_neuron_densities_optimization.py
+++ b/atlas_densities/densities/inhibitory_neuron_densities_optimization.py
@@ -47,7 +47,7 @@
 from atlas_commons.typing import AnnotationT, FloatArray
 from scipy.optimize import linprog
 from tqdm import tqdm
-from voxcell import RegionMap
+from voxcell import RegionMap, voxel_data
 
 from atlas_densities.densities import utils
 from atlas_densities.densities.inhibitory_neuron_densities_helper import (
@@ -199,16 +199,18 @@ def _compute_region_cell_counts(
              ...                             ...
         The index is the sorted list of all region identifiers.
     """
+    vtiv = voxel_data.ValueToIndexVoxels(annotation)
+    density = vtiv.ravel(density)
+
     id_counts = []
-    for id_ in tqdm(hierarchy_info.index):
-        mask = annotation == id_
-        id_counts.append(np.sum(density[mask]) * voxel_volume)
+    for id_ in hierarchy_info.index:
+        indices = vtiv.value_to_1d_indices(id_)
+        id_counts.append(np.sum(density[indices]) * voxel_volume)
 
     result = pd.DataFrame(
         {"brain_region": hierarchy_info["brain_region"], "cell_count": id_counts},
         index=hierarchy_info.index,
     )
-
     return result
 
 
@@ -736,17 +738,17 @@ def create_inhibitory_neuron_densities(  # pylint: disable=too-many-locals
     hierarchy_info = utils.get_hierarchy_info(RegionMap.from_dict(hierarchy), root=region_name)
     average_densities = _resize_average_densities(average_densities, hierarchy_info)
 
-    L.info("Initialization of the linear program: started")
-    region_counts, id_counts = _compute_initial_cell_counts(
-        annotation, voxel_volume, average_densities, hierarchy_info
-    )
-
     L.info("Retrieving overall neuron counts in atomic 3D regions ...")
     neuron_counts = _compute_region_cell_counts(
         annotation, neuron_density, voxel_volume, hierarchy_info
     )
     assert np.all(neuron_counts["cell_count"] >= 0.0)
 
+    L.info("Initialization of the linear program: started")
+    region_counts, id_counts = _compute_initial_cell_counts(
+        annotation, voxel_volume, average_densities, hierarchy_info
+    )
+
     L.info("Setting the known values ...")
     x_result, deltas = set_known_values(region_counts, id_counts, hierarchy_info)
     # We assume that if the cell count of `cell_type` in `region_name` is known with certainty
diff --git a/setup.py b/setup.py
index 9fbd57d..308d6ee 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.0.0",
+        "voxcell>=3.1.7",
     ],
     extras_require={
         "tests": [