[ENH] Add grid serialization and improve binary data handling (#1034)

# Description
Enhanced model serialization to include grid data in binary format. This PR improves the serialization process by:

1. Adding support for serializing and deserializing grid data (custom grid and topography)
2. Improving binary encoding/decoding with proper length tracking for different data sections
3. Adding docstrings to binary encoder functions for better code documentation
4. Fixing type checking in `calculate_line_coordinates_2points` to support tuples in addition to lists
5. Enabling validation of serialization by default in compute functions
6. Updating tests to work with the new serialization format

Relates to #serialization-enhancement

# 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

gempy/API/compute_API.py

@@ -91,7 +91,7 @@ def compute_model_at(gempy_model: GeoModel, at: np.ndarray,
         xyz_coord=at
     )
 
-    sol = compute_model(gempy_model, engine_config, validate_serialization=False)
+    sol = compute_model(gempy_model, engine_config, validate_serialization=True)
     return sol.raw_arrays.custom
 
 

gempy/core/data/core_utils.py

@@ -4,9 +4,9 @@
 
 
 def calculate_line_coordinates_2points(p1, p2, res):
-    if isinstance(p1, list):
+    if isinstance(p1, list) or isinstance(p1, tuple):
         p1 = np.array(p1)
-    if isinstance(p2, list):
+    if isinstance(p2, list) or isinstance(p2, tuple):
         p2 = np.array(p2)
     v = p2 - p1  # vector pointing from p1 to p2
     u = v / np.linalg.norm(v)  # normalize it

gempy/core/data/encoders/binary_encoder.py

@@ -4,12 +4,68 @@
 from ..orientations import OrientationsTable
 
 
-def deserialize_input_data_tables(binary_array: bytes, name_id_map: dict, sp_binary_length_: int) -> tuple[OrientationsTable, SurfacePointsTable]:
+def deserialize_input_data_tables(binary_array: bytes, name_id_map: dict, 
+                                  sp_binary_length_: int, ori_binary_length_: int) -> tuple[OrientationsTable, SurfacePointsTable]:
+    """
+    Deserializes binary data into two tables: OrientationsTable and SurfacePointsTable.
+
+    This function takes a binary array, a mapping of names to IDs, and lengths for
+    specific parts of the binary data to extract and deserialize two distinct data
+    tables: OrientationsTable and SurfacePointsTable. It uses the provided lengths
+    to split the binary data accordingly and reconstructs the table contents from
+    their respective binary representations.
+
+    Args:
+        binary_array (bytes): A bytes array containing the serialized data for
+            both the OrientationsTable and SurfacePointsTable.
+        name_id_map (dict): A dictionary mapping names to IDs which is used to
+            help reconstruct the table objects.
+        sp_binary_length_ (int): The length of the binary segment corresponding
+            to the SurfacePointsTable data.
+        ori_binary_length_ (int): The length of the binary segment corresponding
+            to the OrientationsTable data.
+
+    Returns:
+        tuple[OrientationsTable, SurfacePointsTable]: A tuple containing two table
+        objects: first the OrientationsTable, and second the SurfacePointsTable.
+    """
     sp_binary = binary_array[:sp_binary_length_]
-    ori_binary = binary_array[sp_binary_length_:]
+    ori_binary = binary_array[sp_binary_length_:sp_binary_length_+ori_binary_length_]
     # Reconstruct arrays
     sp_data: np.ndarray = np.frombuffer(sp_binary, dtype=SurfacePointsTable.dt)
     ori_data: np.ndarray = np.frombuffer(ori_binary, dtype=OrientationsTable.dt)
     surface_points_table = SurfacePointsTable(data=sp_data, name_id_map=name_id_map)
     orientations_table = OrientationsTable(data=ori_data, name_id_map=name_id_map)
     return orientations_table, surface_points_table
+
+
+def deserialize_grid(binary_array:bytes, custom_grid_length: int, topography_length: int) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Deserialize binary grid data into two numpy arrays.
+
+    This function takes a binary array representing a grid and splits it into two separate
+    numpy arrays: one for the custom grid and one for the topography. The binary array is
+    segmented based on the provided lengths for the custom grid and topography.
+
+    Args:
+        binary_array: The binary data representing the combined custom grid and topography data.
+        custom_grid_length: The length of the custom grid data segment in bytes.
+        topography_length: The length of the topography data segment in bytes.
+
+    Returns:
+        A tuple where the first element is a numpy array representing the custom grid, and
+        the second element is a numpy array representing the topography data.
+
+    Raises:
+        ValueError: If input lengths do not match the specified boundaries or binary data.
+    """
+
+    
+
+    custom_grid_binary = binary_array[:custom_grid_length]
+    topography_binary = binary_array[custom_grid_length:custom_grid_length + topography_length]
+    custom_grid = np.frombuffer(custom_grid_binary, dtype=np.float64)
+    topography = np.frombuffer(topography_binary)
+    
+    
+    return custom_grid, topography

gempy/core/data/encoders/converters.py

@@ -52,9 +52,10 @@ def instantiate_if_necessary(data: dict, key: str, type: type) -> None:
 loading_model_context = ContextVar('loading_model_context', default={})
 
 @contextmanager
-def loading_model_from_binary(binary_body: bytes):
+def loading_model_from_binary(input_binary: bytes, grid_binary: bytes):
     token = loading_model_context.set({
-            'binary_body': binary_body,
+            'input_binary': input_binary,
+            'grid_binary': grid_binary
     })
     try:
         yield

gempy/core/data/grid.py

@@ -8,6 +8,8 @@
 from gempy_engine.core.data.centered_grid import CenteredGrid
 from gempy_engine.core.data.options import EvaluationOptions
 from gempy_engine.core.data.transforms import Transform
+from .encoders.binary_encoder import deserialize_grid
+from .encoders.converters import loading_model_context
 from .grid_modules import RegularGrid, CustomGrid, Sections
 from .grid_modules.topography import Topography
 
@@ -59,13 +61,49 @@ def deserialize_properties(cls, data: Union["Grid", dict], constructor: ModelWra
                 case dict():
                     grid: Grid = constructor(data)
                     grid._active_grids = Grid.GridTypes(data["active_grids"])
+                    # TODO: Digest binary data
+
+                    metadata = data.get('binary_meta_data', {})
+                    context = loading_model_context.get()
+
+                    if 'grid_binary' not in context:
+                        return grid
+                    
+                    custom_grid_vals, topography_vals = deserialize_grid(
+                        binary_array=context['grid_binary'],
+                        custom_grid_length=metadata["custom_grid_binary_length"],
+                        topography_length=metadata["topography_binary_length"]
+                    )
+                    
+                    if grid.custom_grid is not None:
+                        grid.custom_grid.values = custom_grid_vals.reshape(-1, 3)
+                    
+                    if grid.topography is not None:
+                        grid.topography.set_values2d(values=topography_vals)
+                    
                     grid._update_values()
                     return grid
                 case _:
                     raise ValidationError
         except ValidationError:
             raise
 
+    @property
+    def grid_binary(self):
+        custom_grid_bytes = self._custom_grid.values.astype("float64").tobytes() if self._custom_grid else b''
+        topography_bytes = self._topography.values.astype("float64").tobytes() if self._topography else b''
+        return custom_grid_bytes + topography_bytes
+
+
+    _grid_binary_size: int = 0
+    @computed_field
+    def binary_meta_data(self) -> dict:
+        return {
+                'custom_grid_binary_length': len(self._custom_grid.values.astype("float64").tobytes()) if self._custom_grid else 0,
+                'topography_binary_length': len(self._topography.values.astype("float64").tobytes()) if self._topography else 0,
+                'grid_binary_size': self._grid_binary_size
+        }
+
     @computed_field(alias="active_grids")
     @property
     def active_grids(self) -> GridTypes:

gempy/core/data/grid_modules/topography.py

@@ -29,7 +29,7 @@ class Topography:
 
     # Fields managed internally
     values: short_array_type = field(init=False, default=np.zeros((0, 3)))
-    resolution: Tuple[int, int] = field(init=False, default=(0, 0))
+    resolution: Tuple[int, int] = Field(init=True, default=(0, 0))
     raster_shape: Tuple[int, ...] = field(init=False, default=())
     _mask_topo: Optional[np.ndarray] = field(init=False, default=None, repr=False)
     _x: Optional[np.ndarray] = field(init=False, default=None, repr=False)
@@ -95,7 +95,7 @@ def from_unstructured_mesh(cls, regular_grid, xyz_vertices):
         # Reshape the grid for compatibility with existing structure
         values_2d = np.stack((x_regular, y_regular, z_regular), axis=-1)
 
-        return cls(regular_grid=regular_grid, values_2d=values_2d)
+        return cls(_regular_grid=regular_grid, values_2d=values_2d)
 
 
     @classmethod
@@ -107,7 +107,7 @@ def from_arrays(cls, regular_grid, x_coordinates, y_coordinates, height_values,)
         topography_vals = height_values.values[:, :, np.newaxis]  # shape (73, 34, 1)
         # Stack along the last dimension
         result = np.concatenate([x_vals, y_vals, topography_vals], axis=2)  # shape (73, 34, 3)
-        return cls(regular_grid=regular_grid, values_2d=result)
+        return cls(_regular_grid=regular_grid, values_2d=result)
 
     @property
     def extent(self):
@@ -155,6 +155,34 @@ def set_values(self, values_2d: np.ndarray):
         self.values = values_2d.reshape((-1, 3), order='C')
         return self
 
+    def set_values2d(self, values: np.ndarray) -> "Topography":
+        """
+        Reconstruct the 2D topography (shape = resolution + [3]) from
+        a flat Nx3 array (or from self.values if none is provided).
+        """
+        # default to the already-flattened XYZ array
+        
+        # compute expected size
+        nx, ny = self.resolution
+        expected = nx * ny * 3
+        if values.size != expected:
+            raise ValueError(
+                f"Cannot reshape array of size {values.size} into shape {(nx, ny, 3)}."
+            )
+
+        # reshape in C-order to (nx, ny, 3)
+        self.set_values(
+            values_2d=values.reshape(nx, ny, 3, order="C")
+        )
+
+        # invalidate any cached mask
+        self._mask_topo = None
+        return self
+
+    def set_values2d_(self, values: np.ndarray):
+        resolution = (60, 60)
+        self.values_2d = values.reshape(*resolution, 3)
+
     @property
     def topography_mask(self):
         """This method takes a topography grid of the same extent as the regular

gempy/core/data/structural_frame.py

@@ -476,13 +476,14 @@ def deserialize_binary(cls, data: Union["StructuralFrame", dict], constructor: M
                 metadata = data.get('binary_meta_data', {})
                 context = loading_model_context.get()
 
-                if 'binary_body' not in context:
+                if 'input_binary' not in context:
                     return instance
 
                 instance.orientations, instance.surface_points = deserialize_input_data_tables(
-                    binary_array=context['binary_body'],
+                    binary_array=context['input_binary'],
                     name_id_map=instance.surface_points_copy.name_id_map,
-                    sp_binary_length_=metadata["sp_binary_length"]
+                    sp_binary_length_=metadata["sp_binary_length"],
+                    ori_binary_length_=metadata["ori_binary_length"]
                 )
 
                 return instance
@@ -491,12 +492,13 @@ def deserialize_binary(cls, data: Union["StructuralFrame", dict], constructor: M
 
         # Access the context variable to get injected data
 
-
+    _input_binary_size: int = 0
     @computed_field
     def binary_meta_data(self) -> dict:
         return {
                 'sp_binary_length': len(self.surface_points_copy.data.tobytes()),
-                # 'ori_binary_length': len(self.orientations_copy.data.tobytes()) * (miguel May 2025) This is not necessary at the moment
+                'ori_binary_length': len(self.orientations_copy.data.tobytes()) ,
+                'input_binary_size': self._input_binary_size
         }
 
     # endregion

gempy/modules/serialization/save_load.py

@@ -44,13 +44,7 @@ def save_model(model: GeoModel, path: str | None = None, validate_serialization:
         # If no extension, add the valid extension
         path = str(path_obj) + VALID_EXTENSION
 
-    model_json = model.model_dump_json(by_alias=True, indent=4)
-
-    # Compress the binary data
-    zlib = require_zlib()
-    compressed_binary = zlib.compress(model.structural_frame.input_tables_binary)
-
-    binary_file = _to_binary(model_json, compressed_binary)
+    binary_file = model_to_binary(model)
 
     if validate_serialization:
         model_deserialized = _deserialize_binary_file(binary_file)
@@ -67,6 +61,35 @@ def save_model(model: GeoModel, path: str | None = None, validate_serialization:
     return path  # Return the actual path used (helpful if extension was added)
 
 
+def model_to_binary(model: GeoModel) -> bytes:
+
+    # Compress the binary data
+    zlib = require_zlib()
+    compressed_binary_input = zlib.compress(model.structural_frame.input_tables_binary)
+    compressed_binary_grid = zlib.compress(model.grid.grid_binary)
+
+    compressed_binary_grid = zlib.compress(model.grid.grid_binary, level=6)
+
+    import hashlib
+    print("len raw bytes:", len(model.grid.grid_binary))
+    
+    print("raw bytes hash:", hashlib.sha256(model.grid.grid_binary).hexdigest())
+    print("compressed length:", len(compressed_binary_grid))
+    print("zlib version:", zlib.ZLIB_VERSION)
+
+    # * Add here the serialization meta parameters like: len_bytes
+    model.structural_frame._input_binary_size = len(compressed_binary_input)
+    model.grid._grid_binary_size = len(compressed_binary_grid)
+    
+    model_json = model.model_dump_json(by_alias=True, indent=4)
+    binary_file = _to_binary(
+        header_json=model_json,
+        body_input=compressed_binary_input,
+        body_grid=compressed_binary_grid
+    )
+    return binary_file
+
+
 def load_model(path: str) -> GeoModel:
     """
     Load a GeoModel from a file with extension validation.
@@ -110,39 +133,50 @@ def load_model(path: str) -> GeoModel:
 
 
 def _deserialize_binary_file(binary_file):
+    import json
     # Get header length from first 4 bytes
     header_length = int.from_bytes(binary_file[:4], byteorder='little')
     # Split header and body
-    header_json = binary_file[4:4 + header_length].decode('utf-8')
-    binary_body = binary_file[4 + header_length:]
+    header_json= binary_file[4:4 + header_length].decode('utf-8')
+    header = json.loads(header_json)
+    input_metadata = header["structural_frame"]["binary_meta_data"]
+    input_size = input_metadata["input_binary_size"]
+    
+    grid_metadata = header["grid"]["binary_meta_data"]
+    grid_size = grid_metadata["grid_binary_size"]
+    
+    input_binary = binary_file[4 + header_length: 4 + header_length + input_size]
+    all_sections_length = 4 + header_length + input_size + grid_size
+    if all_sections_length != len(binary_file):
+        raise  ValueError("Binary file is corrupted")
+    
+    grid_binary = binary_file[4 + header_length + input_size: all_sections_length]
     zlib = require_zlib()
-    decompressed_binary = zlib.decompress(binary_body)
+    
     with loading_model_from_binary(
-            binary_body=decompressed_binary,
+            input_binary=(zlib.decompress(input_binary)),
+            grid_binary=(zlib.decompress(grid_binary))
     ):
         model = GeoModel.model_validate_json(header_json)
     return model
 
 
-def _to_binary(header_json, body_) -> bytes:
+def _to_binary(header_json, body_input, body_grid) -> bytes:
     header_json_bytes = header_json.encode('utf-8')
     header_json_length = len(header_json_bytes)
     header_json_length_bytes = header_json_length.to_bytes(4, byteorder='little')
-    file = header_json_length_bytes + header_json_bytes + body_
+    file = header_json_length_bytes + header_json_bytes + body_input + body_grid
     return file
 
 
 def _validate_serialization(original_model, model_deserialized):
-    if False:
-        _verify_models(model_deserialized, original_model)
-
     a = hash(original_model.structural_frame.surface_points_copy.data.tobytes())
     b = hash(model_deserialized.structural_frame.surface_points_copy.data.tobytes())
     o_a = hash(original_model.structural_frame.orientations_copy.data.tobytes())
     o_b = hash(model_deserialized.structural_frame.orientations_copy.data.tobytes())
     assert a == b, "Hashes for surface points are not equal"
     assert o_a == o_b, "Hashes for orientations are not equal"
-    original_model___str__ =  re.sub(r'\s+', ' ', original_model.__str__())
+    original_model___str__ = re.sub(r'\s+', ' ', original_model.__str__())
     deserialized___str__ = re.sub(r'\s+', ' ', model_deserialized.__str__())
     if original_model___str__ != deserialized___str__:
         # Find first char that is not the same

requirements/dev-requirements.txt

@@ -3,6 +3,7 @@
 pytest
 pyvista
 subsurface-terra
+wellpathpy
 
 # Testing
 pytest-approvaltests

test/test_modules/_geophysics_TO_UPDATE/test_gravity.py

@@ -59,7 +59,7 @@ def test_gravity():
         structural_frame=frame,
     )
 
-    gp.compute_model(geo_model, validate_serialization=False)
+    gp.compute_model(geo_model, validate_serialization=True)
 
     import gempy_viewer as gpv
     gpv.plot_2d(geo_model, cell_number=0)