Merge branch 'master' of github.com:Loop3D/LoopStructural

Author: lachlangrose

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,5 +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):
+        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/datatypes/_surface.py

@@ -95,7 +95,7 @@ def save(self, filename):
             with open(filename, 'w') as f:
                 json.dump(self.to_dict(), f)
         elif ext == 'vtk':
-            self.vtk.save(filename)
+            self.vtk().save(filename)
         elif ext == 'obj':
             import meshio
 

LoopStructural/export/geoh5.py

@@ -1,4 +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"):
@@ -27,4 +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
+            )
+        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_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 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 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=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

@@ -7,7 +7,7 @@
 import numpy as np
 import pandas as pd
 from typing import List
-
+import pathlib
 from ...modelling.features.fault import FaultSegment
 
 from ...modelling.features.builders import (
@@ -111,7 +111,6 @@ def __init__(
 
 
         """
-        # print('tet')
         if logfile:
             self.logfile = logfile
             log_to_file(logfile, level=loglevel)
@@ -1540,6 +1539,9 @@ def evaluate_model(self, xyz: np.ndarray, scale: bool = True) -> np.ndarray:
         strat_id = np.zeros(xyz.shape[0], dtype=int)
         # set strat id to -1 to identify which areas of the model aren't covered
         strat_id[:] = -1
+        if self.stratigraphic_column is None:
+            logger.warning("No stratigraphic column defined")
+            return strat_id
         for group in reversed(self.stratigraphic_column.keys()):
             if group == "faults":
                 continue
@@ -1757,6 +1759,9 @@ def stratigraphic_ids(self):
             list of unique stratigraphic ids, featurename, unit name and min and max scalar values
         """
         ids = []
+        if self.stratigraphic_column is None:
+            logger.warning('No stratigraphic column defined')
+            return ids
         for group in self.stratigraphic_column.keys():
             if group == "faults":
                 continue
@@ -1777,6 +1782,8 @@ def get_stratigraphic_surfaces(self, units: List[str] = [], bottoms: bool = True
         ## TODO change the stratigraphic column to its own class and have methods to get the relevant surfaces
         surfaces = []
         units = []
+        if self.stratigraphic_column is None:
+            return []
         for group in self.stratigraphic_column.keys():
             if group == "faults":
                 continue
@@ -1800,10 +1807,60 @@ 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 = True,
+        stratigraphic_surfaces=True,
+        fault_surfaces=True,
+        stratigraphic_data=True,
+        fault_data=True,
+    ):
+        path = pathlib.Path(filename)
+        extension = path.suffix
+        name = path.stem
+        stratigraphic_surfaces = self.get_stratigraphic_surfaces()
+        if fault_surfaces:
+            for s in self.get_fault_surfaces():
+                ## geoh5 can save everything into the same file
+                if extension == ".geoh5":
+                    s.save(filename)
+                else:
+                    s.save(f'{name}_{s.name}.{extension}')
+        if stratigraphic_surfaces:
+            for s in self.get_stratigraphic_surfaces():
+                if extension == ".geoh5":
+                    s.save(filename)
+                else:
+                    s.save(f'{name}_{s.name}.{extension}')
+        if block_model:
+            grid, ids = self.get_block_model()
+            if extension == ".geoh5":
+                grid.save(filename)
+            else:
+                grid.save(f'{name}_block_model.{extension}')
+        if stratigraphic_data:
+            if self.stratigraphic_column is not None:
+                for group in self.stratigraphic_column.keys():
+                    if group == "faults":
+                        continue
+                    for series in self.stratigraphic_column[group].keys():
+                        if extension == ".geoh5":
+                            self.__getitem__(series).save(filename)
+                        else:
+                            self.__getitem__(series).save(f'{name}_{series}.{extension}')
+        if fault_data:
+            for f in self.fault_names():
+                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):