Merge branch 'fix/exporting' of github.com:Loop3D/LoopStructural into…

… fix/exporting

LoopStructural/datatypes/__init__.py

@@ -1,3 +1,4 @@
 from ._surface import Surface
 from ._bounding_box import BoundingBox
 from ._point import ValuePoints, VectorPoints
+from ._structured_grid import StructuredGrid

LoopStructural/datatypes/_bounding_box.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
-from typing import Optional, Union
+from typing import Optional, Union, Dict
 from LoopStructural.utils.exceptions import LoopValueError
 from LoopStructural.utils import rng
+from LoopStructural.datatypes._structured_grid import StructuredGrid
 import numpy as np
 
 from LoopStructural.utils.logging import getLogger
@@ -408,6 +409,14 @@ def vtk(self):
             z,
         )
 
-    @property
-    def structured_grid(self):
-        pass
+    def structured_grid(
+        self, cell_data: Dict[str, np.ndarray] = {}, vertex_data={}, name: str = "bounding_box"
+    ):
+        return StructuredGrid(
+            origin=self.global_origin,
+            step_vector=self.step_vector,
+            nsteps=self.nsteps,
+            properties_cell=cell_data,
+            properties_vertex=vertex_data,
+            name=name,
+        )

LoopStructural/datatypes/_point.py

@@ -1,12 +1,15 @@
 from dataclasses import dataclass
 import numpy as np
 
+from typing import Optional
+
 
 @dataclass
 class ValuePoints:
     locations: np.ndarray
     values: np.ndarray
     name: str
+    properties: Optional[dict] = None
 
     def to_dict(self):
         return {
@@ -22,12 +25,40 @@ def vtk(self):
         points["values"] = self.values
         return points
 
+    def save(self, filename):
+        filename = str(filename)
+        ext = filename.split('.')[-1]
+        if ext == 'json':
+            import json
+
+            with open(filename, 'w') as f:
+                json.dump(self.to_dict(), f)
+        elif ext == 'vtk':
+            self.vtk().save(filename)
+
+        elif ext == 'geoh5':
+            from LoopStructural.export.geoh5 import add_points_to_geoh5
+
+            add_points_to_geoh5(filename, self)
+        elif ext == 'pkl':
+            import pickle
+
+            with open(filename, 'wb') as f:
+                pickle.dump(self, f)
+        elif ext == 'vs':
+            from LoopStructural.export.gocad import _write_pointset
+
+            _write_pointset(self, filename)
+        else:
+            raise ValueError(f'Unknown file extension {ext}')
+
 
 @dataclass
 class VectorPoints:
     locations: np.ndarray
     vectors: np.ndarray
     name: str
+    properties: Optional[dict] = None
 
     def to_dict(self):
         return {

LoopStructural/datatypes/_structured_grid.py

@@ -1,3 +1,4 @@
+from typing import Dict
 import numpy as np
 from dataclasses import dataclass
 
@@ -7,7 +8,8 @@ class StructuredGrid:
     origin: np.ndarray
     step_vector: np.ndarray
     nsteps: np.ndarray
-    data: np.ndarray
+    properties_cell: Dict[str, np.ndarray]
+    properties_vertex: Dict[str, np.ndarray]
     name: str
 
     def to_dict(self):
@@ -16,7 +18,8 @@ def to_dict(self):
             "maximum": self.maximum,
             "step_vector": self.step_vector,
             "nsteps": self.nsteps,
-            "data": self.data,
+            "properties_cell": self.properties_cell,
+            "properties_vertex": self.properties_vertex,
             "name": self.name,
         }
 
@@ -37,8 +40,51 @@ def vtk(self):
             y,
             z,
         )
-        grid[self.name] = self.data.flatten(order="F")
+        for name, data in self.properties_vertex.items():
+            grid[name] = data.flatten(order="F")
+        for name, data in self.properties_cell.items():
+            grid.cell_data[name] = data.flatten(order="F")
         return grid
 
+    def merge(self, other):
+        if not np.all(np.isclose(self.origin, other.origin)):
+            raise ValueError("Origin of grids must be the same")
+        if not np.all(np.isclose(self.step_vector, other.step_vector)):
+            raise ValueError("Step vector of grids must be the same")
+        if not np.all(np.isclose(self.nsteps, other.nsteps)):
+            raise ValueError("Number of steps of grids must be the same")
+
+        for name, data in other.properties_cell.items():
+            self.properties_cell[name] = data
+        for name, data in other.properties_vertex.items():
+            self.properties_vertex[name] = data
+
     def save(self, filename):
