junifer 0.0.5.dev240__py3-none-any.whl → 0.0.6__py3-none-any.whl

junifer/storage/base.py

@@ -5,8 +5,9 @@
 # License: AGPL
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -43,7 +44,7 @@ class BaseFeatureStorage(ABC):
     def __init__(
         self,
         uri: Union[str, Path],
-        storage_types: Union[List[str], str],
+        storage_types: Union[list[str], str],
         single_output: bool = True,
     ) -> None:
         self.uri = uri
@@ -61,7 +62,7 @@ class BaseFeatureStorage(ABC):
         self._valid_inputs = storage_types
         self.single_output = single_output
 
-    def get_valid_inputs(self) -> List[str]:
+    def get_valid_inputs(self) -> list[str]:
         """Get valid storage types for input.
 
         Returns
@@ -76,7 +77,7 @@ class BaseFeatureStorage(ABC):
             klass=NotImplementedError,
         )
 
-    def validate(self, input_: List[str]) -> None:
+    def validate(self, input_: list[str]) -> None:
         """Validate the input to the pipeline step.
 
         Parameters
@@ -98,7 +99,7 @@ class BaseFeatureStorage(ABC):
             )
 
     @abstractmethod
-    def list_features(self) -> Dict[str, Dict[str, Any]]:
+    def list_features(self) -> dict[str, dict[str, Any]]:
         """List the features in the storage.
 
         Returns
@@ -119,8 +120,8 @@ class BaseFeatureStorage(ABC):
         self,
         feature_name: Optional[str] = None,
         feature_md5: Optional[str] = None,
-    ) -> Dict[
-        str, Union[str, List[Union[int, str, Dict[str, str]]], np.ndarray]
+    ) -> dict[
+        str, Union[str, list[Union[int, str, dict[str, str]]], np.ndarray]
     ]:
         """Read stored feature.
 
@@ -169,7 +170,7 @@ class BaseFeatureStorage(ABC):
         )
 
     @abstractmethod
-    def store_metadata(self, meta_md5: str, element: Dict, meta: Dict) -> None:
+    def store_metadata(self, meta_md5: str, element: dict, meta: dict) -> None:
         """Store metadata.
 
         Parameters
@@ -225,11 +226,15 @@ class BaseFeatureStorage(ABC):
             self.store_scalar_table(
                 meta_md5=meta_md5, element=t_element, **kwargs
             )
+        elif kind == "timeseries_2d":
+            self.store_timeseries_2d(
+                meta_md5=meta_md5, element=t_element, **kwargs
+            )
 
     def store_matrix(
         self,
         meta_md5: str,
-        element: Dict,
+        element: dict,
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
         row_names: Optional[Iterable[str]] = None,
@@ -271,8 +276,8 @@ class BaseFeatureStorage(ABC):
     def store_vector(
         self,
         meta_md5: str,
-        element: Dict,
-        data: Union[np.ndarray, List],
+        element: dict,
+        data: Union[np.ndarray, list],
         col_names: Optional[Iterable[str]] = None,
     ) -> None:
         """Store vector.
@@ -297,7 +302,7 @@ class BaseFeatureStorage(ABC):
     def store_timeseries(
         self,
         meta_md5: str,
-        element: Dict,
+        element: dict,
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
     ) -> None:
@@ -320,10 +325,39 @@ class BaseFeatureStorage(ABC):
             klass=NotImplementedError,
         )
 
