pastastore 1.4.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

pastastore/store.py

@@ -1,7 +1,10 @@
+"""Module containing the PastaStore object for managing time series and models."""
+
 import json
+import logging
 import os
 import warnings
-from typing import List, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -14,11 +17,14 @@ 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, PASTAS_LEQ_022
 from pastastore.yaml_interface import PastastoreYAML
 
 FrameorSeriesUnion = Union[pd.DataFrame, pd.Series]
 warnings.showwarning = _custom_warning
 
+logger = logging.getLogger(__name__)
+
 
 class PastaStore:
     """PastaStore object for managing pastas time series and models.
@@ -27,9 +33,8 @@ class PastaStore:
     the database. Different Connectors are available, e.g.:
 
     - PasConnector for storing all data as .pas (JSON) files on disk (recommended)
-    - DictConenctor for storing all data in dictionaries (in-memory)
-    - ArcticConnector for saving data to MongoDB using the Arctic module
-    - PystoreConnector for saving data to disk using the Pystore module
+    - ArcticDBConnector for saving data on disk using arcticdb package
+    - DictConnector for storing all data in dictionaries (in-memory)
 
     Parameters
     ----------
@@ -40,6 +45,8 @@ class PastaStore:
         name of the PastaStore, by default takes the name of the Connector object
     """
 
