[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

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)

tests/test_benchmark/notes.md

@@ -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/test_speed_I.py

@@ -1,17 +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"
 
-    gp.compute_model(two_wells)
+    # 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_prob_model/test_prob_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")