fix: adding add_group and add points from dataframe for geoh5 + some tests

Author: lachlangrose

LoopStructural/export/geoh5.py

@@ -5,14 +5,29 @@
 
 from LoopStructural.datatypes import ValuePoints, VectorPoints
 
-
-def add_surface_to_geoh5(filename, surface, overwrite=True, groupname="Loop"):
+def add_group_to_geoh5(filename, groupname="Loop", parent=None, overwrite=True):
     with geoh5py.workspace.Workspace(filename) as workspace:
+
         group = workspace.get_entity(groupname)[0]
+        if group and overwrite:
+            group.allow_delete = True
+            workspace.remove_entity(group)
         if not group:
             group = geoh5py.groups.ContainerGroup.create(
                 workspace, name=groupname, allow_delete=True
             )
+        if parent is not None:
+            parent = workspace.get_entity(parent)[0]    
+        if parent:
+            parent.add_children(group)
+        return group.uid
+def add_surface_to_geoh5(filename, surface, overwrite=True, group="Loop"):
+    with geoh5py.workspace.Workspace(filename) as workspace:
+        group = workspace.get_entity(group)[0]
+        if not group:
+            group = geoh5py.groups.ContainerGroup.create(
+                workspace, name=group, allow_delete=True
+            )
         if surface.name in workspace.list_entities_name.values():
             existing_surf = workspace.get_entity(surface.name)
             existing_surf[0].allow_delete = True
