pastastore 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

pastastore/store.py

@@ -1,7 +1,9 @@
+"""Module containing the PastaStore object for managing time series and models."""
+
 import json
 import os
 import warnings
-from typing import List, Optional, Tuple, Union
+from typing import List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -14,6 +16,7 @@ from pastastore.base import BaseConnector
 from pastastore.connectors import DictConnector
 from pastastore.plotting import Maps, Plots
 from pastastore.util import _custom_warning
+from pastastore.version import PASTAS_GEQ_150
 from pastastore.yaml_interface import PastastoreYAML
 
 FrameorSeriesUnion = Union[pd.DataFrame, pd.Series]
@@ -72,7 +75,7 @@ class PastaStore:
         self.yaml = PastastoreYAML(self)
 
     def _register_connector_methods(self):
-        """Internal method for registering connector methods."""
+        """Register connector methods (internal method)."""
         methods = [
             func
             for func in dir(self.conn)
@@ -83,26 +86,70 @@ class PastaStore:
 
     @property
     def oseries(self):
+        """
+        Returns the oseries metadata as dataframe.
+
+        Returns
+        -------
+        oseries
+            oseries metadata as dataframe
+        """
         return self.conn.oseries
 
     @property
     def stresses(self):
+        """
+        Returns the stresses metadata as dataframe.
+
+        Returns
+        -------
+        stresses
+            stresses metadata as dataframe
+        """
         return self.conn.stresses
 
     @property
     def models(self):
+        """Return list of model names.
+
+        Returns
+        -------
+        list
+            list of model names
+        """
         return self.conn.models
 
     @property
     def oseries_names(self):
+        """Return list of oseries names.
+
+        Returns
+        -------
+        list
+            list of oseries names
+        """
         return self.conn.oseries_names
 
     @property
     def stresses_names(self):
+        """Return list of streses names.
+
+        Returns
+        -------
+        list
+            list of stresses names
+        """
         return self.conn.stresses_names
 
     @property
     def model_names(self):
+        """Return list of model names.
+
+        Returns
+        -------
+        list
+            list of model names
+        """
         return self.conn.model_names
 
     @property
@@ -111,22 +158,57 @@ class PastaStore:
 
     @property
     def n_oseries(self):
+        """Return number of oseries.
+
+        Returns
+        -------
+        int
+            number of oseries
+        """
         return self.conn.n_oseries
 
     @property
     def n_stresses(self):
+        """Return number of stresses.
+
+        Returns
+        -------
+        int
+            number of stresses
+        """
         return self.conn.n_stresses
 
     @property
     def n_models(self):
+        """Return number of models.
+
+        Returns
+        -------
+        int
+            number of models
+        """
         return self.conn.n_models
 
     @property
     def oseries_models(self):
+        """Return dictionary of models per oseries.
+
+        Returns
+        -------
+        dict
+            dictionary containing list of models (values) for each oseries (keys).
+        """
         return self.conn.oseries_models
 
     @property
     def oseries_with_models(self):
+        """Return list of oseries for which models are contained in the database.
+
+        Returns
+        -------
+        list
+            list of oseries names for which models are contained in the database.
+        """
         return self.conn.oseries_with_models
 
     def __repr__(self):
@@ -136,7 +218,7 @@ class PastaStore:
     def get_oseries_distances(
         self, names: Optional[Union[list, str]] = None
     ) -> FrameorSeriesUnion:
-        """Method to obtain the distances in meters between the oseries.
+        """Get the distances in meters between the oseries.
 
         Parameters
         ----------
@@ -175,7 +257,7 @@ class PastaStore:
         n: int = 1,
         maxdist: Optional[float] = None,
     ) -> FrameorSeriesUnion:
-        """Method to obtain the nearest (n) oseries.
+        """Get the nearest (n) oseries.
 
         Parameters
         ----------
