Skip to content

Commit a4c9ee2

Browse files
committed
refactor: extract ModelSerializer from GeologicalModel
GeologicalModel mixed recipe/pickle I/O in with feature-construction orchestration, evaluation, and export logic (2580 lines total). Move the recipe/JSON and dill/pickle serialization logic into a standalone ModelSerializer, keeping every GeologicalModel public method as a thin unchanged-signature wrapper so the CI-enforced stable API surface (API.md, tests/unit/test_public_api_contract.py) is unaffected. (cherry picked from commit 404ce4b)
1 parent 9c0a278 commit a4c9ee2

2 files changed

Lines changed: 250 additions & 189 deletions

File tree

Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,240 @@
1+
"""Recipe (JSON) and pickle serialization logic for GeologicalModel.
2+
3+
Extracted from GeologicalModel to separate serialization concerns from
4+
feature-container orchestration (see API.md). GeologicalModel's
5+
``@public_api``-decorated methods (``to_dict``, ``to_recipe_dict``,
6+
``from_recipe_dict``, ``to_recipe_json``, ``from_recipe_json``,
7+
``save_recipe``, ``load_recipe``, ``to_file``, ``from_file``) stay defined
8+
directly on the class -- their ``__qualname__`` is part of the CI-checked
9+
stable API surface (``tests/unit/test_public_api_contract.py``) -- and just
10+
delegate to the staticmethods here.
11+
"""
12+
13+
import json
14+
import pathlib
15+
16+
import pandas as pd
17+
18+
from ...geometry import BoundingBox
19+
from ...utils import LoopValueError, getLogger
20+
from ..features import GeologicalFeature, StructuralFrame, UnconformityFeature
21+
from ..features.fault import FaultSegment
22+
from .stratigraphic_column import StratigraphicColumn
23+
24+
logger = getLogger(__name__)
25+
26+
27+
class ModelSerializer:
28+
@staticmethod
29+
def feature_recipe_kind(feature):
30+
if isinstance(feature, GeologicalFeature):
31+
return "foliation"
32+
if isinstance(feature, StructuralFrame):
33+
return "structural_frame"
34+
if isinstance(feature, UnconformityFeature):
35+
return "unconformity"
36+
if isinstance(feature, FaultSegment):
37+
return "fault"
38+
return feature.__class__.__name__.lower()
39+
40+
@staticmethod
41+
def to_dict(model):
42+
result = {}
43+
result["model"] = {}
44+
result["model"]["features"] = [f.name for f in model.features]
45+
result['model']['bounding_box'] = model.bounding_box.to_dict()
46+
result["model"]["stratigraphic_column"] = model.stratigraphic_column
47+
return result
48+
49+
@staticmethod
50+
def to_recipe_dict(model, data_reference=None):
51+
recipe = {
52+
"schema": "LoopStructural.GeologicalModelRecipe",
53+
"version": 1,
54+
"model": {
55+
"bounding_box": model.bounding_box.to_dict(),
56+
"stratigraphic_column": model.stratigraphic_column.to_dict(),
57+
"data_source": None,
58+
"features": [],
59+
},
60+
}
61+
for feature in model.features:
62+
feature_entry = {
63+
"name": feature.name,
64+
"kind": ModelSerializer.feature_recipe_kind(feature),
65+
"faults": [
66+
fault.name
67+
for fault in getattr(feature, "faults", [])
68+
if getattr(fault, "name", None)
69+
],
70+
"regions": [],
71+
}
72+
recipe["model"]["features"].append(feature_entry)
73+
if data_reference is not None:
74+
recipe["model"]["data_source"] = {
75+
"kind": "reference",
76+
"path": str(pathlib.Path(data_reference)),
77+
}
78+
elif not model.data.empty:
79+
recipe["model"]["data_source"] = {
80+
"kind": "inline",
81+
"dataframe": model.data.to_dict(orient="split"),
82+
}
83+
return recipe
84+
85+
@staticmethod
86+
def from_recipe_dict(cls, recipe):
87+
if not isinstance(recipe, dict):
88+
raise TypeError("recipe must be a dictionary")
89+
90+
model_data = recipe.get("model", recipe)
91+
bounding_box = model_data.get("bounding_box")
92+
if isinstance(bounding_box, dict):
93+
bounding_box = BoundingBox.from_dict(bounding_box)
94+
if not isinstance(bounding_box, BoundingBox):
95+
raise TypeError("recipe must include a bounding_box dictionary")
96+
97+
model = cls(bounding_box)
98+
99+
data_source = model_data.get("data_source")
100+
if isinstance(data_source, dict):
101+
kind = data_source.get("kind")
102+
if kind == "reference":
103+
model.data = pd.read_csv(pathlib.Path(data_source["path"]))
104+
elif kind == "inline":
105+
dataframe = data_source.get("dataframe")
106+
if not isinstance(dataframe, dict):
107+
raise TypeError("inline data_source must include a dataframe dictionary")
108+
model.data = pd.DataFrame(**dataframe)
109+
elif kind is not None:
110+
raise ValueError(f"Unsupported data_source kind: {kind}")
111+
elif isinstance(data_source, str):
112+
model.data = pd.read_csv(pathlib.Path(data_source))
113+
elif data_source is not None:
114+
raise TypeError("data_source must be a dictionary, string path, or None")
115+
116+
stratigraphic_column = model_data.get("stratigraphic_column")
117+
if isinstance(stratigraphic_column, dict):
118+
model.stratigraphic_column = StratigraphicColumn.from_dict(stratigraphic_column)
119+
elif stratigraphic_column is not None:
120+
raise TypeError("stratigraphic_column must be a dictionary or None")
121+
122+
features = model_data.get("features", [])
123+
if features is None:
124+
features = []
125+
if not isinstance(features, list):
126+
raise TypeError("features must be a list")
127+
128+
feature_map = {}
129+
for feature_entry in features:
130+
if not isinstance(feature_entry, dict):
131+
raise TypeError("each feature entry must be a dictionary")
132+
feature_name = feature_entry.get("name")
133+
if not isinstance(feature_name, str):
134+
raise TypeError("each feature entry must include a string name")
135+
feature_data = model.data.loc[model.data["feature_name"] == feature_name].copy()
136+
if feature_data.empty:
137+
feature_data = None
138+
feature = model.create_and_add_foliation(feature_name, data=feature_data)
139+
if feature is None:
140+
raise ValueError(f"Could not recreate feature '{feature_name}' from recipe")
141+
feature_map[feature_name] = feature
142+
143+
for feature_entry in features:
144+
feature_name = feature_entry.get("name")
145+
fault_names = feature_entry.get("faults", [])
146+
if not isinstance(fault_names, list):
147+
raise TypeError("faults must be a list")
148+
if fault_names:
149+
feature = feature_map[feature_name]
150+
feature.faults = [feature_map[name] for name in fault_names if name in feature_map]
151+
152+
return model
153+
154+
@staticmethod
155+
def to_recipe_json(model, data_reference=None, indent=2):
156+
recipe = ModelSerializer.to_recipe_dict(model, data_reference=data_reference)
157+
return json.dumps(recipe, indent=indent)
158+
159+
@staticmethod
160+
def from_recipe_json(cls, json_str):
161+
if not isinstance(json_str, str):
162+
raise TypeError("json_str must be a string")
163+
try:
164+
recipe = json.loads(json_str)
165+
except json.JSONDecodeError as e:
166+
raise TypeError(f"json_str is not valid JSON: {e}")
167+
return ModelSerializer.from_recipe_dict(cls, recipe)
168+
169+
@staticmethod
170+
def save_recipe(model, filename, data_reference=None):
171+
filename = pathlib.Path(filename)
172+
recipe = ModelSerializer.to_recipe_dict(model, data_reference=data_reference)
173+
with open(filename, "w") as f:
174+
json.dump(recipe, f, indent=2)
175+
logger.info(f"Recipe saved to {filename}")
176+
177+
@staticmethod
178+
def load_recipe(cls, filename):
179+
filename = pathlib.Path(filename)
180+
if not filename.exists():
181+
raise FileNotFoundError(f"Recipe file not found: {filename}")
182+
with open(filename, "r") as f:
183+
recipe = json.load(f)
184+
logger.info(f"Recipe loaded from {filename}")
185+
return ModelSerializer.from_recipe_dict(cls, recipe)
186+
187+
@staticmethod
188+
def to_file(model, file):
189+
try:
190+
import dill as pickle
191+
except ImportError:
192+
logger.error("Cannot write to file, dill not installed \n" "pip install dill")
193+
return
194+
try:
195+
logger.info(f"Writing GeologicalModel to: {file}")
196+
with open(file, "wb") as handle:
197+
pickle.dump(model, handle)
198+
except pickle.PicklingError:
199+
logger.error("Error saving file")
200+
201+
@staticmethod
202+
def from_file(cls, file, allow_pickle: bool = True):
203+
if not allow_pickle:
204+
raise LoopValueError(
205+
"Pickle-based loading is disabled (allow_pickle=False). "
206+
f"Refusing to unpickle '{file}' because deserialising untrusted "
207+
"pickle/dill data can execute arbitrary code. If you generated "
208+
"this file yourself and trust its contents, call "
209+
"GeologicalModel.from_file(file, allow_pickle=True). Otherwise, "
210+
"use the JSON-based GeologicalModel.from_recipe_dict "
211+
"(paired with GeologicalModel.to_recipe_dict) as a safe "
212+
"alternative serialisation format."
213+
)
214+
logger.warning(
215+
f"Loading GeologicalModel from '{file}' using dill/pickle. "
216+
"Only load model files from trusted sources: deserialising a "
217+
"pickle file can execute arbitrary code. Pass allow_pickle=False "
218+
"to refuse pickle-based loading, or use "
219+
"GeologicalModel.from_recipe_dict for untrusted/JSON-based input."
220+
)
221+
try:
222+
import dill as pickle
223+
except ImportError:
224+
logger.error("Cannot import from file, dill not installed")
225+
return None
226+
path = pathlib.Path(file)
227+
if not path.is_file():
228+
raise LoopValueError(f"Cannot load model, file does not exist: {file}")
229+
try:
230+
with open(path, "rb") as f:
231+
model = pickle.load(f)
232+
except Exception as e:
233+
logger.error(f"Failed to load model from {file}: {e}")
234+
raise LoopValueError(f"Failed to load model from {file}: {e}") from e
235+
if isinstance(model, cls):
236+
logger.info("GeologicalModel initialised from file")
237+
return model
238+
else:
239+
logger.error(f"{file} does not contain a geological model")
240+
return None

0 commit comments

Comments
 (0)