Merge branch 'main' of github.com:Deltares-research/geost

Deltares-research · Sep 4, 2024 · 39869d2 · 39869d2
2 parents 4079adb + b821dc9
commit 39869d2
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 6 deletions.
diff --git a/geost/__init__.py b/geost/__init__.py
@@ -14,4 +14,4 @@
 # read_xml_soil_cores,
 from geost.utils import csv_to_parquet, excel_to_parquet
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
diff --git a/geost/base.py b/geost/base.py
@@ -898,7 +898,6 @@ def get_cumulative_layer_thickness(self, column: str, values: str | List[str]):
             cumulative_thickness
         )
         cum_thickness = cum_thickness.unstack(level=column)
-        cum_thickness[cum_thickness.isna()] = 0
         return cum_thickness
 
     def get_layer_top(self, column: str, values: str | List[str]):
@@ -2225,14 +2224,16 @@ def get_cumulative_layer_thickness(
         cum_thickness = self.data.get_cumulative_layer_thickness(column, values)
 
         if include_in_header:
+            columns = [c + "_thickness" for c in cum_thickness.columns]
             self.header.gdf.drop(
-                columns=[c + "_thickness" for c in cum_thickness.columns],
+                columns=columns,
                 errors="ignore",
                 inplace=True,
             )
             self.header = self.header.gdf.merge(
                 cum_thickness.add_suffix("_thickness"), on="nr", how="left"
             )
+            self.header[columns] = self.header[columns].fillna(0)
         else:
             return cum_thickness
 

diff --git a/tests/test_collections.py b/tests/test_collections.py
@@ -512,7 +512,7 @@ def test_to_csv(self, borehole_collection, tmp_path):
 
 class TestBoreholeCollection:
     @pytest.mark.unittest
-    def test_get_cumulative_layer_thickness(self, borehole_collection):
+    def test_get_cumulative_layer_thickness_multiple(self, borehole_collection):
         borehole_collection.get_cumulative_layer_thickness(
             "lith", ["Z", "K"], include_in_header=True
         )
@@ -526,6 +526,26 @@ def test_get_cumulative_layer_thickness(self, borehole_collection):
             borehole_collection.header["Z_thickness"], expected_sand_thickness
         )
 
+        # Single query
+        borehole_collection.get_cumulative_layer_thickness(
+            "lith", "Z", include_in_header=True
+        )
+        assert_almost_equal(
+            borehole_collection.header["Z_thickness"], expected_sand_thickness
+        )
+
+    @pytest.mark.unittest
+    def test_get_cumulative_layer_thickness_single(self, borehole_collection):
+        expected_sand_thickness = [2.2, 0.0, 2.6, 0.5, 3.0]
+
+        # Single query
+        borehole_collection.get_cumulative_layer_thickness(
+            "lith", "Z", include_in_header=True
+        )
+        assert_almost_equal(
+            borehole_collection.header["Z_thickness"], expected_sand_thickness
+        )
+
     @pytest.mark.unittest
     def test_get_layer_top(self, borehole_collection):
         borehole_collection.get_layer_top("lith", ["Z", "K"], include_in_header=True)

diff --git a/tests/test_data_objects.py b/tests/test_data_objects.py
@@ -78,8 +78,8 @@ def test_get_cumulative_layer_thickness(self, borehole_data):
 
         result = borehole_data.get_cumulative_layer_thickness("lith", ["Z", "K"])
         expected_boreholes_returned = ["A", "B", "C", "D", "E"]
-        expected_sand_thickness = [2.2, 0.0, 2.6, 0.5, 3.0]
-        expected_clay_thickness = [2.0, 2.0, 2.9, 1.1, 0.0]
+        expected_sand_thickness = [2.2, np.nan, 2.6, 0.5, 3.0]
+        expected_clay_thickness = [2.0, 2.0, 2.9, 1.1, np.nan]
 
         assert result.shape == (5, 2)
         assert_array_equal(result.index, expected_boreholes_returned)