@@ -191,7 +273,6 @@ class PastaStore:
         oseries:
             list with the names of the oseries.
         """
-
         distances = self.get_oseries_distances(names)
         if maxdist is not None:
             distances = distances.where(distances <= maxdist, np.nan)
@@ -214,8 +295,7 @@ class PastaStore:
         stresses: Optional[Union[list, str]] = None,
         kind: Optional[Union[str, List[str]]] = None,
     ) -> FrameorSeriesUnion:
-        """Method to obtain the distances in meters between the oseries and
-        stresses.
+        """Get the distances in meters between the oseries and stresses.
 
         Parameters
         ----------
@@ -274,7 +354,7 @@ class PastaStore:
         n: int = 1,
         maxdist: Optional[float] = None,
     ) -> FrameorSeriesUnion:
-        """Method to obtain the nearest (n) stresses of a specific kind.
+        """Get the nearest (n) stresses of a specific kind.
 
         Parameters
         ----------
@@ -295,7 +375,6 @@ class PastaStore:
         stresses:
             list with the names of the stresses.
         """
-
         distances = self.get_distances(oseries, stresses, kind)
         if maxdist is not None:
             distances = distances.where(distances <= maxdist, np.nan)
@@ -317,8 +396,9 @@ class PastaStore:
         progressbar=False,
         ignore_errors=False,
     ):
-        """Get groundwater signatures. NaN-values are returned when the
-        signature could not be computed.
+        """Get groundwater signatures.
+
+        NaN-values are returned when the signature cannot be computed.
 
         Parameters
         ----------
@@ -380,11 +460,16 @@ class PastaStore:
                         i_signatures.append(sign_val)
                 else:
                     raise e
-            signatures_df.loc[name, signatures] = i_signatures
+            signatures_df.loc[name, signatures] = i_signatures.squeeze()
 
         return signatures_df
 
-    def get_tmin_tmax(self, libname, names=None, progressbar=False):
+    def get_tmin_tmax(
+        self,
+        libname: Literal["oseries", "stresses", "models"],
+        names: Union[str, List[str], None] = None,
+        progressbar: bool = False,
+    ):
         """Get tmin and tmax for time series.
 
         Parameters
@@ -403,22 +488,48 @@ class PastaStore:
         tmintmax : pd.dataframe
             Dataframe containing tmin and tmax per time series
         """
-
         names = self.conn._parse_names(names, libname=libname)
         tmintmax = pd.DataFrame(
             index=names, columns=["tmin", "tmax"], dtype="datetime64[ns]"
         )
         desc = f"Get tmin/tmax {libname}"
         for n in tqdm(names, desc=desc) if progressbar else names:
-            if libname == "oseries":
-                s = self.conn.get_oseries(n)
+            if libname == "models":
+                mld = self.conn.get_models(
+                    n,
+                    return_dict=True,
+                )
+                tmintmax.loc[n, "tmin"] = mld["settings"]["tmin"]
+                tmintmax.loc[n, "tmax"] = mld["settings"]["tmax"]
             else:
-                s = self.conn.get_stresses(n)
-            tmintmax.loc[n, "tmin"] = s.first_valid_index()
-            tmintmax.loc[n, "tmax"] = s.last_valid_index()
+                s = (
+                    self.conn.get_oseries(n)
+                    if libname == "oseries"
+                    else self.conn.get_stresses(n)
+                )
+                tmintmax.loc[n, "tmin"] = s.first_valid_index()
+                tmintmax.loc[n, "tmax"] = s.last_valid_index()
+
         return tmintmax
 
     def get_extent(self, libname, names=None, buffer=0.0):
+        """Get extent [xmin, xmax, ymin, ymax] from library.
+
+        Parameters
+        ----------
+        libname : str
+            name of the library containing the time series
+            ('oseries', 'stresses', 'models')
+        names : str, list of str, or None, optional
+            list of names to include for computing the extent
+        buffer : float, optional
+            add this distance to the extent, by default 0.0
+
+        Returns
+        -------
+        extent : list
+            extent [xmin, xmax, ymin, ymax]
+        """
         names = self.conn._parse_names(names, libname=libname)
         if libname in ["oseries", "stresses"]:
             df = getattr(self, libname)
