Handle octree case separately

surface_apps/iso_surfaces/utils.py

@@ -14,15 +14,17 @@
 from geoapps_utils.locations import mask_under_horizon
 from geoapps_utils.numerical import weighted_average
 from geoapps_utils.transformations import rotate_xyz
-from geoh5py.objects import BlockModel, ObjectBase, Surface
+from geoh5py.data import DataAssociationEnum
+from geoh5py.data.data import Data
+from geoh5py.objects import BlockModel, ObjectBase, Octree, Surface
 from scipy.interpolate import interp1d
 from skimage.measure import marching_cubes
 from tqdm import tqdm
 
 
 def entity_to_grid(
     entity: ObjectBase,
-    values: np.ndarray,
+    data: Data,
     resolution: float,
     max_distance: float,
     horizon: Surface | None = None,
@@ -31,7 +33,7 @@ def entity_to_grid(
     Extracts grid and values from a geoh5py object.
 
     :param entity: geoh5py object with locations data.
-    :param values: Data to be reshaped or interpolated to grid ordering.
+    :param data: Data to be reshaped or interpolated to grid ordering.
     :resolution: Grid resolution.
     :param max_distance: Maximum distance used in weighted average.
     :param horizon: Clipping surface to restrict interpolation from bleeding
@@ -44,9 +46,9 @@ def entity_to_grid(
         raise ValueError("Input 'entity' must have a vertices or centroids set.")
 
     if isinstance(entity, BlockModel):
-        grid, values = block_model_to_grid(entity, values)
+        grid, values = block_model_to_grid(entity, data.values)
     else:
-        grid, values = interp_to_grid(entity, values, resolution, max_distance, horizon)
+        grid, values = interp_to_grid(entity, data, resolution, max_distance, horizon)
 
     return grid, values
 
@@ -79,7 +81,7 @@ def block_model_to_grid(
 
 def interp_to_grid(  # pylint: disable=too-many-locals
     entity: ObjectBase,
-    values: np.ndarray,
+    data: Data,
     resolution: float,
     max_distance: float,
     horizon: Surface | None = None,
@@ -88,7 +90,7 @@ def interp_to_grid(  # pylint: disable=too-many-locals
     Interpolate values into a regular grid bounding all finite data points.
 
     :param entity: Geoh5py object with locations data.
-    :param values: Data to be interpolated to grid.
+    :param data: Data to be interpolated to grid.
     :param resolution: Grid resolution
     :param max_distance: Maximum distance used in weighted average.
 
@@ -98,9 +100,16 @@ def interp_to_grid(  # pylint: disable=too-many-locals
     """
 
     grid = []
-    is_finite = np.isfinite(values)
-    locations = entity.locations[is_finite, :]
-    finite_values = values[is_finite]
+    is_finite = np.isfinite(data.values)
+    if isinstance(entity, Octree):
+        locations = entity.centroids[is_finite, :]
+    elif data.association == DataAssociationEnum.CELL:
+        vertices = entity.locations
+        locations = np.mean(vertices[entity.cells], axis=1)[is_finite, :]
+    else:
+        locations = entity.vertices[is_finite, :]
+
+    finite_values = data.values[is_finite]
     for i in np.arange(3):
         grid += [
             np.arange(