junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/stats.py

@@ -4,37 +4,63 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from enum import StrEnum
+from typing import Any
 
 import numpy as np
+import structlog
+from pydantic import validate_call
 from scipy.stats import mode, trim_mean
 from scipy.stats.mstats import winsorize
 
-from .utils import logger, raise_error
+from .utils import raise_error
 
 
-__all__ = ["count", "get_aggfunc_by_name", "select", "winsorized_mean"]
+__all__ = [
+    "AggFunc",
+    "count",
+    "get_aggfunc_by_name",
+    "select",
+    "winsorized_mean",
+]
 
+_log = structlog.get_logger("junifer")
+logger = _log.bind(pkg="stats")
 
+
+class AggFunc(StrEnum):
+    """Accepted aggregation function names.
+
+    * ``mean`` -> :func:`numpy.mean`
+    * ``winsorized_mean`` -> :func:`scipy.stats.mstats.winsorize`
+    * ``trim_mean`` -> :func:`scipy.stats.trim_mean`
+    * ``mode`` -> :func:`scipy.stats.mode`
+    * ``std`` -> :func:`numpy.std`
+    * ``count`` -> :func:`.count`
+    * ``select`` -> :func:`.select`
+
+    """
+
+    Mean = "mean"
+    WinsorizedMean = "winsorized_mean"
+    TrimMean = "trim_mean"
+    Mode = "mode"
+    Std = "std"
+    Count = "count"
+    Select = "select"
+
+
+@validate_call
 def get_aggfunc_by_name(
-    name: str, func_params: Optional[dict[str, Any]] = None
+    name: AggFunc, func_params: dict[str, Any] | None = None
 ) -> Callable:
     """Get an aggregation function by its name.
 
     Parameters
     ----------
-    name : str
-        Name to identify the function. Currently supported names and
-        corresponding functions are:
-
-        * ``mean`` -> :func:`numpy.mean`
-        * ``winsorized_mean`` -> :func:`scipy.stats.mstats.winsorize`
-        * ``trim_mean`` -> :func:`scipy.stats.trim_mean`
-        * ``mode`` -> :func:`scipy.stats.mode`
-        * ``std`` -> :func:`numpy.std`
-        * ``count`` -> :func:`.count`
-        * ``select`` -> :func:`.select`
-
+    name : :enum:`.AggFunc`
+        Aggregation function name.
     func_params : dict, optional
         Parameters to pass to the function.
         E.g. for ``winsorized_mean``: ``func_params = {'limits': [0.1, 0.1]}``
@@ -48,16 +74,6 @@ def get_aggfunc_by_name(
     """
     from functools import partial  # local import to avoid sphinx error
 
-    # check validity of names
-    _valid_func_names = {
-        "winsorized_mean",
-        "mean",
-        "std",
-        "trim_mean",
-        "count",
-        "select",
-        "mode",
-    }
     if func_params is None:
         func_params = {}
     # apply functions