@@ -443,8 +554,10 @@ class PastaStore:
         progressbar: Optional[bool] = False,
         ignore_errors: Optional[bool] = False,
     ) -> FrameorSeriesUnion:
-        """Get model parameters. NaN-values are returned when the parameters
-        are not present in the model or the model is not optimized.
+        """Get model parameters.
+
+        NaN-values are returned when the parameters are not present in the model or the
+        model is not optimized.
 
         Parameters
         ----------
@@ -526,7 +639,6 @@ class PastaStore:
         -------
         s : pandas.DataFrame
         """
-
         modelnames = self.conn._parse_names(modelnames, libname="models")
 
         # if statistics is str
@@ -558,6 +670,7 @@ class PastaStore:
         name: str,
         modelname: str = None,
         add_recharge: bool = True,
+        add_ar_noisemodel: bool = False,
         recharge_name: str = "recharge",
     ) -> ps.Model:
         """Create a pastas Model.
@@ -572,6 +685,8 @@ class PastaStore:
             add recharge to the model by looking for the closest
             precipitation and evaporation time series in the stresses
             library, by default True
+        add_ar1_noisemodel : bool, optional
+            add AR(1) noise model to the model, by default False
         recharge_name : str
             name of the RechargeModel
 
@@ -598,6 +713,8 @@ class PastaStore:
             ml = ps.Model(ts, name=modelname, metadata=meta)
             if add_recharge:
                 self.add_recharge(ml, recharge_name=recharge_name)
+            if add_ar_noisemodel and PASTAS_GEQ_150:
+                ml.add_noisemodel(ps.ArNoiseModel())
             return ml
         else:
             raise ValueError("Empty time series!")
@@ -709,9 +826,9 @@ class PastaStore:
         for var in ("prec", "evap"):
             try:
                 name = self.get_nearest_stresses(ml.oseries.name, kind=var).iloc[0, 0]
-            except AttributeError:
+            except AttributeError as e:
                 msg = "No precipitation or evaporation time series found!"
