added to_kingdom function

geost/base.py

@@ -3,6 +3,7 @@
 from typing import Iterable, List
 
 import geopandas as gpd
+import numpy as np
 import pandas as pd
 from pyproj import CRS
 from shapely.geometry import LineString
@@ -1110,6 +1111,60 @@ def to_qgis3d(
         qgis3d = self._create_geodataframe_3d(relative_to_vertical_reference)
         qgis3d.to_file(outfile, driver="GPKG", **kwargs)
 
+    def to_kingdom(
+        self,
+        outfile: str | WindowsPath,
+        tdstart: int = 1,
+        vw: float = 1500.0,
+        vs: float = 1600.0,
+    ):
+        """
+        Write data to 2 csv files: 1) interval data and 2) time-depth chart. These files
+        can be imported in the Kingdom seismic interpretation software.
+
+        Parameters
+        ----------
+        outfile : str | WindowsPath
+            Path to csv file to be written.
+        tdstart : int
+            startindex for TDchart, default is 1
+        vw : float
+            sound velocity in water in m/s, default is 1500 m/s
+        vs : float
+            sound velocity in sediment in m/s, default is 1600 m/s
+        """
+        # 1. add column needed in kingdom and write interval data
+        kingdom_df = self.df.copy()
+        # Add total depth and rename bottom and top columns to Kingdom requirements
+        kingdom_df.insert(7, "Total depth", (kingdom_df["surface"] - kingdom_df["end"]))
+        kingdom_df.rename(
+            columns={"top": "Start depth", "bottom": "End depth"}, inplace=True
+        )
+        kingdom_df.to_csv(outfile, index=False)
+
+        # 2. create and write time-depth chart
+        tdchart = self[["nr", "surface"]].copy()
+        tdchart.drop_duplicates(inplace=True)
+        tdchart.insert(0, "id", range(tdstart, tdstart + len(tdchart)))
+        # Add measured depth (predefined depths of 0 and 1 m below surface)
+        tdchart = pd.concat(
+            [
+                tdchart.assign(MD=np.zeros(len(tdchart), dtype=np.int64)),
+                tdchart.assign(MD=np.ones(len(tdchart), dtype=np.int64)),
+            ]
+        )
+        # Add two-way travel time
+        tdchart["TWT"] = (-tdchart["surface"] / (vw / 2 / 1000)) + (
+            tdchart["MD"] * 1 / (vs / 2 / 1000)
+        )
+
+        tdchart.drop("surface", axis=1, inplace=True)
+        tdchart.sort_values(by=["id", "MD"], inplace=True)
+        tdchart.to_csv(
+            outfile.parent.joinpath(f"{outfile.stem}_TDCHART{outfile.suffix}"),
+            index=False,
+        )
+
 
 class DiscreteData(AbstractData, PandasExportMixin):
     def __init__(self, df, has_inclined: bool = False):
@@ -2041,6 +2096,30 @@ def to_qgis3d(
         """
         self.data.to_qgis3d(outfile, relative_to_vertical_reference, **kwargs)
 
+    def to_kingdom(
+        self,
+        outfile: str | WindowsPath,
+        tdstart: int = 1,
+        vw: float = 1500.0,
+        vs: float = 1600.0,
+    ):
+        """
+        Write data to 2 csv files: interval data and time-depth chart,
+            for import in Kingdom seismic interpretation software.
+
+        Parameters
+        ----------
+        out_file : str | WindowsPath
+            Path to csv file to be written.
+        tdstart : int
+            startindex for TDchart, default is 1
+        vw : float
+            sound velocity in water in m/s, default is 1500 m/s
+        vs : float
+            sound velocity in sediment in m/s, default is 1600 m/s
+        """
+        self.data.to_kingdom(outfile, tdstart, vw, vs)
+
 
 class CptCollection(Collection):
     pass

tests/test_collections.py

@@ -496,3 +496,13 @@ def test_get_layer_top(self, borehole_collection):
 
         assert_almost_equal(borehole_collection.header["K_top"], expected_clay_top)
         assert_almost_equal(borehole_collection.header["Z_top"], expected_sand_top)
+
+    @pytest.mark.unittest
+    def test_to_kingdom(self, borehole_collection):
+        outfile = Path("temp_kingdom.csv")
+        tdfile = Path(outfile.parent, f"{outfile.stem}_TDCHART{outfile.suffix}")
+        borehole_collection.to_kingdom(outfile)
+        assert outfile.is_file()
+        assert tdfile.is_file()
+        outfile.unlink()
+        tdfile.unlink()

tests/test_data_objects.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import numpy as np
+import pandas as pd
 import pytest
 from numpy.testing import assert_array_almost_equal, assert_array_equal
 from pyvista import MultiBlock
@@ -270,22 +271,80 @@ def test_to_qgis3d(self, borehole_data):
         assert outfile.is_file()
         outfile.unlink()
 
+    @pytest.mark.unittest
+    def test_to_kingdom(self, borehole_data):
+        outfile = Path("temp_kingdom.csv")
+        tdfile = Path(outfile.parent, f"{outfile.stem}_TDCHART{outfile.suffix}")
+        borehole_data.to_kingdom(outfile)
+        assert outfile.is_file()
+        assert tdfile.is_file()
+        out_layerdata = pd.read_csv(outfile)
+        out_tddata = pd.read_csv(tdfile)
+        assert_array_almost_equal(
+            out_layerdata["Start depth"][:6],
+            np.array([0.0, 0.8, 1.5, 2.5, 3.7, 0.0]),
+        )
+        assert_array_almost_equal(
+            out_layerdata["End depth"][:6],
+            np.array([0.8, 1.5, 2.5, 3.7, 4.2, 0.6]),
+        )
+        assert_array_almost_equal(
+            out_layerdata["Total depth"][:6],
+            np.array([4.2, 4.2, 4.2, 4.2, 4.2, 3.9]),
+        )
+        assert_array_almost_equal(
+            out_tddata["MD"],
+            np.array(
+                [
+                    0,
+                    1,
+                    0,
+                    1,
+                    0,
+                    1,
+                    0,
+                    1,
+                    0,
+                    1,
+                ]
+            ),
+        )
+        assert_array_almost_equal(
+            out_tddata["TWT"],
+            np.array(
+                [
+                    -0.26666667,
+                    0.98333333,
+                    -0.4,
+                    0.85,
+                    -0.33333333,
+                    0.91666667,
+                    -0.13333333,
+                    1.11666667,
+                    0.13333333,
+                    1.38333333,
+                ]
+            ),
+        )
+        outfile.unlink()
+        tdfile.unlink()
+
     @pytest.mark.unittest
     def test_create_geodataframe_3d(self, borehole_data):
         relative_to_vertical_reference = True
         gdf = borehole_data._create_geodataframe_3d(relative_to_vertical_reference)
 
-        first_line_coords = get_coordinates(gdf['geometry'].iloc[0], include_z=True)
-        expected_coords = [[2., 3., 0.21], [2., 3., -0.59]]
+        first_line_coords = get_coordinates(gdf["geometry"].iloc[0], include_z=True)
+        expected_coords = [[2.0, 3.0, 0.21], [2.0, 3.0, -0.59]]
 
         assert all(gdf.geom_type == "LineString")
         assert_array_almost_equal(first_line_coords, expected_coords)
 
         relative_to_vertical_reference = False
         gdf = borehole_data._create_geodataframe_3d(relative_to_vertical_reference)
 
-        first_line_coords = get_coordinates(gdf['geometry'].iloc[0], include_z=True)
-        expected_coords = [[2., 3., 0.01], [2., 3., 0.81]]
+        first_line_coords = get_coordinates(gdf["geometry"].iloc[0], include_z=True)
+        expected_coords = [[2.0, 3.0, 0.01], [2.0, 3.0, 0.81]]
 
         assert all(gdf.geom_type == "LineString")
         assert_array_almost_equal(first_line_coords, expected_coords)