[REFACTOR] Use center grid for dual contouring instead of corners grid (#24)

# Combining centers and corners

This PR refactors the dual contouring implementation to use a more efficient approach for handling grid corners:

- Adds `corners_grid` property to `EngineGrid` to store corner values alongside centers
- Modifies `interpolate_on_octree` to compute corners as part of the center grid rather than separately
- Adds corner-related properties to `InterpOutput` for accessing corner data
- Fixes `__len__` in `RegularGrid` to return the correct number of active cells
- Replaces `_get_intersection_on_edges` with direct calls to `find_intersection_on_edge` for improved consistency
- Updates masking logic to work with the new grid structure
- Simplifies octree generation by removing redundant corner-related code

These changes improve the efficiency of the dual contouring algorithm by reducing redundant computations and simplifying the data flow.

Author: Leguark

gempy_engine/API/dual_contouring/_interpolate_on_edges.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple, Optional
+from typing import List, Optional
 
 import numpy as np
 
@@ -24,10 +24,15 @@ def interpolate_on_edges_for_dual_contouring(
         octree_leaves: OctreeLevel,
         mask: Optional[np.ndarray] = None
 ) -> DualContouringData:
-
-    # region define location where we need to interpolate the gradients for dual contouring
-    output_corners: InterpOutput = octree_leaves.outputs_corners[n_scalar_field]
-    intersection_xyz, valid_edges = _get_intersection_on_edges(octree_leaves, output_corners, mask)
+    
+    output: InterpOutput = octree_leaves.outputs_centers[n_scalar_field]
+    intersection_xyz, valid_edges = find_intersection_on_edge(
+        _xyz_corners=octree_leaves.grid_centers.corners_grid.values,
+        scalar_field_on_corners=output.exported_fields.scalar_field[output.grid.corners_grid_slice],
+        scalar_at_sp=output.scalar_field_at_sp,
+        masking=mask
+    )
+    
     interpolation_input.set_temp_grid(EngineGrid(custom_grid=GenericGrid(values=intersection_xyz)))
     # endregion
 
@@ -43,22 +48,7 @@ def interpolate_on_edges_for_dual_contouring(
         xyz_on_centers=octree_leaves.grid_centers.octree_grid.values if mask is None else octree_leaves.grid_centers.octree_grid.values[mask],
         dxdydz=octree_leaves.grid_centers.octree_dxdydz,
         exported_fields_on_edges=output_on_edges[n_scalar_field].exported_fields,
-        n_surfaces_to_export=output_corners.scalar_field_at_sp.shape[0],
+        n_surfaces_to_export=output.scalar_field_at_sp.shape[0],
         tree_depth=options.number_octree_levels,
     )
     return dc_data
-
-
-# TODO: These two functions could be moved to the module
-def _get_intersection_on_edges(octree_level: OctreeLevel, output_corners: InterpOutput,
-                               mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
-    # First find xyz on edges:
-    intersection_xyz, valid_edges = find_intersection_on_edge(
-        _xyz_corners=octree_level.grid_corners.values,
-        scalar_field_on_corners=output_corners.exported_fields.scalar_field,
-        scalar_at_sp=output_corners.scalar_field_at_sp,
-        masking=mask
-    )
-    return intersection_xyz, valid_edges
-
-

gempy_engine/API/dual_contouring/multi_scalar_dual_contouring.py

@@ -97,9 +97,11 @@ def dual_contouring_multi_scalar(data_descriptor: InputDataDescriptor, interpola
 
 
 def _mask_generation(octree_leaves, masking_option: MeshExtractionMaskingOptions) -> np.ndarray | None:
-    all_scalar_fields_outputs: list[InterpOutput] = octree_leaves.outputs_corners
+    all_scalar_fields_outputs: list[InterpOutput] = octree_leaves.outputs_centers
     n_scalar_fields = len(all_scalar_fields_outputs)
-    grid_size = all_scalar_fields_outputs[0].grid_size
+    outputs_ = all_scalar_fields_outputs[0]
+    slice_corners = outputs_.grid.corners_grid_slice
+    grid_size = outputs_.cornersGrid_values.shape[0]
     mask_matrix = BackendTensor.t.zeros((n_scalar_fields, grid_size // 8), dtype=bool)
     onlap_chain_counter = 0
 
@@ -123,15 +125,15 @@ def _mask_generation(octree_leaves, masking_option: MeshExtractionMaskingOptions
             #     raise NotImplementedError("Onlap is not supported yet")
             #     return octree_leaves.outputs_corners[n_scalar_field].squeezed_mask_array.reshape((1, -1, 8)).sum(-1, bool)[0]
             case MeshExtractionMaskingOptions.INTERSECT, StackRelationType.ERODE:
-                x = all_scalar_fields_outputs[i + onlap_chain_counter].squeezed_mask_array.reshape((1, -1, 8))
+                x = all_scalar_fields_outputs[i + onlap_chain_counter].squeezed_mask_array[slice_corners].reshape((1, -1, 8))
                 mask_matrix[i] = BackendTensor.t.sum(x, -1, bool)[0]
                 onlap_chain_counter = 0
             case MeshExtractionMaskingOptions.INTERSECT, StackRelationType.BASEMENT:
-                x = all_scalar_fields_outputs[i].squeezed_mask_array.reshape((1, -1, 8))
+                x = all_scalar_fields_outputs[i].squeezed_mask_array[slice_corners].reshape((1, -1, 8))
                 mask_matrix[i] = BackendTensor.t.sum(x, -1, bool)[0]
                 onlap_chain_counter = 0
             case MeshExtractionMaskingOptions.INTERSECT, StackRelationType.ONLAP:
-                x = all_scalar_fields_outputs[i].squeezed_mask_array.reshape((1, -1, 8))
+                x = all_scalar_fields_outputs[i].squeezed_mask_array[slice_corners].reshape((1, -1, 8))
                 mask_matrix[i] = BackendTensor.t.sum(x, -1, bool)[0]
                 onlap_chain_counter += 1
             case _, StackRelationType.FAULT:

gempy_engine/API/interp_single/_octree_generation.py

@@ -17,7 +17,7 @@
 import numpy as np
 
 
-def interpolate_on_octree(interpolation_input: InterpolationInput, options: InterpolationOptions,
+def interpolate_on_octree_(interpolation_input: InterpolationInput, options: InterpolationOptions,
                           data_shape: InputDataDescriptor) -> OctreeLevel:
     if BackendTensor.engine_backend is not AvailableBackends.PYTORCH and NOT_MAKE_INPUT_DEEP_COPY is False:
         temp_interpolation_input = copy.deepcopy(interpolation_input)
@@ -30,8 +30,9 @@ def interpolate_on_octree(interpolation_input: InterpolationInput, options: Inte
     # * Interpolate - corners
     grid_0_centers: EngineGrid = temp_interpolation_input.grid  # ? This could be moved to the next section
     if options.compute_corners:
-        grid_0_corners: Optional[EngineGrid] = EngineGrid.from_xyz_coords(
-            xyz_coords=_generate_corners(regular_grid=grid_0_centers.octree_grid)
+        xyz_corners = _generate_corners(regular_grid=grid_0_centers.octree_grid)
+        grid_0_corners: EngineGrid = EngineGrid.from_xyz_coords(
+            xyz_coords=xyz_corners
         )
 
         # ! Here we need to swap the grid temporarily but it is
@@ -58,6 +59,51 @@ def interpolate_on_octree(interpolation_input: InterpolationInput, options: Inte
     return next_octree_level
 
 
+def interpolate_on_octree(interpolation_input: InterpolationInput, options: InterpolationOptions,
+                           data_shape: InputDataDescriptor) -> OctreeLevel:
+    if BackendTensor.engine_backend is not AvailableBackends.PYTORCH and NOT_MAKE_INPUT_DEEP_COPY is False:
+        temp_interpolation_input = copy.deepcopy(interpolation_input)
+    else:
+        temp_interpolation_input = interpolation_input
+
+    # * Interpolate - corners
+    if options.compute_corners:
+        grid_0_centers: EngineGrid = temp_interpolation_input.grid  # ? This could be moved to the next section
+        xyz_corners = _generate_corners(regular_grid=grid_0_centers.octree_grid)
+        
+        corner_grid = GenericGrid(values=xyz_corners)
+        grid_0_centers.corners_grid = corner_grid
+        output_0_centers: List[InterpOutput] = interpolate_all_fields(temp_interpolation_input, options, data_shape)  # interpolate - centers
+
+        # * DEP
+        grid_0_corners = None
+        output_0_corners = []
+
+        # * Create next octree level
+        next_octree_level = OctreeLevel(
+            grid_centers=grid_0_centers,
+            grid_corners=grid_0_corners,
+            outputs_centers=output_0_centers,
+            outputs_corners=output_0_corners
+        )
+    else:
+        grid_0_centers: EngineGrid = temp_interpolation_input.grid  # ? This could be moved to the next section
+        output_0_centers: List[InterpOutput] = interpolate_all_fields(temp_interpolation_input, options, data_shape)  # interpolate - centers
+
+        # * DEP
+        output_0_corners = []
+        grid_0_corners = None
+
+        # * Create next octree level
+        next_octree_level = OctreeLevel(
+            grid_centers=grid_0_centers,
+            grid_corners=grid_0_corners,
+            outputs_centers=output_0_centers,
+            outputs_corners=output_0_corners
+        )
+
+    return next_octree_level
+
 def _generate_corners_DEP(regular_grid: RegularGrid, level=1) -> np.ndarray:
     if regular_grid is None: raise ValueError("Regular grid is None")
 

gempy_engine/core/data/engine_grid.py

@@ -22,20 +22,23 @@ class EngineGrid:
     topography: Optional[GenericGrid] = None
     sections: Optional[GenericGrid] = None
     geophysics_grid: Optional[CenteredGrid] = None  # TODO: Not implemented this probably will need something different that the generic grid?
+    corners_grid: Optional[GenericGrid] = None  # TODO: Not implemented this probably will need something different that the generic grid?
 
     debug_vals = None
 
     # ? Should we add the number of octrees here instead of the general options
 
     def __init__(self, octree_grid: Optional[RegularGrid] = None, dense_grid: Optional[RegularGrid] = None,
                  custom_grid: Optional[GenericGrid] = None, topography: Optional[GenericGrid] = None,
-                 sections: Optional[GenericGrid] = None, geophysics_grid: Optional[CenteredGrid] = None):
+                 sections: Optional[GenericGrid] = None, geophysics_grid: Optional[CenteredGrid] = None,
+                 corners_grid: Optional[GenericGrid] = None):
         self.octree_grid = octree_grid
         self.dense_grid = dense_grid
         self.custom_grid = custom_grid
         self.topography = topography
         self.sections = sections
         self.geophysics_grid = geophysics_grid
+        self.corners_grid = corners_grid
 
     @property
     def regular_grid(self):
@@ -78,6 +81,8 @@ def values(self) -> np.ndarray:
             values.append(self.sections.values)
         if self.geophysics_grid is not None:
             values.append(self.geophysics_grid.values)
+        if self.corners_grid is not None:
+            values.append(self.corners_grid.values)
 
         values_array = BackendTensor.t.concatenate(values, dtype=BackendTensor.dtype)
         values_array = BackendTensor.t.array(values_array, dtype=BackendTensor.dtype)
@@ -131,6 +136,14 @@ def geophysics_grid_slice(self) -> slice:
             start + len(self.geophysics_grid) if self.geophysics_grid is not None else start
         )
 
+    @property
+    def corners_grid_slice(self) -> slice:
+        start = self.geophysics_grid_slice.stop
+        return slice(
+            start,
+            start + len(self.corners_grid) if self.corners_grid is not None else start
+        )
+
     @property
     def len_all_grids(self) -> int:
         return self.values.shape[0]

gempy_engine/core/data/interp_output.py

@@ -84,6 +84,14 @@ def custom_grid_values(self):
     def geophysics_grid_values(self):
         return self.block[self.grid.geophysics_grid_slice]
 
+    @property
+    def cornersGrid_values(self):
+        return self.block[self.grid.corners_grid_slice]
+
+    @property
+    def ids_cornersGrid(self):
+        return BackendTensor.t.rint(self.block[self.grid.corners_grid_slice])
+
     @property
     def ids_geophysics_grid(self):
         return BackendTensor.t.rint(self.block[self.grid.geophysics_grid_slice])
@@ -132,6 +140,21 @@ def litho_faults_ids(self):
         # Generate the unique IDs
         unique_ids = litho_ids + faults_ids * multiplier
         return unique_ids
+
+    @property
+    def litho_faults_ids_corners_grid(self):
+        if self.combined_scalar_field is None:  # * This in principle is only used for testing
+            return self.ids_cornersGrid
+
+        litho_ids = BackendTensor.t.rint(self.block[self.grid.corners_grid_slice])
+        faults_ids = BackendTensor.t.rint(self.faults_block[self.grid.corners_grid_slice])
+
+        # Get the number of unique lithology IDs
+        multiplier = len(BackendTensor.t.unique(litho_ids))
+
+        # Generate the unique IDs
+        unique_ids = litho_ids + faults_ids * multiplier
+        return unique_ids
 
     def get_block_from_value_type(self, value_type: ValueType, slice_: slice):
         match value_type:

gempy_engine/core/data/octree_level.py

@@ -48,7 +48,11 @@ def output_corners(self):  # * Alias
     @property
     def last_output_corners(self):
         return self.outputs_corners[-1]
-    
+
+    @property
+    def litho_faults_ids_corners_grid(self):
+        return self.outputs_centers[-1].litho_faults_ids_corners_grid
+
     @property
     def number_of_outputs(self):
         return len(self.outputs_centers)

gempy_engine/core/data/regular_grid.py

@@ -20,7 +20,8 @@ class RegularGrid:
     original_values: np.ndarray = field(default=None, repr=False, init=False)  #: When the regular grid is representing a octree level, only active cells are stored in values. This is the original values of the regular grid.
 
     def __len__(self):
-        return self.regular_grid_shape.prod()
+        # return self.regular_grid_shape.prod()
+        return self.values.shape[0]
 
     def __post_init__(self):
         self.regular_grid_shape = _check_and_convert_list_to_array(self.regular_grid_shape)

gempy_engine/modules/octrees_topology/_octree_internals.py

@@ -14,7 +14,7 @@
 
 def compute_next_octree_locations(prev_octree: OctreeLevel, evaluation_options: EvaluationOptions,
                                   current_octree_level: int) -> EngineGrid:
-    ids = prev_octree.last_output_corners.litho_faults_ids
+    ids = prev_octree.litho_faults_ids_corners_grid
     uv_8 = ids.reshape((-1, 8))
 
     # Old octree

tests/test_common/test_integrations/test_multi_fields.py

@@ -147,9 +147,7 @@ def test_plot_corners(unconformity_complex, n_oct_levels=2):
     interpolation_input, options, structure = unconformity_complex
     options.number_octree_levels = n_oct_levels
     solutions: Solutions = compute_model(interpolation_input, options, structure)
-    output_corners: InterpOutput = solutions.octrees_output[-1].outputs_corners[-1]
-
-    vertices = output_corners.grid.values
+    vertices = solutions.octrees_output[-1].grid_centers.corners_grid.values
     if plot_pyvista or False:
         helper_functions_pyvista.plot_pyvista(solutions.octrees_output, v_just_points=vertices)