pastas-plugins 0.3.0__py3-none-any.whl

pastas_plugins/__init__.py

@@ -0,0 +1,35 @@
+# ruff : noqa: F401
+import pathlib
+from importlib import import_module
+
+# from platform import python_version
+# from packaging.version import parse as parse_version
+from pastas_plugins.version import __version__
+
+
+def list_plugins():
+    plugins = pathlib.Path(__file__).parent.iterdir()
+    plugins = [
+        plugin.stem
+        for plugin in plugins
+        if plugin.is_dir() and not plugin.stem.startswith("_")
+    ]
+    plugins.sort()
+    return plugins
+
+
+def show_plugin_versions():
+    showtip = False
+    plugins = list_plugins()
+    msg = f"pastas_plugins version      : {__version__}\n"
+    for plugin in plugins:
+        try:
+            module = import_module(f"pastas_plugins.{plugin}.version")
+            version = module.__version__
+        except ModuleNotFoundError:
+            showtip = True
+            version = "not available (check dependencies)"
+        msg += f"- {(plugin + ' version'):25s} : {version}\n"
+    if showtip:
+        msg += "\nNote: To install missing dependencies use `pip install pastas-plugins[<plugin-name>]`"
+    print(msg)

pastas_plugins/cross_correlation/__init__.py

@@ -0,0 +1,8 @@
+# ruff: noqa: F401
+from pastas_plugins.cross_correlation.cross_correlation import (
+    ccf,
+    fit_response,
+    prewhiten,
+)
+from pastas_plugins.cross_correlation.plots import plot_ccf_overview, plot_corr
+from pastas_plugins.cross_correlation.version import __version__

pastas_plugins/cross_correlation/cross_correlation.py

