[ENH] Add model serialization validation and refactor grid modules (#1032)

# Description
Refactored grid modules to improve serialization and added validation capabilities to the model computation process. Key changes include:

- Split grid_types.py into separate files for each grid type (RegularGrid, CustomGrid, Sections)
- Converted CustomGrid and Sections to dataclasses with proper field definitions
- Added serialization validation during model computation via environment variable
- Fixed parameter names in grid API functions to match updated class structures
- Added dotenv support to compute_API for configuration
- Updated geophysics_input to be properly serializable
- Added verification tests for model serialization in multiple test cases
- Added approved verification files for serialization tests

Relates to #serialization-improvements

# Checklist
- [x] My code uses type hinting for function and method arguments and return values.
- [x] I have created tests which cover my code.
- [x] The test code either 1. demonstrates at least one valuable use case (e.g. integration tests) 
or 2. verifies that outputs are as expected for given inputs (e.g. unit tests).
- [x] New tests pass locally with my changes.

Author: Leguark

docs/developers_notes/dev_log/2025_05.md

@@ -26,5 +26,8 @@
   - Solutions (Not sure if I want to save them)
       - RawArraysSolutions
       - Solutions
+- Testing 
+  - [ ] 10 Test to go in modules
+  - [ ] test api
 
 ## What do I have in the engine server logic?

gempy/API/compute_API.py

@@ -1,4 +1,7 @@
-from typing import Optional
+import dotenv
+import os
+
+from typing import Optional
 
 import numpy as np
 
@@ -14,8 +17,11 @@
 from ..modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
 from ..optional_dependencies import require_gempy_legacy
 
+dotenv.load_dotenv()
+
 
-def compute_model(gempy_model: GeoModel, engine_config: Optional[GemPyEngineConfig] = None) -> Solutions:
+def compute_model(gempy_model: GeoModel, engine_config: Optional[GemPyEngineConfig] = None,
+                  **kwargs) -> Solutions:
     """
     Compute the geological model given the provided GemPy model.
 
@@ -56,6 +62,12 @@ def compute_model(gempy_model: GeoModel, engine_config: Optional[GemPyEngineConf
         case _:
             raise ValueError(f'Backend {engine_config} not supported')
 
+    if os.getenv("VALIDATE_SERIALIZATION", False) and kwargs.get("validate_serialization", True):
+        from ..modules.serialization.save_load import save_model
+        import tempfile
+        with tempfile.NamedTemporaryFile(mode='w+', delete=True) as tmp:
+            save_model(model=gempy_model, path=tmp.name, validate_serialization=True)
+
     return gempy_model.solutions
 
 
@@ -79,7 +91,7 @@ def compute_model_at(gempy_model: GeoModel, at: np.ndarray,
         xyz_coord=at
     )
 
-    sol = compute_model(gempy_model, engine_config)
+    sol = compute_model(gempy_model, engine_config, validate_serialization=False)
     return sol.raw_arrays.custom
 
 

gempy/API/grid_API.py

@@ -11,7 +11,10 @@
 
 def set_section_grid(grid: Grid, section_dict: dict):
     if grid.sections is None:
-        grid.sections = Sections(regular_grid=grid.regular_grid, section_dict=section_dict)
+        grid.sections = Sections(
+            z_ext=grid.regular_grid.extent[4:],
+            section_dict=section_dict
+        )
     else:
         grid.sections.set_sections(section_dict,
                                    regular_grid=grid.regular_grid)
@@ -54,9 +57,9 @@ def set_topography_from_random(grid: Grid, fractal_dimension: float = 2.0, d_z:
         dz=d_z,
         fractal_dimension=fractal_dimension
     )
-    
+
     grid.topography = Topography(
-        regular_grid=grid.regular_grid,
+        _regular_grid=grid.regular_grid,
         values_2d=random_topography
     )
 
@@ -70,7 +73,7 @@ def set_topography_from_subsurface_structured_grid(grid: Grid, struct: "subsurfa
     return grid.topography
 
 
-def set_topography_from_arrays(grid: Grid,  xyz_vertices: np.ndarray):
+def set_topography_from_arrays(grid: Grid, xyz_vertices: np.ndarray):
     grid.topography = Topography.from_unstructured_mesh(grid.regular_grid, xyz_vertices)
     set_active_grid(grid, [Grid.GridTypes.TOPOGRAPHY])
     return grid.topography
@@ -86,9 +89,9 @@ def set_topography_from_file(grid: Grid, filepath: str, crop_to_extent: Union[Se
 
 
 def set_custom_grid(grid: Grid, xyz_coord: np.ndarray):
-    custom_grid = CustomGrid(xyx_coords=xyz_coord)
+    custom_grid = CustomGrid(values=xyz_coord)
     grid.custom_grid = custom_grid
-    
+
     set_active_grid(grid, [Grid.GridTypes.CUSTOM])
     return grid.custom_grid
 

gempy/core/data/encoders/converters.py

@@ -1,3 +1,5 @@
+from typing import Annotated
+
 from contextlib import contextmanager
 
 from contextvars import ContextVar
@@ -17,6 +19,9 @@ def validate_numpy_array(v):
     return np.array(v) if v is not None else None
 
 
+short_array_type = Annotated[np.ndarray, (BeforeValidator(lambda v: np.array(v) if v is not None else None))]
+
+
 def instantiate_if_necessary(data: dict, key: str, type: type) -> None:
     """
     Creates instances of the specified type for a dictionary key if the key exists and its

gempy/core/data/geo_model.py

@@ -15,6 +15,7 @@
 from gempy_engine.core.data.interpolation_input import InterpolationInput
 from gempy_engine.core.data.raw_arrays_solution import RawArraysSolution
 from gempy_engine.core.data.transforms import Transform, GlobalAnisotropy
+from gempy_engine.modules.geophysics.gravity_gradient import calculate_gravity_gradient
 from .encoders.converters import instantiate_if_necessary
 from .encoders.json_geomodel_encoder import encode_numpy_array
 from .grid import Grid
@@ -23,6 +24,7 @@
 from .surface_points import SurfacePointsTable
 from ...modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
 
+
 """
 TODO:
     - [ ] StructuralFrame will all input points chunked on Elements. Here I will need a property to put all
@@ -62,7 +64,7 @@ class GeoModel(BaseModel):
     # region GemPy engine data types
     _interpolation_options: InterpolationOptions  #: The interpolation options provided by the user.
 
-    geophysics_input: GeophysicsInput = Field(default=None, exclude=True)  #: The geophysics input of the geological model.
+    geophysics_input: GeophysicsInput | None = Field(default=None, exclude=False)  #: The geophysics input of the geological model.
     input_transform: Transform = Field(default=None, exclude=False)  #: The transformation used in the geological model for input points.
 
     interpolation_grid: EngineGrid = Field(default=None, exclude=True)  #: ptional grid used for interpolation. Can be seen as a cache field.
@@ -295,24 +297,30 @@ def add_surface_points(self, X: Sequence[float], Y: Sequence[float], Z: Sequence
         arbitrary_types_allowed=True,
         use_enum_values=False,
         json_encoders={
-                np.ndarray: encode_numpy_array
+                np.ndarray: encode_numpy_array,
         }
     )
 
     @model_validator(mode='wrap')
     @classmethod
     def deserialize_properties(cls, data: Union["GeoModel", dict], constructor: ModelWrapValidatorHandler["GeoModel"]) -> "GeoModel":
         match data:
-            case GeoModel():
+            case GeoModel(): 
                 return data
-            case dict():
+            case dict(): # 
                 instance: GeoModel = constructor(data)
                 instantiate_if_necessary(
                     data=data,
                     key="_interpolation_options",
                     type=InterpolationOptions
                 )
                 instance._interpolation_options = data.get("_interpolation_options")
+                
+                # * Reset geophysics if necessary
+                centered_grid = instance.grid.centered_grid
+                if centered_grid is not None and instance.geophysics_input is not None:
+                    instance.geophysics_input.tz = calculate_gravity_gradient(centered_grid)
+                
                 return instance
             case _:
                 raise ValidationError

gempy/core/data/grid_modules/__init__.py

@@ -1,2 +1,4 @@
-from .grid_types import Sections, RegularGrid, CustomGrid
+from .regular_grid import RegularGrid
+from .custom_grid import CustomGrid
+from .sections_grid import Sections
 from .topography import Topography

gempy/core/data/grid_modules/custom_grid.py

@@ -0,0 +1,33 @@
+import dataclasses
+import numpy as np
+from pydantic import Field
+
+
+@dataclasses.dataclass
+class CustomGrid:
+    """Object that contains arbitrary XYZ coordinates.
+
+    Args:
+        xyx_coords (numpy.ndarray like): XYZ (in columns) of the desired coordinates
+
+    Attributes:
+        values (np.ndarray): XYZ coordinates
+    """
+
+    values: np.ndarray = Field(
+        exclude=True, 
+        default_factory=lambda: np.zeros((0, 3)),
+        repr=False
+    )
+    
+       
+    def __post_init__(self):
+        custom_grid = np.atleast_2d(self.values)
+        assert type(custom_grid) is np.ndarray and custom_grid.shape[1] == 3, \
+            'The shape of new grid must be (n,3)  where n is the number of' \
+            ' points of the grid'
+
+    
+    @property
+    def length(self):
+        return self.values.shape[0]

gempy/core/data/grid_modules/grid_types.py renamed to gempy/core/data/grid_modules/regular_grid.py

@@ -5,11 +5,9 @@
 
 import numpy as np
 
-from ..core_utils import calculate_line_coordinates_2points
 from ..encoders.converters import numpy_array_short_validator
 from .... import optional_dependencies
-from ....optional_dependencies import require_pandas
-from gempy_engine.core.data.transforms import Transform, TransformOpsOrder
+from gempy_engine.core.data.transforms import Transform
 
 
 @dataclasses.dataclass
@@ -255,145 +253,3 @@ def plot_rotation(regular_grid, pivot, point_x_axis, point_y_axis):
         plt.show()
 
 
-class Sections:
-    """
-    Object that creates a grid of cross sections between two points.
-
-    Args:
-        regular_grid: Model.grid.regular_grid
-        section_dict: {'section name': ([p1_x, p1_y], [p2_x, p2_y], [xyres, zres])}
-    """
-
-    def __init__(self, regular_grid=None, z_ext=None, section_dict=None):
-        pd = require_pandas()
-        if regular_grid is not None:
-            self.z_ext = regular_grid.extent[4:]
-        else:
-            self.z_ext = z_ext
-
-        self.section_dict = section_dict
-        self.names = []
-        self.points = []
-        self.resolution = []
-        self.length = [0]
-        self.dist = []
-        self.df = pd.DataFrame()
-        self.df['dist'] = self.dist
-        self.values = np.empty((0, 3))
-        self.extent = None
-
-        if section_dict is not None:
-            self.set_sections(section_dict)
-
-    def _repr_html_(self):
-        return self.df.to_html()
-
-    def __repr__(self):
-        return self.df.to_string()
-
-    def show(self):
-        pass
-
-    def set_sections(self, section_dict, regular_grid=None, z_ext=None):
-        pd = require_pandas()
-        self.section_dict = section_dict
-        if regular_grid is not None:
-            self.z_ext = regular_grid.extent[4:]
-
-        self.names = np.array(list(self.section_dict.keys()))
-
-        self.get_section_params()
-        self.calculate_all_distances()
-        self.df = pd.DataFrame.from_dict(self.section_dict, orient='index', columns=['start', 'stop', 'resolution'])
-        self.df['dist'] = self.dist
-
-        self.compute_section_coordinates()
-
-    def get_section_params(self):
-        self.points = []
-        self.resolution = []
-        self.length = [0]
-
-        for i, section in enumerate(self.names):
-            points = [self.section_dict[section][0], self.section_dict[section][1]]
-            assert points[0] != points[
-                1], 'The start and end points of the section must not be identical.'
-
-            self.points.append(points)
-            self.resolution.append(self.section_dict[section][2])
-            self.length = np.append(self.length, self.section_dict[section][2][0] *
-                                    self.section_dict[section][2][1])
-        self.length = np.array(self.length).cumsum()
-
-    def calculate_all_distances(self):
-        self.coordinates = np.array(self.points).ravel().reshape(-1,
-                                                                 4)  # axis are x1,y1,x2,y2
-        self.dist = np.sqrt(np.diff(self.coordinates[:, [0, 2]]) ** 2 + np.diff(
-            self.coordinates[:, [1, 3]]) ** 2)
-
-    def compute_section_coordinates(self):
-        for i in range(len(self.names)):
-            xy = calculate_line_coordinates_2points(self.coordinates[i, :2],
-                                                    self.coordinates[i, 2:],
-                                                    self.resolution[i][0])
-            zaxis = np.linspace(self.z_ext[0], self.z_ext[1], self.resolution[i][1],
-                                dtype="float64")
-            X, Z = np.meshgrid(xy[:, 0], zaxis, indexing='ij')
-            Y, _ = np.meshgrid(xy[:, 1], zaxis, indexing='ij')
-            xyz = np.vstack((X.flatten(), Y.flatten(), Z.flatten())).T
-            if i == 0:
-                self.values = xyz
-            else:
-                self.values = np.vstack((self.values, xyz))
-
-    def generate_axis_coord(self):
-        for i, name in enumerate(self.names):
-            xy = calculate_line_coordinates_2points(
-                self.coordinates[i, :2],
-                self.coordinates[i, 2:],
-                self.resolution[i][0]
-            )
-            yield name, xy
-
-    def get_section_args(self, section_name: str):
-        where = np.where(self.names == section_name)[0][0]
-        return self.length[where], self.length[where + 1]
-
-    def get_section_grid(self, section_name: str):
-        l0, l1 = self.get_section_args(section_name)
-        return self.values[l0:l1]
-
-
-class CustomGrid:
-    """Object that contains arbitrary XYZ coordinates.
-
-    Args:
-        xyx_coords (numpy.ndarray like): XYZ (in columns) of the desired coordinates
-
-    Attributes:
-        values (np.ndarray): XYZ coordinates
-    """
-
-    def __init__(self, xyx_coords: np.ndarray):
-        self.values = np.zeros((0, 3))
-        self.set_custom_grid(xyx_coords)
-
-    def set_custom_grid(self, custom_grid: np.ndarray):
-        """
-        Give the coordinates of an external generated grid
-
-        Args:
-            custom_grid (numpy.ndarray like): XYZ (in columns) of the desired coordinates
-
-        Returns:
-              numpy.ndarray: Unraveled 3D numpy array where every row correspond to the xyz coordinates of a regular
-               grid
-        """
-        custom_grid = np.atleast_2d(custom_grid)
-        assert type(custom_grid) is np.ndarray and custom_grid.shape[1] == 3, \
-            'The shape of new grid must be (n,3)  where n is the number of' \
-            ' points of the grid'
-
-        self.values = custom_grid
-        self.length = self.values.shape[0]
-        return self.values