switch to ValueToIndexVoxels in voxcell

atlas_densities/densities/fitting.py

@@ -21,21 +21,19 @@
 import itertools as it
 import logging
 import warnings
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 import pandas as pd
 from atlas_commons.typing import AnnotationT, BoolArray, FloatArray
 from scipy.optimize import curve_fit
 from tqdm import tqdm
+from voxcell import voxel_data, RegionMap
 
 from atlas_densities.densities import utils
 from atlas_densities.densities.measurement_to_density import remove_unknown_regions
 from atlas_densities.exceptions import AtlasDensitiesError, AtlasDensitiesWarning
 
-if TYPE_CHECKING:  # pragma: no cover
-    from voxcell import RegionMap
-
 L = logging.getLogger(__name__)
 
 MarkerVolumes = Dict[str, Dict[str, Union[FloatArray, List[int]]]]
@@ -294,62 +292,6 @@ def _compute_average_intensities_helper(index, gene_marker_volumes, id_):
     return res
 
 
-class ValueToIndexVoxels:
-    """Efficient access to indices of unique values of the values array.
-
-    Useful for when one has an "annotations volume" or "brain region volume" that has
-    regions indicated by unique values, and these are used to create masks.  Often,
-    it's faster to avoid mask creation, and use indices directly
-
-    Example:
-        vtiv = ValueToIndexVoxels(br.raw)
-        values, funcs = zip(*((i, np.sum if i % 2 else np.mean) for i in vtiv.values[:10]))
-        list(vtiv.apply(values, funcs, density.raw))
-    """
-
-    def __init__(self, values):
-        """Initialize.
-
-        Args:
-            values(np.array): volume with each voxel marked with a value; usually to group regions
-        """
-        self._order = "C" if values.flags["C_CONTIGUOUS"] else "F"
-
-        values = values.ravel(order=self._order)
-        uniques, codes, counts = np.unique(values, return_inverse=True, return_counts=True)
-
-        offsets = np.empty(len(counts) + 1, dtype=np.uint64)
-        offsets[0] = 0
-        offsets[1:] = np.cumsum(counts)
-
-        self._codes = codes
-        self._offsets = offsets
-        self._indices = np.argsort(values, kind="stable")
-        self._mapping = {v: i for i, v in enumerate(uniques)}
-        self.index_dtype = values.dtype
-
-    def ravel(self, voxel_data):
-        """Ensures `voxel_data` matches the layout that the 1D indices can be used."""
-        return voxel_data.ravel(order=self._order)
-
-    @property
-    def values(self):
-        """Unique values that are found in the original volume."""
-        return np.fromiter(self._mapping, dtype=self.index_dtype)
-
-    def value_to_1d_indices(self, value):
-        """Return the indices array indices corresponding to the 'value'.
-
-        Note: These are 1D indices, so the assumption is they are applied to a volume
-        who has been ValueToIndexVoxels::ravel(volume)
-        """
-        if value not in self._mapping:
-            return np.array([], dtype=np.uint64)
-
-        group_index = self._mapping[value]
-        return self._indices[self._offsets[group_index] : self._offsets[group_index + 1]]
-
-
 def compute_average_intensities(
     annotation: AnnotationT,
     gene_marker_volumes: MarkerVolumes,
@@ -388,7 +330,7 @@ def compute_average_intensities(
     """
     gene_marker_volumes = _apply_density_slices(gene_marker_volumes)
 
-    index = ValueToIndexVoxels(annotation)
+    index = voxel_data.ValueToIndexVoxels(annotation)
 
     _helper = _compute_average_intensities_helper
     work = []
@@ -732,7 +674,7 @@ def _check_average_densities_sanity(average_densities: pd.DataFrame) -> None:
         )
 
 
-def _get_group_names(region_map: "RegionMap", group_ids_config: dict) -> dict[str, set[str]]:
+def _get_group_names(region_map: RegionMap, group_ids_config: dict) -> dict[str, set[str]]:
     """
     Get AIBS names for regions in several region groups of interest.
 
@@ -780,7 +722,7 @@ def _get_group_region_names(groups):
 
 
 def linear_fitting(  # pylint: disable=too-many-arguments
-    region_map: "RegionMap",
+    region_map: RegionMap,
     annotation: AnnotationT,
     neuron_density: FloatArray,
     gene_marker_volumes: MarkerVolumes,