[CLN] Tidying up GeoModel

Author: Leguark

gempy/core/data/geo_model.py

@@ -49,39 +49,19 @@ class GeoModelMeta:
     owner: str | None
 
 
-# @dataclass(init=False)
 class GeoModel(BaseModel):
     """
     Class representing a geological model.
 
     """
 
-    model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-        use_enum_values=False,
-        json_encoders={
-                np.ndarray: encode_numpy_array
-        }
-    )
-
     meta: GeoModelMeta = Field(exclude=False)  #: Meta-information about the geological model, like its name, creation and modification dates, and owner.
     structural_frame: StructuralFrame = Field(exclude=False)  #: The structural information of the geological model.
     grid: Grid = Field(exclude=False, default=None)  #: The general grid used in the geological model.
 
     # region GemPy engine data types
     _interpolation_options: InterpolationOptions  #: The interpolation options provided by the user.
 
-    @computed_field(alias="_interpolation_options")
-    @property
-    def interpolation_options(self) -> InterpolationOptions:
-        self._infer_dense_grid_solution()
-        self._interpolation_options.cache_model_name = self.meta.name
-        return self._interpolation_options
-
-    @interpolation_options.setter
-    def interpolation_options(self, value):
-        self._interpolation_options = value
-
     geophysics_input: GeophysicsInput = Field(default=None, exclude=True)  #: 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.
 
@@ -92,78 +72,22 @@ def interpolation_options(self, value):
     # endregion
     _solutions: Solutions = PrivateAttr(init=False, default=None)  #: The computed solutions of the geological model. 
 
-    @model_validator(mode='wrap')
-    @classmethod
-    def deserialize_properties(cls, data: Union["GeoModel", dict], constructor: ModelWrapValidatorHandler["GeoModel"]) -> "GeoModel":
-        try:
-            match data:
-                case GeoModel():
-                    return data
-                case dict():
-                    instance: GeoModel = constructor(data)
-                    instantiate_if_necessary(
-                        data=data,
-                        key="_interpolation_options",
-                        type=InterpolationOptions
-                    )
-                    instance._interpolation_options = data.get("_interpolation_options")
-                    return instance
-                case _:
-                    raise ValidationError
-        except ValidationError:
-            raise
-
-    @classmethod
-    def from_args(cls, name: str, structural_frame: StructuralFrame, grid: Grid, interpolation_options: InterpolationOptions):
-        # TODO: Fill the arguments properly
-        meta = GeoModelMeta(
-            name=name,
-            creation_date=datetime.datetime.now().isoformat(),
-            last_modification_date=None,
-            owner=None
-        )
-
-        structural_frame = structural_frame  # ? This could be Optional
-
-        grid = grid
-        _interpolation_options = interpolation_options
-        input_transform = Transform.from_input_points(
-            surface_points=structural_frame.surface_points_copy,
-            orientations=structural_frame.orientations_copy
-        )
-
-        model = GeoModel(
-            meta=meta,
-            structural_frame=structural_frame,
-            grid=grid,
-            input_transform=input_transform,
-            _interpolation_options=_interpolation_options
-        )
-
-        return model
-
     def __repr__(self):
         # TODO: Improve this
         return pprint.pformat(self.__dict__)
 
-    def update_transform(self, auto_anisotropy: GlobalAnisotropy = GlobalAnisotropy.NONE, anisotropy_limit: Optional[np.ndarray] = None):
-        """Update the transformation of the geological model.
-
-            This function updates the transformation of the geological model using the provided surface points and orientations.
-            It also applies anisotropy based on the specified type and limit.
-
-            Args:
-                auto_anisotropy (GlobalAnisotropy): The type of anisotropy to apply. Defaults to GlobalAnisotropy.NONE.
-                anisotropy_limit (Optional[np.ndarray]): Anisotropy limit values. If None, no limit is applied.
-
-        """
+    # region Properties
 
-        self.input_transform = Transform.from_input_points(
-            surface_points=self.surface_points_copy,
-            orientations=self.orientations_copy
-        )
+    @computed_field(alias="_interpolation_options")
+    @property
+    def interpolation_options(self) -> InterpolationOptions:
+        self._infer_dense_grid_solution()
+        self._interpolation_options.cache_model_name = self.meta.name
+        return self._interpolation_options
 
-        self.input_transform.apply_anisotropy(anisotropy_type=auto_anisotropy, anisotropy_limit=anisotropy_limit)
+    @interpolation_options.setter
+    def interpolation_options(self, value):
+        self._interpolation_options = value
 
     @property
     def solutions(self) -> Solutions:
