diff --git a/src/nhflotools/pwnlayers/layers.py b/src/nhflotools/pwnlayers/layers.py index ec68adf..e18cc82 100644 --- a/src/nhflotools/pwnlayers/layers.py +++ b/src/nhflotools/pwnlayers/layers.py @@ -8,6 +8,7 @@ from scipy.interpolate import griddata from nhflotools.pwnlayers.io import read_pwn_data2 +from nhflotools.pwnlayers.utils import fix_missings_botms_and_min_layer_thickness from nhflotools.pwnlayers2.layers import get_mensink_layer_model as get_mensink_layer_model2 logger = logging.getLogger(__name__) @@ -112,7 +113,7 @@ def get_pwn_layer_model( layer_model_regis["top"] = top if fix_min_layer_thickness: - _fix_missings_botms_and_min_layer_thickness(layer_model_regis) + fix_missings_botms_and_min_layer_thickness(layer_model_regis) # Get PWN layer models ds_pwn_data = read_pwn_data2( @@ -166,7 +167,7 @@ def get_pwn_layer_model( koppeltabel_header_other="ASSUMPTION1", ) if fix_min_layer_thickness: - _fix_missings_botms_and_min_layer_thickness(layer_model_bergen_regis) + fix_missings_botms_and_min_layer_thickness(layer_model_bergen_regis) # Combine PWN layer model with REGIS layer model layer_model_mensink_bergen_regis, _cat_mensink_bergen_regis = combine_two_layer_models( @@ -179,7 +180,7 @@ def get_pwn_layer_model( koppeltabel_header_other="ASSUMPTION1", ) if fix_min_layer_thickness: - _fix_missings_botms_and_min_layer_thickness(layer_model_mensink_bergen_regis) + fix_missings_botms_and_min_layer_thickness(layer_model_mensink_bergen_regis) # Remove inactive layers and set kh and kv of non-existing cells to default values layer_model_mensink_bergen_regis["kh"] = layer_model_mensink_bergen_regis.kh.where( @@ -287,41 +288,6 @@ def get_top_from_ahn( return top -def _fix_missings_botms_and_min_layer_thickness(ds): - """ - Fix missing botms and ensure all layers have a positive thickness. - - Parameters - ---------- - ds : xarray Dataset - Dataset containing the layer model with top and botm. - - Raises - ------ - ValueError - If top contains nan values. - """ - if ds["top"].isnull().any(): - msg = "Top should not contain nan values" - raise ValueError(msg) - - out = xr.concat((ds["top"].expand_dims(dim={"layer": ["mv"]}, axis=0), ds["botm"]), dim="layer") - - # Use ffill here to fill the nan's with the previous layer. Layer thickness is zero for non existing layers - out = out.ffill(dim="layer") - out.values = np.minimum.accumulate(out.values, axis=out.dims.index("layer")) - topbotm_fixed = out.isel(layer=slice(1, None)).transpose("layer", "icell2d") - - # inform - ncell, nisnull = ds["botm"].size, ds["botm"].isnull().sum() - nfixed = (~np.isclose(ds.botm, topbotm_fixed)).sum() - - logger.info( - f"Fixed {nisnull / ncell * 100.0:.1f}% missing botms using downward fill. Shifted {(nfixed - nisnull) / ncell * 100.0:.1f}% botms to ensure all layers have a positive thickness, assuming more info is in the upper layer." - ) - ds["botm"].values = topbotm_fixed.values - - def combine_two_layer_models( df_koppeltabel, layer_model_regis, @@ -1235,7 +1201,7 @@ def n(s): return transition if fix_min_layer_thickness: ds = xr.Dataset({"botm": out, "top": data["top"]}) - _fix_missings_botms_and_min_layer_thickness(ds) + fix_missings_botms_and_min_layer_thickness(ds) out = ds["botm"] return out @@ -1594,7 +1560,7 @@ def n(s): return transition if fix_min_layer_thickness: ds = xr.Dataset({"botm": out, "top": data["top"]}) - _fix_missings_botms_and_min_layer_thickness(ds) + fix_missings_botms_and_min_layer_thickness(ds) out = ds["botm"] return out diff --git a/src/nhflotools/pwnlayers/utils.py b/src/nhflotools/pwnlayers/utils.py new file mode 100644 index 0000000..085d2db --- /dev/null +++ b/src/nhflotools/pwnlayers/utils.py @@ -0,0 +1,42 @@ +import logging + +import numpy as np +import xarray as xr + +logger = logging.getLogger(__name__) + +def fix_missings_botms_and_min_layer_thickness(ds): + """ + Fix missing botms and ensure all layers have a positive thickness. + + Parameters + ---------- + ds : xarray Dataset + Dataset containing the layer model with top and botm. + + Raises + ------ + ValueError + If top contains nan values. + """ + if ds["top"].isnull().any(): + msg = "Top should not contain nan values" + raise ValueError(msg) + + out = xr.concat((ds["top"].expand_dims(dim={"layer": ["mv"]}, axis=0), ds["botm"]), dim="layer") + + # Use ffill here to fill the nan's with the previous layer. Layer thickness is zero for non existing layers + out = out.ffill(dim="layer") + out.values = np.minimum.accumulate(out.values, axis=out.dims.index("layer")) + topbotm_fixed = out.isel(layer=slice(1, None)).transpose("layer", "icell2d") + ds["botm"].values = topbotm_fixed.values + + # inform + ncell, nisnull = ds["botm"].size, ds["botm"].isnull().sum() + nfixed = (~np.isclose(ds.botm, topbotm_fixed)).sum() + + logger.info( + "Fixed %.1f%% missing botms using downward fill. Shifted %.1f%% botms to ensure all layers have a positive thickness, assuming more info is in the upper layer.", + nisnull / ncell * 100.0, + (nfixed - nisnull) / ncell * 100.0 + ) diff --git a/src/nhflotools/pwnlayers2/layers.py b/src/nhflotools/pwnlayers2/layers.py index c69a8f0..56a8ca8 100644 --- a/src/nhflotools/pwnlayers2/layers.py +++ b/src/nhflotools/pwnlayers2/layers.py @@ -16,7 +16,7 @@ from packaging import version from shapely.ops import unary_union -from nhflotools.pwnlayers.layers import _fix_missings_botms_and_min_layer_thickness +from nhflotools.pwnlayers.utils import fix_missings_botms_and_min_layer_thickness logger = logging.getLogger(__name__) @@ -277,7 +277,7 @@ def get_mensink_botm( if fix_min_layer_thickness: ds = xr.Dataset({"botm": out, "top": a["top"]}) - _fix_missings_botms_and_min_layer_thickness(ds) + fix_missings_botms_and_min_layer_thickness(ds) out = ds["botm"] return out @@ -545,4 +545,4 @@ def n(s): check = mask.astype(int) + transition.astype(int) assert (check <= 1).all(), "Transition cells should not overlap with mask." return transition - return out \ No newline at end of file + return out