From c0683a831c29c45ad66442ff088885acd5d13e1a Mon Sep 17 00:00:00 2001 From: Adrien Berchet Date: Fri, 5 Apr 2024 10:56:59 +0200 Subject: [PATCH] Add the VoxelData.value_to_indices() method (#33) * Add the VoxelData.value_to_indices() method * Avoid copy for scalar values * Add edge case in test * Add example of use in dosctring --- tests/test_voxel_data.py | 32 ++++++++++++++++++++++++++++++++ voxcell/voxel_data.py | 24 +++++++++++++++++++++++- 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/tests/test_voxel_data.py b/tests/test_voxel_data.py index 2e96db4..b1acaaa 100644 --- a/tests/test_voxel_data.py +++ b/tests/test_voxel_data.py @@ -421,6 +421,22 @@ def test_ValueToIndexVoxels(): npt.assert_array_equal(vtiv.value_to_1d_indices(3), [6, 7, 8]) npt.assert_array_equal(vtiv.value_to_1d_indices(4), []) + npt.assert_array_equal(vtiv.value_to_indices(1), [[0, 0], [0, 1], [0, 2], [1, 0]]) + npt.assert_array_equal(vtiv.value_to_indices(2), [[1, 1], [1, 2]]) + npt.assert_array_equal(vtiv.value_to_indices(3), [[2, 0], [2, 1], [2, 2]]) + npt.assert_array_equal( + vtiv.value_to_indices(4), + np.zeros_like([], shape=(0, 2), dtype=vtiv._indices.dtype), + ) + npt.assert_array_equal( + vtiv.value_to_indices(range(1, 5)), + [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]], + ) + npt.assert_array_equal( + vtiv.value_to_indices(range(10, 15)), # All values are missing in the input array + np.zeros_like([], shape=(0, 2), dtype=vtiv._indices.dtype), + ) + assert vtiv.index_size == 3 assert vtiv.index_dtype == np.int64 assert list(vtiv.values) == [1, 2, 3] @@ -437,6 +453,22 @@ def test_ValueToIndexVoxels(): npt.assert_array_equal(vtiv.value_to_1d_indices(3), [2, 5, 8]) npt.assert_array_equal(vtiv.value_to_1d_indices(4), []) + npt.assert_array_equal(vtiv.value_to_indices(1), [[0, 0], [1, 0], [0, 1], [0, 2]]) + npt.assert_array_equal(vtiv.value_to_indices(2), [[1, 1], [1, 2]]) + npt.assert_array_equal(vtiv.value_to_indices(3), [[2, 0], [2, 1], [2, 2]]) + npt.assert_array_equal( + vtiv.value_to_indices(4), + np.zeros_like([], shape=(0, 2), dtype=vtiv._indices.dtype), + ) + npt.assert_array_equal( + vtiv.value_to_indices(range(1, 5)), + [[0, 0], [1, 0], [0, 1], [0, 2], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]], + ) + npt.assert_array_equal( + vtiv.value_to_indices(range(10, 15)), # All values are missing in the input array + np.zeros_like([], shape=(0, 2), dtype=vtiv._indices.dtype), + ) + 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.]) diff --git a/voxcell/voxel_data.py b/voxcell/voxel_data.py index 5ea8c48..4a9fe92 100644 --- a/voxcell/voxel_data.py +++ b/voxcell/voxel_data.py @@ -461,7 +461,7 @@ def values(self): 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'. + """Return the indices array 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) @@ -472,6 +472,28 @@ def value_to_1d_indices(self, value): group_index = self._mapping[value] return self._indices[self._offsets[group_index]:self._offsets[group_index + 1]] + def value_to_indices(self, values): + """Return the ND-indices array corresponding to the 'values'. + + This can be convenient to get the positions of the given values in the VoxelData space: + raw = np.array([[11, 12], [21, 22]]) + v = VoxelData(raw, voxel_dimensions=(2, 3), offset=np.array([2, 2])) + vtiv = ValueToIndexVoxels(v.raw) + positions = v.indices_to_positions(vtiv.value_to_indices(11)) + + Note: The given 'values' can be given as one scalar value or as a list of values. In both + case a list of ND-indices will be returned. + """ + if np.isscalar(values): + flat_indices = self.value_to_1d_indices(values) + else: + flat_indices = np.concatenate( + [self.value_to_1d_indices(i) for i in values] + ) + return np.array( + np.unravel_index(flat_indices, self._shape, order=self._order) + ).T + def ravel(self, voxel_data): """Ensure `voxel_data` matches the layout that the 1D indices can be used.""" if voxel_data.shape != self._shape: