fix: adding gocad export for structured grid support

lachlangrose · lachlangrose · commit 0892dccb463e · 2026-04-10T10:55:29.000+09:30
diff --git a/LoopStructural/datatypes/_structured_grid.py b/LoopStructural/datatypes/_structured_grid.py
@@ -28,6 +28,7 @@ class StructuredGrid:
     name : str, optional
         Name of the grid, by default "default_grid"
     """
+
     origin: np.ndarray = field(default_factory=lambda: np.array([0, 0, 0]))
     step_vector: np.ndarray = field(default_factory=lambda: np.array([1, 1, 1]))
     nsteps: np.ndarray = field(default_factory=lambda: np.array([10, 10, 10]))
@@ -134,15 +135,15 @@ def cell_centres(self):
             self.nsteps[2] - 1,
         )
         x, y, z = np.meshgrid(x, y, z, indexing="ij")
-        return np.vstack([x.flatten(order='f'), y.flatten(order='f'), z.flatten(order='f')]).T
+        return np.vstack([x.flatten(order="f"), y.flatten(order="f"), z.flatten(order="f")]).T
 
     @property
     def nodes(self):
         x = np.linspace(self.origin[0], self.maximum[0], self.nsteps[0])
         y = np.linspace(self.origin[1], self.maximum[1], self.nsteps[1])
         z = np.linspace(self.origin[2], self.maximum[2], self.nsteps[2])
         x, y, z = np.meshgrid(x, y, z, indexing="ij")
-        return np.vstack([x.flatten(order='f'), y.flatten(order='f'), z.flatten(order='f')]).T
+        return np.vstack([x.flatten(order="f"), y.flatten(order="f"), z.flatten(order="f")]).T
 
     def merge(self, other):
         if not np.all(np.isclose(self.origin, other.origin)):
@@ -157,36 +158,33 @@ def merge(self, other):
         for name, data in other.properties.items():
             self.properties[name] = data
 
-    def save(self, filename, *,group='Loop'):
+    def save(self, filename, *, group="Loop"):
         filename = str(filename)
-        ext = filename.split('.')[-1]
-        if ext == 'json':
+        ext = filename.split(".")[-1].lower()
+        if ext == "json":
             import json
 
-            with open(filename, 'w') as f:
+            with open(filename, "w") as f:
                 json.dump(self.to_dict(), f)
-        elif ext == 'vtk':
+        elif ext == "vtk":
             self.vtk().save(filename)
 
-        elif ext == 'geoh5':
+        elif ext == "geoh5":
             from LoopStructural.export.geoh5 import add_structured_grid_to_geoh5
 
             add_structured_grid_to_geoh5(filename, self, groupname=group)
-        elif ext == 'pkl':
+        elif ext == "pkl":
             import pickle
 
-            with open(filename, 'wb') as f:
+            with open(filename, "wb") as f:
                 pickle.dump(self, f)
-        elif ext == 'omf':
+        elif ext == "omf":
             from LoopStructural.export.omf_wrapper import add_structured_grid_to_omf
 
             add_structured_grid_to_omf(self, filename)
-        elif ext == 'vs':
-            raise NotImplementedError(
-                "Saving structured grids in gocad format is not yet implemented"
-            )
-            # from LoopStructural.export.gocad import _write_structued_grid
+        elif ext == "vo" or ext == "gocad":
+            from LoopStructural.export.gocad import _write_structured_grid_gocad
 
-            # _write_pointset(self, filename)
+            _write_structured_grid_gocad(self, filename)
         else:
-            raise ValueError(f'Unknown file extension {ext}')
+            raise ValueError(f"Unknown file extension {ext}")
diff --git a/LoopStructural/export/gocad.py b/LoopStructural/export/gocad.py
@@ -1,5 +1,153 @@
+import re
+from pathlib import Path
+
 import numpy as np
 
+from LoopStructural.utils import getLogger
+
+logger = getLogger(__name__)
+
+
+def _normalise_voxet_property(values, property_name, nsteps):
+    array = np.asarray(values)
+    expected_shape = tuple(int(step) for step in nsteps)
+    expected_size = int(np.prod(expected_shape))
+
+    if array.shape == expected_shape:
+        flat_values = array.reshape(-1, order="F")
+    else:
+        flat_values = np.squeeze(array)
+        if flat_values.shape == expected_shape:
+            flat_values = flat_values.reshape(-1, order="F")
+        elif flat_values.ndim == 1 and flat_values.size == expected_size:
+            flat_values = flat_values
+        else:
+            raise ValueError(
+                f"Property '{property_name}' must have shape {expected_shape} or size {expected_size}"
+            )
+
+    if np.issubdtype(flat_values.dtype, np.integer):
+        if flat_values.size == 0:
+            export_dtype = np.int8
+            storage_type = "Octet"
+            element_size = 1
+        elif flat_values.min() >= np.iinfo(np.int8).min and flat_values.max() <= np.iinfo(np.int8).max:
+            export_dtype = np.int8
+            storage_type = "Octet"
+            element_size = 1
+        else:
+            export_dtype = np.dtype(">i4")
+            storage_type = "Integer"
+            element_size = 4
+        no_data_value = None
+    elif np.issubdtype(flat_values.dtype, np.floating):
+        export_dtype = np.dtype(">f4")
+        storage_type = "Float"
+        element_size = 4
+        no_data_value = -999999.0
+        flat_values = np.nan_to_num(flat_values, nan=no_data_value)
+    else:
+        raise ValueError(f"Property '{property_name}' has unsupported dtype {flat_values.dtype}")
+
+    return {
+        "values": np.asarray(flat_values, dtype=export_dtype),
+        "storage_type": storage_type,
+        "element_size": element_size,
+        "no_data_value": no_data_value,
+    }
+
+
+def _write_structured_grid_gocad(grid, file_name):
+    """Write a StructuredGrid to GOCAD VOXET format."""
+    vo_path = Path(file_name).with_suffix(".vo")
+    axis_n = np.asarray(grid.nsteps, dtype=int)
+    property_source = grid.properties
+    axis_min = np.min(grid.nodes, axis=0)
+    axis_max = np.max(grid.nodes, axis=0)
+
+    if property_source:
+        if grid.cell_properties:
+            logger.warning(
+                "StructuredGrid GOCAD export uses point properties; cell_properties were not exported"
+            )
+    elif grid.cell_properties:
+        axis_n = np.asarray(grid.nsteps, dtype=int) - 1
+        if np.any(axis_n <= 0):
+            raise ValueError("StructuredGrid cell_properties require at least two grid nodes per axis")
+        property_source = grid.cell_properties
+        axis_min = np.min(grid.cell_centres, axis=0)
+        axis_max = np.max(grid.cell_centres, axis=0)
+    else:
+        raise ValueError("StructuredGrid has no properties to export to GOCAD")
+
+    export_properties = []
+    for index, (property_name, values) in enumerate(property_source.items(), start=1):
+        export_info = _normalise_voxet_property(values, property_name, axis_n)
+        safe_name = re.sub(r"[^0-9A-Za-z_-]+", "_", property_name).strip("_") or f"property_{index}"
+        data_path = vo_path.with_name(f"{vo_path.stem}_{safe_name}@@")
+        export_properties.append(
+            {
+                "index": index,
+                "name": property_name,
+                "data_path": data_path,
+                **export_info,
+            }
+        )
+
+    with open(vo_path, "w") as fp:
+        fp.write(
+            f"""GOCAD Voxet 1
+HEADER {{
+name: {grid.name}
+}}
+GOCAD_ORIGINAL_COORDINATE_SYSTEM
+NAME Default
+AXIS_NAME \"X\" \"Y\" \"Z\"
+AXIS_UNIT \"m\" \"m\" \"m\"
+ZPOSITIVE Elevation
+END_ORIGINAL_COORDINATE_SYSTEM
+AXIS_O 0.000000 0.000000 0.000000
+AXIS_U 1.000000 0.000000 0.000000
+AXIS_V 0.000000 1.000000 0.000000
+AXIS_W 0.000000 0.000000 1.000000
+AXIS_MIN {axis_min[0]} {axis_min[1]} {axis_min[2]}
+AXIS_MAX {axis_max[0]} {axis_max[1]} {axis_max[2]}
+AXIS_N {axis_n[0]} {axis_n[1]} {axis_n[2]}
+AXIS_NAME \"X\" \"Y\" \"Z\"
+AXIS_UNIT \"m\" \"m\" \"m\"
+AXIS_TYPE even even even
+"""
+        )
+        for export_property in export_properties:
+            fp.write(
+                f"""PROPERTY {export_property['index']} {export_property['name']}
+PROPERTY_CLASS {export_property['index']} {export_property['name']}
+PROP_UNIT {export_property['index']} {export_property['name']}
+PROPERTY_CLASS_HEADER {export_property['index']} {export_property['name']} {{
+}}
+PROPERTY_SUBCLASS {export_property['index']} QUANTITY {export_property['storage_type']}
+"""
+            )
+            if export_property["no_data_value"] is not None:
+                fp.write(
+                    f"PROP_NO_DATA_VALUE {export_property['index']} {export_property['no_data_value']}\n"
+                )
+            fp.write(
+                f"""PROP_ETYPE {export_property['index']} IEEE
+PROP_FORMAT {export_property['index']} RAW
+PROP_ESIZE {export_property['index']} {export_property['element_size']}
+PROP_OFFSET {export_property['index']} 0
+PROP_FILE {export_property['index']} {export_property['data_path'].name}
+"""
+            )
+        fp.write("END\n")
+
+    for export_property in export_properties:
+        with open(export_property["data_path"], "wb") as fp:
+            export_property["values"].tofile(fp)
+
+    return True
+
 
 def _write_feat_surfs_gocad(surf, file_name):
     """
@@ -23,8 +171,6 @@ def _write_feat_surfs_gocad(surf, file_name):
     True if successful
 
     """
-    from pathlib import Path
-
     properties_header = None
     if surf.properties:
 
@@ -112,8 +258,6 @@ def _write_feat_surfs_gocad(surf, file_name):
 #     True if successful
 
 #     """
-#     from pathlib import Path
-
 #     file_name = Path(file_name).with_suffix(".vs")
 #     with open(f"{file_name}", "w") as fd:
 #         fd.write(
@@ -123,4 +267,4 @@ def _write_feat_surfs_gocad(surf, file_name):
 # }}
 # GOCAD_ORIGINAL_COORDINATE_SYSTEM
 # NAME Default
-# PROJECTION Unknown
+# PROJECTION Unknown
diff --git a/tests/unit/datatypes/test__structured_grid.py b/tests/unit/datatypes/test__structured_grid.py
@@ -8,7 +8,7 @@ def test_structured_grid_to_dict():
     origin = np.array([0, 0, 0])
     nsteps = np.array([10, 10, 10])
     step_vector = np.array([1, 1, 1])
-    data = {'rng': rng.random(size=(10, 10, 10))}
+    data = {"rng": rng.random(size=(10, 10, 10))}
     name = "grid_data"
 
     grid = StructuredGrid(
@@ -55,7 +55,7 @@ def test_structured_grid_vtk():
     origin = np.array([0, 0, 0])
     nsteps = np.array([10, 10, 10])
     step_vector = np.array([1, 1, 1])
-    data = {'rng': rng.random(size=(10, 10, 10))}
+    data = {"rng": rng.random(size=(10, 10, 10))}
     name = "grid_data"
 
     grid = StructuredGrid(
@@ -74,7 +74,65 @@ def test_structured_grid_vtk():
     assert isinstance(vtk_grid, pv.RectilinearGrid)
     assert np.array_equal(vtk_grid.dimensions, nsteps)
     assert np.array_equal(grid_origin, origin)
-    assert np.array_equal(vtk_grid['rng'], data['rng'].flatten(order="F"))
+    assert np.array_equal(vtk_grid["rng"], data["rng"].flatten(order="F"))
+
+
+def test_structured_grid_save_gocad_voxet(tmp_path):
+    origin = np.array([10.0, 20.0, 30.0])
+    nsteps = np.array([2, 3, 4])
+    step_vector = np.array([1.0, 2.0, 3.0])
+    values = np.arange(np.prod(nsteps), dtype=np.float32).reshape(tuple(nsteps), order="F")
+    grid = StructuredGrid(
+        origin=origin,
+        step_vector=step_vector,
+        nsteps=nsteps,
+        properties={"scalar": values},
+        cell_properties={},
+        name="grid_data",
+    )
+
+    output = tmp_path / "structured_grid.vo"
+    grid.save(output)
+
+    header = output.read_text()
+    binary_path = tmp_path / "structured_grid_scalar@@"
+
+    assert "GOCAD Voxet 1" in header
+    assert "PROPERTY 1 scalar" in header
+    assert f"AXIS_N {nsteps[0]} {nsteps[1]} {nsteps[2]}" in header
+    assert "PROP_FILE 1 structured_grid_scalar@@" in header
+    assert binary_path.exists()
+
+    exported = np.fromfile(binary_path, dtype=">f4")
+    np.testing.assert_array_equal(exported, values.flatten(order="F"))
+
+
+def test_structured_grid_save_gocad_voxet_uses_cell_properties_when_needed(tmp_path):
+    origin = np.array([0.0, 0.0, 0.0])
+    nsteps = np.array([3, 3, 3])
+    step_vector = np.array([1.0, 1.0, 1.0])
+    values = np.arange(np.prod(nsteps - 1), dtype=np.float32).reshape(tuple(nsteps - 1), order="F")
+    grid = StructuredGrid(
+        origin=origin,
+        step_vector=step_vector,
+        nsteps=nsteps,
+        properties={},
+        cell_properties={"cells": values},
+        name="cell_grid",
+    )
+
+    output = tmp_path / "cell_grid.vo"
+    grid.save(output)
+
+    header = output.read_text()
+    binary_path = tmp_path / "cell_grid_cells@@"
+
+    assert "PROPERTY 1 cells" in header
+    assert "AXIS_N 2 2 2" in header
+    assert binary_path.exists()
+
+    exported = np.fromfile(binary_path, dtype=">f4")
+    np.testing.assert_array_equal(exported, values.flatten(order="F"))
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/interpolator/test_2d_discrete_support.py b/tests/unit/interpolator/test_2d_discrete_support.py
@@ -1,5 +1,6 @@
 from LoopStructural.interpolators import StructuredGrid2D
 import numpy as np
+import pyvista as pv
 
 
 ## structured grid 2d tests
@@ -59,3 +60,18 @@ def test_get_element_outside2d():
     point = np.array([grid.origin - np.ones(2)])
     idc, inside = grid.position_to_cell_corners(point)
     assert not inside[0]
+
+
+def test_structured_grid2d_vtk_assigns_quad_cell_types():
+    grid = StructuredGrid2D(origin=np.zeros(2), nsteps=np.array([4, 4]))
+
+    vtk_grid = grid.vtk(z=0.0)
+
+    assert vtk_grid.n_cells == grid.n_elements
+    assert vtk_grid.celltypes.shape == (grid.n_elements,)
+    assert np.all(vtk_grid.celltypes == pv.CellType.QUAD)
+    assert vtk_grid.get_cell(0).point_ids == [0, 1, 5, 4]
+
+    triangulated = vtk_grid.triangulate()
+
+    assert triangulated.n_cells == grid.n_elements * 2
