[FEATURE] Add mesh extraction module using marching cubes algorithm

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

@@ -0,0 +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 set_meshes_with_marching_cubes(model: GeoModel) -> None:
+    """Extract meshes for all structural elements using the marching cubes algorithm.
+
+    Parameters
+    ----------
+    model : GeoModel
+        The geological model containing solutions and structural elements.
+
+    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 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]])
+
+    # Store mesh in the structural element
+    structural_element.vertices = verts
+    structural_element.edges = faces

test/test_modules/test_marching_cubes.py

@@ -4,35 +4,11 @@
 import gempy as gp
 from gempy.core.data.enumerators import ExampleModel
 from gempy.core.data.grid_modules import RegularGrid
+from gempy.modules.mesh_extranction import marching_cubes
 from gempy.optional_dependencies import require_gempy_viewer
 
-from skimage import measure
-
 PLOT = True
 
-def marching_cubes(block, elements, spacing, extent):
-    """
-        Extract the surface meshes using marching cubes
-        Args:
-            block (np.array): The block to extract the surface meshes from.
-            elements (list): IDs of unique structural elements in model
-            spacing (tuple): The spacing between grid points in the block.
-
-        Returns:
-            mc_vertices (list): Vertices of the surface meshes.
-            mc_edges (list): Edges of the surface meshes.
-        """
-
-    # Extract the surface meshes using marching cubes
-    mc_vertices = []
-    mc_edges = []
-    for i in range(0, len(elements)):
-        verts, faces, _, _ = measure.marching_cubes(block, i,
-                                                    spacing=spacing)
-        mc_vertices.append(verts + [extent[0], extent[2], extent[4]])
-        mc_edges.append(faces)
-    return mc_vertices, mc_edges
-
 
 def test_marching_cubes_implementation():
     model = gp.generate_example_model(ExampleModel.COMBINATION, compute_model=False)
@@ -50,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)
 
@@ -60,121 +37,13 @@ def test_marching_cubes_implementation():
 
     assert arrays.scalar_field_matrix.shape == (3, 8_000)  # * 3 surfaces, 8000 points
 
-    # TODO: Maybe to complicated because it includes accounting for faults, multiple elements in groups
-    #  and transformation to real coordinates
-
-    # 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:
-        # 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)
-    else:
-        pass
-
-    # Extract scalar field values for elements
-    scalar_values = model.solutions.raw_arrays.scalar_field_at_surface_points
-
-    # Get indices of non fault elements
-    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))
-
-    # Extract marching cubes for non fault elements
-    for idx in false_indices:
-
-        # 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)
-
-    # Reorder everything correctly if faults exist
-    # TODO: All of the following is just complicated code to reorder the elements to match the order of the elements
-    #  in the structural frame, probably unnecessary in gempy strucuture
-    #
-    # if faults is not None:
-    #
-    #     # TODO: This is a very convoluted way to get a boolean list of faults per element
-    #     bool_list = np.zeros(4, dtype=bool)
-    #     for i in range(len(model.structural_frame.structural_groups)):
-    #         print(i)
-    #         if model.structural_frame.group_is_fault[i]:
-    #             for j in range(len(model.structural_frame.structural_groups[i].elements)):
-    #                 bool_list[i + j] = True
-    #         if not model.structural_frame.group_is_fault[i]:
-    #             for k in range(len(model.structural_frame.structural_groups[i].elements)):
-    #                 bool_list[i + k] = False
-    #
-    #     true_count = sum(bool_list)
-    #
-    #     # Split arr_list into two parts
-    #     true_elements_vertices = mc_vertices[:true_count]
-    #     false_elements_vertices = mc_vertices[true_count:]
-    #     true_elements_edges = mc_edges[:true_count]
-    #     false_elements_edges = mc_edges[true_count:]
-    #
-    #     # Create a new list to store reordered elements
-    #     mc_vertices = []
-    #     mc_edges = []
-    #
-    #     # Iterator for both true and false elements
-    #     true_idx, false_idx = 0, 0
-    #
-    #     # Populate reordered_list based on bool_list
-    #     for is_true in bool_list:
-    #         if is_true:
-    #             mc_vertices.append(true_elements_vertices[true_idx] + [model.grid.regular_grid.extent[0],
-    #                                                                    model.grid.regular_grid.extent[2],
-    #                                                                    model.grid.regular_grid.extent[4]])
-    #             mc_edges.append(true_elements_edges[true_idx])
-    #             true_idx += 1
-    #         else:
-    #             mc_vertices.append(false_elements_vertices[false_idx] + [model.grid.regular_grid.extent[0],
-    #                                                                      model.grid.regular_grid.extent[2],
-    #                                                                      model.grid.regular_grid.extent[4]])
-    #             mc_edges.append(false_elements_edges[false_idx])
-    #             false_idx += 1
+    marching_cubes.set_meshes_with_marching_cubes(model)
 
     if PLOT:
         gpv = require_gempy_viewer()
-        # gtv: gpv.GemPyToVista = gpv.plot_3d(model, show_data=True, image=True)
-        import pyvista as pv
-        # pyvista_plotter: pv.Plotter = gtv.p
-
-        # TODO: This opens interactive window as of now
-        pyvista_plotter = pv.Plotter()
-
-        # 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()
+        gtv: gpv.GemPyToVista = gpv.plot_3d(
+            model=model,
+            show_data=True,
+            image=True,
+            show=True
+        )