@@ -0,0 +1,188 @@
+from typing import Tuple, Union
+
+import numpy as np
+import pandas as pd
+import pastas as ps
+import scipy as sc
+from statsmodels.tsa.ar_model import AutoReg
+from statsmodels.tsa.arima.model import ARIMA
+from statsmodels.tsa.filters.filtertools import convolution_filter
+
+
+def ccf(
+    x: pd.Series,
+    y: pd.Series,
+    nlags: Union[int, None] = None,
+    adjusted: bool = True,
+    alpha: Union[float, None] = None,
+) -> Union[pd.Series, pd.DataFrame]:
+    """Cross-correlation of two time series.
+
+    Parameters
+    ----------
+    x : pd.Series
+        Time series
+    y : pd.Series
+        Time series, len(y) should be equal to len(x).
+    nlags : int or None, optional
+        Number of lags to return cross-correlations for, by default None which
+        uses nlags equal to len(x).
+    adjusted : bool, optional
+        If True, denominators for cross-correlation are len(x)-k, otherwise
+        len(x), by default True
+    alpha : float or None, optional
+        If a float between 0 and 1 is given, the confidence intervals for the
+        given level are returned in a DataFrame. For instance if alpha=0.05,
+        95% confidence intervals are returned where the standard deviation is
+        computed according to 1/sqrt(len(x)).
+
+    Returns
+    -------
+    pandas Series or DataFrame
+    """
+    # check if lengths are equal
+    assert len(x) == len(y), "Length of series x and y should be equal"
+    # check if series are equidistant
+    for series in (x, y):
+        if pd.infer_freq(series.index) is None:
+            msg = (
+                "The frequency of the index of time series %s could not be "
+                "inferred. Please provide a time series with a equidistant time step."
+            )
+            raise ValueError(msg % series.name)
+
+    n = len(x)
+
+    xbar = x - x.mean()
+    ybar = y - y.mean()
+
+    d = np.arange(n, 0, -1) if adjusted else n
+    cc = sc.signal.correlate(xbar, ybar, mode="full", method="fft")
+    cvf = cc[n - 1 :] / (np.std(x) * np.std(y) * d)
+
+    nlags = n if nlags is None else nlags
+    index = pd.Index(np.arange(nlags), name="Lags")
+    ret = cvf[:nlags]
+
+    if alpha is not None:
+        interval = sc.stats.norm.ppf(1.0 - alpha / 2.0) / np.sqrt(n)
+        crosscorr = pd.DataFrame(
+            data=np.vstack([ret, ret - interval, ret + interval]).T,
+            index=index,
+            columns=["Cross-correlation", f"CI {alpha / 2}", f"CI {1 - alpha / 2}"],
+        )
+    else:
+        crosscorr = pd.Series(
+            data=ret,
+            index=index,
+            name="Cross-correlation",
+        )
+    return crosscorr
+
+
+def prewhiten(
+    x: pd.Series, y: Union[pd.Series, None] = None, ar: int = 20, arima: bool = False
+) -> Union[pd.Series, Tuple[pd.Series]]:
+    """Prewhiten time series using AR(ar) model.
+
+    An AR(ar) model is fitted on time series x. The goal is to obtain residuals that
+    adhere to a white noise process. Next, the AR(ar) model is applied to time series Y.
+
+    Note
+    ----
+    If prewhitened time series for x still shows significant autocorrelation or partial
+    autocorrelation, try increasing the number of autoregressive parameters.
+
+    Parameters
+    ----------
+    x : pd.Series
+        time series on which AR(ar) model will be fitted
+    y : pd.Series, optional
+        time series that will be filtered using the AR(ar) model fitted on x
+    ar : int, optional
+        number of autoregressive parameters (sometimes called `p`), by default 20
+    arima: bool, optional
+        use an ARIMA(ar,0,0) model instead of an AR(ar) model, by default False
+        which causes a significant speedup at the cost of a very small accuracy
+        penalty
+
+    Returns
+    -------
+    pwx : pd.Series
+        prewhitened time series for x (should no longer show significant
+        autocorrelation or partial autocorrelation)
+    pwy : pd.Series, optional
+        prewhitened time series for y, if y is provided
+    """
+
+    # fit AR model on x
+    if arima:
+        ml = ARIMA(x.values, order=(ar, 0, 0), trend="c").fit()
+    else:
+        ml = AutoReg(x.values, lags=ar, trend="c").fit()
+
+    # get model filtered model residuals
+    residuals = ml.resid[ar:] if arima else ml.resid
+    pwx = pd.Series(residuals, index=x.index[ar:])
+
+    if y is not None:
+        # apply same filter on y
+        arparams = ml.arparams if arima else ml.params[1:]
+        filt = np.append(1.0, -arparams)
+        pwy = convolution_filter(y.values, filt=filt, nsides=1)
+        pwy = pd.Series(data=pwy[ar:], index=y.index[ar:])
+        return pwx, pwy
+    else:
+        return pwx
+
+
+def fit_response(
+    ccf: pd.Series,
+    rfunc: ps.typing.RFunc,
+    scale_factor: float = 1.0,
+    dt: float = 1.0,
+) -> np.ndarray[float]:
+    """Fit the response function to the cross-correlation function using least
+    squares optimization.
+
+    Parameters:
+    -----------
+    ccf : pd.Series
+        The cross-correlation function.
+    rfunc : ps.typing.RFunc
+        The response function to fit on the impulse response.
+    scale_factor : float, optional
+        Scale factor applied to the cross-correlation function to obtain the
+        impulse response, by default 1.0.
+    dt : float, optional
+        Time step of the response function, by default 1.0.
+
+    Returns:
+    --------
+    np.ndarray[float]
+        The optimized parameters of the response function.
+
+    """
+
+    def obj_func(p):
+        """Objective function for least squares optimization."""
+        impulse_response = (ccf * scale_factor).values
+        blockr = rfunc.block(p, dt=dt, cutoff=rfunc.cutoff)
+
+        # make sure length of residuals is constant
+        if len(blockr) > len(impulse_response):
+            blockr = blockr[: len(impulse_response)]
+        elif len(blockr) < len(impulse_response):
+            blockr = np.append(blockr, np.zeros(len(impulse_response) - len(blockr)))
+
+        return impulse_response - blockr
+
+    params = rfunc.get_init_parameters(rfunc._name)
+    pini = params["initial"].values
+    bounds = (
+        params["pmin"].fillna(-np.inf).values,
+        params["pmax"].fillna(np.inf).values,
+    )
+
+    res = sc.optimize.least_squares(obj_func, x0=pini, bounds=bounds)
+    return res.x

