[ENH] Adding first version of marching cubes

Author: Leguark

gempy/core/data/geo_model.py

@@ -147,6 +147,8 @@ def solutions(self) -> Solutions:
 
     @solutions.setter
     def solutions(self, value):
+        # * This is set  from the gempy engine
+        
         self._solutions = value
 
         # * Set solutions per group

gempy/modules/mesh_extranction/marching_cubes.py

@@ -1,64 +1,88 @@
 import numpy as np
+from typing import Optional
 from skimage import measure
+from gempy_engine.core.data.interp_output import InterpOutput
+from gempy_engine.core.data.raw_arrays_solution import RawArraysSolution
 
+from gempy.core.data import GeoModel, StructuralElement, StructuralGroup
+from gempy.core.data.grid_modules import RegularGrid
 
-def compute_marching_cubes(model):
-    # Empty lists to store vertices and edges
-    mc_vertices = []
-    mc_edges = []
-    # Boolean list of fault groups
-    faults = model.structural_frame.group_is_fault
-    # MC for faults, directly on fault block not on scalar field
-    if faults is not None:
-        _extract_fault_mesh(mc_edges, mc_vertices, model)
-    else:
-        pass
+
+def set_meshes_with_marching_cubes(model: GeoModel) -> None:
+    """Extract meshes for all structural elements using the marching cubes algorithm.
     
-    # Extract scalar field values for elements
-    scalar_values = model.solutions.raw_arrays.scalar_field_at_surface_points
+    Parameters
+    ----------
+    model : GeoModel
+        The geological model containing solutions and structural elements.
     
-    # Get indices of non fault elements
-    false_indices = _get_lithology_idx(faults, model)
-    # Extract marching cubes for non fault elements
-    for idx in false_indices:
-        _extract_meshes_for_lithologies(idx, mc_edges, mc_vertices, model, scalar_values)
-    return mc_edges, mc_vertices
-
-
-def _extract_meshes_for_lithologies(idx, mc_edges, mc_vertices, model, scalar_values):
-    # Get correct scalar field for structural group
-    scalar_field = model.solutions.raw_arrays.scalar_field_matrix[idx].reshape(model.grid.regular_grid.resolution)
-    # Extract marching cubes for each scalar value for all elements of a group
-    for i in range(len(scalar_values[idx])):
-        verts, faces, _, _ = measure.marching_cubes(scalar_field, scalar_values[idx][i],
-                                                    spacing=(model.grid.regular_grid.dx,
-                                                             model.grid.regular_grid.dy,
-                                                             model.grid.regular_grid.dz))
-
-        mc_vertices.append(verts + [model.grid.regular_grid.extent[0],
-                                    model.grid.regular_grid.extent[2],
-                                    model.grid.regular_grid.extent[4]])
-        mc_edges.append(faces)
+    Raises
+    ------
+    ValueError
+        If the model solutions do not contain dense grid data.
+    """
+    # Verify that solutions contain dense grid data
+    if (model.solutions is None or 
+            model.solutions.block_solution_type != RawArraysSolution.BlockSolutionType.DENSE_GRID):
+        raise ValueError("Model solutions must contain dense grid data for mesh extraction.")
+    
+    regular_grid: RegularGrid = model.grid.regular_grid
+    structural_groups: list[StructuralGroup] = model.structural_frame.structural_groups
+    
+    if not model.solutions.octrees_output or not model.solutions.octrees_output[0].outputs_centers:
+        raise ValueError("No interpolation outputs available for mesh extraction.")
+    
+    output_lvl0: list[InterpOutput] = model.solutions.octrees_output[0].outputs_centers
 
+    for e, structural_group in enumerate(structural_groups):
+        if e >= len(output_lvl0):
+            continue
+            
+        output_group: InterpOutput = output_lvl0[e]
+        scalar_field_matrix = output_group.exported_fields_dense_grid.scalar_field
+        
+        for element in structural_group.elements:
+            extract_mesh_for_element(
+                structural_element=element,
+                regular_grid=regular_grid,
+                scalar_field=scalar_field_matrix
+            )
 
