junifer/external/h5io

@@ -1 +1 @@
-Subproject commit 4413a9fc3db2ca8fd0982f86e28b326eee6e0b28
+Subproject commit 600d9e25c625b11520d86b08f9f0b761f71b0542

junifer/storage/hdf5.py

@@ -6,7 +6,6 @@
 
 
 from collections import defaultdict
-from functools import reduce
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Optional, Union
 
@@ -15,12 +14,82 @@ import pandas as pd
 from tqdm import tqdm
 
 from ..api.decorators import register_storage
-from ..external.h5io.h5io import ChunkedArray, has_hdf5, read_hdf5, write_hdf5
+from ..external.h5io.h5io import (
+    ChunkedArray,
+    ChunkedList,
+    has_hdf5,
+    read_hdf5,
+    write_hdf5,
+)
 from ..utils import logger, raise_error
 from .base import BaseFeatureStorage
 from .utils import element_to_prefix, matrix_to_vector, store_matrix_checks
 
 
+def _create_chunk(
+    chunk_data: List[np.ndarray],
+    kind: str,
+    element_count: int,
+    chunk_size: int,
+    i_chunk: int,
+) -> Union[ChunkedArray, ChunkedList]:
+    """Create chunked array or list.
+
+    Parameters
+    ----------
+    chunk_data : list of numpy.ndarray
+        The data to be chunked.
+    kind : str
+        The kind of data to be chunked.
+    element_count : int
+        The total number of elements.
+    chunk_size : int
+        The chunk size.
+    i_chunk : int
+        The chunk index.
+
+    Returns
+    -------
+    ChunkedArray or ChunkedList
+        The chunked array or list.
+
+    Raises
+    ------
+    ValueError
+        If `kind` is not one of ['vector', 'matrix', 'timeseries'].
+    """
+    if kind in ["vector", "matrix"]:
+        features_data = np.concatenate(chunk_data, axis=-1)
+        array_shape = [features_data.shape[0]]
+        array_chunk_size = [features_data.shape[0]]
+        # Append second dimension for 3D
+        if features_data.ndim == 3:
+            array_shape.append(features_data.shape[1])
+            array_chunk_size.append(features_data.shape[1])
+        # Append final dimension of element count
+        array_shape.append(element_count)
+        # Append final dimension of chunk size
+        array_chunk_size.append(chunk_size)
+        out = ChunkedArray(
+            data=features_data,
+            shape=tuple(array_shape),
+            chunk_size=tuple(array_chunk_size),
+            n_chunk=i_chunk,
+        )
+    elif kind == "timeseries":
+        out = ChunkedList(
+            data=chunk_data,
+            size=element_count,
+            offset=i_chunk * chunk_size,
+        )
+    else:
+        raise_error(
+            f"Invalid kind: {kind}. "
+            "Must be one of ['vector', 'matrix', 'timeseries']."
+        )
+    return out
+
+
 @register_storage
 class HDF5FeatureStorage(BaseFeatureStorage):
     """Concrete implementation for feature storage via HDF5.
@@ -176,13 +245,13 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             )
 
         # Read metadata
-        logger.info(f"Loading HDF5 metadata from: {uri}")
+        logger.debug(f"Loading HDF5 metadata from: {uri}")
         metadata = read_hdf5(
             fname=uri,
             title="meta",
             slash="ignore",
         )
-        logger.info(f"Loaded HDF5 metadata from: {uri}")
+        logger.debug(f"Loaded HDF5 metadata from: {uri}")
 
         return metadata
 
@@ -249,13 +318,13 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             )
 
         # Read data
-        logger.info(f"Loading HDF5 data for {md5} from: {uri}")
+        logger.debug(f"Loading HDF5 data for {md5} from: {uri}")
         data = read_hdf5(
             fname=uri,
             title=md5,
             slash="ignore",
         )
-        logger.info(f"Loaded HDF5 data for {md5} from: {uri}")
+        logger.debug(f"Loaded HDF5 data for {md5} from: {uri}")
 
         return data
 
@@ -386,9 +455,12 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         elif hdf_data["kind"] == "timeseries":
             # 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
-                n_rows, _ = hdf_data["data"][:, :, idx].shape
+                t_data = hdf_data["data"][idx]
+                all_data.append(t_data)
+                n_rows, _ = t_data.shape
                 # Set rows for the index
                 for key, val in element.items():
                     element_idx[key].extend([val] * n_rows)
@@ -399,7 +471,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             # Set column headers for dataframe
             columns = hdf_data["column_headers"]
             # Convert data from 3D to 2D
-            reshaped_data = hdf_data["data"].reshape(-1, 1)
+            reshaped_data = np.concatenate(all_data, axis=0)
 
         # Create dataframe for index
         idx_df = pd.DataFrame(data=element_idx)  # type: ignore
@@ -413,7 +485,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             data=reshaped_data,
             index=hdf_data_idx,
             columns=columns,  # type: ignore
-            dtype=hdf_data["data"].dtype,
+            dtype=reshaped_data.dtype,
         )
         logger.debug(f"Converted HDF5 data for {md5} to pandas.DataFrame ...")
 
@@ -556,9 +628,19 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         data_to_write = kwargs
 
         # Optional casting of float64 values to float32 for numpy.ndarray
+        if isinstance(data, np.ndarray):
             if data.dtype == np.dtype("float64") and self.force_float32:
-            data = data.astype(dtype=np.dtype("float32"), casting="same_kind")
-
+                data = data.astype(
+                    dtype=np.dtype("float32"), casting="same_kind"
+                )
+        elif isinstance(data, list):
+            if self.force_float32:
+                data = [
+                    x.astype(dtype=np.dtype("float32"), casting="same_kind")
+                    if x.dtype == np.dtype("float64")
+                    else x
+                    for x in data
+                ]
         # Handle cases for existing and new entry
         if not stored_data:
             logger.debug(f"Writing new data for {meta_md5} ...")
@@ -608,14 +690,18 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             logger.debug(
                 f"Existing data found for {meta_md5}, appending to it ..."
             )
+
+            t_data = stored_data["data"]
+            if kind == "timeseries":
+                t_data.append(data)
+            else:
+                t_data = np.concatenate((t_data, data), axis=-1)
             # Existing entry; append to existing
             # "element" and "data"
             data_to_write.update(
                 {
                     "element": stored_data["element"] + element,
-                    "data": np.concatenate(
-                        (stored_data["data"], data), axis=-1
-                    ),
+                    "data": t_data,
                     # for serialization / deserialization of storage type
                     "kind": kind,
                 }
@@ -783,7 +869,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             kind="timeseries",
             meta_md5=meta_md5,
             element=[element],  # convert to list
-            data=data[:, :, np.newaxis],  # convert to 3D
+            data=[data],  # convert to list
             column_headers=col_names,
             row_header_column_name="timepoint",
         )
@@ -809,8 +895,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             )
 
         # Glob files
-        globbed_files = self.uri.parent.glob(  # type: ignore
-            f"*{self.uri.name}"  # type: ignore
+        globbed_files = list(
+            self.uri.parent.glob(f"*_{self.uri.name}")  # type: ignore
         )
 
         # Create new storage instance
@@ -819,10 +905,11 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         # Run loop to collect metadata
         logger.info(
             "Collecting metadata from "
-            f"{self.uri.parent}/*{self.uri.name}"  # type: ignore
+            f"{self.uri.parent}/*_{self.uri.name}"  # type: ignore
         )
         # Collect element files per feature MD5
         elements_per_feature_md5 = defaultdict(list)
+        out_metadata = {}
         for file_ in tqdm(globbed_files, desc="file-metadata"):
             logger.debug(f"Reading HDF5 file: {file_} ...")
             # Create new storage instance to load metadata
@@ -830,105 +917,101 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             # Load metadata from new instance
             in_metadata = in_storage._read_metadata()
 
-            logger.info(f"Updating HDF5 metadata with metadata from: {file_}")
-            # Load metadata; empty dictionary if first entry;
-            # can be replaced with store_metadata() if run on a loop
-            # for the metadata entries from in_storage
-            try:
-                out_metadata = out_storage._read_metadata()
-            except IOError:
-                out_metadata = {}
-            # Update metadata
+            logger.debug(f"Updating HDF5 metadata with metadata from: {file_}")
+
+            # Update metadata to store
             out_metadata.update(in_metadata)
-            # Save metadata
+
+            # Update element files for found MD5s
+            for feature_md5 in in_metadata.keys():
+                elements_per_feature_md5[feature_md5].append(file_)
+
+        logger.info("Writing metadata to HDF5 file ...")
+        # Save metadata out metadata
         out_storage._write_processed_data(
             fname=str(self.uri.resolve()),  # type: ignore
             processed_data=out_metadata,
             title="meta",
         )
-            # Update element files for found MD5s
-            for feature_md5 in in_metadata.keys():
-                elements_per_feature_md5[feature_md5].append(file_)
 
         # Run loop to collect data per feature per file
         logger.info(
             "Collecting data from "
-            f"{self.uri.parent}/*{self.uri.name}"  # type: ignore
+            f"{self.uri.parent}/*_{self.uri.name}"  # type: ignore
         )
+        logger.info(f"Will collect {len(elements_per_feature_md5)} features.")
         for feature_md5, element_files in tqdm(
             elements_per_feature_md5.items(), desc="feature"
         ):
             element_count = len(element_files)
-            # Chunk size for collecting
-            chunk_size = min(self.chunk_size, element_count)
-            # Operate on chunks
-            for chunk_idx, chunk_start in tqdm(
-                enumerate(range(0, element_count, chunk_size)), desc="chunk"
-            ):
-                # Store the chunk files' data
-                stored_data_for_chunk: List[Dict[str, Any]] = []
-                # Read the files of a chunk
-                for i in tqdm(
-                    range(chunk_start, chunk_start + chunk_size),
-                    desc="file-data",
-                ):
-                    file_ = element_files[i]
+
+            i_file = 0
+            i_chunk = 0
+            t_chunk_size = min(self.chunk_size, element_count)
+            chunk_data = []
+            elements = []
+            static_data = None
+            kind = None
+            for file_ in tqdm(element_files, desc="file-data"):
                 logger.debug(
                     f"Reading feature MD5: '{feature_md5}' "
                     f"from HDF5 file: {file_} ..."
                 )
-                    # Read from HDF5 and collect data
-                    stored_data_for_chunk.append(
-                        read_hdf5(
+
+                # Read the data
+                t_data = read_hdf5(
                     fname=str(file_),
                     title=feature_md5,
                     slash="ignore",
                 )
-                    )
-
-                # Concatenate the features data for a chunk
-                features_data = np.concatenate(
-                    [x["data"] for x in stored_data_for_chunk], axis=-1
-                )
-                # Make dictionary to write the collected data;
-                # first the static data then the dynamic data
-                data_to_write = {
-                    key: val
-                    for key, val in stored_data_for_chunk[0].items()
-                    if key not in ("data", "element")
+                if i_file == 0:
+                    # Store the "static" data
+                    static_data = {
+                        k: v
+                        for k, v in t_data.items()
+                        if k not in ["data", "element"]
                     }
-                # Join the features element for a chunk
-                data_to_write["element"] = reduce(
-                    lambda acc, elem: acc + elem,
-                    [x["element"] for x in stored_data_for_chunk],
-                    [],
-                )
+                    kind = static_data["kind"]
+
+                # Append the "dynamic" data
+                if kind == "timeseries":
+                    chunk_data.extend(t_data["data"])
+                else:
+                    chunk_data.append(t_data["data"])
+                elements.extend(t_data["element"])
+
+                i_file += 1
+                if (i_file % t_chunk_size == 0) or i_file == element_count:
+                    # If we have reached the chunk size or the end of the
+                    # elements, write the data
+
+                    # Store one chunk of data
+                    to_write = static_data.copy()
+                    to_write["element"] = []
                     # Write data in chunks to avoid memory usage spikes
                     # Start with the case for 2D
-                array_shape = [features_data.shape[0]]
-                array_chunk_size = [features_data.shape[0]]
-                # Append second dimension for 3D
-                if features_data.ndim == 3:
-                    array_shape.append(features_data.shape[1])
-                    array_chunk_size.append(features_data.shape[1])
-                # Append final dimension of element count
-                array_shape.append(element_count)
-                # Append final dimension of chunk size
-                array_chunk_size.append(chunk_size)
                     # Write chunked array
-                data_to_write["data"] = ChunkedArray(
-                    data=features_data,
-                    shape=tuple(array_shape),
-                    chunk_size=tuple(array_chunk_size),
-                    n_chunk=chunk_idx,
+                    to_write["data"] = _create_chunk(
+                        chunk_data=chunk_data,
+                        kind=kind,
+                        element_count=element_count,
+                        chunk_size=t_chunk_size,
+                        i_chunk=i_chunk,
                     )
+                    if i_file == element_count:
+                        to_write["element"] = elements
+
                     # Write to HDF5
                     write_hdf5(
                         fname=str(self.uri.resolve()),  # type: ignore
-                    data=data_to_write,
-                    overwrite=self.overwrite,  # type: ignore
+                        data=to_write,
+                        overwrite="update",  # type: ignore
                         compression=0,
                         title=feature_md5,
                         slash="error",
                         use_json=False,
                     )
+
+                    # Increment counters and data
+                    i_chunk += 1
+                    chunk_data = []

junifer/storage/tests/test_hdf5.py

@@ -4,7 +4,9 @@
 #          Federico Raimondo <f.raimondo@fz-juelich.de>
 # License: AGPL
 
+from copy import deepcopy
 from pathlib import Path
+from typing import Dict, Tuple
 
 import h5py
 import numpy as np
@@ -13,7 +15,11 @@ from numpy.testing import assert_array_equal
 from pandas.testing import assert_frame_equal
 
 from junifer.storage import HDF5FeatureStorage
-from junifer.storage.utils import element_to_prefix, process_meta
+from junifer.storage.utils import (
+    element_to_prefix,
+    matrix_to_vector,
+    process_meta,
+)
 
 
 def test_get_valid_inputs() -> None:
@@ -785,7 +791,91 @@ def test_store_timeseries(tmp_path: Path) -> None:
     assert_array_equal(read_df.values, data)
 
 
-def test_multi_output_store_and_collect(tmp_path: Path):
+def _create_data_to_store(n_elements: int, kind: str) -> Tuple[str, Dict]:
+    """Create data to store.
+
+    Parameters
+    ----------
+    n_elements : int
+        The number of elements to create.
+    kind : str
+        The kind of data to create.
+
+    Returns
+    -------
+    str
+        The meta md5.
+    dict
+        The data to store.
+    """
+    all_data = []
+    t_md5 = None
+    if kind == "vector":
+        data_to_store = {
+            "data": np.arange(10),
+            "col_names": [f"col-{i}" for i in range(10)],
+        }
+    elif kind == "matrix":
+        data_to_store = {
+            "data": np.arange(100).reshape(10, 10),
+            "row_names": [f"row-{i}" for i in range(10)],
+            "col_names": [f"col-{i}" for i in range(10)],
+            "matrix_kind": "full",
+        }
+    elif kind == "timeseries":
+        data_to_store = {
+            "data": np.arange(20).reshape(2, 10),
+            "col_names": [f"col-{i}" for i in range(10)],
+        }
+    else:
+        raise ValueError(f"Unknown kind {kind}.")
+    for i in range(n_elements):
+        element = {"subject": f"sub-{i // 2}", "session": f"ses-{i % 2}"}
+        meta = {
+            "element": element,
+            "dependencies": ["numpy"],
+            "marker": {"name": f"test-{kind}"},
+            "type": "BOLD",
+        }
+        # Process the metadata
+        meta_md5, meta_to_store, element_to_store = process_meta(meta)
+        if kind == "timeseries":
+            t_data = data_to_store["data"]
+            data_to_store["data"] = np.r_[
+                t_data, (t_data[-1, :] + t_data[-1, :])[None]
+            ]
+        else:
+            data_to_store["data"] = data_to_store["data"] + i
+
+        if t_md5 is None:
+            t_md5 = meta_md5
+        else:
+            assert t_md5 == meta_md5
+
+        # Store metadata
+        all_data.append(
+            {
+                "element": element_to_store,
+                "meta": meta_to_store,
+                "data": deepcopy(data_to_store),
+            }
+        )
+    return t_md5, all_data
+
+
+@pytest.mark.parametrize(
+    "n_elements, chunk_size, kind",
+    [
+        (10, 3, "vector"),
+        (10, 5, "vector"),
+        (10, 3, "matrix"),
+        (10, 5, "matrix"),
+        (10, 5, "timeseries"),
+    ],
+)
+def test_multi_output_store_and_collect(
+    tmp_path: Path, n_elements: int, chunk_size: int, kind: str
+) -> None:
     """Test multi output storing and collection.
 
     Parameters
