GEOPY-2786: add a mthode to access the reshaped values in geoh5py for 2d grid and 3dgrid

geoh5py/objects/block_model.py

@@ -24,6 +24,7 @@
 
 import numpy as np
 
+from ..data import Data
 from ..shared.utils import xy_rotation_matrix
 from .grid_object import GridObject
 
@@ -124,6 +125,30 @@ def shape(self) -> tuple[int, int, int]:
         """
         return self.u_cells.shape[0], self.v_cells.shape[0], self.z_cells.shape[0]
 
+    def shaped_data_values(
+        self, data: str | uuid.UUID | Data | np.ndarray
+    ) -> np.ndarray:
+        """
+        Get the values of a data entity as a 3D array with the same shape as the grid.
+
+        Data values are stored under the hood as a flatten array in Fortran order;
+        first, the values are reshaped to the grid shape with Fortran order,
+        then transposed to match the (n_v, n_u, n_z) shape.
+
+        :param data: The data to get the values from.
+
+        :return: The shaped values of the data entity.
+        """
+        values = (
+            data
+            if isinstance(data, np.ndarray)
+            else self._get_data_to_reshape(data).values
+        )
+
+        return values.reshape(
+            (self.shape[2], self.shape[0], self.shape[1]), order="F"
+        ).transpose(2, 1, 0)
+
     @property
     def u_cell_delimiters(self) -> np.ndarray:
         """

geoh5py/objects/grid2d.py

@@ -25,6 +25,7 @@
 
 import numpy as np
 
+from ..data import Data
 from ..objects import GeoImage
 from ..shared.conversion import Grid2DConversion
 from ..shared.utils import mask_by_extent, xy_rotation_matrix, yz_rotation_matrix
@@ -247,6 +248,24 @@ def shape(self) -> tuple[np.int32, np.int32]:
         """
         return self.u_count, self.v_count
 
+    def shaped_data_values(
+        self, data: str | uuid.UUID | Data | np.ndarray
+    ) -> np.ndarray:
+        """
+        Get the values of a data entity as a 2D array with the same shape as the grid.
+
+        :param data: The data to get the values from.
+
+        :return: The shaped values of the data entity.
+        """
+        values = (
+            data
+            if isinstance(data, np.ndarray)
+            else self._get_data_to_reshape(data).values
+        )
+
+        return values.reshape(self.v_count, self.u_count)
+
     @property
     def u_cell_size(self) -> float:
         """

geoh5py/objects/grid_object.py

@@ -20,11 +20,14 @@
 
 from __future__ import annotations
 
+import uuid
 from abc import ABC, abstractmethod
 from numbers import Real
 
 import numpy as np
 
+from geoh5py.data import Data, DataAssociationEnum
+
 from .object_base import ObjectBase
 
 
@@ -160,3 +163,34 @@ def validate_mask(self, mask: np.ndarray | None):
             )
 
         return mask
+
+    def _get_data_to_reshape(self, data: str | uuid.UUID | Data) -> Data:
+        """
+        Get a unique data entity with association 'CELL' from the data name, uid or object.
+
+        :raises ValueError: if no data are found.
+        :raises ValueError: if multiple data are found.
+        :raises ValueError: if the data association is not 'CELL'.
+
+        :param data: The data to get the values from.
+
+        :return: The unique data.
+        """
+        if not isinstance(data, Data):
+            data_list = self.get_data(data)
+
+            if len(data_list) == 0:
+                raise ValueError(f"No data '{data}' found.")
+            if len(data_list) > 1:
+                raise ValueError(
+                    f"Multiple data '{data}' found. Please specify a unique data name or uid."
+                )
+
+            data = data_list[0]
+
+        if data.association != DataAssociationEnum.CELL:
+            raise ValueError(
+                f"Data '{data.name}' has association '{data.association}'"
+                ", expected 'CELL'."
+            )
+        return data

tests/block_model_test.py

@@ -151,3 +151,25 @@ def test_create_block_model_data(tmp_path):
             grid_copy.mask_by_extent(np.vstack([[-100, -100], [1, 100]]), inverse=True)
             == ~mask
         )
+
+
+def test_shaped_data_values():
+    with Workspace() as workspace:
+        block_model = BlockModel.create(
+            workspace,
+            u_cell_delimiters=np.array([0.0, 1.0, 2.0]),
+            v_cell_delimiters=np.array([0.0, 1.0, 2.0]),
+            z_cell_delimiters=np.array([0.0, -1.0, -2.0]),
+        )
+        values = np.arange(block_model.n_cells, dtype=float)
+        block_model.add_data({"DataValues": {"association": "CELL", "values": values}})
+        shaped = block_model.shaped_data_values("DataValues")
+        assert shaped.shape == block_model.shape
+        assert np.allclose(shaped.ravel(), values)
+        assert np.allclose(shaped, block_model.shaped_data_values(values))
+
+        block_model.add_data({"DataValues2": {"association": "CELL", "values": shaped}})
+
+        data1 = block_model.get_data("DataValues")[0].values
+        data2 = block_model.get_data("DataValues2")[0].values
+        assert np.allclose(data1, data2)

tests/grid_2d_test.py

@@ -184,3 +184,42 @@ def test_grid2d_to_geoimage(tmp_path):
 
         with pytest.raises(TypeError, match="The dtype of the keys must be"):
             converter.key_to_data(grid, [0, 1])
+
+
+def test_shaped_data_values():
+    n_u, n_v = 5, 4
+    with Workspace() as workspace:
+        grid = Grid2D.create(
+            workspace,
+            u_cell_size=1.0,
+            v_cell_size=1.0,
+            u_count=n_u,
+            v_count=n_v,
+        )
+        values = np.arange(grid.n_cells, dtype=float)
+        grid.add_data({"DataValues": {"association": "CELL", "values": values}})
+        shaped = grid.shaped_data_values("DataValues")
+        assert shaped.shape == (n_v, n_u)
+        assert np.allclose(shaped.flatten(), values)
+        assert np.allclose(shaped, grid.shaped_data_values(values))
+
+        # check for data restoration
+        grid.add_data({"DataValues2": {"association": "CELL", "values": shaped}})
+
+        data1 = grid.get_data("DataValues")[0].values
+        data2 = grid.get_data("DataValues2")[0].values
+        assert np.allclose(data1, data2)
+
+
+def test_get_unique_data_errors():
+    with Workspace() as workspace:
+        grid = Grid2D.create(workspace, u_count=2, v_count=2)
+
+        with pytest.raises(ValueError, match="No data"):
+            grid.shaped_data_values("NonExistent")
+
+        grid.add_data(
+            {"ObjectData": {"association": "OBJECT", "values": np.array([1.0])}}
+        )
+        with pytest.raises(ValueError, match="expected 'CELL'"):
+            grid.shaped_data_values("ObjectData")