-def _get_lithology_idx(faults, model):
-    if faults is not None:
-        false_indices = [i for i, fault in enumerate(faults) if not fault]
-    else:
-        false_indices = np.arange(len(model.structural_frame.structural_groups))
-    return false_indices
 
+def extract_mesh_for_element(structural_element: StructuralElement, 
+                             regular_grid: RegularGrid, 
+                             scalar_field: np.ndarray, 
+                             mask: Optional[np.ndarray] = None) -> None:
+    """Extract a mesh for a single structural element using marching cubes.
+    
+    Parameters
+    ----------
+    structural_element : StructuralElement
+        The structural element for which to extract a mesh.
+    regular_grid : RegularGrid
+        The regular grid defining the spatial discretization.
+    scalar_field : np.ndarray
+        The scalar field used for isosurface extraction.
+    mask : np.ndarray, optional
+        Optional mask to restrict the mesh extraction to specific regions.
+    """
+    # Apply mask if provided
+    volume = scalar_field.reshape(regular_grid.resolution)
+    if mask is not None:
+        volume = volume * mask
+    
+    # Extract mesh using marching cubes
+    verts, faces, _, _ = measure.marching_cubes(
+        volume=volume,
+        level=structural_element.scalar_field,
+        spacing=(regular_grid.dx, regular_grid.dy, regular_grid.dz)
+    )
+    
+    # Adjust vertices to correct coordinates in the model's extent
+    verts = (verts + [regular_grid.extent[0],
+                      regular_grid.extent[2],
+                      regular_grid.extent[4]])
 
-def _extract_fault_mesh(mc_edges, mc_vertices, model):
-    # TODO: This should also use the scalar fields probably
-    for i in np.unique(model.solutions.raw_arrays.fault_block)[:-1]:
-        fault_block = model.solutions.raw_arrays.fault_block.reshape(model.grid.regular_grid.resolution)
-        verts, faces, _, _ = measure.marching_cubes(fault_block,
-                                                    i,
-                                                    spacing=(model.grid.regular_grid.dx,
-                                                             model.grid.regular_grid.dy,
-                                                             model.grid.regular_grid.dz))
-        mc_vertices.append(verts + [model.grid.regular_grid.extent[0],
-                                    model.grid.regular_grid.extent[2],
-                                    model.grid.regular_grid.extent[4]])
-        mc_edges.append(faces)
+    # Store mesh in the structural element
+    structural_element.vertices = verts
+    structural_element.edges = faces

test/test_modules/test_marching_cubes.py

@@ -7,9 +7,9 @@
 from gempy.modules.mesh_extranction import marching_cubes
 from gempy.optional_dependencies import require_gempy_viewer
 
-
 PLOT = True
 
+
 def test_marching_cubes_implementation():
     model = gp.generate_example_model(ExampleModel.COMBINATION, compute_model=False)
 
@@ -26,6 +26,7 @@ def test_marching_cubes_implementation():
         reset=True
     )
 
+    model.interpolation_options.evaluation_options.number_octree_levels = 1
     model.interpolation_options.evaluation_options.mesh_extraction = False  # * Not extracting the mesh with dual contouring
     gp.compute_model(model)
 
@@ -36,28 +37,13 @@ def test_marching_cubes_implementation():
 
     assert arrays.scalar_field_matrix.shape == (3, 8_000)  # * 3 surfaces, 8000 points
 
-    mc_edges, mc_vertices = marching_cubes.compute_marching_cubes(model)
+    marching_cubes.set_meshes_with_marching_cubes(model)
 
     if PLOT:
         gpv = require_gempy_viewer()
         gtv: gpv.GemPyToVista = gpv.plot_3d(
             model=model,
             show_data=True,
-            image=False,
-            show=False
+            image=True,
+            show=True
         )
-        import pyvista as pv
-
-        # TODO: This opens interactive window as of now
-        pyvista_plotter: pv.Pltter = gtv.p
-
-        # Add the meshes to the plot
-        for i in range(len(mc_vertices)):
-            pyvista_plotter.add_mesh(
-                pv.PolyData(mc_vertices[i],
-                            np.insert(mc_edges[i], 0, 3, axis=1).ravel()),
-                color='blue')
-
-        pyvista_plotter.show()
-
-