pastas_plugins/cross_correlation/plots.py

@@ -0,0 +1,245 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from statsmodels.graphics.tsaplots import _plot_corr, plot_acf, plot_pacf
+
+from pastas_plugins.cross_correlation.cross_correlation import ccf
+
+
+def plot_corr(
+    corr: pd.Series | pd.DataFrame,
+    ax: plt.Axes | None = None,
+    vlines_kwargs: dict | None = None,
+    **kwargs,
+):
+    """Helper function for the statsmodels _plot_corr function.
+
+    Parameters
+    ----------
+    corr : pd.Series or pd.DataFrame
+        the correlation result to plot
+    ax : plt.Axes, optional
+        axes to plot on, by default None
+    vlines_kwargs : dict, optional
+        keyword arguments for the vlines function, by default None
+
+    Returns
+    -------
+    plt.Axes
+        axes with the plot
+    """
+
+    if ax is None:
+        _, ax = plt.subplots(**kwargs)
+
+    acf_x = corr.values if isinstance(corr, pd.Series) else corr.iloc[:, 0].values
+    confint = corr.iloc[:, 1:3].values if isinstance(corr, pd.DataFrame) else None
+    lags = corr.index.values
+    vlines_kwargs = {} if vlines_kwargs is None else vlines_kwargs
+    _plot_corr(
+        ax=ax,
+        title="",
+        acf_x=acf_x,
+        confint=confint,
+        lags=lags,
+        irregular=False,
+        use_vlines=True,
+        vlines_kwargs=vlines_kwargs,
+    )
+    return ax
+
+
+def plot_ccf_overview(x, y, nlags=None, tmin=None, tmax=None, axes=None):
+    """Plot an overview of the cross-correlation between two time series.
+
+    Parameters
+    ----------
+    x : pd.Series
+        Time series 1
+    y : pd.Series
+        Time series 2
+    nlags : int, optional
+        number of lags to return cross-correlations for, by default None which
+        uses number of lags equal to len(x).
+    tmin : str or pd.Timestamp, optional
+        tmin for both time series, by default None
+    tmax : str or pd.Timestamp, optional
+        tmax for both time series, by default None
+    axes : Axes mosaic, optional
+        if provided, use axes from previous plot
+
+    Returns
+    -------
+    axes : Axes mosaic
+        return axes of subplots mosaic
+    """
+    if tmin is None:
+        tmin = np.min([x.index[0], y.index[0]])
+    if tmax is None:
+        tmax = np.max([x.index[-1], y.index[-1]])
+
+    x = x.loc[tmin:tmax]
+    y = y.loc[tmin:tmax]
+
+    if axes is None:
+        mosaic = [
+            ["x", "x", "norm", "norm"],
+            ["y", "y", "norm", "norm"],
+            ["x-acf", "y-acf", "ccf", "ccf"],
+            ["x-pacf", "y-pacf", "ccf", "ccf"],
+        ]
+
+        fig, axes = plt.subplot_mosaic(mosaic, figsize=(16, 8))
+        rescale_axes = False
+        newaxes = True
+        color1 = "C0"
+        color2 = "C1"
+    else:
+        fig = axes["x"].figure
+        rescale_axes = True
+        newaxes = False
+        color1 = "C2"
+        color2 = "C3"
+
+    # set names if not provided
+    if x.name is None:
+        x.name = "x"
+    if y.name is None:
+        y.name = "y"
+
+    # plot time series
+    axes["x"].plot(x.index, x, label=x.name, color=color1)
+    axes["x"].legend(loc=(0, 1), frameon=False)
+    axes["x"].set_ylabel("x")
+    axes["x"].set_xlim(pd.Timestamp(tmin), pd.Timestamp(tmax))
+    axes["y"].plot(y.index, y, label=y.name, c=color2)
+    axes["y"].legend(loc=(0, 1), frameon=False)
+    axes["y"].set_ylabel("y")
+    axes["y"].set_xlim(pd.Timestamp(tmin), pd.Timestamp(tmax))
+
+    # plot normalized series
+    xnorm = (x - x.mean()) / x.std()
+    ynorm = (y - y.mean()) / y.std()
+    axes["norm"].plot(
+        xnorm.index, xnorm, label=x.name + " (normalized)", alpha=0.7, color=color1
+    )
+    axes["norm"].plot(
+        ynorm.index, ynorm, label=y.name + " (normalized)", alpha=0.7, color=color2
+    )
+    axes["norm"].legend(loc=(0, 1), frameon=False, ncol=2)
+    axes["norm"].set_ylabel("normalized [-]")
+    axes["norm"].set_xlim(pd.Timestamp(tmin), pd.Timestamp(tmax))
+    handles, _ = axes["norm"].get_legend_handles_labels()
+
+    # plot acf, pacf
+    plot_acf(
+        xnorm,
+        ax=axes["x-acf"],
+        color=color1,
+        alpha=0.05,
+        title="",
+        zero=False,
+        auto_ylims=True,
+        vlines_kwargs={"color": "k"},
+    )
+    plot_acf(
+        ynorm,
+        ax=axes["y-acf"],
+        color=color2,
+        alpha=0.05,
+        title="",
+        zero=False,
+        auto_ylims=True,
+        vlines_kwargs={"color": "k"},
+    )
+    axes["x-acf"].set_xlim(left=0.0)
+    axes["x-acf"].set_ylabel("ACF [-]")
+    (p1,) = axes["x-acf"].plot([], [], marker="o", ls="none", color=color1)
+    if not newaxes:
+        leg = axes["x-acf"].get_legend()
+        handles = leg.legend_handles
+        labels = [t.get_text() for t in leg.get_texts()]
+        handles += [p1]
+        labels += [x.name]
+    else:
+        handles = [p1]
+        labels = [x.name]
+
+    axes["x-acf"].legend(handles, labels, loc=(0, 1), frameon=False)
+    axes["y-acf"].set_xlim(left=0.0)
+    axes["y-acf"].get_children()[3].set_facecolor(color2)
+    (p2,) = axes["y-acf"].plot([], [], marker="o", ls="none", color=color2)
+    if not newaxes:
+        leg = axes["y-acf"].get_legend()
+        handles = leg.legend_handles
+        labels = [t.get_text() for t in leg.get_texts()]
+        handles += [p2]
+        labels += [y.name]
+    else:
+        handles = [p2]
+        labels = [y.name]
+    axes["y-acf"].legend(handles, labels, loc=(0, 1), frameon=False)
+
+    plot_pacf(
+        xnorm,
+        method="ywm",
+        ax=axes["x-pacf"],
+        color=color1,
+        alpha=0.05,
+        title="",
+        zero=False,
+        auto_ylims=True,
+        vlines_kwargs={"color": "k"},
+    )
+    plot_pacf(
+        ynorm,
+        method="ywm",
+        ax=axes["y-pacf"],
+        color=color2,
+        alpha=0.05,
+        title="",
+        zero=False,
+        auto_ylims=True,
+        vlines_kwargs={"color": "k"},
+    )
+    axes["x-pacf"].set_xlim(left=0.0)
+    axes["x-pacf"].set_ylabel("PACF [-]")
+    axes["x-pacf"].set_xlabel("Lags")
+
+    axes["y-pacf"].set_xlim(left=0.0)
+    axes["y-pacf"].get_children()[3].set_facecolor("C1")
+    axes["y-pacf"].set_xlabel("Lags")
+
+    # ccf
+    cc = ccf(x, y, nlags=nlags)
+    axes["ccf"].bar(
+        cc.index,
+        cc,
+        width=1.0,
+        linewidth=0.5,
+        alpha=0.5,
+        label=f"CCF ({x.name}|{y.name})",
+    )
+    axes["ccf"].set_ylabel("CCF [-]")
+    axes["ccf"].set_xlabel("Lags")
+    axes["ccf"].legend(loc=(0, 1), frameon=False)
+    axes["ccf"].set_xlim(left=0.0)
+
+    share_x = [axes["x"], axes["y"], axes["norm"]]
+    for i, iax in enumerate(share_x):
+        if i < (len(share_x) - 1):
+            iax.sharex(share_x[-1])
+
+    # share_x = [axes["x-acf"], axes["x-pacf"], axes["y-acf"], axes["y-pacf"]]
+    # for i, iax in enumerate(share_x):
+    #     if i < (len(share_x) - 1):
+    #         iax.sharex(share_x[-1])
+
+    fig.tight_layout()
+    fig.align_ylabels()
+
+    if rescale_axes:
+        for iax in axes.values():
+            iax.autoscale()
+
+    return axes