@@ -101,11 +117,7 @@ def get_aggfunc_by_name(
         func = partial(select, **func_params)
     elif name == "mode":
         func = partial(mode, **func_params)
-    else:
-        raise_error(
-            f"Function {name} unknown. Please provide any of "
-            f"{_valid_func_names}"
-        )
+
     return func
 
 
@@ -134,7 +146,7 @@ def count(data: np.ndarray, axis: int = 0) -> np.ndarray:
 
 
 def winsorized_mean(
-    data: np.ndarray, axis: Optional[int] = None, **win_params
+    data: np.ndarray, axis: int | None = None, **win_params
 ) -> np.ndarray:
     """Compute a winsorized mean by chaining winsorization and mean.
 
@@ -169,8 +181,8 @@ def winsorized_mean(
 def select(
     data: np.ndarray,
     axis: int = 0,
-    pick: Optional[list[int]] = None,
-    drop: Optional[list[int]] = None,
+    pick: list[int] | None = None,
+    drop: list[int] | None = None,
 ) -> np.ndarray:
     """Select a subset of the data.
 

junifer/storage/__init__.pyi

@@ -1,11 +1,15 @@
 __all__ = [
     "BaseFeatureStorage",
+    "MatrixKind",
+    "StorageType",
+    "HDF5FeatureStorage",
     "PandasBaseFeatureStorage",
     "SQLiteFeatureStorage",
-    "HDF5FeatureStorage",
+    "Upsert",
+    "logger",
 ]
 
-from .base import BaseFeatureStorage
-from .pandas_base import PandasBaseFeatureStorage
-from .sqlite import SQLiteFeatureStorage
+from .base import BaseFeatureStorage, MatrixKind, StorageType, logger
 from .hdf5 import HDF5FeatureStorage
+from .pandas_base import PandasBaseFeatureStorage
+from .sqlite import SQLiteFeatureStorage, Upsert

junifer/storage/base.py

@@ -6,20 +6,44 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
+from enum import StrEnum
 from pathlib import Path
-from typing import Any, ClassVar, Optional, Union
+from typing import Any, ClassVar
 
 import numpy as np
 import pandas as pd
+import structlog
+from pydantic import BaseModel, ConfigDict
 
-from ..utils import logger, raise_error
-from .utils import process_meta
+from ..utils import raise_error
 
 
-__all__ = ["BaseFeatureStorage"]
+__all__ = ["BaseFeatureStorage", "MatrixKind", "StorageType"]
 
 
-class BaseFeatureStorage(ABC):
+class MatrixKind(StrEnum):
+    """Accepted matrix kind value."""
+
+    UpperTriangle = "triu"
+    LowerTriangle = "tril"
+    Full = "full"
+
+
+_log = structlog.get_logger("junifer")
+logger = _log.bind(pkg="storage", step="storage")
+
+
+class StorageType(StrEnum):
+    """Accepted storage type."""
+
+    Vector = "vector"
+    Matrix = "matrix"
+    Timeseries = "timeseries"
+    Timeseries2D = "timeseries_2d"
+    ScalarTable = "scalar_table"
+
+
+class BaseFeatureStorage(BaseModel, ABC):
     """Abstract base class for feature storage.
 
     For every storage, one needs to provide a concrete
@@ -27,7 +51,7 @@ class BaseFeatureStorage(ABC):
 
     Parameters
     ----------
-    uri : str or pathlib.Path
+    uri : pathlib.Path
         The path to the storage.
     single_output : bool, optional
         Whether to have single output (default True).
@@ -35,27 +59,26 @@ class BaseFeatureStorage(ABC):
     Raises
     ------
     AttributeError
-        If the storage does not have `_STORAGE_TYPES` attribute.
+        If the storage does not have ``_STORAGE_TYPES`` attribute.
 
     """
 
-    _STORAGE_TYPES: ClassVar[Sequence[str]]
+    _STORAGE_TYPES: ClassVar[Sequence[StorageType]]
 
-    def __init__(
-        self,
-        uri: Union[str, Path],
-        single_output: bool = True,
-    ) -> None:
+    model_config = ConfigDict(
+        extra="forbid", frozen=True, use_enum_values=True
+    )
+
+    uri: Path
+    single_output: bool = True
+
+    def model_post_init(self, context: Any):  # noqa: D102
         # Check for missing storage types attribute
         if not hasattr(self, "_STORAGE_TYPES"):
             raise_error(
                 msg="Missing `_STORAGE_TYPES` for the storage",
                 klass=AttributeError,
             )
-        # Convert str to Path
-        if not isinstance(uri, Path):
-            uri = Path(uri)
-        self.uri = uri
         # Create parent directories if not present
         if not self.uri.parent.exists():
             logger.info(
@@ -63,21 +86,12 @@ class BaseFeatureStorage(ABC):
                 "does not exist, creating now"
             )
             self.uri.parent.mkdir(parents=True, exist_ok=True)
-        self.single_output = single_output
-
-    def get_valid_inputs(self) -> list[str]:
-        """Get valid storage types for input.
-
-        Returns
-        -------
-        list of str
-            The list of storage types that can be used as input for this
-            storage.
-
-        """
-        return list(self._STORAGE_TYPES)
+        logger.info(
+            f"Parameters: {self.model_dump(mode='json')}",
+            component=self.__class__.__name__,
+        )
 
-    def validate(self, input_: list[str]) -> None:
+    def validate_input(self, input_: list[str]) -> None:
         """Validate the input to the pipeline step.
 
         Parameters
@@ -119,11 +133,9 @@ class BaseFeatureStorage(ABC):
     @abstractmethod
     def read(
         self,
-        feature_name: Optional[str] = None,
-        feature_md5: Optional[str] = None,
-    ) -> dict[
-        str, Union[str, list[Union[int, str, dict[str, str]]], np.ndarray]
-    ]:
+        feature_name: str | None = None,
+        feature_md5: str | None = None,
+    ) -> dict[str, str | list[int | str | dict[str, str]] | np.ndarray]:
         """Read stored feature.
 
         Parameters
@@ -147,8 +159,8 @@ class BaseFeatureStorage(ABC):
     @abstractmethod
     def read_df(
         self,
-        feature_name: Optional[str] = None,
-        feature_md5: Optional[str] = None,
+        feature_name: str | None = None,
+        feature_md5: str | None = None,
     ) -> pd.DataFrame:
         """Read feature into a pandas DataFrame.
 
@@ -171,7 +183,13 @@ class BaseFeatureStorage(ABC):
         )  # pragma: no cover
 
     @abstractmethod
-    def store_metadata(self, meta_md5: str, element: dict, meta: dict) -> None:
+    def store_metadata(
+        self,
+        meta_md5: str,
+        element: dict,
+        meta: dict,
+        processed_element: dict | None = None,
+    ) -> None:
         """Store metadata.
 
         Parameters
@@ -182,6 +200,11 @@ class BaseFeatureStorage(ABC):
             The element as a dictionary.
         meta : dict
             The metadata as a dictionary.
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, with all
+            the keys of the element, which can be used to choose where to
+            store the data (e.g., the file of each element being processed).
+            If None, ``element`` is used (default None).
 
         """
         raise_error(
@@ -189,12 +212,12 @@ class BaseFeatureStorage(ABC):
             klass=NotImplementedError,
         )  # pragma: no cover
 
-    def store(self, kind: str, **kwargs) -> None:
+    def store(self, kind: StorageType, **kwargs) -> None:
         """Store extracted features data.
 
         Parameters
         ----------
-        kind : {"matrix", "timeseries", "vector", "scalar_table"}
+        kind : :enum:`.StorageType`
             The storage kind.
         **kwargs
             The keyword arguments.
@@ -212,11 +235,24 @@ class BaseFeatureStorage(ABC):
                 msg=f"I don't know how to store {kind}.",
                 klass=ValueError,
             )