-                raise Exception(msg)
+                raise Exception(msg) from e
             if isinstance(name, float):
                 if np.isnan(name):
                     raise ValueError(
@@ -828,8 +945,8 @@ class PastaStore:
         """
         try:
             from art_tools import pastas_get_model_results
-        except Exception:
-            raise ModuleNotFoundError("You need 'art_tools' to use this method!")
+        except Exception as e:
+            raise ModuleNotFoundError("You need 'art_tools' to use this method!") from e
 
         if mls is None:
             mls = self.conn.models
@@ -870,7 +987,7 @@ class PastaStore:
                 "File already exists! " "Use 'overwrite=True' to " "force writing file."
             )
         elif os.path.exists(fname):
-            warnings.warn(f"Overwriting file '{os.path.basename(fname)}'")
+            warnings.warn(f"Overwriting file '{os.path.basename(fname)}'", stacklevel=1)
 
         with ZipFile(fname, "w", compression=ZIP_DEFLATED) as archive:
             # oseries
@@ -991,8 +1108,9 @@ class PastaStore:
         libname: str,
         s: Optional[Union[list, str]] = None,
         case_sensitive: bool = True,
+        sort=True,
     ):
-        """Search for names of time series or models starting with s.
+        """Search for names of time series or models starting with `s`.
 
         Parameters
         ----------
@@ -1002,19 +1120,20 @@ class PastaStore:
             find names with part of this string or strings in list
         case_sensitive : bool, optional
             whether search should be case sensitive, by default True
+        sort : bool, optional
+            sort list of names
 
         Returns
         -------
         matches : list
             list of names that match search result
         """
-
         if libname == "models":
-            lib_names = getattr(self, "model_names")
+            lib_names = self.model_names
         elif libname == "stresses":
-            lib_names = getattr(self, "stresses_names")
+            lib_names = self.stresses_names
         elif libname == "oseries":
-            lib_names = getattr(self, "oseries_names")
+            lib_names = self.oseries_names
         else:
             raise ValueError("Provide valid libname: 'models', 'stresses' or 'oseries'")
 
@@ -1031,7 +1150,8 @@ class PastaStore:
                 else:
                     m = np.append(m, [n for n in lib_names if sub.lower() in n.lower()])
             matches = list(np.unique(m))
-
+        if sort:
+            matches.sort()
         return matches
 
     def get_model_timeseries_names(
@@ -1060,7 +1180,6 @@ class PastaStore:
             indicating whether a stress is contained within a time series
             model.
         """
-
         model_names = self.conn._parse_names(modelnames, libname="models")
         structure = pd.DataFrame(
             index=model_names, columns=["oseries"] + self.stresses_names
@@ -1127,6 +1246,53 @@ class PastaStore:
                 "'libname' must be one of ['oseries', 'stresses', 'models']!"
             )
         getter = getattr(self.conn, f"get_{libname}")
-        for n in tqdm(names) if progressbar else names:
+        for n in (
+            tqdm(names, desc=f"Applying {func.__name__}") if progressbar else names
+        ):
             result[n] = func(getter(n))
         return result
+
+    def within(self, extent, names=None, libname="oseries"):
+        """Get names of items within extent.
+
+        Parameters
+        ----------
+        extent : list
+            list with [xmin, xmax, ymin, ymax]
+        names : str, list of str, optional
+            list of names to include, by default None
+        libname : str, optional
+            name of library, must be one of ('oseries', 'stresses', 'models'), by
+            default "oseries"
+
+        Returns
+        -------
+        list
+            list of items within extent
+        """
+        xmin, xmax, ymin, ymax = extent
+        names = self.conn._parse_names(names, libname)
+        if libname == "oseries":
+            df = self.oseries.loc[names]
+        elif libname == "stresses":
+            df = self.stresses.loc[names]
+        elif libname == "models":
+            onames = np.unique(
+                [
+                    self.get_models(modelname, return_dict=True)["oseries"]["name"]
+                    for modelname in names
+                ]
+            )
+            df = self.oseries.loc[onames]
+        else:
+            raise ValueError(
+                "libname must be one of ['oseries', 'stresses', 'models']"
+                f", got '{libname}'"
+            )
+        mask = (
+            (df["x"] <= xmax)
+            & (df["x"] >= xmin)
+            & (df["y"] >= ymin)
+            & (df["y"] <= ymax)
+        )
+        return df.loc[mask].index.tolist()

pastastore/styling.py

@@ -0,0 +1,67 @@
+"""Module containing dataframe styling functions."""
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def float_styler(val, norm, cmap=None):
+    """Style float values in DataFrame.
+
+    Parameters
+    ----------
+    val : float
+        value in cell
+    norm : matplotlib.colors.Normalize
+        normalizer to map values to range(0, 1)
+    cmap : colormap, optional
+        colormap to use, by default None, which uses RdYlBu
+
+    Returns
+    -------
+    str
+        css value pairs for styling dataframe
+
+    Usage
+    -----
+    Given some dataframe
+
+    >>> df.map(float_styler, subset=["some column"], norm=norm, cmap=cmap)
+    """
+    if cmap is None:
+        cmap = plt.get_cmap("RdYlBu")
+    bg = cmap(norm(val))
+    color = mpl.colors.rgb2hex(bg)
+    c = "White" if np.mean(bg[:3]) < 0.4 else "Black"
+    return f"background-color: {color}; color: {c}"
+
+
+def boolean_styler(b):
+    """Style boolean values in DataFrame.
+
+    Parameters
+    ----------
+    b : bool
+        value in cell
+
+    Returns
+    -------
+    str
+        css value pairs for styling dataframe
+
+    Usage
+    -----
+    Given some dataframe
+
+    >>> df.map(boolean_styler, subset=["some column"])
+    """
+    if b:
+        return (
+            f"background-color: {mpl.colors.rgb2hex((231/255, 255/255, 239/255))}; "
+            "color: darkgreen"
+        )
+    else:
+        return (
+            f"background-color: {mpl.colors.rgb2hex((255/255, 238/255, 238/255))}; "
+            "color: darkred"
+        )