-        raise NotImplementedError("Saving structured grids is not yet implemented")
+        filename = str(filename)
+        ext = filename.split('.')[-1]
+        if ext == 'json':
+            import json
+
+            with open(filename, 'w') as f:
+                json.dump(self.to_dict(), f)
+        elif ext == 'vtk':
+            self.vtk().save(filename)
+
+        elif ext == 'geoh5':
+            from LoopStructural.export.geoh5 import add_structured_grid_to_geoh5
+
+            add_structured_grid_to_geoh5(filename, self)
+        elif ext == 'pkl':
+            import pickle
+
+            with open(filename, 'wb') as f:
+                pickle.dump(self, f)
+        elif ext == 'vs':
+            raise NotImplementedError(
+                "Saving structured grids in gocad format is not yet implemented"
+            )
+            # from LoopStructural.export.gocad import _write_structued_grid
+
+            # _write_pointset(self, filename)
+        else:
+            raise ValueError(f'Unknown file extension {ext}')

LoopStructural/export/geoh5.py

@@ -1,5 +1,7 @@
 import geoh5py
+import geoh5py.workspace
 import numpy as np
+from LoopStructural.datatypes import ValuePoints, VectorPoints
 
 
 def add_surface_to_geoh5(filename, surface, overwrite=True, groupname="Loop"):
@@ -28,56 +30,71 @@ def add_surface_to_geoh5(filename, surface, overwrite=True, groupname="Loop"):
         surface.add_data(data)
 
 
-def add_points_to_geoh5(filename, points, overwrite=True, groupname="Loop"):
+def add_points_to_geoh5(filename, point, overwrite=True, groupname="Loop"):
     with geoh5py.workspace.Workspace(filename) as workspace:
         group = workspace.get_entity(groupname)[0]
         if not group:
             group = geoh5py.groups.ContainerGroup.create(
                 workspace, name=groupname, allow_delete=True
             )
-        for point in points:
-            if point.name in workspace.list_entities_name.values():
-                existing_point = workspace.get_entity(point.name)
-                existing_point[0].allow_delete = True
-                if overwrite:
-                    workspace.remove_entity(existing_point[0])
-            data = {}
-            if point.properties is not None:
-                for k, v in point.properties.items():
-                    data[k] = {'association': "VERTEX", "values": v}
-            point = geoh5py.objects.Points.create(
-                workspace,
-                name=point.name,
-                vertices=point.vertices,
-                parent=group,
-            )
-            point.add_data(data)
+        if point.name in workspace.list_entities_name.values():
+            existing_point = workspace.get_entity(point.name)
+            existing_point[0].allow_delete = True
+            if overwrite:
+                workspace.remove_entity(existing_point[0])
+        data = {}
+        if point.properties is not None:
+            for k, v in point.properties.items():
+                data[k] = {'association': "VERTEX", "values": v}
+        if isinstance(point, VectorPoints):
+            data['vx'] = {'association': "VERTEX", "values": point.vectors[:, 0]}
+            data['vy'] = {'association': "VERTEX", "values": point.vectors[:, 1]}
+            data['vz'] = {'association': "VERTEX", "values": point.vectors[:, 2]}
+
+        if isinstance(point, ValuePoints):
+            data['val'] = {'association': "VERTEX", "values": point.values}
+        point = geoh5py.objects.Points.create(
+            workspace,
+            name=point.name,
+            vertices=point.locations,
+            parent=group,
+        )
+        point.add_data(data)
 
 
-def add_block_model_to_geoh5py(filename, block_model, overwrite=True, groupname="Loop"):
+def add_structured_grid_to_geoh5(filename, structured_grid, overwrite=True, groupname="Loop"):
     with geoh5py.workspace.Workspace(filename) as workspace:
         group = workspace.get_entity(groupname)[0]
         if not group:
             group = geoh5py.groups.ContainerGroup.create(
                 workspace, name=groupname, allow_delete=True
             )
-        if block_model.name in workspace.list_entities_name.values():
-            existing_block = workspace.get_entity(block_model.name)
+        if structured_grid.name in workspace.list_entities_name.values():
+            existing_block = workspace.get_entity(structured_grid.name)
             existing_block[0].allow_delete = True
             if overwrite:
                 workspace.remove_entity(existing_block[0])
         data = {}
