Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
[CLN] Remove all the unnecessary code during serialization
  • Loading branch information
Leguark committed May 23, 2025
commit 0308f5a8d01a106c22e99cb9296dde0c516f47e6
12 changes: 0 additions & 12 deletions gempy/core/data/encoders/converters.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,18 +46,6 @@ def instantiate_if_necessary(data: dict, key: str, type: type) -> None:
# First, create a context variable
loading_model_context = ContextVar('loading_model_context', default={})

@contextmanager
def loading_model_injection(surface_points_binary: np.ndarray, orientations_binary: np.ndarray):
token = loading_model_context.set({
'surface_points_binary': surface_points_binary,
'orientations_binary' : orientations_binary
})
try:
yield
finally:
loading_model_context.reset(token)


@contextmanager
def loading_model_from_binary(binary_body: bytes):
token = loading_model_context.set({
Expand Down
47 changes: 4 additions & 43 deletions gempy/core/data/structural_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,10 @@ def orientations(self, modified_orientations: OrientationsTable) -> None:
"""Distributes the modified orientations back to the structural elements."""
for element in self.structural_elements:
element.orientations.data = modified_orientations.get_orientations_by_id(element.id).data

@property
def input_tables_binary(self):
return self.surface_points_copy.data.tobytes() + self.orientations_copy.data.tobytes()

@property
def element_id_name_map(self) -> dict[int, str]:
Expand Down Expand Up @@ -489,49 +493,6 @@ def deserialize_binary(cls, data: Union["StructuralFrame", dict], constructor: M
# Access the context variable to get injected data


@model_validator(mode="after")
def deserialize_surface_points(self: "StructuralFrame"):
# Access the context variable to get injected data
context = loading_model_context.get()

if 'surface_points_binary' not in context:
return self

# Check if we have a binary payload to digest
binary_array = context['surface_points_binary']
if not isinstance(binary_array, np.ndarray):
return self
if binary_array.shape[0] < 1:
return self

self.surface_points = SurfacePointsTable(
data=binary_array,
name_id_map=self.surface_points_copy.name_id_map
)

return self

@model_validator(mode="after")
def deserialize_orientations(self: "StructuralFrame"):
# TODO: Check here the binary size of surface_points_binary

# Access the context variable to get injected data
context = loading_model_context.get()
if 'orientations_binary' not in context:
return self

# Check if we have a binary payload to digest
binary_array = context['orientations_binary']
if not isinstance(binary_array, np.ndarray):
return self

self.orientations = OrientationsTable(
data=binary_array,
name_id_map=self.orientations_copy.name_id_map
)

return self

@computed_field
def binary_meta_data(self) -> dict:
return {
Expand Down
54 changes: 9 additions & 45 deletions gempy/modules/serialization/save_load.py
Original file line number Diff line number Diff line change
@@ -1,42 +1,25 @@
import json

import numpy as np

from ...core.data import GeoModel
from ...core.data.encoders.converters import loading_model_injection, loading_model_from_binary
from ...core.data.encoders.converters import loading_model_from_binary
from ...optional_dependencies import require_zlib


def save_model(model: GeoModel, path: str):
import zlib

# TODO: Serialize to json
model_json = model.model_dump_json(by_alias=True, indent=4)

# TODO: Serialize to binary
data: np.ndarray = model.structural_frame.surface_points_copy.data
sp_binary = data.tobytes()
ori_binary = model.structural_frame.orientations_copy.data.tobytes()

# Compress the binary data
compressed_binary = zlib.compress(sp_binary + ori_binary)
zlib = require_zlib()
compressed_binary = zlib.compress(model.structural_frame.input_tables_binary)

# Add compression info to metadata
# model_dict = model.model_dump(by_alias=True)
# model_dict["_binary_metadata"] = {
# "sp_shape" : model.structural_frame.surface_points_copy.data.shape,
# "sp_dtype" : str(model.structural_frame.surface_points_copy.data.dtype),
# "ori_shape" : model.structural_frame.orientations_copy.data.shape,
# "ori_dtype" : str(model.structural_frame.orientations_copy.data.dtype),
# "compression": "zlib",
# "sp_length" : len(sp_binary) # Need this to split the arrays after decompression
# }

# TODO: Putting both together
binary_file = _to_binary(model_json, compressed_binary)

# TODO: Add validation

with open(path, 'wb') as f:
f.write(binary_file)



def load_model(path: str) -> GeoModel:
with open(path, 'rb') as f:
binary_file = f.read()
Expand All @@ -46,14 +29,8 @@ def load_model(path: str) -> GeoModel:

# Split header and body
header_json = binary_file[4:4 + header_length].decode('utf-8')
header_dict = json.loads(header_json)

# metadata = header_dict.pop("_binary_metadata")

# Decompress the binary data
# ori_data, sp_data = _foo(binary_file, header_length, metadata)

binary_body = binary_file[4 + header_length:]

zlib = require_zlib()
decompressed_binary = zlib.decompress(binary_body)
with loading_model_from_binary(
Expand All @@ -64,19 +41,6 @@ def load_model(path: str) -> GeoModel:
return model


def _foo(binary_file, header_length, metadata):
zlib = require_zlib()
body = binary_file[4 + header_length:]
decompressed_binary = zlib.decompress(body)
# Split the decompressed data using the stored length
sp_binary = decompressed_binary[:metadata["sp_length"]]
ori_binary = decompressed_binary[metadata["sp_length"]:]
# Reconstruct arrays
sp_data = np.frombuffer(sp_binary, dtype=np.dtype(metadata["sp_dtype"])).reshape(metadata["sp_shape"])
ori_data = np.frombuffer(ori_binary, dtype=np.dtype(metadata["ori_dtype"])).reshape(metadata["ori_shape"])
return ori_data, sp_data


def _to_binary(header_json, body_) -> bytes:
header_json_bytes = header_json.encode('utf-8')
header_json_length = len(header_json_bytes)
Expand Down
11 changes: 5 additions & 6 deletions test/test_modules/test_serialize_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,12 @@
import os
import pprint

from gempy_engine.core.data import InterpolationOptions

import gempy as gp
from gempy.core.data.encoders.converters import loading_model_injection
from gempy.core.data.encoders.converters import loading_model_from_binary
from gempy.core.data.enumerators import ExampleModel
from gempy.modules.serialization.save_load import save_model, load_model
from gempy_engine.core.data import InterpolationOptions
from test.verify_helper import verify_json


Expand All @@ -20,10 +21,8 @@ def test_generate_horizontal_stratigraphic_model():
with open(file_path, "w") as f:
f.write(model_json)

# TODO: modify this for the binary
with loading_model_injection(
surface_points_binary=model.structural_frame.surface_points_copy.data, # TODO: Here we need to pass the binary array
orientations_binary=model.structural_frame.orientations_copy.data
with loading_model_from_binary(
binary_body=model.structural_frame.input_tables_binary
):
model_deserialized = gp.data.GeoModel.model_validate_json(model_json)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -249,8 +249,9 @@
},
"is_dirty": true,
"basement_color": "#ffbe00",
"serialize_sp": 3507338795,
"serialize_orientations": 553806131
"binary_meta_data": {
"sp_binary_length": 432
}
},
"grid": {
"_octree_grid": null,
Expand Down