pastastore 1.7.2__py3-none-any.whl → 1.8.0__py3-none-any.whl

pastastore/base.py

@@ -2,21 +2,15 @@
 """Base classes for PastaStore Connectors."""
 
 import functools
-import json
 import warnings
 
 # import weakref
 from abc import ABC, abstractmethod
-from collections.abc import Iterable
 from itertools import chain
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
 import pastas as ps
-from numpy import isin
-from packaging.version import parse as parse_version
-from pandas.testing import assert_series_equal
-from pastas.io.pas import PastasEncoder
 from tqdm.auto import tqdm
 
 from pastastore.util import ItemInLibraryException, _custom_warning, validate_names
@@ -70,7 +64,7 @@ class BaseConnector(ABC):
     def _get_library(self, libname: str):
         """Get library handle.
 
-        Must be overriden by subclass.
+        Must be overridden by subclass.
 
         Parameters
         ----------
@@ -94,7 +88,7 @@ class BaseConnector(ABC):
     ) -> None:
         """Add item for both time series and pastas.Models (internal method).
 
-        Must be overriden by subclass.
+        Must be overridden by subclass.
 
         Parameters
         ----------
@@ -112,7 +106,7 @@ class BaseConnector(ABC):
     def _get_item(self, libname: str, name: str) -> Union[FrameorSeriesUnion, Dict]:
         """Get item (series or pastas.Models) (internal method).
 
-        Must be overriden by subclass.
+        Must be overridden by subclass.
 
         Parameters
         ----------
@@ -131,7 +125,7 @@ class BaseConnector(ABC):
     def _del_item(self, libname: str, name: str) -> None:
         """Delete items (series or models) (internal method).
 
-        Must be overriden by subclass.
+        Must be overridden by subclass.
 
         Parameters
         ----------
@@ -145,7 +139,7 @@ class BaseConnector(ABC):
     def _get_metadata(self, libname: str, name: str) -> Dict:
         """Get metadata (internal method).
 
-        Must be overriden by subclass.
+        Must be overridden by subclass.
 
         Parameters
         ----------
@@ -165,7 +159,7 @@ class BaseConnector(ABC):
     def oseries_names(self):
         """List of oseries names.
 
-        Property must be overriden by subclass.
+        Property must be overridden by subclass.
         """
 
     @property
@@ -173,7 +167,7 @@ class BaseConnector(ABC):
     def stresses_names(self):
         """List of stresses names.
 
-        Property must be overriden by subclass.
+        Property must be overridden by subclass.
         """
 
     @property
@@ -181,7 +175,37 @@ class BaseConnector(ABC):
     def model_names(self):
         """List of model names.
 
-        Property must be overriden by subclass.
+        Property must be overridden by subclass.
+        """
+
+    @abstractmethod
+    def _parallel(
+        self,
+        func: Callable,
+        names: List[str],
+        progressbar: Optional[bool] = True,
+        max_workers: Optional[int] = None,
+        chunksize: Optional[int] = None,
+        desc: str = "",
+    ) -> None:
+        """Parallel processing of function.
+
+        Must be overridden by subclass.
+
+        Parameters
+        ----------
+        func : function
+            function to apply in parallel
+        names : list
+            list of names to apply function to
+        progressbar : bool, optional
+            show progressbar, by default True
+        max_workers : int, optional
+            maximum number of workers, by default None
+        chunksize : int, optional
+            chunksize for parallel processing, by default None
+        desc : str, optional
+            description for progressbar, by default ""
         """
 
     def set_check_model_series_values(self, b: bool):
@@ -1111,8 +1135,8 @@ class BaseConnector(ABC):
             time series contained in library
         """
         names = self._parse_names(names, libname)
