Add benchmarking tests and examples generator for performance analysis (#8)

# Test speed

[ENH] Add example generator and initial speed test

Introduce `examples_generator.py` to create sample geological models, including a "two wells" example. Add `test_speed_I` for benchmarking model computation and validation of interpolation options.

[ENH][TEST] Add custom grid setups and profiling for performance

Introduce custom grid configurations to improve control over model evaluations. Add profiling in test_benchmark to measure and analyze computation speeds and iterations. Update tests to track interpolation counts and backend changes.

Author: Leguark

docs/dev_logs/2025_May.md

@@ -5,7 +5,8 @@
   - [ ] Keep extracting code into modules and functions 
   - [ ] Posterior analysis for gempy 
 - [ ] Investigate speed
-- [ ] Investigate segmentation function 
+- [ ] Investigate segmentation function
+  https://chatgpt.com/share/68172f53-4f7c-8000-8f42-205a033ed899
 - [ ] Create a robust api to make sure we do not mess up assigning priors to gempy parameters
 - [ ] do we need to use index_put or just normal indexing?
 --- Vector examples
@@ -18,9 +19,13 @@
 ### Doing:
 - [x] Compressing code
 - [x] GemPy posterior visualization
-- [ ] Add proper thickness likelihood
+- [x] Add proper thickness likelihood
 - [ ] Speed
-
+  - [ ] Time test for torch CPU
+  - [ ] Time test for keops CPU
+  - [ ] Time test for keops GPU
+  - [ ] Options and config presets
+   
 ### Possible Design:
 - Geological Model
 - Priors

gempy_probability/modules/model_definition/model_examples.py

@@ -43,6 +43,8 @@ def model(geo_model: gempy.core.data.GeoModel, normal, y_obs_list):
     # endregion
 
     # region Forward model computation
+    
+    geo_model.counter +=1
 
     # * Compute the geological model
     geo_model.solutions = gempy_engine.compute_model(
@@ -51,7 +53,9 @@ def model(geo_model: gempy.core.data.GeoModel, normal, y_obs_list):
         data_descriptor=geo_model.input_data_descriptor,
         geophysics_input=geo_model.geophysics_input,
     )
-
+    # if i does not exist init
+    
+    
     # Compute and observe the thickness of the geological layer
     model_solutions: gp.data.Solutions = geo_model.solutions
     thickness = apparent_thickness_likelihood(model_solutions)

gempy_probability/plugins/examples_generator.py

@@ -0,0 +1,41 @@
+from enum import Enum, auto
+import gempy as gp
+import os
+from gempy_engine.core.backend_tensor import BackendTensor
+
+
+class ExampleProbModel(Enum):
+    TWO_WELLS = auto()
+
+
+def generate_example_model(example_model: ExampleProbModel, compute_model: bool = True) -> gp.data.GeoModel:
+    match example_model:
+        case ExampleProbModel.TWO_WELLS:
+            return _generate_two_wells_model(compute_model)
+        case _:
+            raise ValueError(f"Example model {example_model} not found.")
+
+
+def _generate_two_wells_model(compute_model: bool) -> gp.data.GeoModel:
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    data_path = os.path.abspath(os.path.join(current_dir, '..', '..', 'examples', 'tutorials', 'data'))
+    geo_model = gp.create_geomodel(
+        project_name='Wells',
+        extent=[0, 12000, -500, 500, 0, 4000],
+        refinement=4,
+        importer_helper=gp.data.ImporterHelper(
+            path_to_orientations=os.path.join(data_path, "2-layers", "2-layers_orientations.csv"),
+            path_to_surface_points=os.path.join(data_path, "2-layers", "2-layers_surface_points.csv")
+        )
+    )
+
+    BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
+
+    geo_model.interpolation_options.uni_degree = 0
+    geo_model.interpolation_options.mesh_extraction = False
+
+    if compute_model:
+        # Compute the geological model
+        gp.compute_model(geo_model)
+
+    return geo_model

tests/test_benchmark/__init__.py

@@ -0,0 +1,7 @@
+- Running 250 iterations took 28 sec. This is 8.81 it/s
+- The number of gempy evaluations are:
+  - 53 evals/s or 18 ms per eval
+  - 6 evals per iteration
+- More than 53 gempy evaluations per second seem hard to improve
+- Wait till I have a likelihood with a bigger model and a likelihood function that needs
+ big grids (i.e.) gravity

tests/test_benchmark/notes.md

@@ -0,0 +1,49 @@
+import pytest
+
+import gempy as gp
+import gempy_viewer as gpv
+import numpy as np
+
+from gempy_probability.plugins.examples_generator import generate_example_model, ExampleProbModel
+
+
+def test_speed_I():
+    
+    two_wells: gp.data.GeoModel = generate_example_model(
+        example_model=ExampleProbModel.TWO_WELLS,
+        compute_model=False
+    )
+
+    assert two_wells.interpolation_options.number_octree_levels == 4, "Number of octrees should be 4"
+
+    # region Minimal grid for the specific likelihood function
+    x_loc = 6000
+    y_loc = 0
+    z_loc = np.linspace(0, 4000, 100)
+    xyz_coord = np.array([[x_loc, y_loc, z] for z in z_loc])
+    gp.set_custom_grid(two_wells.grid, xyz_coord=xyz_coord)
+    # endregion
+
+    two_wells.grid.active_grids = gp.data.Grid.GridTypes.CUSTOM
+    
+    profiler = cProfile.Profile()
+    profiler.enable()
+    iterations = 100
+    for _ in range(iterations):
+        gp.compute_model(
+            gempy_model=two_wells,
+            engine_config=gp.data.GemPyEngineConfig(
+                backend=gp.data.AvailableBackends.numpy
+            )
+        )
+    profiler.disable()
+    stats = pstats.Stats(profiler).sort_stats('cumtime')
+    stats.print_stats(10)
+
+    if PLOT := False:
+        gpv.plot_2d(two_wells, show_scalar=False)
+
+
+@pytest.mark.skip(reason="Not implemented yet")
+def test_speed_on_gravity_likelihood():
+    pass

tests/test_benchmark/test_speed_I.py

@@ -38,8 +38,9 @@ def test_basic_gempy_I() -> None:
 
     geo_model.grid.active_grids = gp.data.Grid.GridTypes.CUSTOM
     assert geo_model.grid.values.shape[0] == 100, "Custom grid should have 100 cells"
+    geo_model.counter = 0
     gp.compute_model(gempy_model=geo_model)
-
+    BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
 
     normal = dist.Normal(
         loc=(geo_model.surface_points_copy_transformed.xyz[0, 2]),
@@ -73,6 +74,7 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable,
         num_samples=50
     )
     prior = predictive(geo_model, normal, y_obs_list)
+    print("Number of interpolations: ", geo_model.counter)
 
     data = az.from_pyro(prior=prior)
     az.plot_trace(data.prior)
@@ -103,9 +105,12 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable,
         posterior_samples=posterior_samples
     )
     posterior_predictive = posterior_predictive_fn(geo_model, normal, y_obs_list)
+
     data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive)
     # endregion
 
+    print("Number of interpolations: ", geo_model.counter)
+
     if True: # * Save the arviz data
         data.to_netcdf("arviz_data.nc")