Issue #1952 improve timeseries resampling further

docs/api/changelog.rst

@@ -46,6 +46,14 @@ Added
 - Added :class:`imod.mf6.AdvectionUTVD` to the transport packages. This is the
   most accurate advection scheme in most cases, and is recommended for when
   sharp fronts need to be preserved or when computing on an unstructured grid.
+- Added ``skipna: bool = False`` to :meth:`imod.mf6.Well.from_imod5_data` and
+  :meth:`imod.mf6.LayeredWell.from_imod5_data`. NaN rates in the timeseries of
+  associated wells are forward filled by default, as before. With
+  ``skipna=True`` they are instead left out of the time weighted mean when
+  resampling to the simulation times. The last stress period now also gets the
+  last valid rate instead of NaN when the last rate is NaN, or zero if there
+  are no valid rates.
+
 
 Fixed
 ~~~~~

imod/mf6/wel.py

@@ -120,6 +120,7 @@ def _prepare_well_rates_from_groups(
     pkg_data: dict[str, Any],
     unique_well_groups: Sequence[pd.api.typing.DataFrameGroupBy],
     start_times: StressPeriodTimesType,
+    skipna: bool = False,
 ) -> xr.DataArray:
     """
     Prepare well rates from dataframe groups, grouped by unique well locations.
@@ -129,19 +130,21 @@ def _prepare_well_rates_from_groups(
     if has_associated:
         # Resample times per group
         unique_well_groups = [
-            _process_timeseries(df_group, start_times)
+            _process_timeseries(df_group, start_times, skipna)
             for df_group in unique_well_groups
         ]
     return _df_groups_to_da_rates(unique_well_groups)
 
 
 def _process_timeseries(
-    df_group: pd.api.typing.DataFrameGroupBy, start_times: StressPeriodTimesType
+    df_group: pd.api.typing.DataFrameGroupBy,
+    start_times: StressPeriodTimesType,
+    skipna: bool = False,
 ) -> pd.DataFrame:
     if _is_steady_state(start_times):
         return average_timeseries(df_group)
     else:
-        return resample_timeseries(df_group, start_times)
+        return resample_timeseries(df_group, start_times, skipna=skipna)
 
 
 def _prepare_df_ipf_associated(
@@ -616,6 +619,7 @@ def from_imod5_data(
         times: StressPeriodTimesType,
         minimum_k: float = 0.0,
         minimum_thickness: float = 0.0,
+        skipna: bool = False,
     ) -> "GridAgnosticWell":
         """
         Convert wells to imod5 data, loaded with
