[TEST] Adding test to visualize gempy posterior models

Author: Leguark

tests/test_posterior_analysis/__init__.py

@@ -0,0 +1,54 @@
+import numpy as np
+import os
+import pyro.distributions as dist
+import torch
+
+import gempy as gp
+from gempy_engine.core.backend_tensor import BackendTensor
+import arviz as az
+import matplotlib.pyplot as plt
+
+
+def test_gempy_posterior_I():
+    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=3,
+        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")
+        )
+    )
+
+    geo_model.interpolation_options.uni_degree = 0
+    geo_model.interpolation_options.mesh_extraction = False
+    geo_model.interpolation_options.sigmoid_slope = 1100.
+
+    data = az.from_netcdf("../arviz_data.nc")
+
+    if PLOT:= False:
+        _plot_posterior(data)
+
+
+def _plot_posterior(data):
+    az.plot_trace(data)
+    plt.show()
+    from gempy_probability.modules.plot.plot_posterior import default_red, default_blue
+    az.plot_density(
+        data=[data.posterior_predictive, data.prior_predictive],
+        shade=.9,
+        var_names=[r'$\mu_{thickness}$'],
+        data_labels=["Posterior Predictive", "Prior Predictive"],
+        colors=[default_red, default_blue],
+    )
+    plt.show()
+    az.plot_density(
+        data=[data, data.prior],
+        shade=.9,
+        hdi_prob=.99,
+        data_labels=["Posterior", "Prior"],
+        colors=[default_red, default_blue],
+    )
+    plt.show()

tests/test_posterior_analysis/test_gempy_posterior_I.py

@@ -20,6 +20,9 @@ def test_basic_gempy_I() -> None:
         )
     )
 
+    # TODO: Make this a more elegant way 
+    BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
+    
     # TODO: Convert this into an options preset
     geo_model.interpolation_options.uni_degree = 0
     geo_model.interpolation_options.mesh_extraction = False
@@ -37,8 +40,6 @@ def test_basic_gempy_I() -> None:
     assert geo_model.grid.values.shape[0] == 100, "Custom grid should have 100 cells"
     gp.compute_model(gempy_model=geo_model)
 
-    # TODO: Make this a more elegant way 
-    BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
 
     normal = dist.Normal(
         loc=(geo_model.surface_points_copy_transformed.xyz[0, 2]),
@@ -105,6 +106,9 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable,
     data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive)
     # endregion
 
+    if True: # * Save the arviz data
+        data.to_netcdf("arviz_data.nc")
+        
     az.plot_trace(data)
     plt.show()
     from gempy_probability.modules.plot.plot_posterior import default_red, default_blue
@@ -140,5 +144,3 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable,
     print(f"Thickness mean: {posterior_thickness_mean}")
     print("MCMC convergence diagnostics:")
     print(f"Divergences: {float(data.sample_stats.diverging.sum())}")
-    print(f"Acceptance rate: {float(data.sample_stats.acceptance_rate.mean())}")
-    print(f"Mean tree depth: {float(data.sample_stats.mean_tree_depth.mean())}")