@@ -34,6 +49,7 @@ def add_surface_to_geoh5(filename, surface, 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(
@@ -54,56 +70,73 @@ def add_points_to_geoh5(filename, point, overwrite=True, groupname="Loop"):
             data['vz'] = {'association': "VERTEX", "values": point.vectors[:, 2]}
 
         if isinstance(point, ValuePoints):
-            data['val'] = {'association': "VERTEX", "values": point.values}
+            data['values'] = {'association': "VERTEX", "values": point.values}
         point = geoh5py.objects.Points.create(
             workspace,
             name=point.name,
             vertices=point.locations,
             parent=group,
         )
         point.add_data(data)
+    
+def overwrite_object(workspace, name, overwrite):
+    if name in workspace.list_entities_name.values():
+        existing_entity = workspace.get_entity(name)
+        existing_entity[0].allow_delete = True
+        if overwrite:
+            workspace.remove_entity(existing_entity[0])
 
-def add_points_from_df(filename, df, overwrite=True, child = None, parent = None, 
-                        normal_cols=['nx','ny','nz']):
+def add_points_from_df(filename, df, name='pointset', overwrite=True,                        columns=None, groupname="Loop", x_col='X', y_col='Y', z_col='Z'):
+        """
+        Add points to a geoh5 file from a pandas DataFrame. The DataFrame must have columns 'name', 'X', 'Y', 'Z' for the point locations.
+        Additional columns can be added as data associated with the points.
+        Parameters
+        ----------
+        filename: str
+            Path to the geoh5 file.
+        df: pandas.DataFrame
+            DataFrame containing point data. Must have columns 'name', 'X', 'Y', 'Z'. Additional columns will be added as data.
+        overwrite: bool, optional
+            Whether to overwrite existing points with the same name. Default is True.
+        columns: list of str, optional
+            List of columns in the DataFrame to add as data. If None, all columns except 'name', 'X', 'Y', 'Z' will be added. Default is None.
+                
+        """
+        if columns is None:
+            columns = df.columns.tolist()
+        if x_col not in columns or y_col not in columns or z_col not in columns:
+            raise ValueError("DataFrame must contain 'name', 'X', 'Y', 'Z' columns. " \
+            "Specify the column names using x_col, y_col, z_col parameters if they are different.")
         with geoh5py.workspace.Workspace(filename) as workspace:
-            entities = workspace.get_entity(child)
-            child_name = child
-            child = entities[0] if entities else None
-            if not child:
-                child = geoh5py.groups.ContainerGroup.create(
-                    workspace, name=child_name, allow_delete=True,
-                )
-            if parent:
-                parent.add_children(child)
-
-            for _, row in df.iterrows():
-                name = row['name']
-                loc = np.array([[row['X'], row['Y'], row['Z']]])  # shape (1,3)
-
-                # remove existing entity if present and overwrite requested
-                if name in workspace.list_entities_name.values():
-                    existing = workspace.get_entity(name)
-                    if existing:
-                        existing[0].allow_delete = True
-                        if overwrite:
-                            workspace.remove_entity(existing[0])
+            if groupname:
+                group = workspace.get_entity(groupname)
+                group = group[0] if group else None
+                if not group:
+                    group = geoh5py.groups.ContainerGroup.create(
+                        workspace, name=groupname, allow_delete=True,
+                    )
+            
+            location = np.array(df[[x_col, y_col, z_col]].values)  # shape (n,3)
+            
+            overwrite_object(workspace, name, overwrite)
+            
 
-                pts = geoh5py.objects.Points.create(
-                    workspace,
-                    name=name,
-                    vertices=loc,
-                    parent=child,
-                )
-
-                # build data dict from normal_cols (and any other columns you want)
-                data = {}
-                for col in normal_cols:
-                    if col in row and not pd.isna(row[col]):
-                        # association must be "VERTEX" and values length must match vertices (1)
-                        data[col] = {"association": "VERTEX", "values": np.array([row[col]])}
+            pts = geoh5py.objects.Points.create(
+                workspace,
+                name=name,
+                vertices=location,
+                parent=group,
+            )
+            data = {}
+            for col in columns:
+                if col in ['name', x_col, y_col, z_col]:
+                    continue
+                data[col] = {"association": "VERTEX", "values": np.array(df[col]).flatten()}
+                
 
-                if data:
-                    pts.add_data(data)
+            if data:
+                pts.add_data(data)
+                
 
 def add_structured_grid_to_geoh5(filename, structured_grid, overwrite=True, groupname="Loop"):
     with geoh5py.workspace.Workspace(filename) as workspace:

tests/unit/io/test_geoh5.py

@@ -0,0 +1,90 @@
+from LoopStructural.export.geoh5 import add_group_to_geoh5, add_points_to_geoh5, add_points_from_df
+import geoh5py
+import pytest
+from pathlib import Path
+from LoopStructural.datatypes import ValuePoints, VectorPoints
+import numpy as np
+@pytest.fixture
+def tmp_path():
+    import tempfile
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+@pytest.fixture
+def test_setup(tmp_path):
+    filename = tmp_path / "test.geoh5"
+    with geoh5py.workspace.Workspace.create(filename) as workspace:
+        yield filename
+        workspace.close()
+
+def test_add_group_to_geoh5(test_setup):
+    filename = test_setup
+    group_uid = add_group_to_geoh5(filename, groupname="TestGroup")
+
+    with geoh5py.workspace.Workspace(filename) as workspace:
+        assert workspace.get_entity(group_uid)[0].name == "TestGroup"
+
+def test_add_points_to_geoh5(test_setup):
+    filename = test_setup
+    group_uid = add_group_to_geoh5(filename, groupname="TestGroup")
+    points = ValuePoints(
+        name="TestPoints",
+        locations=[[0, 0, 0], [1, 1, 1], [2, 2, 2]],
+        values=[10., 20, 30],
+    )
+    add_points_to_geoh5(filename, points, groupname=group_uid)
+    with geoh5py.workspace.Workspace(filename) as workspace:
+        point_entity = workspace.get_entity("TestPoints")[0]
+        assert point_entity.name == "TestPoints"
+        assert point_entity.vertices.tolist() == [[0, 0, 0], [1, 1, 1], [2, 2, 2]]
+        assert np.sum(point_entity.get_data("values")[0].values-np.array([10., 20., 30.])) == 0
+
+def test_add_vector_points_to_geoh5(test_setup):
+    filename = test_setup
+    group_uid = add_group_to_geoh5(filename, groupname="TestGroup")
+    points = VectorPoints(
+        name="TestVectorPoints",
+        locations=[[0, 0, 0], [1, 1, 1], [2, 2, 2]],
+        vectors=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+    )
+    add_points_to_geoh5(filename, points, groupname=group_uid)
+    with geoh5py.workspace.Workspace(filename) as workspace:
+        point_entity = workspace.get_entity("TestVectorPoints")[0]
+        assert point_entity.name == "TestVectorPoints"
+        assert point_entity.vertices.tolist() == [[0, 0, 0], [1, 1, 1], [2, 2, 2]]
+        assert np.sum(point_entity.get_data("vx")[0].values-np.array([1., 0., 0.])) == 0
+        assert np.sum(point_entity.get_data("vy")[0].values-np.array([0., 1., 0.])) == 0
+        assert np.sum(point_entity.get_data("vz")[0].values-np.array([0., 0., 1.])) == 0
+
+def test_add_df_to_geoh5(test_setup):
+    import pandas as pd
+    filename = test_setup
+    group_uid = add_group_to_geoh5(filename, groupname="TestGroup")
+    df = pd.DataFrame({
+        'X': [0, 1, 2],
+        'Y': [0, 1, 2],
+        'Z': [0, 1, 2],
+        'value': [10., 20., 30.],
+    })
+    add_points_from_df(filename, df, name='df_points', groupname=group_uid)
+    with geoh5py.workspace.Workspace(filename) as workspace:
+        point_entity = workspace.get_entity("TestGroup")[0].children[0]
+        assert point_entity.name == "df_points"
+        assert point_entity.vertices.tolist() == [[0, 0, 0], [1, 1, 1], [2, 2, 2]]
+        assert np.sum(point_entity.get_data("value")[0].values-np.array([10., 20., 30.])) == 0
+
+def test_add_df_with_alternate_xyz_to_geoh5(test_setup):
+    import pandas as pd
+    filename = test_setup
+    group_uid = add_group_to_geoh5(filename, groupname="TestGroup")
+    df = pd.DataFrame({
+        'EAST': [0, 1, 2],
+        'NORTH': [0, 1, 2],
+        'RL': [0, 1, 2],
+        'value': [10., 20., 30.],
+    })
+    add_points_from_df(filename, df, name='df_points',groupname=group_uid, x_col='EAST', y_col='NORTH', z_col='RL')
+    with geoh5py.workspace.Workspace(filename) as workspace:
+        point_entity = workspace.get_entity("TestGroup")[0].children[0]
+        assert point_entity.name == "df_points"
+        assert point_entity.vertices.tolist() == [[0, 0, 0], [1, 1, 1], [2, 2, 2]]
+        assert np.sum(point_entity.get_data("value")[0].values-np.array([10., 20., 30.])) == 0