+    _accessors = set()
+
     def __init__(
         self,
         connector: Optional[BaseConnector] = None,
@@ -72,7 +79,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 +90,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 +162,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 +222,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 +261,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 +277,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 +299,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 +358,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 +379,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 +400,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
         ----------
@@ -384,7 +468,12 @@ class PastaStore:
 
         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 +492,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 +558,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 +643,6 @@ class PastaStore:
         -------
         s : pandas.DataFrame
         """
-
         modelnames = self.conn._parse_names(modelnames, libname="models")
 
         # if statistics is str
@@ -556,8 +672,9 @@ class PastaStore:
     def create_model(
         self,
         name: str,
-        modelname: str = None,
+        modelname: Optional[str] = None,
         add_recharge: bool = True,
+        add_ar_noisemodel: bool = False,
         recharge_name: str = "recharge",
     ) -> ps.Model:
         """Create a pastas Model.
@@ -572,6 +689,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_ar_noisemodel : bool, optional
+            add AR(1) noise model to the model, by default False
         recharge_name : str
             name of the RechargeModel
 
@@ -591,13 +710,15 @@ class PastaStore:
         meta = self.conn.get_metadata("oseries", name, as_frame=False)
         ts = self.conn.get_oseries(name)
 
-        # convert to Timeseries and create model
+        # convert to time series and create model
         if not ts.dropna().empty:
             if modelname is None:
                 modelname = name
             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!")
@@ -694,50 +815,363 @@ class PastaStore:
         recharge_name : str
             name of the RechargeModel
         """
-        # get nearest prec and evap stns
-        if "prec" not in self.stresses.kind.values:
-            raise ValueError(
-                "No stresses with kind='prec' found in store. "
-                "add_recharge() requires stresses with kind='prec'!"
-            )
-        if "evap" not in self.stresses.kind.values:
-            raise ValueError(
-                "No stresses with kind='evap' found in store. "
-                "add_recharge() requires stresses with kind='evap'!"
-            )
-        names = []
-        for var in ("prec", "evap"):
-            try:
-                name = self.get_nearest_stresses(ml.oseries.name, kind=var).iloc[0, 0]
-            except AttributeError:
-                msg = "No precipitation or evaporation time series found!"
-                raise Exception(msg)
-            if isinstance(name, float):
-                if np.isnan(name):
+        if recharge is None:
+            recharge = ps.rch.Linear()
+        if rfunc is None:
+            rfunc = ps.Exponential
+
+        self.add_stressmodel(
+            ml,
+            stresses={"prec": "nearest", "evap": "nearest"},
+            rfunc=rfunc,
+            stressmodel=ps.RechargeModel,
+            stressmodel_name=recharge_name,
+            recharge=recharge,
+        )
+
+    def _parse_stresses(
+        self,
+        stresses: Union[str, List[str], Dict[str, str]],
+        kind: Optional[str],
+        stressmodel,
+        oseries: Optional[str] = None,
+    ):
+        # parse stresses for RechargeModel, allow list of len 2 or 3 and
+        # set correct kwarg names
+        if stressmodel._name == "RechargeModel":
+            if isinstance(stresses, list):
+                if len(stresses) == 2:
+                    stresses = {
+                        "prec": stresses[0],
+                        "evap": stresses[1],
+                    }
+                elif len(stresses) == 3:
+                    stresses = {
+                        "prec": stresses[0],
+                        "evap": stresses[1],
+                        "temp": stresses[2],
+                    }
+                else:
                     raise ValueError(
-                        f"Unable to find nearest '{var}' stress! "
-                        "Check x and y coordinates."
+                        "RechargeModel requires 2 or 3 stress names, "
+                        f"got: {len(stresses)}!"
                     )
+        # if stresses is list, create dictionary normally
+        elif isinstance(stresses, list):
+            stresses = {"stress": stresses}
+        # if stresses is str, make it a list of len 1
+        elif isinstance(stresses, str):
+            stresses = {"stress": [stresses]}
+
+        # check if stresses is a dictionary, else raise TypeError
+        if not isinstance(stresses, dict):
+            raise TypeError("stresses must be a list, string or dictionary!")
+
+        # if no kind specified, set to well for WellModel
+        if stressmodel._name == "WellModel":
+            if kind is None:
+                kind = "well"
+
+        # store a copy of the user input for kind
+        if isinstance(kind, list):
+            _kind = kind.copy()
+        else:
+            _kind = kind
+
+        # create empty list for gathering metadata
+        metadata = []
+        # loop over stresses keys/values
+        for i, (k, v) in enumerate(stresses.items()):
+            # if entry in dictionary is str, make it list of len 1
+            if isinstance(v, str):
+                v = [v]
+            # parse value
+            if isinstance(v, list):
+                for item in v:
+                    names = []  # empty list for names
+                    # parse nearest
+                    if item.startswith("nearest"):
+                        # check oseries defined if nearest option is used
+                        if not oseries:
+                            raise ValueError(
+                                "Getting nearest stress(es) requires oseries name!"
+                            )
+                        try:
+                            if len(item.split()) == 3:  # nearest <n> <kind>
+                                n = int(item.split()[1])
+                                kind = item.split()[2]
+                            elif len(item.split()) == 2:  # nearest <n> | <kind>
+                                try:
+                                    n = int(item.split()[1])  # try converting to <n>
+                                except ValueError:
+                                    n = 1
+                                    kind = item.split()[1]  # interpret as <kind>
+                            else:  # nearest
+                                n = 1
+                                # if RechargeModel, we can infer kind
+                                if (
+                                    _kind is None
+                                    and stressmodel._name == "RechargeModel"
+                                ):
+                                    kind = k
+                                elif _kind is None:  # catch no kind with bare nearest
+                                    raise ValueError(
+                                        "Bare 'nearest' found but no kind specified."
+                                    )
+                                elif isinstance(_kind, list):
+                                    kind = _kind[i]  # if multiple kind, select i-th
+                        except Exception as e:
+                            # raise if nearest parsing failed
+                            raise ValueError(
+                                f"Could not parse stresses: '{item}'! "
+                                "When using option 'nearest', use 'nearest' and specify"
+                                " kind, or 'nearest <kind>' or 'nearest <n> <kind>'!"
+                            ) from e
+                        # check if kind exists at all
+                        if kind not in self.stresses.kind.values:
+                            raise ValueError(
+                                f"Could not find stresses with kind='{kind}'!"
+                            )
+                        # get stress names of <n> nearest <kind> stresses
+                        inames = self.get_nearest_stresses(
+                            oseries, kind=kind, n=n
+                        ).iloc[0]
+                        # check if any NaNs in result
+                        if inames.isna().any():
+                            nkind = (self.stresses.kind == kind).sum()
+                            raise ValueError(
+                                f"Could not find {n} nearest stress(es) for '{kind}'! "
+                                f"There are only {nkind} '{kind}' stresses."
+                            )
+                        # append names
+                        names += inames.tolist()
+                    else:
+                        # assume name is name of stress
+                        names.append(item)
+                # get stresses and metadata
+                stress_series, imeta = self.get_stresses(
+                    names, return_metadata=True, squeeze=True
+                )
+                # replace stress name(s) with time series
+                if len(names) > 1:
+                    stresses[k] = list(stress_series.values())
+                else:
+                    stresses[k] = stress_series
+                # gather metadata
+                if isinstance(imeta, list):
+                    metadata += imeta
+                else:
+                    metadata.append(imeta)
+
+        return stresses, metadata
+
+    def get_stressmodel(
+        self,
+        stresses: Union[str, List[str], Dict[str, str]],
+        stressmodel=ps.StressModel,
+        stressmodel_name: Optional[str] = None,
+        rfunc=ps.Exponential,
+        rfunc_kwargs: Optional[dict] = None,
+        kind: Optional[Union[List[str], str]] = None,
+        oseries: Optional[str] = None,
+        **kwargs,
+    ):
+        """Get a Pastas stressmodel from stresses time series in Pastastore.
+
+        Supports "nearest" selection. Any stress name can be replaced by
+        "nearest [<n>] <kind>" where <n> is optional and represents the number of
+        nearest stresses and <kind> and represents the kind of stress to
+        consider. <kind> can also be specified directly with the `kind` kwarg.
+
+        Note: the 'nearest' option requires the oseries name to be provided.
+        Additionally, 'x' and 'y' metadata must be stored for oseries and stresses.
+
+        Parameters
+        ----------
+        stresses : str, list of str, or dict
+            name(s) of the time series to use for the stressmodel, or dictionary
+            with key(s) and value(s) as time series name(s). Options include:
+               - name of stress: `"prec_stn"`
+               - list of stress names: `["prec_stn", "evap_stn"]`
+               - dict for RechargeModel: `{"prec": "prec_stn", "evap": "evap_stn"}`
+               - dict for StressModel: `{"stress": "well1"}`
+               - nearest, specifying kind: `"nearest well"`
+               - nearest specifying number and kind: `"nearest 2 well"`
+        stressmodel : str or class
+            stressmodel class to use, by default ps.StressModel
+        stressmodel_name : str, optional
+            name of the stressmodel, by default None, which uses the stress name,
+            if there is 1 stress otherwise the name of the stressmodel type. For
+            RechargeModels, the name defaults to 'recharge'.
+        rfunc : str or class
+            response function class to use, by default ps.Exponential
+        rfunc_kwargs : dict, optional
+            keyword arguments to pass to the response function, by default None
+        kind : str or list of str, optional
+            specify kind of stress(es) to use, by default None, useful in combination
+            with 'nearest' option for defining stresses
+        oseries : str, optional
+            name of the oseries to use for the stressmodel, by default None, used when
+            'nearest' option is used for defining stresses.
+        **kwargs
+            additional keyword arguments to pass to the stressmodel
+
+        Returns
+        -------
+        stressmodel : pastas.StressModel
+            pastas StressModel that can be added to pastas Model.
+        """
+        # get stressmodel class, if str was provided
+        if isinstance(stressmodel, str):
+            stressmodel = getattr(ps, stressmodel)
+
+        # parse stresses names to get time series and metadata
+        stresses, metadata = self._parse_stresses(
+            stresses=stresses, stressmodel=stressmodel, kind=kind, oseries=oseries
+        )
+
+        # get stressmodel name if not provided
+        if stressmodel_name is None:
+            if stressmodel._name == "RechargeModel":
+                stressmodel_name = "recharge"
+            elif len(metadata) == 1:
+                stressmodel_name = stresses["stress"].squeeze().name
+            else:
+                stressmodel_name = stressmodel._name
+
+        # check if metadata is list of len 1 and unpack
+        if isinstance(metadata, list) and len(metadata) == 1:
+            metadata = metadata[0]
+
+        # get stressmodel time series settings
+        if kind and "settings" not in kwargs:
+            # try using kind to get predefined settings options
+            if isinstance(kind, str):
+                kwargs["settings"] = ps.rcParams["timeseries"].get(kind, None)
             else:
-                names.append(name)
-        if len(names) == 0:
-            msg = "No precipitation or evaporation time series found!"
-            raise Exception(msg)
-
-        # get data
-        tsdict = self.conn.get_stresses(names)
-        metadata = self.conn.get_metadata("stresses", names, as_frame=False)
-        # add recharge to model
-        rch = ps.RechargeModel(
-            tsdict[names[0]],
-            tsdict[names[1]],
+                kwargs["settings"] = [
+                    ps.rcParams["timeseries"].get(ikind, None) for ikind in kind
+                ]
+        elif kind is None and "settings" not in kwargs:
+            # try using kind stored in metadata to get predefined settings options
+            if isinstance(metadata, list):
+                kwargs["settings"] = [
+                    ps.rcParams["timeseries"].get(imeta.get("kind", None), None)
+                    for imeta in metadata
+                ]
+            elif isinstance(metadata, dict):
+                kwargs["settings"] = ps.rcParams["timeseries"].get(
+                    metadata.get("kind", None), None
+                )
+
+        # get rfunc class if str was provided
+        if isinstance(rfunc, str):
+            rfunc = getattr(ps, rfunc)
+
+        # create empty rfunc_kwargs if not provided
+        if rfunc_kwargs is None:
+            rfunc_kwargs = {}
+
+        # special for WellModels
+        if stressmodel._name == "WellModel":
+            names = [s.squeeze().name for s in stresses["stress"]]
+            # check oseries is provided
+            if oseries is None:
+                raise ValueError("WellModel requires 'oseries' to compute distances!")
+            # compute distances and add to kwargs
+            distances = (
+                self.get_distances(oseries=oseries, stresses=names).T.squeeze().values
+            )
+            kwargs["distances"] = distances
+            # set settings to well
+            if "settings" not in kwargs:
+                kwargs["settings"] = "well"
+            # override rfunc and set to HantushWellModel
+            rfunc = ps.HantushWellModel
+
+        # do not add metadata for pastas 0.22 and WellModel
+        if not PASTAS_LEQ_022 and (stressmodel._name != "WellModel"):
+            kwargs["metadata"] = metadata
+
+        return stressmodel(
+            **stresses,
+            rfunc=rfunc(**rfunc_kwargs),
+            name=stressmodel_name,
+            **kwargs,
+        )
+
+    def add_stressmodel(
+        self,
+        ml: Union[ps.Model, str],
+        stresses: Union[str, List[str], Dict[str, str]],
+        stressmodel=ps.StressModel,
+        stressmodel_name: Optional[str] = None,
+        rfunc=ps.Exponential,
+        rfunc_kwargs: Optional[dict] = None,
+        kind: Optional[Union[List[str], str]] = None,
+        **kwargs,
+    ):
+        """Add a pastas StressModel from stresses time series in Pastastore.
+
+        Supports "nearest" selection. Any stress name can be replaced by
+        "nearest [<n>] <kind>" where <n> is optional and represents the number of
+        nearest stresses and <kind> and represents the kind of stress to
+        consider. <kind> can also be specified directly with the `kind` kwarg.
+
+        Note: the 'nearest' option requires the oseries name to be provided.
+        Additionally, 'x' and 'y' metadata must be stored for oseries and stresses.
+
+        Parameters
+        ----------
+        ml : pastas.Model or str
+            pastas.Model object to add StressModel to, if passed as string,
+            model is loaded from store, the stressmodel is added and then written
+            back to the store.
+        stresses : str, list of str, or dict
+            name(s) of the time series to use for the stressmodel, or dictionary
+            with key(s) and value(s) as time series name(s). Options include:
+               - name of stress: `"prec_stn"`
+               - list of stress names: `["prec_stn", "evap_stn"]`
+               - dict for RechargeModel: `{"prec": "prec_stn", "evap": "evap_stn"}`
+               - dict for StressModel: `{"stress": "well1"}`
+               - nearest, specifying kind: `"nearest well"`
+               - nearest specifying number and kind: `"nearest 2 well"`
+        stressmodel : str or class
+            stressmodel class to use, by default ps.StressModel
+        stressmodel_name : str, optional
+            name of the stressmodel, by default None, which uses the stress name,
+            if there is 1 stress otherwise the name of the stressmodel type. For
+            RechargeModels, the name defaults to 'recharge'.
+        rfunc : str or class
+            response function class to use, by default ps.Exponential
+        rfunc_kwargs : dict, optional
+            keyword arguments to pass to the response function, by default None
+        kind : str or list of str, optional
+            specify kind of stress(es) to use, by default None, useful in combination
+            with 'nearest' option for defining stresses
+        **kwargs
+            additional keyword arguments to pass to the stressmodel
+        """
+        sm = self.get_stressmodel(
+            stresses=stresses,
+            stressmodel=stressmodel,
+            stressmodel_name=stressmodel_name,
             rfunc=rfunc,
-            name=recharge_name,
-            recharge=recharge,
-            settings=("prec", "evap"),
-            metadata=metadata,
+            rfunc_kwargs=rfunc_kwargs,
+            kind=kind,
+            oseries=ml if isinstance(ml, str) else ml.oseries.name,
+            **kwargs,
         )
-        ml.add_stressmodel(rch)
+        if isinstance(ml, str):
+            ml = self.get_model(ml)
+            ml.add_stressmodel(sm)
+            self.conn.add_model(ml, overwrite=True)
+            logger.info(
+                f"Stressmodel '{sm.name}' added to model '{ml.name}' "
+                "and stored in database."
+            )
+        else:
+            ml.add_stressmodel(sm)
 
     def solve_models(
         self,
@@ -828,8 +1262,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 +1304,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
@@ -1011,13 +1445,12 @@ class PastaStore:
         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'")
 
@@ -1064,7 +1497,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
@@ -1138,6 +1570,23 @@ class PastaStore:
         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":