diff --git a/geost/__init__.py b/geost/__init__.py index 77d5191..39dad0d 100644 --- 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 index 9a1ff77..cd0b03c 100644 --- 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 index 4f57c89..bb0d884 100644 --- 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 index c925379..43611e4 100644 --- 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)