pastas_plugins/cross_correlation/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pastas_plugins/modflow/__init__.py

@@ -0,0 +1,4 @@
+# ruff: noqa: F401
+from pastas_plugins.modflow.modflow import ModflowRch
+from pastas_plugins.modflow.stressmodels import ModflowModel
+from pastas_plugins.modflow.version import __version__

pastas_plugins/modflow/modflow.py

@@ -0,0 +1,186 @@
+import functools
+import logging
+from typing import List, Protocol
+
+import flopy
+import numpy as np
+from pandas import DataFrame, Series
+from pastas.typing import ArrayLike
+
+logger = logging.getLogger(__name__)
+
+
+class Modflow(Protocol):
+    def __init__(self) -> None: ...
+
+    def get_init_parameters(self) -> DataFrame: ...
+
+    def create_model(self) -> None: ...
+
+    def simulate(self) -> ArrayLike: ...
+
+
+class ModflowRch:
+    def __init__(
+        self, exe_name: str, sim_ws: str, raise_on_modflow_error: bool = False
+    ):
+        self.exe_name = exe_name
+        self.sim_ws = sim_ws
+        self._name = "mf_rch"
+        self._stress = None
+        self._simulation = None
+        self._gwf = None
+        self._changing_packages = (
+            "STO",
+            "GHB",
+            "RCH",
+        )
+        self.raise_on_modflow_error = raise_on_modflow_error
+
+    def get_init_parameters(self, name: str) -> DataFrame:
+        parameters = DataFrame(columns=["initial", "pmin", "pmax", "vary", "name"])
+        parameters.loc[name + "_sy"] = (0.05, 0.001, 0.5, True, name)
+        parameters.loc[name + "_c"] = (220, 1e1, 1e8, True, name)
+        parameters.loc[name + "_f"] = (-1.0, -2.0, 0.0, True, name)
+        return parameters
+
+    def create_model(self) -> None:
+        sim = flopy.mf6.MFSimulation(
+            sim_name=self._name,
+            version="mf6",
+            exe_name=self.exe_name,
+            sim_ws=self.sim_ws,
+            lazy_io=True,
+        )
+
+        _ = flopy.mf6.ModflowTdis(
+            sim,
+            time_units="DAYS",
+            nper=self._nper,
+            perioddata=[(1, 1, 1) for _ in range(self._nper)],
+        )
+
+        gwf = flopy.mf6.ModflowGwf(
+            sim,
+            modelname=self._name,
+        )
+
+        _ = flopy.mf6.ModflowIms(
+            sim,
+            complexity="SIMPLE",
+            outer_dvclose=1e-2,
+            inner_dvclose=1e-2,
+            rcloserecord=1e-1,
+            linear_acceleration="BICGSTAB",
+            pname=None,
+        )
+        # sim.register_ims_package(imsgwf, [self._name])
+
+        _ = flopy.mf6.ModflowGwfdis(
+            gwf,
+            length_units="METERS",
+            nlay=1,
+            nrow=1,
+            ncol=1,
+            delr=1,
+            delc=1,
+            top=1.0,
+            botm=0.0,
+            idomain=1,
+            pname=None,
+        )
+
+        _ = flopy.mf6.ModflowGwfnpf(
+            gwf, save_flows=False, icelltype=0, k=1.0, pname="npf"
+        )
+
+        _ = flopy.mf6.ModflowGwfic(gwf, strt=0.0, pname="ic")
+
+        _ = flopy.mf6.ModflowGwfoc(
+            gwf,
+            head_filerecord=f"{gwf.name}.hds",
+            saverecord=[("HEAD", "ALL")],
+            pname=None,
+        )
+
+        sim.write_simulation(silent=True)
+        self._simulation = sim
+        self._gwf = gwf
+
+    def update_model(self, p: ArrayLike):
+        sy, c, f = p[0:3]
+
+        d = 0.0
+        r = self._stress[0] + f * self._stress[1]
+
+        # remove existing packages
+        if all(
+            [True for x in self._gwf.get_package_list() if x in self._changing_packages]
+        ):
+            [self._gwf.remove_package(x) for x in self._changing_packages]
+
+        haq = (self._gwf.dis.top.array - self._gwf.dis.botm.array)[0]
+        sto = flopy.mf6.ModflowGwfsto(
+            self._gwf,
+            save_flows=False,
+            iconvert=0,
+            ss=sy / haq,
+            transient=True,
+            pname="sto",
+        )
+        sto.write()
+
+        # ghb
+        ghb = flopy.mf6.ModflowGwfghb(
+            self._gwf,
+            maxbound=1,
+            stress_period_data={0: [[(0, 0, 0), d, 1.0 / c]]},
+            pname="ghb",
+        )
+        ghb.write()
+
+        rts = [(i, x) for i, x in zip(range(self._nper + 1), np.append(r, 0.0))]
+
+        ts_dict = {
+            "filename": "recharge.ts",
+            "timeseries": rts,
+            "time_series_namerecord": ["recharge"],
+            "interpolation_methodrecord": ["stepwise"],
+        }
+
+        rch = flopy.mf6.ModflowGwfrch(
+            self._gwf,
+            maxbound=1,
+            pname="rch",
+            stress_period_data={0: [[(0, 0, 0), "recharge"]]},
+            timeseries=ts_dict,
+        )
+        rch.write()
+        rch.ts.write()
+
+        self._gwf.name_file.write()
+
+    @functools.lru_cache(maxsize=5)
+    def _get_head(self, p):
+        self.update_model(p=p)
+        success, _ = self._simulation.run_simulation(silent=True)
+        if success:
+            return self._gwf.output.head().get_ts((0, 0, 0))[:, 1]
+        else:
+            logger.error(
+                "ModflowError: model run failed with parameters: "
+                f"sy={p[0]}, c={p[1]}, f={p[2]}"
+            )
+            if self.raise_on_modflow_error:
+                raise Exception(
+                    "Modflow run failed. Check the LIST file for more information."
+                )
+            else:
+                return np.zeros(self._nper)
+
+    def simulate(self, p: ArrayLike, stress: List[Series]) -> ArrayLike:
+        if self._simulation is None:
+            self._stress = stress
+            self._nper = len(self._stress[0])
+            self.create_model()
+        return self._get_head(tuple(p))