@@ -707,6 +711,11 @@ def from_imod5_data(
             On creating point wells, no point wells will be placed in cells with
             a lower thickness than this. Wells are placed when ``to_mf6_pkg`` is
             called.
+        skipna: bool, default False
+            How to treat NaN rates in the timeseries of associated wells, when
+            ``times`` is a list of datetimes. By default, NaN rates are forward
+            filled. If True, NaN rates are instead left out of the time
+            weighted mean.
         """
 
         pkg_data = imod5_data[key]
@@ -735,7 +744,7 @@ def from_imod5_data(
             for (var, dtype), value in zip(varnames, index_values)
         }
         cls_input["rate"] = _prepare_well_rates_from_groups(
-            pkg_data, unique_well_groups, start_times
+            pkg_data, unique_well_groups, start_times, skipna
         )
         cls_input["minimum_k"] = minimum_k
         cls_input["minimum_thickness"] = minimum_thickness

imod/tests/test_mf6/test_mf6_wel.py

@@ -1,3 +1,4 @@
+import copy
 import sys
 from contextlib import nullcontext as does_not_raise
 from datetime import datetime
@@ -920,6 +921,24 @@ def test_import__as_steady_state(imod5_dataset, wel_class):
     np.testing.assert_almost_equal(wel.dataset["rate"].values, -323.89361702)
 
 
+@parametrize("wel_class", [Well, LayeredWell])
+def test_import__skipna(imod5_dataset, wel_class):
+    data = copy.deepcopy(imod5_dataset[0])
+    times = list(pd.date_range(datetime(1989, 1, 1), datetime(2013, 1, 1), 8400))
+    # Set a rate in the middle of the associated timeseries to NaN
+    df = data["wel-1-WELLS_L3"]["dataframe"][0]
+    df.loc[df.index[len(df) // 2], "rate"] = np.nan
+
+    wel_ffill = wel_class.from_imod5_data("wel-1-WELLS_L3", data, times)
+    wel_skipna = wel_class.from_imod5_data("wel-1-WELLS_L3", data, times, skipna=True)
+
+    rate_ffill = wel_ffill.dataset["rate"].values
+    rate_skipna = wel_skipna.dataset["rate"].values
+    assert np.isfinite(rate_ffill).all()
+    assert np.isfinite(rate_skipna).all()
+    assert not np.array_equal(rate_ffill, rate_skipna)
+
+
 @pytest.mark.unittest_jit
 @parametrize("wel_class", [Well])
 def test_import_and_cleanup(imod5_dataset, wel_class: Well):

imod/tests/test_mf6/test_utilities/test_resampling.py

@@ -166,6 +166,48 @@ def test_timeseries_resampling_varying_rates(output_days, expected_rates):
     )
 
 
+@pytest.mark.parametrize(
+    "record_days, rates, output_days, skipna, expected_rates",
+    [
+        # Inner NaN: forward filled by default, left out of the weights with skipna.
+        ([1, 4, 9], [-10.0, np.nan, 0.0], [1, 11], False, [-8.0, 0.0]),
+        ([1, 4, 9], [-10.0, np.nan, 0.0], [1, 11], True, [-6.0, 0.0]),
+        # Trailing NaN: left out of the weights, last row gets the last valid rate.
+        ([1, 4, 9], [-10.0, -20.0, np.nan], [1, 11], False, [-17.0, -20.0]),
+        ([1, 4, 9], [-10.0, -20.0, np.nan], [1, 11], True, [-16.25, -20.0]),
+        # Period with only NaN rates gets zero.
+        (
+            [1, 4, 9],
+            [-10.0, np.nan, -30.0],
+            [1, 4, 9, 11],
+            True,
+            [-10.0, 0.0, -30.0, -30.0],
+        ),
+        # Leading NaN: zero by default, left out of the weights with skipna.
+        ([1, 4], [np.nan, -20.0], [1, 11], False, [-14.0, -20.0]),
+        ([1, 4], [np.nan, -20.0], [1, 11], True, [-20.0, -20.0]),
+        # Only NaN rates.
+        ([1, 4], [np.nan, np.nan], [1, 11], False, [0.0, 0.0]),
+        ([1, 4], [np.nan, np.nan], [1, 11], True, [0.0, 0.0]),
+    ],
+)
+def test_timeseries_resampling_nan_rates(
+    record_days, rates, output_days, skipna, expected_rates
+):
+    # Tests resampling of timeseries with NaN rates, checking both skipna options.
+    timeseries = initialize_timeseries(
+        [datetime(2020, 1, day) for day in record_days], rates
+    )
+    new_dates = pd.DatetimeIndex([datetime(2020, 1, day) for day in output_days])
+
+    new_timeseries = resample_timeseries(timeseries, new_dates, skipna=skipna)
+
+    expected_timeseries = initialize_timeseries(new_dates, expected_rates)
+    pd.testing.assert_frame_equal(
+        new_timeseries, expected_timeseries, check_dtype=False
+    )
+
+
 def test_timeseries_resampling_6():
     # In this test, we resample a timeseries after the last timestep of
     # timeseries, should be forward filled.

imod/util/expand_repetitions.py

@@ -4,6 +4,8 @@
 import numpy as np
 import pandas as pd
 
+LOCATION_COLUMNS = ["x", "y", "id", "filt_top", "filt_bot", "index"]
+
 
 def expand_repetitions(
     repeat_stress: list[datetime], time_min: datetime, time_max: datetime
@@ -46,118 +48,171 @@ def expand_repetitions(
 def average_timeseries(well_rate: pd.DataFrame) -> pd.DataFrame:
     """
     Compute the mean value of the timeseries for a single well. Time column is
-    dropped.
+    dropped. If no records are valid, the rate is set to 0.0.
 
     Parameters
     ----------
     well_rate:  pd.DataFrame
         input timeseries for a single well
+
+    Returns
+    -------
+    pd.DataFrame
+        A single-row DataFrame with the mean rate and the location columns
+        from the first record.
     """
     # Take first item
     output_frame = well_rate.iloc[[0]].drop(columns=["time", "rate"])
-    # Compute mean
+    # Compute mean.
+    rate = well_rate["rate"].mean()
     col_index = output_frame.shape[1]
-    output_frame.insert(col_index, "rate", well_rate["rate"].mean())
+    output_frame.insert(col_index, "rate", 0.0 if pd.isna(rate) else rate)
     return output_frame
 
 
-def resample_timeseries(
-    well_rate: pd.DataFrame, times: list[datetime] | pd.DatetimeIndex
+def _align_records_to_output_times(
+    well_rate: pd.DataFrame,
+    output_frame: pd.DataFrame,
+    times: list[datetime] | pd.DatetimeIndex,
 ) -> pd.DataFrame:
     """
-    On input, well_rate is a dataframe containing a timeseries for rate for one well
-    while "times" is a list of datetimes.
-    This function creates a new dataframe containing a timeseries defined on the
-    times in "times".
-
-    Parameters
-    ----------
-    well_rate:  pd.DataFrame
-        input timeseries for a single well
-    times: datetime
-        list of times on which the output dataframe should have entries.
-
-    returns:
-    --------
-    a new dataframe containing a timeseries defined on the times in "times".
+    Merge well records and output times, so that every row holds the record in
+    effect at that time, and clip records before the first output time. NaN
+    rates in the records are kept.
     """
-    is_steady_state = len(times) == 0
-    # Fill NaT for the last timestep in well_rate if it is outside
-    # pandas' Timestamp nanosecond range (2262), for example 2999-12-31.
-    # FUTURE: Can be removed when we require pandas 3.0
-    tc_index = well_rate.columns.get_loc("time")
-    well_rate.iloc[-1:, tc_index] = well_rate.iloc[-1:, tc_index].fillna(
-        pd.Timestamp.max
-    )
-
-    output_frame = pd.DataFrame(times)
-    output_frame = output_frame.rename(columns={0: "time"})
+    # Look up the rate of each row from the record in effect, so a NaN rate stays
+    # a property of its record instead of being filled by the forward fill.
+    well_rate_numbered = well_rate.assign(record_nr=np.arange(len(well_rate)))
     intermediate_df = pd.merge(
         output_frame,
-        well_rate,
+        well_rate_numbered,
         how="outer",
         on="time",
         validate="m:m",
     ).ffill()
-    # Catch case where multiple well rates precede first timestep on multiple
-    # occasions, these have to be clipped to avoid including them in the
-    # computation of weighted averages later in the function. We only want to
-    # include the last well timestep BEFORE the first output timestep.
+    rates = pd.Series(well_rate["rate"].to_numpy(dtype=float, na_value=np.nan))
+    intermediate_df["rate"] = intermediate_df["record_nr"].map(rates)
+    intermediate_df = intermediate_df.drop(columns="record_nr")
+    # Only keep the last well record BEFORE the first output time, earlier
+    # records would distort the weighted averages.
     n_before_first_timestep = (intermediate_df["time"] < times[0]).cumsum()
     to_preserve = n_before_first_timestep == n_before_first_timestep.max()
-    intermediate_df = intermediate_df[to_preserve]
+    return intermediate_df[to_preserve]
+
 
-    # The entries before the start of the well timeseries do not have data yet,
-    # so we fill them in here. Keep rate to zero and pad the location columns with
-    # the first entry.
-    location_columns = ["x", "y", "id", "filt_top", "filt_bot", "index"]
+def _zero_fill_before_first_record(
+    intermediate_df: pd.DataFrame, well_rate: pd.DataFrame
+) -> None:
+    """
+    The entries before the start of the well timeseries do not have data yet,
+    so we fill them in here. Keep rate to zero and pad the location columns with
+    the first entry. Modifies intermediate_df in place.
+    """
     time_before_start_input = (
         intermediate_df["time"].values < well_rate["time"].values[0]
     )
     if time_before_start_input[0]:
         intermediate_df.loc[time_before_start_input, "rate"] = 0.0
-        first_well_location_values = well_rate.iloc[0][location_columns].values
-        intermediate_df.loc[time_before_start_input, location_columns] = (
+        first_well_location_values = well_rate.iloc[0][LOCATION_COLUMNS].values
+        intermediate_df.loc[time_before_start_input, LOCATION_COLUMNS] = (
             first_well_location_values
         )
 
-    # compute time difference from perious to current row
+
+def _add_durations_and_volumes(intermediate_df: pd.DataFrame) -> None:
+    """Add the duration until the next row and the rate volume"""
+    # Compute time difference from previous to current row
     time_diff_col = intermediate_df["time"].diff()
     col_index = intermediate_df.shape[1] - 1
     intermediate_df.insert(col_index, "time_to_next", time_diff_col.values)
 
-    # shift the new column 1 place down so that they become the time to the next row
+    # Shift the column up so each row holds the time until the next row
     intermediate_df["time_to_next"] = intermediate_df["time_to_next"].shift(-1)
 
-    # Integrate by grouping by the period number
     intermediate_df["duration_sec"] = intermediate_df["time_to_next"].dt.total_seconds()
+    # A NaN duration drops the row from both the volume and duration sums.
+    intermediate_df.loc[intermediate_df["rate"].isna(), "duration_sec"] = np.nan
     intermediate_df["volume"] = (
         intermediate_df["rate"] * intermediate_df["duration_sec"]
     )
+
+
+def resample_timeseries(
+    well_rate: pd.DataFrame,
+    times: list[datetime] | pd.DatetimeIndex,
+    skipna: bool = False,
+) -> pd.DataFrame:
+    """
+    Resample the rate timeseries of a single well to the given output times.
+
+    The rate at each output time is the average rate over the period until the
+    next output time, weighted by duration. The last record before the first
+    output time sets the rate at the start, the rate is zero before the first
+    record, and the last output time gets the last well rate.
+
+    NaN rates are forward filled by default, or set to zero if they come before
+    the first valid rate. With ``skipna=True`` they are left out of the weights
+    instead. Periods with only NaN rates count as zero, and the last output
+    time gets the last valid rate.
+
+    Parameters
+    ----------
+    well_rate: pd.DataFrame
+        Timeseries of a single well, with columns "time", "rate", "x", "y",
+        "id", "filt_top", "filt_bot" and "index".
+    times: list of datetime or pd.DatetimeIndex
+        Times on which the output dataframe should have entries.
+    skipna: bool, default False
+        Leave NaN rates out of the weights instead of forward filling them.
+
+    Returns
+    -------
+    pd.DataFrame
+        Timeseries with the same columns as ``well_rate``, with one row per
+        entry in ``times``.
+    """
+    is_steady_state = len(times) == 0
+    if not skipna:
+        # NaN rates before the first valid rate have nothing to fill from.
+        well_rate = well_rate.assign(rate=well_rate["rate"].ffill().fillna(0.0))
+
+    output_frame = pd.DataFrame(times)
+    output_frame = output_frame.rename(columns={0: "time"})
+    intermediate_df = _align_records_to_output_times(well_rate, output_frame, times)
+    _zero_fill_before_first_record(intermediate_df, well_rate)
+    _add_durations_and_volumes(intermediate_df)
+
+    # Integrate by grouping by the period number
     intermediate_df["period_nr"] = intermediate_df["time"].isin(times).cumsum()
-    # Drop the kept record before the first output time (period_nr 0). Its rate
-    # is already forward filled to the first output time.
+    # Drop the record before the first output time (period_nr 0), as it is
+    # already forward filled onto the first output time. Keeping it would shift
+    # all averages one row.
     intermediate_df = intermediate_df[intermediate_df["period_nr"] > 0]
     gb = intermediate_df.groupby("period_nr")
 
     output_frame["rate"] = (gb["volume"].sum() / gb["duration_sec"].sum()).reset_index(
         drop=True
     )
-    # If last value is nan (fell outside range), pad with last well rate.
+    # If last value is nan (fell outside range), pad with last valid well rate.
     last_index = output_frame.index[-1]
     if output_frame.isna().loc[last_index, "rate"]:
-        output_frame.loc[last_index, "rate"] = well_rate["rate"].iloc[-1]
+        valid_rates = well_rate["rate"].dropna()
+        output_frame.loc[last_index, "rate"] = (
+            valid_rates.iloc[-1] if len(valid_rates) > 0 else 0.0
+        )
+    # Periods where all rates are NaN.
+    output_frame["rate"] = output_frame["rate"].fillna(0.0)
 
     if is_steady_state:
         # Take first element, the slice is to force pandas to return it as
         # dataframe instead of series.
-        location_dataframe = intermediate_df[location_columns].iloc[slice(0, 1), :]
+        location_dataframe = intermediate_df[LOCATION_COLUMNS].iloc[slice(0, 1), :]
         # Concat along columns and drop time column
         return pd.concat([output_frame, location_dataframe], axis=1).drop(
             columns="time"
         )
     else:
-        columns_to_merge = ["time"] + location_columns
+        columns_to_merge = ["time"] + LOCATION_COLUMNS
         return pd.merge(
             output_frame,
             intermediate_df[columns_to_merge],