@@ -297,10 +221,98 @@ def input_data_descriptor(self) -> InputDataDescriptor:
         # TODO: This should have the exact same dirty logic as interpolation_input
         return self.structural_frame.input_data_descriptor
 
+    # endregion
+    # region Constructors
+    @classmethod
+    def from_args(cls, name: str, structural_frame: StructuralFrame, grid: Grid, interpolation_options: InterpolationOptions):
+        # TODO: Fill the arguments properly
+        meta = GeoModelMeta(
+            name=name,
+            creation_date=datetime.datetime.now().isoformat(),
+            last_modification_date=None,
+            owner=None
+        )
+
+        structural_frame = structural_frame  # ? This could be Optional
+
+        grid = grid
+        _interpolation_options = interpolation_options
+        input_transform = Transform.from_input_points(
+            surface_points=structural_frame.surface_points_copy,
+            orientations=structural_frame.orientations_copy
+        )
+
+        model = GeoModel(
+            meta=meta,
+            structural_frame=structural_frame,
+            grid=grid,
+            input_transform=input_transform,
+            _interpolation_options=_interpolation_options
+        )
+
+        return model
+
+    # endregion
+
+    # region Methods
+    def update_transform(self, auto_anisotropy: GlobalAnisotropy = GlobalAnisotropy.NONE, anisotropy_limit: Optional[np.ndarray] = None):
+        """Update the transformation of the geological model.
+
+            This function updates the transformation of the geological model using the provided surface points and orientations.
+            It also applies anisotropy based on the specified type and limit.
+
+            Args:
+                auto_anisotropy (GlobalAnisotropy): The type of anisotropy to apply. Defaults to GlobalAnisotropy.NONE.
+                anisotropy_limit (Optional[np.ndarray]): Anisotropy limit values. If None, no limit is applied.
+
+        """
+
+        self.input_transform = Transform.from_input_points(
+            surface_points=self.surface_points_copy,
+            orientations=self.orientations_copy
+        )
+
+        self.input_transform.apply_anisotropy(anisotropy_type=auto_anisotropy, anisotropy_limit=anisotropy_limit)
+
     def add_surface_points(self, X: Sequence[float], Y: Sequence[float], Z: Sequence[float],
                            surface: Sequence[str], nugget: Optional[Sequence[float]] = None) -> None:
         raise NotImplementedError("This method is deprecated. Use `gp.add_surface_points` instead")
 
+    # endregion
+
+    # region Pydantic serialization
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        use_enum_values=False,
+        json_encoders={
+                np.ndarray: encode_numpy_array
+        }
+    )
+
+    @model_validator(mode='wrap')
+    @classmethod
+    def deserialize_properties(cls, data: Union["GeoModel", dict], constructor: ModelWrapValidatorHandler["GeoModel"]) -> "GeoModel":
+        try:
+            match data:
+                case GeoModel():
+                    return data
+                case dict():
+                    instance: GeoModel = constructor(data)
+                    instantiate_if_necessary(
+                        data=data,
+                        key="_interpolation_options",
+                        type=InterpolationOptions
+                    )
+                    instance._interpolation_options = data.get("_interpolation_options")
+                    return instance
+                case _:
+                    raise ValidationError
+        except ValidationError:
+            raise
+
+    # endregion
+
     def _infer_dense_grid_solution(self):
         octrees_set: bool = Grid.GridTypes.OCTREE in self.grid.active_grids
         dense_set: bool = Grid.GridTypes.DENSE in self.grid.active_grids

test/test_modules/test_serialize_model.py

@@ -10,15 +10,8 @@
 
 def test_generate_horizontal_stratigraphic_model():
     model: gp.data.GeoModel = gp.generate_example_model(ExampleModel.HORIZONTAL_STRAT, compute_model=False)
-
     model_json = model.model_dump_json(by_alias=True)
 
-    # Pretty print JSON
-    pprint.pp(model_json)
-
-    # Ensure the 'verify/' directory exists
-    os.makedirs("verify", exist_ok=True)
-
     # Write the JSON to disk
     file_path = os.path.join("temp", "horizontal_stratigraphic_model.json")
     with open(file_path, "w") as f:
@@ -42,6 +35,8 @@ def test_generate_horizontal_stratigraphic_model():
 
     # # Validate json against schema
     if False:
+        # Ensure the 'verify/' directory exists
+        os.makedirs("verify", exist_ok=True)
         verify_json(model_json, name="verify/Horizontal Stratigraphic Model serialization")