Add ValueToIndexVoxels to simplify using indices instead of masks (#30)

Co-authored-by: Eleftherios Zisis <[email protected]>

tests/test_voxel_data.py

@@ -409,3 +409,34 @@ def test_offset_and_voxel_dimensions_type():
 
     voxel_data = test_module.VoxelData(np.ones((2, 2, 2)), voxel_dimensions=(3, 4, 5))
     assert voxel_data.offset.dtype == np.float32
+
+
+def test_ValueToIndexVoxels():
+    values = np.array([[1., 1., 1.], [1., 2., 2.], [3., 3., 3.]])
+    br = np.array([[1, 1, 1], [1, 2, 2], [3, 3, 3]], order='C')
+    vtiv = test_module.ValueToIndexVoxels(br)
+
+    npt.assert_array_equal(vtiv.value_to_1d_indices(1), [0, 1, 2, 3])
+    npt.assert_array_equal(vtiv.value_to_1d_indices(2), [4, 5])
+    npt.assert_array_equal(vtiv.value_to_1d_indices(3), [6, 7, 8])
+    npt.assert_array_equal(vtiv.value_to_1d_indices(4), [])
+
+    assert vtiv.index_size == 3
+    assert vtiv.index_dtype == np.int64
+    assert list(vtiv.values) == [1, 2, 3]
+
+    for order in ('K', 'A', 'C', 'F'):
+        r = vtiv.ravel(np.array(values, order=order))
+        npt.assert_array_equal(r[vtiv.value_to_1d_indices(1)], [1., 1., 1., 1.])
+
+    br = np.array([[1, 1, 1], [1, 2, 2], [3, 3, 3]], order='F')
+    vtiv = test_module.ValueToIndexVoxels(br)
+
+    npt.assert_array_equal(vtiv.value_to_1d_indices(1), [0, 1, 3, 6])
+    npt.assert_array_equal(vtiv.value_to_1d_indices(2), [4, 7])
+    npt.assert_array_equal(vtiv.value_to_1d_indices(3), [2, 5, 8])
+    npt.assert_array_equal(vtiv.value_to_1d_indices(4), [])
+
+    for order in ('K', 'A', 'C', 'F'):
+        r = vtiv.ravel(np.array(values, order=order))
+        npt.assert_array_equal(r[vtiv.value_to_1d_indices(1)], [1., 1., 1., 1.])

voxcell/voxel_data.py

@@ -409,6 +409,73 @@ def __init__(self, *args, **kwargs):
         self.raw = self.raw.astype(bool)
 
 
+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:
+        # To calculate the cell count based on densities of a certain ID in the brain_regions volume
+        vtiv = ValueToIndexVoxels(brain_regions.raw)
+        density_copy = vtiv.ravel(density.raw.copy())
+        indices = vtiv.value_to_1d_indices(value=id_)
+        cell_count = np.sum(density_copy[indices]) * voxel_volume)
+    """
+
+    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, counts = np.unique(values, return_counts=True)
+
+        offsets = np.empty(len(counts) + 1, dtype=np.uint64)
+        offsets[0] = 0
+        offsets[1:] = np.cumsum(counts)
+
+        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
+
+    @property
+    def index_size(self):
+        """Return the size of the unique index values."""
+        return len(self._mapping)
+
+    @property
+    def index_dtype(self):
+        """Return the dytpe of the index values."""
+        return self._index_dtype
+
+    @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=self._indices.dtype)
+
+        group_index = self._mapping[value]
+        return self._indices[self._offsets[group_index]:self._offsets[group_index + 1]]
+
+    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)
+
+
 def values_to_region_attribute(values, region_map, attr="acronym"):
     """Convert region ids to the corresponding region attribute.