-        if block_model.properties is not None:
-            for k, v in block_model.properties.items():
-                data[k] = {'association': "CELL", "values": v}
+        if structured_grid.properties_cell is not None:
+            for k, v in structured_grid.properties_cell.items():
+                data[k] = {
+                    'association': "CELL",
+                    "values": np.rot90(v.reshape(structured_grid.nsteps - 1, order="F")).flatten(),
+                }
         block = geoh5py.objects.BlockModel.create(
             workspace,
-            name=block_model.name,
-            origin=[25, -100, 50],
-            u_cell_delimiters=np.cumsum(np.ones(16) * 5),  # Offsets along u
-            v_cell_delimiters=np.cumsum(np.ones(32) * 5),  # Offsets along v
-            z_cell_delimiters=np.cumsum(np.ones(16) * -2.5),  # Offsets along z (down)
-            rotation=30.0,
+            name=structured_grid.name,
+            origin=structured_grid.origin,
+            u_cell_delimiters=np.cumsum(
+                np.ones(structured_grid.nsteps[0]) * structured_grid.step_vector[0]
+            ),  # Offsets along u
+            v_cell_delimiters=np.cumsum(
+                np.ones(structured_grid.nsteps[1]) * structured_grid.step_vector[1]
+            ),  # Offsets along v
+            z_cell_delimiters=np.cumsum(
+                np.ones(structured_grid.nsteps[2]) * structured_grid.step_vector[2]
+            ),  # Offsets along z (down)
+            rotation=0.0,
             parent=group,
         )
         block.add_data(data)

LoopStructural/modelling/core/geological_model.py

@@ -37,7 +37,6 @@
 from ...modelling.intrusions import IntrusionBuilder
 
 from ...modelling.intrusions import IntrusionFrameBuilder
-from LoopStructural.modelling.features import fault
 
 
 logger = getLogger(__name__)
@@ -112,7 +111,6 @@ def __init__(
 
 
         """
-        # print('tet')
         if logfile:
             self.logfile = logfile
             log_to_file(logfile, level=loglevel)
@@ -1803,18 +1801,18 @@ def get_stratigraphic_surfaces(self, units: List[str] = [], bottoms: bool = True
 
         return surfaces
 
-    def get_block_model(self):
-        grid = self.bounding_box.vtk()
+    def get_block_model(self, name='block model'):
+        grid = self.bounding_box.structured_grid(name=name)
 
-        grid.cell_data['stratigraphy'] = self.evaluate_model(
+        grid.properties_cell['stratigraphy'] = self.evaluate_model(
             self.bounding_box.cell_centers(), scale=False
         )
         return grid, self.stratigraphic_ids()
 
     def save(
         self,
         filename: str,
-        block_model: bool = False,
+        block_model: bool = True,
         stratigraphic_surfaces=True,
         fault_surfaces=True,
         stratigraphic_data=True,
@@ -1855,7 +1853,8 @@ def save(
                             self.__getitem__(series).save(f'{name}_{series}.{extension}')
         if fault_data:
             for f in self.fault_names():
-                if extension == ".geoh5":
-                    self.__getitem__(f).save(filename)
-                else:
-                    self.__getitem__(f).save(f'{name}_{f}.{extension}')
+                for d in self.__getitem__(f).get_data():
+                    if extension == ".geoh5":
+                        d.save(filename)
+                    else:
+                        d.save(f'{name}_{group}.{extension}')

LoopStructural/modelling/features/_base_geological_feature.py

@@ -314,11 +314,14 @@ def scalar_field(self, bounding_box=None):
             if self.model is None:
                 raise ValueError("Must specify bounding box")
             bounding_box = self.model.bounding_box
-        grid = bounding_box.vtk()
+        grid = bounding_box.structured_grid(name=self.name)
         value = self.evaluate_value(
             self.model.scale(bounding_box.regular_grid(local=False, order='F'))
         )
-        grid[self.name] = value
+        grid.properties_vertex[self.name] = value
+
+        value = self.evaluate_value(bounding_box.cell_centers(order='F'))
+        grid.properties_cell[self.name] = value
         return grid
 
     def vector_field(self, bounding_box=None, tolerance=0.05, scale=1.0):