@@ -795,121 +885,96 @@ def test_multi_output_store_and_collect(tmp_path: Path):
 
     """
     uri = tmp_path / "test_multi_output_store_and_collect.hdf5"
-    storage = HDF5FeatureStorage(uri=uri, single_output=False)
+    storage = HDF5FeatureStorage(
+        uri=uri, single_output=False, chunk_size=chunk_size
+    )
 
-    # Metadata to store
-    meta_1 = {
-        "element": {"subject": "test-01", "session": "ses-01"},
-        "dependencies": ["numpy"],
-        "marker": {"name": "fc"},
-        "type": "BOLD",
-    }
-    meta_2 = {
-        "element": {"subject": "test-02", "session": "ses-01"},
-        "dependencies": ["numpy"],
-        "marker": {"name": "fc"},
-        "type": "BOLD",
-    }
-    meta_3 = {
-        "element": {"subject": "test-01", "session": "ses-02"},
-        "dependencies": ["numpy"],
-        "marker": {"name": "fc"},
-        "type": "BOLD",
-    }
-
-    # Data to store
-    data_1 = np.array([10, 20, 30, 40, 50])
-    data_2 = data_1 * 10
-    data_3 = data_1 * 20
-    col_headers = ["f1", "f2", "f3", "f4", "f5"]
-
-    # Process metadata for storage
-    hash_1, meta_to_store_1, element_to_store_1 = process_meta(meta_1)
-
-    # Process metadata for storage
-    hash_2, meta_to_store_2, element_to_store_2 = process_meta(meta_2)
-
-    # Process metadata for storage
-    hash_3, meta_to_store_3, element_to_store_3 = process_meta(meta_3)
-
-    # Check hash equality as element is not considered for hash
-    assert hash_1 == hash_2
-    assert hash_2 == hash_3
+    meta_md5, all_data = _create_data_to_store(n_elements, kind)
 
+    for t_data in all_data:
         # Store metadata for tables
         storage.store_metadata(
-        meta_md5=hash_1,
-        element=element_to_store_1,
-        meta=meta_to_store_1,
+            meta_md5=meta_md5,
+            element=t_data["element"],
+            meta=t_data["meta"],
         )
-    storage.store_metadata(
-        meta_md5=hash_2,
-        element=element_to_store_2,
-        meta=meta_to_store_2,
-    )
-    storage.store_metadata(
-        meta_md5=hash_3,
-        element=element_to_store_3,
-        meta=meta_to_store_3,
-    )
-
+        if kind == "vector":
             # Store tables
             storage.store_vector(
-        meta_md5=hash_1,
-        element=element_to_store_1,
-        data=data_1,
-        col_names=col_headers,
+                meta_md5=meta_md5,
+                element=t_data["element"],
+                **t_data["data"],
             )
-    storage.store_vector(
-        meta_md5=hash_2,
-        element=element_to_store_2,
-        data=data_2,
-        col_names=col_headers,
+        elif kind == "matrix":
+            # Store tables
+            storage.store_matrix(
+                meta_md5=meta_md5,
+                element=t_data["element"],
+                **t_data["data"],
             )
-    storage.store_vector(
-        meta_md5=hash_3,
-        element=element_to_store_3,
-        data=data_3,
-        col_names=col_headers,
+        elif kind == "timeseries":
+            storage.store_timeseries(
+                meta_md5=meta_md5,
+                element=t_data["element"],
+                **t_data["data"],
             )
-
     # Check that base URI does not exist yet
     assert not uri.exists()
 
+    for t_data in all_data:
         # Convert element to preifx
-    prefix_1 = element_to_prefix(meta_1["element"])  # type: ignore
-    prefix_2 = element_to_prefix(meta_2["element"])  # type: ignore
-    prefix_3 = element_to_prefix(meta_3["element"])  # type: ignore
-
+        prefix = element_to_prefix(t_data["element"])
         # URIs for data storage
-    uri_1 = uri.parent / f"{prefix_1}{uri.name}"
-    uri_2 = uri.parent / f"{prefix_2}{uri.name}"
-    uri_3 = uri.parent / f"{prefix_3}{uri.name}"
-
+        elem_uri = uri.parent / f"{prefix}{uri.name}"
         # Check URIs for data storage exist
-    assert uri_1.exists()
-    assert uri_2.exists()
-    assert uri_3.exists()
+        assert elem_uri.exists()
 
         # Read stored metadata from different files using element
-    read_meta_1 = storage._read_metadata(element=meta_1["element"])
-    read_meta_2 = storage._read_metadata(element=meta_2["element"])
-    read_meta_3 = storage._read_metadata(element=meta_3["element"])
-
+        read_meta = storage._read_metadata(element=t_data["element"])
         # Check if metadata are equal
-    assert read_meta_1 == read_meta_2
-    assert read_meta_2 == read_meta_3
+        assert read_meta == {meta_md5: t_data["meta"]}
 
     # Collect data
     storage.collect()
     # Check that base URI exists now
     assert uri.exists()
 
-    # Read unified metadata
+    # # Read unified metadata
     read_unified_meta = storage.list_features()
+    assert meta_md5 in read_unified_meta
 
     # Check if aggregated metadata are equal
-    assert read_unified_meta == {**read_meta_1, **read_meta_2, **read_meta_3}
+    assert all(x["meta"] == read_unified_meta[meta_md5] for x in all_data)
+
+    all_df = storage.read_df(feature_md5=meta_md5)
+    if kind == "timeseries":
+        data_size = np.sum([x["data"]["data"].shape[0] for x in all_data])
+        assert len(all_df) == data_size
+        idx_names = [x for x in all_df.index.names if x != "timepoint"]
+    else:
+        assert len(all_df) == len(all_data)
+        idx_names = all_df.index.names
+    for t_data in all_data:
+        t_series = all_df.loc[tuple([t_data["element"][v] for v in idx_names])]
+        if kind == "vector":
+            assert_array_equal(t_series.values, t_data["data"]["data"])
+            series_names = t_series.index.values.tolist()
+            assert series_names == t_data["data"]["col_names"]
+        elif kind == "matrix":
+            flat_data, columns = matrix_to_vector(
+                t_data["data"]["data"],
+                col_names=t_data["data"]["col_names"],
+                row_names=t_data["data"]["row_names"],
+                matrix_kind="full",
+                diagonal=True,
+            )
+            assert_array_equal(t_series.values, flat_data)
+            series_names = t_series.index.values.tolist()
+            assert series_names == columns
+        elif kind == "timeseries":
+            assert_array_equal(t_series.values, t_data["data"]["data"])
+            series_names = t_series.columns.values.tolist()
+            assert series_names == t_data["data"]["col_names"]
 
 
 def test_collect_error_single_output() -> None: