[CLN] Small cleaning while reviewing new code

Author: Leguark

gempy/modules/mesh_extranction/marching_cubes.py

@@ -22,16 +22,16 @@ def set_meshes_with_marching_cubes(model: GeoModel) -> None:
         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):
+    solution_not_having_dense: bool = model.solutions.block_solution_type != RawArraysSolution.BlockSolutionType.DENSE_GRID
+    if model.solutions is None or solution_not_having_dense:
         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
 
     # TODO: How to get this properly in gempy
@@ -50,26 +50,12 @@ def set_meshes_with_marching_cubes(model: GeoModel) -> None:
     scalar_values = model.solutions.raw_arrays.scalar_field_at_surface_points
 
     # TODO: Here I just get my own masks, cause the gempy masks dont work as expected
-    masks = []
-    masks.append(np.ones_like(model.solutions.raw_arrays.scalar_field_matrix[0].reshape(model.grid.regular_grid.resolution),
-                              dtype=bool))
-    for idx in lith_group_indices:
-        mask = model.solutions.raw_arrays.scalar_field_matrix[idx].reshape(model.grid.regular_grid.resolution) <= \
-               scalar_values[idx][-1]
-
-        masks.append(mask)
+    masks = _get_masking_arrays(lith_group_indices, model, scalar_values)
 
     # TODO: Attribute of element.scalar_field was None, changed it to scalar field value of that element
     #  This should probably be done somewhere else and maybe renamed to scalar_field_value?
     #  This is just the most basic solution to be clear what I did
-    counter = 0
-    for e, structural_group in enumerate(structural_groups):
-        if e >= len(output_lvl0):
-            continue
-
-        for element in structural_group.elements:
-            element.scalar_field = model.solutions.scalar_field_at_surface_points[counter]
-            counter += 1
+    _set_scalar_field_to_element(model, output_lvl0, structural_groups)
 
     # Trying to use the exiting gempy masks
     # masks = []
@@ -96,7 +82,7 @@ def set_meshes_with_marching_cubes(model: GeoModel) -> None:
         if structural_group.is_fault:
             mask = np.ones_like(scalar_field, dtype=bool)
         else:
-            mask = masks[non_fault_counter] # TODO: I need the entry without faults here
+            mask = masks[non_fault_counter]  # TODO: I need the entry without faults here
             non_fault_counter += 1
 
         for element in structural_group.elements:
@@ -108,9 +94,34 @@ def set_meshes_with_marching_cubes(model: GeoModel) -> None:
             )
 
 
-def extract_mesh_for_element(structural_element: StructuralElement, 
-                             regular_grid: RegularGrid, 
-                             scalar_field: np.ndarray, 
+# TODO: This should be set somewhere else
+def _set_scalar_field_to_element(model, output_lvl0, structural_groups):
+    counter = 0
+    for e, structural_group in enumerate(structural_groups):
+        if e >= len(output_lvl0):
+            continue
+
+        for element in structural_group.elements:
+            element.scalar_field = model.solutions.scalar_field_at_surface_points[counter]
+            counter += 1
+
+
+# TODO: This should be set somewhere else
+def _get_masking_arrays(lith_group_indices, model, scalar_values):
+    masks = []
+    masks.append(np.ones_like(model.solutions.raw_arrays.scalar_field_matrix[0].reshape(model.grid.regular_grid.resolution),
+                              dtype=bool))
+    for idx in lith_group_indices:
+        mask = model.solutions.raw_arrays.scalar_field_matrix[idx].reshape(model.grid.regular_grid.resolution) <= \
+               scalar_values[idx][-1]
+
+        masks.append(mask)
+    return masks
+
+
+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.
     
@@ -132,12 +143,12 @@ def extract_mesh_for_element(structural_element: StructuralElement,
         spacing=(regular_grid.dx, regular_grid.dy, regular_grid.dz),
         mask=mask
     )
-    
+
     # 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
+    structural_element.edges = faces