[ENH/TEST] Add support for Greenstone model generation and corresponding test

gempy/API/compute_API.py

@@ -4,7 +4,6 @@
 from typing import Optional
 
 import gempy_engine
-from gempy.API.gp2_gp3_compatibility.gp3_to_gp2_input import gempy3_to_gempy2
 from gempy_engine.config import AvailableBackends
 from gempy_engine.core.backend_tensor import BackendTensor
 from gempy_engine.core.data import Solutions
@@ -14,7 +13,6 @@
 from ..core.data.geo_model import GeoModel
 from ..modules.data_manipulation import interpolation_input_from_structural_frame
 from ..modules.optimize_nuggets import nugget_optimizer
-from ..optional_dependencies import require_gempy_legacy
 
 dotenv.load_dotenv()
 

gempy/API/examples_generator.py

@@ -25,6 +25,8 @@ def generate_example_model(example_model: ExampleModel, compute_model: bool = Tr
             return _generate_one_fault_model_gravity(compute_model)
         case ExampleModel.GRABEN:
             return _generate_graben_model(compute_model)
+        case ExampleModel.GREENSTONE:
+            return _generate_greenstone_model(compute_model)
         case _:
             raise NotImplementedError(f"Example model {example_model} not implemented.")
 
@@ -482,3 +484,26 @@ def _generate_graben_model(compute_model: bool) -> gp.data.GeoModel:
 
     return geo_data
 
+
+def _generate_greenstone_model(compute_model: bool) -> gp.data.GeoModel:
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # Build the path relative to the test file location
+    path = os.path.join(test_dir, '..', '..', 'examples', 'data', 'gempy_models', 'Greenstone.gempy')
+    with open(path, 'rb') as f:
+        binary_file = f.read()
+
+    from gempy.modules.serialization.save_load import _load_model_from_bytes
+    geo_model: gp.data.GeoModel = _load_model_from_bytes(binary_file)
+
+    if compute_model:
+        sol = gp.compute_model(
+            gempy_model=geo_model,
+            engine_config=gp.data.GemPyEngineConfig(
+                backend=gp.data.AvailableBackends.numpy,
+                dtype='float32'
+            )
+    )
+    return geo_model
+
+

gempy/core/data/enumerators.py

@@ -9,3 +9,4 @@ class ExampleModel(Enum):
     COMBINATION = auto()
     ONE_FAULT_GRAVITY = auto()
     GRABEN = auto()
+    GREENSTONE = auto()

gempy/modules/data_manipulation/_engine_factory.py

@@ -42,11 +42,18 @@ def interpolation_input_from_structural_frame(geo_model: "gempy.data.GeoModel")
         extent_transformed=geo_model.extent_transformed_transformed_by_input
     )
 
+    weights = []
+    if geo_model.solutions is not None:
+        for stack_sol in geo_model.solutions.root_output.outputs_centers:
+            weights.append(stack_sol.weights)
+
+
     interpolation_input: InterpolationInput = InterpolationInput(
         surface_points=surface_points,
         orientations=orientations,
         grid=grid,
-        unit_values=structural_frame.elements_ids  # TODO: Here we will need to pass densities etc.
+        unit_values=structural_frame.elements_ids,  # TODO: Here we will need to pass densities etc.
+        weights=weights
     )
 
     return interpolation_input

test/test_model_types/test_example_models_I.py

@@ -80,4 +80,13 @@ def test_generate_combination_model():
         model=model,
         name="Combination Scalar Field"
     )
-
+
+
+def test_generate_greenstone_model():
+    model = gp.generate_example_model(ExampleModel.GREENSTONE, compute_model=True)
+    print(model.structural_frame)
+
+    if PLOT:
+        gpv = require_gempy_viewer()
+        gpv.plot_3d(model, image=True)
+

test/test_modules/test_cg/test_cg_solver.py

@@ -0,0 +1,82 @@
+import gempy as gp
+from gempy.core.data.enumerators import ExampleModel
+from gempy.optional_dependencies import require_gempy_viewer
+import gempy_engine.core.backend_tensor as BackendTensor
+
+from gempy_engine.modules.weights_cache.weights_cache_interface import WeightCache
+
+PLOT = True
+
+
+def test_solve_with_cg():
+    model = gp.generate_example_model(ExampleModel.GREENSTONE, compute_model=False)
+    print(model.structural_frame)
+
+    WeightCache.clear_cache()
+    BackendTensor.PYKEOPS = True
+
+    sol = gp.compute_model(
+        gempy_model=model,
+        engine_config=gp.data.GemPyEngineConfig(
+            backend=gp.data.AvailableBackends.PYTORCH,
+            use_gpu=True,
+            dtype='float64'
+        )
+    )
+
+    if PLOT:
+        gpv = require_gempy_viewer()
+        gpv.plot_3d(model, image=True)
+
+
+def test_save_weights():
+    model = gp.generate_example_model(ExampleModel.GREENSTONE, compute_model=False)
+    print(model.structural_frame)
+
+    sol = gp.compute_model(
+        gempy_model=model,
+        engine_config=gp.data.GemPyEngineConfig(
+            backend=gp.data.AvailableBackends.PYTORCH,
+            use_gpu=False,
+            dtype='float32'
+        )
+    )
+    weights1 = sol.octrees_output[0].outputs_centers[0].weights
+    weights2 = sol.octrees_output[0].outputs_centers[1].weights
+    weights3 = sol.octrees_output[0].outputs_centers[2].weights
+
+    WeightCache.clear_cache()
+    BackendTensor.PYKEOPS = True
+
+    sol = gp.compute_model(
+        gempy_model=model,
+        engine_config=gp.data.GemPyEngineConfig(
+            backend=gp.data.AvailableBackends.PYTORCH,
+            use_gpu=False,
+            dtype='float32'
+        )
+    )
+
+    if PLOT:
+        gpv = require_gempy_viewer()
+        gpv.plot_3d(model, image=True)
+
+
+def test_keops_x_torch():
+    model = gp.generate_example_model(ExampleModel.GREENSTONE, compute_model=False)
+    print(model.structural_frame)
+
+    WeightCache.clear_cache()
+    BackendTensor.PYKEOPS = True
+    sol = gp.compute_model(
+        gempy_model=model,
+        engine_config=gp.data.GemPyEngineConfig(
+            backend=gp.data.AvailableBackends.PYTORCH,
+            use_gpu=False,
+            dtype='float32'
+        )
+    )
+
+    if PLOT:
+        gpv = require_gempy_viewer()
+        gpv.plot_3d(model, image=True)