+    def store_timeseries_2d(
+        self,
+        meta_md5: str,
+        element: dict,
+        data: np.ndarray,
+        col_names: Optional[Iterable[str]] = None,
+        row_names: Optional[Iterable[str]] = None,
+    ) -> None:
+        """Store 2D timeseries.
+
+        Parameters
+        ----------
+        meta_md5 : str
+            The metadata MD5 hash.
+        element : dict
+            The element as a dictionary.
+        data : numpy.ndarray
+            The timeseries data to store.
+        col_names : list or tuple of str, optional
+            The column labels (default None).
+        row_names : list or tuple of str, optional
+            The row labels (default None).
+
+        """
+        raise_error(
+            msg="Concrete classes need to implement store_timeseries_2d().",
+            klass=NotImplementedError,
+        )
+
     def store_scalar_table(
         self,
         meta_md5: str,
-        element: Dict,
+        element: dict,
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
         row_names: Optional[Iterable[str]] = None,

junifer/storage/hdf5.py

@@ -4,10 +4,10 @@
 #          Federico Raimondo <f.raimondo@fz-juelich.de>
 # License: AGPL
 
-
 from collections import defaultdict
+from collections.abc import Iterable
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -23,14 +23,20 @@ from ..external.h5io.h5io import (
 )
 from ..utils import logger, raise_error
 from .base import BaseFeatureStorage
-from .utils import element_to_prefix, matrix_to_vector, store_matrix_checks
+from .utils import (
+    element_to_prefix,
+    matrix_to_vector,
+    store_matrix_checks,
+    store_timeseries_2d_checks,
+    timeseries2d_to_vector,
+)
 
 
 __all__ = ["HDF5FeatureStorage"]
 
 
 def _create_chunk(
-    chunk_data: List[np.ndarray],
+    chunk_data: list[np.ndarray],
     kind: str,
     element_count: int,
     chunk_size: int,
@@ -81,7 +87,7 @@ def _create_chunk(
             chunk_size=tuple(array_chunk_size),
             n_chunk=i_chunk,
         )
-    elif kind in ["timeseries", "scalar_table"]:
+    elif kind in ["timeseries", "scalar_table", "timeseries_2d"]:
         out = ChunkedList(
             data=chunk_data,
             size=element_count,
@@ -90,8 +96,8 @@ def _create_chunk(
     else:
         raise_error(
             f"Invalid kind: {kind}. "
-            "Must be one of ['vector', 'matrix', 'timeseries',"
-            "'scalar_table']."
+            "Must be one of ['vector', 'matrix', 'timeseries', "
+            "'timeseries_2d', 'scalar_table']."
         )
     return out
 
@@ -164,7 +170,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         self.force_float32 = force_float32
         self.chunk_size = chunk_size
 
-    def get_valid_inputs(self) -> List[str]:
+    def get_valid_inputs(self) -> list[str]:
         """Get valid storage types for input.
 
         Returns
@@ -176,7 +182,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         """
         return ["matrix", "vector", "timeseries", "scalar_table"]
 
-    def _fetch_correct_uri_for_io(self, element: Optional[Dict]) -> str:
+    def _fetch_correct_uri_for_io(self, element: Optional[dict]) -> str:
         """Return proper URI for I/O based on `element`.
 
         If `element` is None, will return `self.uri`.
@@ -195,8 +201,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         if not self.single_output and not element:
             raise_error(
                 msg=(
-                    "`element` must be provided when `single_output` "
-                    "is False"
+                    "`element` must be provided when `single_output` is False"
                 ),
                 klass=RuntimeError,
             )
@@ -210,8 +215,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         return f"{self.uri.parent}/{prefix}{self.uri.name}"  # type: ignore
 
     def _read_metadata(
-        self, element: Optional[Dict[str, str]] = None
-    ) -> Dict[str, Dict[str, Any]]:
+        self, element: Optional[dict[str, str]] = None
+    ) -> dict[str, dict[str, Any]]:
         """Read metadata (should not be called directly).
 
         Parameters
@@ -261,7 +266,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         return metadata
 
-    def list_features(self) -> Dict[str, Dict[str, Any]]:
+    def list_features(self) -> dict[str, dict[str, Any]]:
         """List the features in the storage.
 
         Returns
@@ -281,8 +286,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         return metadata
 
     def _read_data(
-        self, md5: str, element: Optional[Dict[str, str]] = None
-    ) -> Dict[str, Any]:
+        self, md5: str, element: Optional[dict[str, str]] = None
+    ) -> dict[str, Any]:
         """Read data (should not be called directly).
 
         Parameters
@@ -338,8 +343,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         self,
         feature_name: Optional[str] = None,
         feature_md5: Optional[str] = None,
-    ) -> Dict[
-        str, Union[str, List[Union[int, str, Dict[str, str]]], np.ndarray]
+    ) -> dict[
+        str, Union[str, list[Union[int, str, dict[str, str]]], np.ndarray]
     ]:
         """Read stored feature.
 
@@ -513,6 +518,27 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             columns = hdf_data["column_headers"]
             # Convert data from 3D to 2D
             reshaped_data = np.concatenate(all_data, axis=0)
+        elif hdf_data["kind"] == "timeseries_2d":
+            # Create dictionary for aggregating index data
+            element_idx = defaultdict(list)
+            all_data = []
+            for idx, element in enumerate(hdf_data["element"]):
+                # Get row count for the element
+                t_data = hdf_data["data"][idx]
+                flat_data, columns = timeseries2d_to_vector(
+                    data=t_data,
+                    col_names=hdf_data["column_headers"],
+                    row_names=hdf_data["row_headers"],
+                )
+                all_data.append(flat_data)
+                n_timepoints = flat_data.shape[0]
+                # Set rows for the index
+                for key, val in element.items():
+                    element_idx[key].extend([val] * n_timepoints)
+                # Add extra column for timepoints
+                element_idx["timepoint"].extend(np.arange(n_timepoints))
+            # Convert data from 3D to 2D
+            reshaped_data = np.concatenate(all_data, axis=0)
         elif hdf_data["kind"] == "scalar_table":
             # Create dictionary for aggregating index data
             element_idx = defaultdict(list)
@@ -562,7 +588,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         return df
 
     def _write_processed_data(
-        self, fname: str, processed_data: Dict[str, Any], title: str
+        self, fname: str, processed_data: dict[str, Any], title: str
     ) -> None:
         """Write processed data to HDF5 (should not be called directly).
 
@@ -594,8 +620,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     def store_metadata(
         self,
         meta_md5: str,
-        element: Dict[str, str],
-        meta: Dict[str, Any],
+        element: dict[str, str],
+        meta: dict[str, Any],
     ) -> None:
         """Store metadata.
 
@@ -655,7 +681,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         self,
         kind: str,
         meta_md5: str,
-        element: List[Dict[str, str]],
+        element: list[dict[str, str]],
         data: np.ndarray,
         **kwargs: Any,
     ) -> None:
@@ -764,7 +790,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             )
 
             t_data = stored_data["data"]
-            if kind in ["timeseries", "scalar_table"]:
+            if kind in ["timeseries", "scalar_table", "timeseries_2d"]:
                 t_data += data
             else:
                 t_data = np.concatenate((t_data, data), axis=-1)
@@ -797,7 +823,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     def store_matrix(
         self,
         meta_md5: str,
-        element: Dict[str, str],
+        element: dict[str, str],
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
         row_names: Optional[Iterable[str]] = None,
@@ -876,8 +902,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     def store_vector(
         self,
         meta_md5: str,
-        element: Dict[str, str],
-        data: Union[np.ndarray, List],
+        element: dict[str, str],
+        data: Union[np.ndarray, list],
         col_names: Optional[Iterable[str]] = None,
     ) -> None:
         """Store vector.
@@ -919,7 +945,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     def store_timeseries(
         self,
         meta_md5: str,
-        element: Dict[str, str],
+        element: dict[str, str],
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
     ) -> None:
@@ -946,10 +972,48 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             row_header_column_name="timepoint",
         )
 
+    def store_timeseries_2d(
+        self,
+        meta_md5: str,
+        element: dict[str, str],
+        data: np.ndarray,
+        col_names: Optional[Iterable[str]] = None,
+        row_names: Optional[Iterable[str]] = None,
+    ) -> None:
+        """Store a 2D timeseries.
+
+        Parameters
+        ----------
+        meta_md5 : str
+            The metadata MD5 hash.
+        element : dict
+            The element as dictionary.
+        data : numpy.ndarray
+            The 2D timeseries data to store.
+        col_names : list or tuple of str, optional
+            The column labels (default None).
+        row_names : list or tuple of str, optional
+            The row labels (default None).
+
+        """
+        store_timeseries_2d_checks(
+            data_shape=data.shape,
+            row_names_len=len(row_names),  # type: ignore
+            col_names_len=len(col_names),  # type: ignore
+        )
+        self._store_data(
+            kind="timeseries_2d",
+            meta_md5=meta_md5,
+            element=[element],  # convert to list
+            data=[data],  # convert to list
+            column_headers=col_names,
+            row_headers=row_names,
+        )
+
     def store_scalar_table(
         self,
         meta_md5: str,
-        element: Dict,
+        element: dict,
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
         row_names: Optional[Iterable[str]] = None,
@@ -1013,8 +1077,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         # Run loop to collect metadata
         logger.info(
-            "Collecting metadata from "
-            f"{self.uri.parent}/*_{self.uri.name}"  # type: ignore
+            f"Collecting metadata from {self.uri.parent}/*_{self.uri.name}"  # type: ignore
         )
         # Collect element files per feature MD5
         elements_per_feature_md5 = defaultdict(list)
@@ -1045,8 +1108,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         # Run loop to collect data per feature per file
         logger.info(
-            "Collecting data from "
-            f"{self.uri.parent}/*_{self.uri.name}"  # type: ignore
+            f"Collecting data from {self.uri.parent}/*_{self.uri.name}"  # type: ignore
         )
         logger.info(f"Will collect {len(elements_per_feature_md5)} features.")
 
@@ -1091,7 +1153,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
                     kind = static_data["kind"]
 
                 # Append the "dynamic" data
-                if kind in ["timeseries", "scalar_table"]:
+                if kind in ["timeseries", "scalar_table", "timeseries_2d"]:
                     chunk_data.extend(t_data["data"])
                 else:
                     chunk_data.append(t_data["data"])

junifer/storage/pandas_base.py

@@ -5,8 +5,9 @@
 # License: AGPL
 
 import json
+from collections.abc import Iterable
 from pathlib import Path
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -44,7 +45,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
     ) -> None:
         super().__init__(uri=uri, single_output=single_output, **kwargs)
 
-    def get_valid_inputs(self) -> List[str]:
+    def get_valid_inputs(self) -> list[str]:
         """Get valid storage types for input.
 
         Returns
@@ -56,7 +57,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
         """
         return ["matrix", "vector", "timeseries"]
 
-    def _meta_row(self, meta: Dict, meta_md5: str) -> pd.DataFrame:
+    def _meta_row(self, meta: dict, meta_md5: str) -> pd.DataFrame:
         """Convert the metadata to a pandas DataFrame.
 
         Parameters
@@ -80,7 +81,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
 
     @staticmethod
     def element_to_index(
-        element: Dict, n_rows: int = 1, rows_col_name: Optional[str] = None
+        element: dict, n_rows: int = 1, rows_col_name: Optional[str] = None
     ) -> Union[pd.Index, pd.MultiIndex]:
         """Convert the element metadata to index.
 
@@ -101,7 +102,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
 
         """
         # Make mapping between element access keys and values
-        elem_idx: Dict[str, Iterable[str]] = {
+        elem_idx: dict[str, Iterable[str]] = {
             k: [v] * n_rows for k, v in element.items()
         }
 
@@ -129,7 +130,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
         return index
 
     def store_df(
-        self, meta_md5: str, element: Dict, df: Union[pd.DataFrame, pd.Series]
+        self, meta_md5: str, element: dict, df: Union[pd.DataFrame, pd.Series]
     ) -> None:
         """Implement pandas DataFrame storing.
 
@@ -157,8 +158,8 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
     def _store_2d(
         self,
         meta_md5: str,
-        element: Dict,
-        data: Union[np.ndarray, List],
+        element: dict,
+        data: Union[np.ndarray, list],
         col_names: Optional[Iterable[str]] = None,
         rows_col_name: Optional[str] = None,
     ) -> None:
@@ -194,8 +195,8 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
     def store_vector(
         self,
         meta_md5: str,
-        element: Dict,
-        data: Union[np.ndarray, List],
+        element: dict,
+        data: Union[np.ndarray, list],
         col_names: Optional[Iterable[str]] = None,
     ) -> None:
         """Store vector.
@@ -232,7 +233,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
     def store_timeseries(
         self,
         meta_md5: str,
-        element: Dict,
+        element: dict,
         data: np.ndarray,
         col_names: Optional[Iterable[str]] = None,
     ) -> None:

junifer/storage/sqlite.py

@@ -6,7 +6,7 @@
 
 import json
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -92,7 +92,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         # Set upsert
         self._upsert = upsert
 
-    def get_engine(self, element: Optional[Dict] = None) -> "Engine":
+    def get_engine(self, element: Optional[dict] = None) -> "Engine":
         """Get engine.
 
         Parameters
@@ -209,7 +209,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
                         msg=f"Invalid option {if_exists} for if_exists."
                     )
 
-    def list_features(self) -> Dict[str, Dict[str, Any]]:
+    def list_features(self) -> dict[str, dict[str, Any]]:
         """List the features in the storage.
 
         Returns
@@ -229,7 +229,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         # Format index names for retrieved data
         meta_df.index = meta_df.index.str.replace(r"meta_", "")
         # Convert dataframe to dictionary
-        out: Dict[str, Dict[str, str]] = meta_df.to_dict(
+        out: dict[str, dict[str, str]] = meta_df.to_dict(
             orient="index"
         )  # type: ignore
         # Format output
@@ -242,8 +242,8 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         self,
         feature_name: Optional[str] = None,
         feature_md5: Optional[str] = None,
-    ) -> Dict[
-        str, Union[str, List[Union[int, str, Dict[str, str]]], np.ndarray]
+    ) -> dict[
+        str, Union[str, list[Union[int, str, dict[str, str]]], np.ndarray]
     ]:
         """Read stored feature.
 
@@ -358,7 +358,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         df = df.set_index(index_names)
         return df
 
-    def store_metadata(self, meta_md5: str, element: Dict, meta: Dict) -> None:
+    def store_metadata(self, meta_md5: str, element: dict, meta: dict) -> None:
         """Implement metadata storing in the storage.
 
         Parameters
@@ -381,7 +381,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             self._save_upsert(meta_df, "meta", engine)
 
     def store_df(
-        self, meta_md5: str, element: Dict, df: Union[pd.DataFrame, pd.Series]
+        self, meta_md5: str, element: dict, df: Union[pd.DataFrame, pd.Series]
     ) -> None:
         """Implement pandas DataFrame storing.
 
@@ -434,10 +434,10 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
     def store_matrix(
         self,
         meta_md5: str,
-        element: Dict,
+        element: dict,
         data: np.ndarray,
-        col_names: Optional[List[str]] = None,
-        row_names: Optional[List[str]] = None,
+        col_names: Optional[list[str]] = None,
+        row_names: Optional[list[str]] = None,
         matrix_kind: str = "full",
         diagonal: bool = True,
     ) -> None: