[ENH] Refactor dual contouring code with improved documentation and modular functions (#25)

Leguark · web-flow · commit bcac51db0f73 · 2025-09-25T10:45:27.000+02:00
# Refactor dual contouring implementation for improved clarity and maintainability

This PR refactors the `dual_contouring_multi_scalar` function to improve code organization, readability, and maintainability. The changes include:

- Added comprehensive docstrings to explain function purpose and parameters
- Extracted helper functions to break down complex logic into smaller, focused components:
  - `_get_triangulation_codes`: Determines appropriate triangulation strategy
  - `_validate_stack_relations`: Validates stack relation configurations
  - `_get_masked_codes`: Applies masks to triangulation codes
  - `_interp_on_edges`: Handles interpolation on edge intersection points
  - Enhanced `_mask_generation` with better documentation

The refactoring preserves all existing functionality while making the code more modular and easier to understand. Edge intersection detection and processing are now more clearly separated, improving the overall architecture.
diff --git a/gempy_engine/API/dual_contouring/multi_scalar_dual_contouring.py b/gempy_engine/API/dual_contouring/multi_scalar_dual_contouring.py
@@ -1,6 +1,6 @@
 import copy
 import warnings
-from typing import List
+from typing import List, Any
 
 import numpy as np
 
@@ -19,12 +19,29 @@
 from ...core.data.options import MeshExtractionMaskingOptions
 from ...core.data.stack_relation_type import StackRelationType
 from ...core.utils import gempy_profiler_decorator
+from ...modules.dual_contouring.dual_contouring_interface import find_intersection_on_edge
 from ...modules.dual_contouring.fancy_triangulation import get_left_right_array
 
 
 @gempy_profiler_decorator
-def dual_contouring_multi_scalar(data_descriptor: InputDataDescriptor, interpolation_input: InterpolationInput,
-                                 options: InterpolationOptions, octree_list: list[OctreeLevel]) -> List[DualContouringMesh]:
+def dual_contouring_multi_scalar(
+    data_descriptor: InputDataDescriptor,
+    interpolation_input: InterpolationInput,
+    options: InterpolationOptions,
+    octree_list: List[OctreeLevel]
+) -> List[DualContouringMesh]:
+    """
+    Perform dual contouring for multiple scalar fields.
+    
+    Args:
+        data_descriptor: Input data descriptor containing stack structure information
+        interpolation_input: Input data for interpolation
+        options: Interpolation options including debug and extraction settings
+        octree_list: List of octree levels with the last being the leaf level
+        
+    Returns:
+        List of dual contouring meshes for all processed scalar fields
+    """
     # Dual Contouring prep:
     MaskBuffer.clean()
 
@@ -35,110 +52,251 @@ def dual_contouring_multi_scalar(data_descriptor: InputDataDescriptor, interpola
     dual_contouring_options.evaluation_options.compute_scalar_gradient = True
 
     if options.debug_water_tight:
-        _experimental_water_tight(all_meshes, data_descriptor, interpolation_input, octree_leaves, dual_contouring_options)
+        _experimental_water_tight(
+            all_meshes, data_descriptor, interpolation_input, octree_leaves, dual_contouring_options
+        )
         return all_meshes
 
-    # region new triangulations
-    is_pure_octree = bool(np.all(octree_list[0].grid_centers.octree_grid_shape == 2))
-    match (options.evaluation_options.mesh_extraction_fancy, is_pure_octree):
-        case (True, True):
-            left_right_codes = get_left_right_array(octree_list)
-        case (True, False):
-            left_right_codes = None
-            warnings.warn("Fancy triangulation only works with regular grid of resolution [2,2,2]. Defaulting to regular triangulation")
-        case (False, _):
-            left_right_codes = None
-        case _:
-            raise ValueError("Invalid combination of options")
-    # endregion
+    # Determine triangulation strategy
+    left_right_codes = _get_triangulation_codes(octree_list, options)
 
+    # Generate masks for all scalar fields
     all_mask_arrays: np.ndarray = _mask_generation(
         octree_leaves=octree_leaves,
         masking_option=options.evaluation_options.mesh_extraction_masking_options
     )
 
+    # Process each scalar field
+    all_stack_intersection = []
+    all_valid_edges = []
+    all_left_right_codes = []
+
     for n_scalar_field in range(data_descriptor.stack_structure.n_stacks):
-        previous_stack_is_onlap = data_descriptor.stack_relation[n_scalar_field - 1] == 'Onlap'
-        was_erosion_before = data_descriptor.stack_relation[n_scalar_field - 1] == 'Erosion'
-        if previous_stack_is_onlap and was_erosion_before:  # ? (July, 2023) Is this still valid? I thought we have all the combinations
-            raise NotImplementedError("Erosion and Onlap are not supported yet")
-            pass
+        _validate_stack_relations(data_descriptor, n_scalar_field)
 
         mask: np.ndarray = all_mask_arrays[n_scalar_field]
+        left_right_codes_per_stack = _get_masked_codes(left_right_codes, mask)
 
-        if mask is not None and left_right_codes is not None:
-            left_right_codes_per_stack = left_right_codes[mask]
-        else:
-            left_right_codes_per_stack = left_right_codes
-
-        # @off
-        dc_data: DualContouringData = interpolate_on_edges_for_dual_contouring(
-            data_descriptor=data_descriptor,
-            interpolation_input=interpolation_input,
-            options=dual_contouring_options,
-            n_scalar_field=n_scalar_field,
-            octree_leaves=octree_leaves,
-            mask=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
         )
 
+        all_stack_intersection.append(intersection_xyz)
+        all_valid_edges.append(valid_edges)
+        all_left_right_codes.append(left_right_codes_per_stack)
+
+    # Interpolate on edges for all stacks
+    output_on_edges = _interp_on_edges(
+        all_stack_intersection, data_descriptor, dual_contouring_options, interpolation_input
+    )
+
+    # Generate meshes for each scalar field
+    for n_scalar_field in range(data_descriptor.stack_structure.n_stacks):
+        output: InterpOutput = octree_leaves.outputs_centers[n_scalar_field]
+        mask = all_mask_arrays[n_scalar_field]
+        
+        dc_data = DualContouringData(
+            xyz_on_edge=all_stack_intersection[n_scalar_field],
+            valid_edges=all_valid_edges[n_scalar_field],
+            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.scalar_field_at_sp.shape[0],
+            tree_depth=options.number_octree_levels,
+        )
+        
         meshes: List[DualContouringMesh] = compute_dual_contouring(
             dc_data_per_stack=dc_data,
-            left_right_codes=left_right_codes_per_stack,
+            left_right_codes=all_left_right_codes[n_scalar_field],
             debug=options.debug
         )
 
-        # ! If the order of the meshes does not match the order of scalar_field_at_surface points we need to reorder them HERE
-
+        # TODO: If the order of the meshes does not match the order of scalar_field_at_surface points, reorder them here
         if meshes is not None:
             all_meshes.extend(meshes)
-        # @on
 
     return all_meshes
 
 
-def _mask_generation(octree_leaves, masking_option: MeshExtractionMaskingOptions) -> np.ndarray | None:
-    all_scalar_fields_outputs: list[InterpOutput] = octree_leaves.outputs_centers
+def _get_triangulation_codes(octree_list: List[OctreeLevel], options: InterpolationOptions) -> np.ndarray | None:
+    """
+    Determine the appropriate triangulation codes based on options and octree structure.
+    
+    Args:
+        octree_list: List of octree levels
+        options: Interpolation options
+        
+    Returns:
+        Left-right codes array if fancy triangulation is enabled and supported, None otherwise
+    """
+    is_pure_octree = bool(np.all(octree_list[0].grid_centers.octree_grid_shape == 2))
+    
+    match (options.evaluation_options.mesh_extraction_fancy, is_pure_octree):
+        case (True, True):
+            return get_left_right_array(octree_list)
+        case (True, False):
+            warnings.warn(
+                "Fancy triangulation only works with regular grid of resolution [2,2,2]. "
+                "Defaulting to regular triangulation"
+            )
+            return None
+        case (False, _):
+            return None
+        case _:
+            raise ValueError("Invalid combination of options")
+
+
+def _validate_stack_relations(data_descriptor: InputDataDescriptor, n_scalar_field: int) -> None:
+    """
+    Validate stack relations for the given scalar field.
+    
+    Args:
+        data_descriptor: Input data descriptor containing stack relations
+        n_scalar_field: Current scalar field index
+        
+    Raises:
+        NotImplementedError: If unsupported combination of Erosion and Onlap is detected
+    """
+    if n_scalar_field == 0:
+        return
+        
+    previous_stack_is_onlap = data_descriptor.stack_relation[n_scalar_field - 1] == 'Onlap'
+    was_erosion_before = data_descriptor.stack_relation[n_scalar_field - 1] == 'Erosion'
+    
+    if previous_stack_is_onlap and was_erosion_before:
+        # TODO (July, 2023): Is this still valid? I thought we have all the combinations
+        raise NotImplementedError("Erosion and Onlap are not supported yet")
+
+
+def _get_masked_codes(left_right_codes: np.ndarray | None, mask: np.ndarray | None) -> np.ndarray | None:
+    """
+    Apply mask to left-right codes if both are available.
+    
+    Args:
+        left_right_codes: Original left-right codes array
+        mask: Boolean mask array
+        
+    Returns:
+        Masked codes if both inputs are not None, otherwise original codes
+    """
+    if mask is not None and left_right_codes is not None:
+        return left_right_codes[mask]
+    return left_right_codes
+
+
+def _interp_on_edges(
+    all_stack_intersection: List[Any],
+    data_descriptor: InputDataDescriptor,
+    dual_contouring_options: InterpolationOptions,
+    interpolation_input: InterpolationInput
+) -> List[InterpOutput]:
+    """
+    Interpolate scalar fields on edge intersection points.
+    
+    Args:
+        all_stack_intersection: List of intersection points for all stacks
+        data_descriptor: Input data descriptor
+        dual_contouring_options: Dual contouring specific options
+        interpolation_input: Interpolation input data
+        
+    Returns:
+        List of interpolation outputs for each stack
+    """
+    from ...core.data.engine_grid import EngineGrid
+    from ...core.data.generic_grid import GenericGrid
+    from ..interp_single.interp_features import interpolate_all_fields_no_octree
+    
+    # Set temporary grid with concatenated intersection points
+    interpolation_input.set_temp_grid(
+        EngineGrid(
+            custom_grid=GenericGrid(
+                values=BackendTensor.t.concatenate(all_stack_intersection, axis=0)
+            )
+        )
+    )
+
+    # TODO (@miguel 21 June): By definition in `interpolate_all_fields_no_octree`
+    # we just need to interpolate up to the n_scalar_field, but need to test this
+    # This should be done with buffer weights to avoid waste
+    output_on_edges: List[InterpOutput] = interpolate_all_fields_no_octree(
+        interpolation_input=interpolation_input,
+        options=dual_contouring_options,
+        data_descriptor=data_descriptor
+    )
+    
+    # Restore original grid
+    interpolation_input.set_grid_to_original()
+    return output_on_edges
+
+
+def _mask_generation(
+    octree_leaves: OctreeLevel,
+    masking_option: MeshExtractionMaskingOptions
+) -> np.ndarray | None:
+    """
+    Generate masks for mesh extraction based on masking options and stack relations.
+    
+    Args:
+        octree_leaves: Octree leaf level containing scalar field outputs
+        masking_option: Mesh extraction masking configuration
+        
+    Returns:
+        Matrix of boolean masks for each scalar field
+        
+    Raises:
+        NotImplementedError: For unsupported masking options
+        ValueError: For invalid option combinations
+    """
+    all_scalar_fields_outputs: List[InterpOutput] = octree_leaves.outputs_centers
     n_scalar_fields = len(all_scalar_fields_outputs)
     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
 
     for i in range(n_scalar_fields):
         stack_relation = all_scalar_fields_outputs[i].scalar_fields.stack_relation
+        
         match (masking_option, stack_relation):
             case MeshExtractionMaskingOptions.RAW, _:
                 mask_matrix[i] = BackendTensor.t.ones(grid_size // 8, dtype=bool)
+                
             case MeshExtractionMaskingOptions.DISJOINT, _:
-                raise NotImplementedError("Disjoint is not supported yet. Not even sure if there is anything to support")
-            # case (MeshExtractionMaskingOptions.DISJOINT | MeshExtractionMaskingOptions.INTERSECT, StackRelationType.FAULT):
-            #     mask_matrix[i] = np.ones(grid_size//8, dtype=bool)
-            # case MeshExtractionMaskingOptions.DISJOINT, StackRelationType.ERODE | StackRelationType.BASEMENT:
-            #     mask_scalar = all_scalar_fields_outputs[i - 1].squeezed_mask_array.reshape((1, -1, 8)).sum(-1, bool)[0]
-            #     if MaskBuffer.previous_mask is None:
-            #         mask = mask_scalar
-            #     else:
-            #         mask = (MaskBuffer.previous_mask ^ mask_scalar) * mask_scalar
-            #     MaskBuffer.previous_mask = mask
-            # case MeshExtractionMaskingOptions.DISJOINT, StackRelationType.ONLAP:
-            #     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]
+                raise NotImplementedError(
+                    "Disjoint is not supported yet. Not even sure if there is anything to support"
+                )
+                
             case MeshExtractionMaskingOptions.INTERSECT, StackRelationType.ERODE:
-                x = all_scalar_fields_outputs[i + onlap_chain_counter].squeezed_mask_array[slice_corners].reshape((1, -1, 8))
+                mask_array = all_scalar_fields_outputs[i + onlap_chain_counter].squeezed_mask_array
+                x = 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[slice_corners].reshape((1, -1, 8))
+                mask_array = all_scalar_fields_outputs[i].squeezed_mask_array
+                x = 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[slice_corners].reshape((1, -1, 8))
+                mask_array = all_scalar_fields_outputs[i].squeezed_mask_array
+                x = 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:
                 mask_matrix[i] = BackendTensor.t.ones(grid_size // 8, dtype=bool)
+                
             case _:
                 raise ValueError("Invalid combination of options")
 
-    return mask_matrix
+    return mask_matrix