-        # Process and store metadata
+        # Imported here to avoid circular import
+        from .utils import process_meta
+
+        # Process and store metadata: the data is stored for the element with
+        # the keys that it depends on, and the element being processed is
+        # passed as ``processed_element`` (e.g., to choose the file to store
+        # the data of each element being processed)
         t_meta = kwargs.pop("meta")
+        processed_element = t_meta.get("element")
         meta_md5, t_meta, t_element = process_meta(t_meta)
-        self.store_metadata(meta_md5=meta_md5, element=t_element, meta=t_meta)
+        self.store_metadata(
+            meta_md5=meta_md5,
+            element=t_element,
+            meta=t_meta,
+            processed_element=processed_element,
+        )
         # Store data
+        kwargs["processed_element"] = processed_element
         if kind == "matrix":
             self.store_matrix(meta_md5=meta_md5, element=t_element, **kwargs)
         elif kind == "timeseries":
@@ -239,10 +275,11 @@ class BaseFeatureStorage(ABC):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
-        row_names: Optional[Sequence[str]] = None,
-        matrix_kind: str = "full",
+        col_names: Sequence[str] | None = None,
+        row_names: Sequence[str] | None = None,
+        matrix_kind: MatrixKind = MatrixKind.Full,
         diagonal: bool = True,
+        processed_element: dict | None = None,
     ) -> None:
         """Store matrix.
 
@@ -258,17 +295,17 @@ class BaseFeatureStorage(ABC):
             The column labels (default None).
         row_names : list-like of str, optional
             The row labels (default None).
-        matrix_kind : str, optional
-            The kind of matrix:
-
-            * ``triu`` : store upper triangular only
-            * ``tril`` : store lower triangular
-            * ``full`` : full matrix
-
-            (default "full").
+        matrix_kind : :enum:`.MatrixKind`, optional
+            The matrix kind (default ``MatrixKind.Full``).
         diagonal : bool, optional
-            Whether to store the diagonal. If ``matrix_kind = full``, setting
-            this to False will raise an error (default True).
+            Whether to store the diagonal. If ``matrix_kind=MatrixKind.Full``,
+            setting this to False will raise an error (default True).
+
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, with all
+            the keys of the element, which can be used to choose where to
+            store the data (e.g., the file of each element being processed).
+            If None, ``element`` is used (default None).
 
         """
         raise_error(
@@ -280,8 +317,9 @@ class BaseFeatureStorage(ABC):
         self,
         meta_md5: str,
         element: dict,
-        data: Union[np.ndarray, list],
-        col_names: Optional[Sequence[str]] = None,
+        data: np.ndarray | list,
+        col_names: Sequence[str] | None = None,
+        processed_element: dict | None = None,
     ) -> None:
         """Store vector.
 
@@ -296,6 +334,12 @@ class BaseFeatureStorage(ABC):
         col_names : list-like of str, optional
             The column labels (default None).
 
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, with all
+            the keys of the element, which can be used to choose where to
+            store the data (e.g., the file of each element being processed).
+            If None, ``element`` is used (default None).
+
         """
         raise_error(
             msg="Concrete classes need to implement store_vector().",
@@ -307,7 +351,8 @@ class BaseFeatureStorage(ABC):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
+        col_names: Sequence[str] | None = None,
+        processed_element: dict | None = None,
     ) -> None:
         """Store timeseries.
 
@@ -322,6 +367,12 @@ class BaseFeatureStorage(ABC):
         col_names : list-like of str, optional
             The column labels (default None).
 
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, with all
+            the keys of the element, which can be used to choose where to
+            store the data (e.g., the file of each element being processed).
+            If None, ``element`` is used (default None).
+
         """
         raise_error(
             msg="Concrete classes need to implement store_timeseries().",
@@ -333,8 +384,9 @@ class BaseFeatureStorage(ABC):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
-        row_names: Optional[Sequence[str]] = None,
+        col_names: Sequence[str] | None = None,
+        row_names: Sequence[str] | None = None,
+        processed_element: dict | None = None,
     ) -> None:
         """Store 2D timeseries.
 
@@ -351,6 +403,12 @@ class BaseFeatureStorage(ABC):
         row_names : list-like of str, optional
             The row labels (default None).
 
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, with all
+            the keys of the element, which can be used to choose where to
+            store the data (e.g., the file of each element being processed).
+            If None, ``element`` is used (default None).
+
         """
         raise_error(
             msg="Concrete classes need to implement store_timeseries_2d().",
@@ -362,9 +420,10 @@ class BaseFeatureStorage(ABC):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
-        row_names: Optional[Sequence[str]] = None,
-        row_header_col_name: Optional[str] = "feature",
+        col_names: Sequence[str] | None = None,
+        row_names: Sequence[str] | None = None,
+        row_header_col_name: str | None = "feature",
+        processed_element: dict | None = None,
     ) -> None:
         """Store table with scalar values.
 
@@ -383,6 +442,12 @@ class BaseFeatureStorage(ABC):
         row_header_col_name : str, optional
             The column name for the row header column (default "feature").
 
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, with all
+            the keys of the element, which can be used to choose where to
+            store the data (e.g., the file of each element being processed).
+            If None, ``element`` is used (default None).
+
         """
         raise_error(
             msg="Concrete classes need to implement store_scalar_table().",