[CLN] Cleaning the tests

Author: Leguark

gempy_engine/modules/activator/activator_interface.py

@@ -124,7 +124,11 @@ def soft_segment_unbounded(Z, edges, ids, sigmoid_slope):
 
     # weighted sum by the ids
     ids__sum = (membership * ids).sum(dim=-1)
-    return np.atleast_2d(ids__sum.numpy())
+    
+    # make it at least 2d
+    ids__sum = ids__sum[None, :]
+
+    return ids__sum
 
 
 import numpy as np

tests/test_common/test_modules/test_activator_fns.py

@@ -17,83 +17,66 @@
 
 
 def test_activator_3_layers_segmentation_function(simple_model_3_layers, simple_grid_3d_more_points_grid):
-    interpolation_input = simple_model_3_layers[0]
-    options = simple_model_3_layers[1]
-    data_shape = simple_model_3_layers[2].tensors_structure
-    grid = dataclasses.replace(simple_grid_3d_more_points_grid)
-    interpolation_input.set_temp_grid(grid)
-
-    ids = np.array([1, 2, 3, 4])
+    Z_x, grid, ids_block, interpolation_input = _run_test(
+        backend=AvailableBackends.numpy,
+        ids=np.array([1, 20, 3, 4]),
+        simple_grid_3d_more_points_grid=simple_grid_3d_more_points_grid,
+        simple_model_3_layers=simple_model_3_layers
+    )
 
-    interp_input: SolverInput = input_preprocess(data_shape, interpolation_input)
-    weights = _solve_interpolation(interp_input, options.kernel_options)
+    if plot:
+        _plot_continious(grid, ids_block, interpolation_input)
 
-    exported_fields = _evaluate_sys_eq(interp_input, weights, options)
-    exported_fields.set_structure_values(
-        reference_sp_position=data_shape.reference_sp_position,
-        slice_feature=interpolation_input.slice_feature,
-        grid_size=interpolation_input.grid.len_all_grids)
 
-    Z_x: np.ndarray = exported_fields.scalar_field
-    sasp = exported_fields.scalar_field_at_surface_points
-    ids = np.array([1, 20, 3, 4])
+def test_activator_3_layers_segmentation_function_II(simple_model_3_layers, simple_grid_3d_more_points_grid):
+    Z_x, grid, ids_block, interpolation_input = _run_test(
+        backend=AvailableBackends.numpy,
+        ids=np.array([1, 2, 3, 4]),
+        simple_grid_3d_more_points_grid=simple_grid_3d_more_points_grid,
+        simple_model_3_layers=simple_model_3_layers
+    )
 
-    print(Z_x, Z_x.shape[0])
-    print(sasp)
+    BackendTensor.change_backend_gempy(AvailableBackends.numpy)
 
-    ids_block = activate_formation_block(
-        exported_fields=exported_fields,
-        ids=ids,
-        sigmoid_slope = 500*4
-    )[0, :-7]
+    if plot:
+        _plot_continious(grid, ids_block, interpolation_input)
 
-    if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
-        ids_block = ids_block.detach().numpy()
-        Z_x = Z_x.detach().numpy()
-        interpolation_input.surface_points.sp_coords = interpolation_input.surface_points.sp_coords.detach().numpy()
 
+def test_activator_3_layers_segmentation_function_torch(simple_model_3_layers, simple_grid_3d_more_points_grid):
+    Z_x, grid, ids_block, interpolation_input = _run_test(
+        backend=AvailableBackends.PYTORCH,
+        ids=np.array([1, 2, 3, 4]),
+        simple_grid_3d_more_points_grid=simple_grid_3d_more_points_grid,
+        simple_model_3_layers=simple_model_3_layers
+    )
     if plot:
         _plot_continious(grid, ids_block, interpolation_input)
 
 
-def test_activator_3_layers_segmentation_function_II(simple_model_3_layers, simple_grid_3d_more_points_grid):
+def _run_test(backend, ids, simple_grid_3d_more_points_grid, simple_model_3_layers):
     interpolation_input = simple_model_3_layers[0]
     options = simple_model_3_layers[1]
     data_shape = simple_model_3_layers[2].tensors_structure
     grid = dataclasses.replace(simple_grid_3d_more_points_grid)
     interpolation_input.set_temp_grid(grid)
-
     interp_input: SolverInput = input_preprocess(data_shape, interpolation_input)
     weights = _solve_interpolation(interp_input, options.kernel_options)
-
     exported_fields = _evaluate_sys_eq(interp_input, weights, options)
     exported_fields.set_structure_values(
         reference_sp_position=data_shape.reference_sp_position,
         slice_feature=interpolation_input.slice_feature,
         grid_size=interpolation_input.grid.len_all_grids)
-
     Z_x: np.ndarray = exported_fields.scalar_field
     sasp = exported_fields.scalar_field_at_surface_points
-    ids = np.array([1, 2, 3, 4])
-
     print(Z_x, Z_x.shape[0])
     print(sasp)
-
-    BackendTensor.change_backend_gempy(AvailableBackends.numpy)
+    BackendTensor.change_backend_gempy(backend)
     ids_block = activate_formation_block(
         exported_fields=exported_fields,
         ids=ids,
         sigmoid_slope=500 * 4
     )[0, :-7]
-
-    BackendTensor.change_backend_gempy(AvailableBackends.numpy)
-    if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
-        ids_block = ids_block.detach().numpy()
-        Z_x = Z_x.detach().numpy()
-        interpolation_input.surface_points.sp_coords = interpolation_input.surface_points.sp_coords.detach().numpy()
-
-    if plot:
-        _plot_continious(grid, ids_block, interpolation_input)
+    return Z_x, grid, ids_block, interpolation_input
 
 
 def _plot_continious(grid, ids_block, interpolation_input):
@@ -113,5 +96,3 @@ def _plot_continious(grid, ids_block, interpolation_input):
     plt.plot(xyz[:, 0], xyz[:, 2], "o")
     plt.colorbar()
     plt.show()
-
-