-        for nam in names:
-            yield self._get_series(libname, nam, progressbar=False)
+        for name in names:
+            yield self._get_series(libname, name, progressbar=False)
 
     def iter_oseries(self, names: Optional[List[str]] = None):
         """Iterate over oseries in library.
@@ -1338,626 +1362,6 @@ class BaseConnector(ABC):
         return d
 
 
-class ConnectorUtil:
-    """Mix-in class for general Connector helper functions.
-
-    Only for internal methods, and not methods that are related to CRUD operations on
-    database.
-    """
-
-    def _parse_names(
-        self,
-        names: Optional[Union[list, str]] = None,
-        libname: Optional[str] = "oseries",
-    ) -> list:
-        """Parse names kwarg, returns iterable with name(s) (internal method).
-
-        Parameters
-        ----------
-        names : Union[list, str], optional
-            str or list of str or None or 'all' (last two options
-            retrieves all names)
-        libname : str, optional
-            name of library, default is 'oseries'
-
-        Returns
-        -------
-        list
-            list of names
-        """
-        if not isinstance(names, str) and isinstance(names, Iterable):
-            return names
-        elif isinstance(names, str) and names != "all":
-            return [names]
-        elif names is None or names == "all":
-            if libname == "oseries":
-                return self.oseries_names
-            elif libname == "stresses":
-                return self.stresses_names
-            elif libname == "models":
-                return self.model_names
-            elif libname == "oseries_models":
-                return self.oseries_with_models
-            else:
-                raise ValueError(f"No library '{libname}'!")
-        else:
-            raise NotImplementedError(f"Cannot parse 'names': {names}")
-
-    @staticmethod
-    def _meta_list_to_frame(metalist: list, names: list):
-        """Convert list of metadata dictionaries to DataFrame.
-
-        Parameters
-        ----------
-        metalist : list
-            list of metadata dictionaries
-        names : list
-            list of names corresponding to data in metalist
-
-        Returns
-        -------
-        pandas.DataFrame
-            DataFrame containing overview of metadata
-        """
-        # convert to dataframe
-        if len(metalist) > 1:
-            meta = pd.DataFrame(metalist)
-            if len({"x", "y"}.difference(meta.columns)) == 0:
-                meta["x"] = meta["x"].astype(float)
-                meta["y"] = meta["y"].astype(float)
-        elif len(metalist) == 1:
-            meta = pd.DataFrame(metalist)
-        elif len(metalist) == 0:
-            meta = pd.DataFrame()
-
-        meta.index = names
-        meta.index.name = "name"
-        return meta
-
-    def _parse_model_dict(self, mdict: dict, update_ts_settings: bool = False):
-        """Parse dictionary describing pastas models (internal method).
-
-        Parameters
-        ----------
-        mdict : dict
-            dictionary describing pastas.Model
-        update_ts_settings : bool, optional
-            update stored tmin and tmax in time series settings
-            based on time series loaded from store.
-
-        Returns
-        -------
-        ml : pastas.Model
-            time series analysis model
-        """
-        PASFILE_LEQ_022 = parse_version(
-            mdict["file_info"]["pastas_version"]
-        ) <= parse_version("0.22.0")
-
-        # oseries
-        if "series" not in mdict["oseries"]:
-            name = str(mdict["oseries"]["name"])
-            if name not in self.oseries.index:
-                msg = "oseries '{}' not present in library".format(name)
-                raise LookupError(msg)
-            mdict["oseries"]["series"] = self.get_oseries(name).squeeze()
-            # update tmin/tmax from time series
-            if update_ts_settings:
-                mdict["oseries"]["settings"]["tmin"] = mdict["oseries"]["series"].index[
-                    0
-                ]
-                mdict["oseries"]["settings"]["tmax"] = mdict["oseries"]["series"].index[
-                    -1
-                ]
-
-        # StressModel, WellModel
-        for ts in mdict["stressmodels"].values():
-            if "stress" in ts.keys():
-                # WellModel
-                classkey = "stressmodel" if PASFILE_LEQ_022 else "class"
-                if ts[classkey] == "WellModel":
-                    for stress in ts["stress"]:
-                        if "series" not in stress:
-                            name = str(stress["name"])
-                            if name in self.stresses.index:
-                                stress["series"] = self.get_stresses(name).squeeze()
-                                # update tmin/tmax from time series
-                                if update_ts_settings:
-                                    stress["settings"]["tmin"] = stress["series"].index[
-                                        0
-                                    ]
-                                    stress["settings"]["tmax"] = stress["series"].index[
-                                        -1
-                                    ]
-                # StressModel
-                else:
-                    for stress in ts["stress"] if PASFILE_LEQ_022 else [ts["stress"]]:
-                        if "series" not in stress:
-                            name = str(stress["name"])
-                            if name in self.stresses.index:
-                                stress["series"] = self.get_stresses(name).squeeze()
-                                # update tmin/tmax from time series
-                                if update_ts_settings:
-                                    stress["settings"]["tmin"] = stress["series"].index[
-                                        0
-                                    ]
-                                    stress["settings"]["tmax"] = stress["series"].index[
-                                        -1
-                                    ]
-
-            # RechargeModel, TarsoModel
-            if ("prec" in ts.keys()) and ("evap" in ts.keys()):
-                for stress in [ts["prec"], ts["evap"]]:
-                    if "series" not in stress:
-                        name = str(stress["name"])
-                        if name in self.stresses.index:
-                            stress["series"] = self.get_stresses(name).squeeze()
-                            # update tmin/tmax from time series
-                            if update_ts_settings:
-                                stress["settings"]["tmin"] = stress["series"].index[0]
-                                stress["settings"]["tmax"] = stress["series"].index[-1]
-                        else:
-                            msg = "stress '{}' not present in library".format(name)
-                            raise KeyError(msg)
-
-        # hack for pcov w dtype object (when filled with NaNs on store?)
-        if "fit" in mdict:
-            if "pcov" in mdict["fit"]:
-                pcov = mdict["fit"]["pcov"]
-                if pcov.dtypes.apply(lambda dtyp: isinstance(dtyp, object)).any():
-                    mdict["fit"]["pcov"] = pcov.astype(float)
-
-        # check pastas version vs pas-file version
-        file_version = mdict["file_info"]["pastas_version"]
-
-        # check file version and pastas version
-        # if file<0.23  and pastas>=1.0 --> error
-        PASTAS_GT_023 = parse_version(ps.__version__) > parse_version("0.23.1")
-        if PASFILE_LEQ_022 and PASTAS_GT_023:
-            raise UserWarning(
-                f"This file was created with Pastas v{file_version} "
-                f"and cannot be loaded with Pastas v{ps.__version__} Please load and "
-                "save the file with Pastas 0.23 first to update the file "
-                "format."
-            )
-
-        try:
-            # pastas>=0.15.0
-            ml = ps.io.base._load_model(mdict)
-        except AttributeError:
-            # pastas<0.15.0
-            ml = ps.io.base.load_model(mdict)
-        return ml
-
-    @staticmethod
-    def _validate_input_series(series):
-        """Check if series is pandas.DataFrame or pandas.Series.
-
-        Parameters
-        ----------
-        series : object
-            object to validate
-
-        Raises
-        ------
-        TypeError
-            if object is not of type pandas.DataFrame or pandas.Series
-        """
-        if not (isinstance(series, pd.DataFrame) or isinstance(series, pd.Series)):
-            raise TypeError("Please provide pandas.DataFrame or pandas.Series!")
-        if isinstance(series, pd.DataFrame):
-            if series.columns.size > 1:
-                raise ValueError("Only DataFrames with one column are supported!")
-
-    @staticmethod
-    def _set_series_name(series, name):
-        """Set series name to match user defined name in store.
-
-        Parameters
-        ----------
-        series : pandas.Series or pandas.DataFrame
-            set name for this time series
-        name : str
-            name of the time series (used in the pastastore)
-        """
-        if isinstance(series, pd.Series):
-            series.name = name
-            # empty string on index name causes trouble when reading
-            # data from ArcticDB: TODO: check if still an issue?
-            if series.index.name == "":
-                series.index.name = None
-
-        if isinstance(series, pd.DataFrame):
-            series.columns = [name]
-            # check for hydropandas objects which are instances of DataFrame but
-            # do have a name attribute
-            if hasattr(series, "name"):
-                series.name = name
-        return series
-
-    @staticmethod
-    def _check_stressmodels_supported(ml):
-        supported_stressmodels = [
-            "StressModel",
-            "StressModel2",
-            "RechargeModel",
-            "WellModel",
-            "TarsoModel",
-            "Constant",
-            "LinearTrend",
-            "StepModel",
-        ]
-        if isinstance(ml, ps.Model):
-            smtyps = [sm._name for sm in ml.stressmodels.values()]
-        elif isinstance(ml, dict):
-            classkey = "stressmodel" if PASTAS_LEQ_022 else "class"
-            smtyps = [sm[classkey] for sm in ml["stressmodels"].values()]
-        check = isin(smtyps, supported_stressmodels)
-        if not all(check):
-            unsupported = set(smtyps) - set(supported_stressmodels)
-            raise NotImplementedError(
-                "PastaStore does not support storing models with the "
-                f"following stressmodels: {unsupported}"
-            )
-
-    @staticmethod
-    def _check_model_series_names_for_store(ml):
-        prec_evap_model = ["RechargeModel", "TarsoModel"]
-
-        if isinstance(ml, ps.Model):
-            series_names = [
-                istress.series.name
-                for sm in ml.stressmodels.values()
-                for istress in sm.stress
-            ]
-
-        elif isinstance(ml, dict):
-            # non RechargeModel, Tarsomodel, WellModel stressmodels
-            classkey = "stressmodel" if PASTAS_LEQ_022 else "class"
-            if PASTAS_LEQ_022:
-                series_names = [
-                    istress["name"]
-                    for sm in ml["stressmodels"].values()
-                    if sm[classkey] not in (prec_evap_model + ["WellModel"])
-                    for istress in sm["stress"]
-                ]
-            else:
-                series_names = [
-                    sm["stress"]["name"]
-                    for sm in ml["stressmodels"].values()
-                    if sm[classkey] not in (prec_evap_model + ["WellModel"])
-                ]
-
-            # WellModel
-            if isin(
-                ["WellModel"],
-                [i[classkey] for i in ml["stressmodels"].values()],
-            ).any():
-                series_names += [
-                    istress["name"]
-                    for sm in ml["stressmodels"].values()
-                    if sm[classkey] in ["WellModel"]
-                    for istress in sm["stress"]
-                ]
-
-            # RechargeModel, TarsoModel
-            if isin(
-                prec_evap_model,
-                [i[classkey] for i in ml["stressmodels"].values()],
-            ).any():
-                series_names += [
-                    istress["name"]
-                    for sm in ml["stressmodels"].values()
-                    if sm[classkey] in prec_evap_model
-                    for istress in [sm["prec"], sm["evap"]]
-                ]
-
-        else:
-            raise TypeError("Expected pastas.Model or dict!")
-        if len(series_names) - len(set(series_names)) > 0:
-            msg = (
-                "There are multiple stresses series with the same name! "
-                "Each series name must be unique for the PastaStore!"
-            )
-            raise ValueError(msg)
-
-    def _check_oseries_in_store(self, ml: Union[ps.Model, dict]):
-        """Check if Model oseries are contained in PastaStore (internal method).
-
-        Parameters
-        ----------
-        ml : Union[ps.Model, dict]
-            pastas Model
-        """
-        if isinstance(ml, ps.Model):
-            name = ml.oseries.name
-        elif isinstance(ml, dict):
-            name = str(ml["oseries"]["name"])
-        else:
-            raise TypeError("Expected pastas.Model or dict!")
-        if name not in self.oseries.index:
-            msg = (
-                f"Cannot add model because oseries '{name}' "
-                "is not contained in store."
-            )
-            raise LookupError(msg)
-        # expensive check
-        if self.CHECK_MODEL_SERIES_VALUES and isinstance(ml, ps.Model):
-            s_org = self.get_oseries(name).squeeze().dropna()
-            if PASTAS_LEQ_022:
-                so = ml.oseries.series_original
-            else:
-                so = ml.oseries._series_original
-            try:
-                assert_series_equal(
-                    so.dropna(),
-                    s_org,
-                    atol=self.SERIES_EQUALITY_ABSOLUTE_TOLERANCE,
-                    rtol=self.SERIES_EQUALITY_RELATIVE_TOLERANCE,
-                )
-            except AssertionError as e:
-                raise ValueError(
-                    f"Cannot add model because model oseries '{name}'"
-                    " is different from stored oseries! See stacktrace for differences."
-                ) from e
-
-    def _check_stresses_in_store(self, ml: Union[ps.Model, dict]):
-        """Check if stresses time series are contained in PastaStore (internal method).
-
-        Parameters
-        ----------
-        ml : Union[ps.Model, dict]
-            pastas Model
-        """
-        prec_evap_model = ["RechargeModel", "TarsoModel"]
-        if isinstance(ml, ps.Model):
-            for sm in ml.stressmodels.values():
-                if sm._name in prec_evap_model:
-                    stresses = [sm.prec, sm.evap]
-                else:
-                    stresses = sm.stress
-                for s in stresses:
-                    if str(s.name) not in self.stresses.index:
-                        msg = (
-                            f"Cannot add model because stress '{s.name}' "
-                            "is not contained in store."
-                        )
-                        raise LookupError(msg)
-                    if self.CHECK_MODEL_SERIES_VALUES:
-                        s_org = self.get_stresses(s.name).squeeze()
-                        if PASTAS_LEQ_022:
-                            so = s.series_original
-                        else:
-                            so = s._series_original
-                        try:
-                            assert_series_equal(
-                                so,
-                                s_org,
-                                atol=self.SERIES_EQUALITY_ABSOLUTE_TOLERANCE,
-                                rtol=self.SERIES_EQUALITY_RELATIVE_TOLERANCE,
-                            )
-                        except AssertionError as e:
-                            raise ValueError(
-                                f"Cannot add model because model stress "
-                                f"'{s.name}' is different from stored stress! "
-                                "See stacktrace for differences."
-                            ) from e
-        elif isinstance(ml, dict):
-            for sm in ml["stressmodels"].values():
-                classkey = "stressmodel" if PASTAS_LEQ_022 else "class"
-                if sm[classkey] in prec_evap_model:
-                    stresses = [sm["prec"], sm["evap"]]
-                elif sm[classkey] in ["WellModel"]:
-                    stresses = sm["stress"]
-                else:
-                    stresses = sm["stress"] if PASTAS_LEQ_022 else [sm["stress"]]
-                for s in stresses:
-                    if str(s["name"]) not in self.stresses.index:
-                        msg = (
-                            f"Cannot add model because stress '{s['name']}' "
-                            "is not contained in store."
-                        )
-                        raise LookupError(msg)
-        else:
-            raise TypeError("Expected pastas.Model or dict!")
-
-    def _stored_series_to_json(
-        self,
-        libname: str,
-        names: Optional[Union[list, str]] = None,
-        squeeze: bool = True,
-        progressbar: bool = False,
-    ):
-        """Write stored series to JSON.
-
-        Parameters
-        ----------
-        libname : str
-            library name
-        names : Optional[Union[list, str]], optional
-            names of series, by default None
-        squeeze : bool, optional
-            return single entry as json string instead
-            of list, by default True
-        progressbar : bool, optional
-            show progressbar, by default False
-
-        Returns
-        -------
-        files : list or str
-            list of series converted to JSON string or single string
-            if single entry is returned and squeeze is True
-        """
-        names = self._parse_names(names, libname=libname)
-        files = []
-        for n in tqdm(names, desc=libname) if progressbar else names:
-            s = self._get_series(libname, n, progressbar=False)
-            if isinstance(s, pd.Series):
-                s = s.to_frame()
-            try:
-                sjson = s.to_json(orient="columns")
-            except ValueError as e:
-                msg = (
-                    f"DatetimeIndex of '{n}' probably contains NaT "
-                    "or duplicate timestamps!"
-                )
-                raise ValueError(msg) from e
-            files.append(sjson)
-        if len(files) == 1 and squeeze:
-            return files[0]
-        else:
-            return files
-
-    def _stored_metadata_to_json(
-        self,
-        libname: str,
-        names: Optional[Union[list, str]] = None,
-        squeeze: bool = True,
-        progressbar: bool = False,
-    ):
-        """Write metadata from stored series to JSON.
-
-        Parameters
-        ----------
-        libname : str
-            library containing series
-        names : Optional[Union[list, str]], optional
-            names to parse, by default None
-        squeeze : bool, optional
-            return single entry as json string instead of list, by default True
-        progressbar : bool, optional
-            show progressbar, by default False
-
-        Returns
-        -------
-        files : list or str
-            list of json string
-        """
-        names = self._parse_names(names, libname=libname)
-        files = []
-        for n in tqdm(names, desc=libname) if progressbar else names:
-            meta = self.get_metadata(libname, n, as_frame=False)
-            meta_json = json.dumps(meta, cls=PastasEncoder, indent=4)
-            files.append(meta_json)
-        if len(files) == 1 and squeeze:
-            return files[0]
-        else:
-            return files
-
-    def _series_to_archive(
-        self,
-        archive,
-        libname: str,
-        names: Optional[Union[list, str]] = None,
-        progressbar: bool = True,
-    ):
-        """Write DataFrame or Series to zipfile (internal method).
-
-        Parameters
-        ----------
-        archive : zipfile.ZipFile
-            reference to an archive to write data to
-        libname : str
-            name of the library to write to zipfile
-        names : str or list of str, optional
-            names of the time series to write to archive, by default None,
-            which writes all time series to archive
-        progressbar : bool, optional
-            show progressbar, by default True
-        """
-        names = self._parse_names(names, libname=libname)
-        for n in tqdm(names, desc=libname) if progressbar else names:
-            sjson = self._stored_series_to_json(
-                libname, names=n, progressbar=False, squeeze=True
-            )
-            meta_json = self._stored_metadata_to_json(
-                libname, names=n, progressbar=False, squeeze=True
-            )
-            archive.writestr(f"{libname}/{n}.json", sjson)
-            archive.writestr(f"{libname}/{n}_meta.json", meta_json)
-
-    def _models_to_archive(self, archive, names=None, progressbar=True):
-        """Write pastas.Model to zipfile (internal method).
-
-        Parameters
-        ----------
-        archive : zipfile.ZipFile
-            reference to an archive to write data to
-        names : str or list of str, optional
-            names of the models to write to archive, by default None,
-            which writes all models to archive
-        progressbar : bool, optional
-            show progressbar, by default True
-        """
-        names = self._parse_names(names, libname="models")
-        for n in tqdm(names, desc="models") if progressbar else names:
-            m = self.get_models(n, return_dict=True)
-            jsondict = json.dumps(m, cls=PastasEncoder, indent=4)
-            archive.writestr(f"models/{n}.pas", jsondict)
-
-    @staticmethod
-    def _series_from_json(fjson: str, squeeze: bool = True):
-        """Load time series from JSON.
-
-        Parameters
-        ----------
-        fjson : str
-            path to file
-        squeeze : bool, optional
-            squeeze time series object to obtain pandas Series
-
-        Returns
-        -------
-        s : pd.DataFrame
-            DataFrame containing time series
-        """
-        s = pd.read_json(fjson, orient="columns", precise_float=True, dtype=False)
-        if not isinstance(s.index, pd.DatetimeIndex):
-            s.index = pd.to_datetime(s.index, unit="ms")
-        s = s.sort_index()  # needed for some reason ...
-        if squeeze:
-            return s.squeeze()
-        return s
-
-    @staticmethod
-    def _metadata_from_json(fjson: str):
-        """Load metadata dictionary from JSON.
-
-        Parameters
-        ----------
-        fjson : str
-            path to file
-
-        Returns
-        -------
-        meta : dict
-            dictionary containing metadata
-        """
-        with open(fjson, "r") as f:
-            meta = json.load(f)
-        return meta
-
-    def _get_model_orphans(self):
-        """Get models whose oseries no longer exist in database.
-
-        Returns
-        -------
-        dict
-            dictionary with oseries names as keys and lists of model names
-            as values
-        """
-        d = {}
-        for mlnam in tqdm(self.model_names, desc="Identifying model orphans"):
-            mdict = self.get_models(mlnam, return_dict=True)
-            onam = mdict["oseries"]["name"]
-            if onam not in self.oseries_names:
-                if onam in d:
-                    d[onam] = d[onam].append(mlnam)
-                else:
-                    d[onam] = [mlnam]
-        return d
-
-
 class ModelAccessor:
     """Object for managing access to stored models.