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

junifer/storage/hdf5.py

@@ -4,25 +4,22 @@
 #          Federico Raimondo <f.raimondo@fz-juelich.de>
 # License: AGPL
 
+import io
 from collections import defaultdict
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, ClassVar, Optional, Union
+from typing import Any, ClassVar, Literal
 
+import h5py
 import numpy as np
 import pandas as pd
+from h5io import read_hdf5, write_hdf5
+from pydantic import PositiveInt
 from tqdm import tqdm
 
 from ..api.decorators import register_storage
-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 raise_error
+from .base import BaseFeatureStorage, MatrixKind, StorageType, logger
 from .utils import (
     element_to_prefix,
     matrix_to_vector,
@@ -35,21 +32,135 @@ from .utils import (
 __all__ = ["HDF5FeatureStorage"]
 
 
-def _create_chunk(
+# Types of the data of collected features in files written with junifer's
+# fork of h5io, and the h5io types they are stored as
+_LEGACY_TYPES = {"chunkedarray": "ndarray", "chunkedlist": "list"}
+
+
+def _node_type(node: h5py.Group | h5py.Dataset) -> str:
+    """Get the h5io type of a node.
+
+    Parameters
+    ----------
+    node : h5py.Group or h5py.Dataset
+        The node.
+
+    Returns
+    -------
+    str
+        The type.
+
+    """
+    title = node.attrs["TITLE"]
+    return title.decode() if isinstance(title, bytes) else title
+
+
+def _has_hdf5(fname: str, title: str) -> bool:
+    """Check whether a title is in an HDF5 file.
+
+    Parameters
+    ----------
+    fname : str
+        The HDF5 file.
+    title : str
+        The title.
+
+    Returns
+    -------
+    bool
+        Whether the title is in the file.
+
+    """
+    with h5py.File(fname, mode="r") as fid:
+        return title in fid
+
+
+def _read_hdf5(fname: str, title: str) -> Any:
+    """Read a title of an HDF5 file.
+
+    Features collected with junifer's fork of h5io have their data stored with
+    types that h5io does not know (see ``_LEGACY_TYPES``). To read them, the
+    feature is copied in memory with the h5io types.
+
+    Parameters
+    ----------
+    fname : str
+        The HDF5 file.
+    title : str
+        The title.
+
+    Returns
+    -------
+    Any
+        The data.
+
+    """
+    with h5py.File(fname, mode="r") as fid:
+        node = fid.get(title)
+        data = node.get("key_data") if isinstance(node, h5py.Group) else None
+        if data is not None and _node_type(data) in _LEGACY_TYPES:
+            with h5py.File(io.BytesIO(), mode="w") as copy:
+                fid.copy(node, copy, name=title)
+                copy_data = copy[title]["key_data"]
+                copy_data.attrs["TITLE"] = _LEGACY_TYPES[_node_type(data)]
+                return read_hdf5(copy, title=title, slash="ignore")
+    return read_hdf5(fname, title=title, slash="ignore")
+
+
+def _read_elements(fname: str, feature_md5s: list[str]) -> dict[str, list]:
+    """Get the elements that a file has for each feature, without the data.
+
+    Parameters
+    ----------
+    fname : str
+        The HDF5 file.
+    feature_md5s : list of str
+        The MD5 of the feature to check.
+
+    Returns
+    -------
+    dict
+        The elements present in the file for each feature (MD5 -> list of
+        elements). Each element is a tuple of its values, e.g.,
+        ``("sub-01",)``, as all the elements of a feature have the same keys.
+
+    """
+    out = {}
+    with h5py.File(fname, mode="r") as fid:
+        for feature_md5 in feature_md5s:
+            # Copy only the elements of the feature, not its data
+            with h5py.File(io.BytesIO(), mode="w") as copy:
+                fid.copy(fid[feature_md5]["key_element"], copy, name="element")
+                elements = read_hdf5(copy, title="element", slash="ignore")
+            out[feature_md5] = [tuple(e.values()) for e in elements]
+    return out
+
+
+def _write_chunk(
+    fname: str,
+    key: str,
     chunk_data: list[np.ndarray],
-    kind: str,
+    kind: StorageType,
     element_count: int,
     chunk_size: int,
     i_chunk: int,
-) -> Union[ChunkedArray, ChunkedList]:
-    """Create chunked array or list.
+) -> None:
+    """Write a chunk of the data of the elements.
+
+    The data is written in a dataset (``ndarray`` in h5io) for vectors and
+    matrices, with the elements in the last dimension, or in a group of
+    datasets (``list`` in h5io) for the other kinds.
 
     Parameters
     ----------
+    fname : str
+        The HDF5 file.
+    key : str
+        The name of the dataset or group to write to.
     chunk_data : list of numpy.ndarray
-        The data to be chunked.
-    kind : str
-        The kind of data to be chunked.
+        The data of the elements in the chunk.
+    kind : :enum:`.StorageType`
+        The kind of data.
     element_count : int
         The total number of elements.
     chunk_size : int
@@ -57,49 +168,29 @@ def _create_chunk(
     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',
-        'scalar_table'].
-
     """
-    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 in ["timeseries", "scalar_table", "timeseries_2d"]:
-        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', "
-            "'timeseries_2d', 'scalar_table']."
-        )
-    return out
+    offset = i_chunk * chunk_size
+    with h5py.File(fname, mode="a") as fid:
+        if kind in ["vector", "matrix"]:
+            features_data = np.concatenate(chunk_data, axis=-1)
+            if key not in fid:
+                dataset = fid.create_dataset(
+                    key,
+                    shape=(*features_data.shape[:-1], element_count),
+                    dtype=features_data.dtype,
+                    chunks=(*features_data.shape[:-1], chunk_size),
+                )
+                dataset.attrs["TITLE"] = "ndarray"
+            n_elements = features_data.shape[-1]
+            fid[key][..., offset : offset + n_elements] = features_data
+        elif kind in ["timeseries", "scalar_table", "timeseries_2d"]:
+            group = fid.require_group(key)
+            group.attrs["TITLE"] = "list"
+            for i, element_data in enumerate(chunk_data):
+                dataset = group.create_dataset(
+                    f"idx_{offset + i}", data=element_data
+                )
+                dataset.attrs["TITLE"] = "ndarray"
 
 
 @register_storage
@@ -108,7 +199,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
     Parameters
     ----------
-    uri : str or pathlib.Path
+    uri : pathlib.Path
         The path to the file to be used.
     single_output : bool, optional
         If False, will create one HDF5 file per element. The name
@@ -124,7 +215,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     force_float32 : bool, optional
         Whether to force casting of numpy.ndarray values to float32 if float64
         values are found (default True).
-    chunk_size : int, optional
+    chunk_size : positive int, optional
         The chunk size to use when collecting data from element files in
         :meth:`.collect`. If the file count is smaller than the value, the
         minimum is used (default 100).
@@ -135,34 +226,21 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
     """
 
-    _STORAGE_TYPES: ClassVar[Sequence[str]] = [
-        "vector",
-        "timeseries",
-        "matrix",
-        "scalar_table",
-        "timeseries_2d",
+    _STORAGE_TYPES: ClassVar[Sequence[StorageType]] = [
+        StorageType.Vector,
+        StorageType.Timeseries,
+        StorageType.Matrix,
+        StorageType.ScalarTable,
+        StorageType.Timeseries2D,
     ]
 
-    def __init__(
-        self,
-        uri: Union[str, Path],
-        single_output: bool = True,
-        overwrite: Union[bool, str] = "update",
-        compression: int = 7,
-        force_float32: bool = True,
-        chunk_size: int = 100,
-    ) -> None:
-        self.overwrite = overwrite
-        self.compression = compression
-        self.force_float32 = force_float32
-        self.chunk_size = chunk_size
-        super().__init__(
-            uri=uri,
-            single_output=single_output,
-        )
+    overwrite: bool | str = "update"
+    compression: Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] = 7
+    force_float32: bool = True
+    chunk_size: PositiveInt = 100
 
-    def _fetch_correct_uri_for_io(self, element: Optional[dict]) -> str:
-        """Return proper URI for I/O based on ``element``.
+    def _fetch_correct_uri_for_io(self, element: dict | None) -> str:
+        """Return proper URI for I/O based on `element`.
 
         If ``element`` is None, will return ``self.uri``.
 
@@ -197,7 +275,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         return f"{self.uri.parent}/{prefix}{self.uri.name}"
 
     def _read_metadata(
-        self, element: Optional[dict[str, str]] = None
+        self, element: dict[str, str] | None = None
     ) -> dict[str, dict[str, Any]]:
         """Read metadata (should not be called directly).
 
@@ -231,7 +309,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             )
 
         # Check if group is found in the storage
-        if not has_hdf5(fname=uri, title="meta"):
+        if not _has_hdf5(fname=uri, title="meta"):
             raise_error(
                 f"Invalid junifer HDF5 file at: {uri}",
                 klass=RuntimeError,
@@ -239,11 +317,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         # Read metadata
         logger.debug(f"Loading HDF5 metadata from: {uri}")
-        metadata = read_hdf5(
-            fname=uri,
-            title="meta",
-            slash="ignore",
-        )
+        metadata = _read_hdf5(fname=uri, title="meta")
         logger.debug(f"Loaded HDF5 metadata from: {uri}")
 
         return metadata
@@ -260,15 +334,11 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         """
         # Read metadata
-        metadata = read_hdf5(
-            fname=str(self.uri.resolve()),
-            title="meta",
-            slash="ignore",
-        )
+        metadata = _read_hdf5(fname=str(self.uri.resolve()), title="meta")
         return metadata
 
     def _read_data(
-        self, md5: str, element: Optional[dict[str, str]] = None
+        self, md5: str, element: dict[str, str] | None = None
     ) -> dict[str, Any]:
         """Read data (should not be called directly).
 
@@ -276,7 +346,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         ----------
         md5 : str
             The MD5 used as the HDF5 group name.
-        element : dict, optional
+        element : dict or None, optional
             The element as dictionary (default None).
 
         Returns
@@ -304,7 +374,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             )
 
         # Check if group is found in the storage
-        if not has_hdf5(fname=uri, title=md5):
+        if not _has_hdf5(fname=uri, title=md5):
             raise_error(
                 f"{md5} not found in HDF5 file at: {uri}",
                 klass=RuntimeError,
@@ -312,22 +382,16 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         # Read data
         logger.debug(f"Loading HDF5 data for {md5} from: {uri}")
-        data = read_hdf5(
-            fname=uri,
-            title=md5,
-            slash="ignore",
-        )
+        data = _read_hdf5(fname=uri, title=md5)
         logger.debug(f"Loaded HDF5 data for {md5} from: {uri}")
 
         return data
 
     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
@@ -371,11 +435,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
                 )
             )
         # Parameter check pass; read metadata
-        metadata = read_hdf5(
-            fname=str(self.uri.resolve()),
-            title="meta",
-            slash="ignore",
-        )
+        metadata = _read_hdf5(fname=str(self.uri.resolve()), title="meta")
         # Initialize MD5 variable
         md5: str = ""
 
@@ -383,7 +443,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         if feature_md5:
             logger.debug(
                 f"Validating feature MD5 '{feature_md5}' in metadata "
-                f"for: {self.uri.resolve()} ..."  # type: ignore
+                f"for: {self.uri.resolve()} ..."
             )
             # Validate MD5
             if feature_md5 in metadata:
@@ -423,17 +483,13 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             md5 = feature_name_duplicates_with_different_md5[0]
 
         # Read data from HDF5
-        hdf_data = read_hdf5(
-            fname=str(self.uri.resolve()),
-            title=md5,
-            slash="ignore",
-        )
+        hdf_data = _read_hdf5(fname=str(self.uri.resolve()), title=md5)
         return hdf_data
 
     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.
 
@@ -613,6 +669,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         meta_md5: str,
         element: dict[str, str],
         meta: dict[str, Any],
+        processed_element: dict[str, str] | None = None,
     ) -> None:
         """Store metadata.
 
@@ -629,15 +686,20 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             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, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
-        # Get correct URI for element;
+        # Get correct URI for the element being processed;
         # is different from uri if single_output is False
-        uri = self._fetch_correct_uri_for_io(element=element)
+        file_element = processed_element or element
+        uri = self._fetch_correct_uri_for_io(element=file_element)
 
         # Check if file exists, then read metadata else create empty dictionary
         if Path(uri).exists():
-            metadata = self._read_metadata(element=element)
+            metadata = self._read_metadata(element=file_element)
         else:
             logger.debug(f"Creating new file at {uri} ...")
             metadata = {}
@@ -670,10 +732,11 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
     def _store_data(
         self,
-        kind: str,
+        kind: StorageType,
         meta_md5: str,
         element: list[dict[str, str]],
         data: np.ndarray,
+        processed_element: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> None:
         """Store data.
@@ -685,7 +748,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         Parameters
         ----------
-        kind : {"matrix", "vector", "timeseries", "scalar_table"}
+        kind : :enum:`.StorageType`
             The storage kind.
         meta_md5 : str
             The metadata MD5 hash.
@@ -693,18 +756,23 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             The element as list of dictionary.
         data : numpy.ndarray
             The data to store.
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
         **kwargs : dict
             Keyword arguments passed from the calling method.
 
         """
-        # Get correct URI for element;
+        # Get correct URI for the element being processed;
         # is different from uri if single_output is False
-        uri = self._fetch_correct_uri_for_io(element=element[0])
+        file_element = processed_element or element[0]
+        uri = self._fetch_correct_uri_for_io(element=file_element)
 
         # Check if MD5 exists, then read data else create empty dictionary
         # File should be present here already
-        if has_hdf5(fname=uri, title=meta_md5):
-            stored_data = self._read_data(md5=meta_md5, element=element[0])
+        if _has_hdf5(fname=uri, title=meta_md5):
+            stored_data = self._read_data(md5=meta_md5, element=file_element)
         else:
             logger.debug(f"Creating new data map for {meta_md5} ...")
             stored_data = {}
@@ -814,13 +882,14 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     def store_matrix(
         self,
         meta_md5: str,
-        element: dict[str, 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,
         row_header_col_name: str = "ROI",
+        processed_element: dict[str, str] | None = None,
     ) -> None:
         """Store matrix.
 
@@ -839,28 +908,17 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             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`` is "full",
+            Whether to store the diagonal. If ``matrix_kind=MatrixKind.Full``,
             setting this to False will raise an error (default True).
         row_header_col_name : str, optional
             The column name for the row header column (default "ROI").
-
-        Raises
-        ------
-        ValueError
-            If invalid ``matrix_kind`` is provided, ``diagonal = False``
-            for ``matrix_kind = "full"``, non-square data is provided
-            for ``matrix_kind = {"triu", "tril"}``, length of ``row_names``
-            do not match data row count, or length of ``col_names`` do not
-            match data column count.
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         # Row data validation
@@ -879,9 +937,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         )
         # Store
         self._store_data(
-            kind="matrix",
+            kind=StorageType.Matrix,
             meta_md5=meta_md5,
             element=[element],  # convert to list
+            processed_element=processed_element,
             data=data[:, :, np.newaxis],  # convert to 3D
             column_headers=col_names,
             row_headers=row_names,
@@ -893,9 +952,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
     def store_vector(
         self,
         meta_md5: str,
-        element: dict,
-        data: Union[np.ndarray, list],
-        col_names: Optional[Sequence[str]] = None,
+        element: dict[str, str],
+        data: np.ndarray | list,
+        col_names: Sequence[str] | None = None,
+        processed_element: dict[str, str] | None = None,
     ) -> None:
         """Store vector.
 
@@ -909,6 +969,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             The vector data to store.
         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, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         if isinstance(data, list):
@@ -926,9 +990,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             processed_data = data.ravel()
 
         self._store_data(
-            kind="vector",
+            kind=StorageType.Vector,
             meta_md5=meta_md5,
             element=[element],  # convert to list
+            processed_element=processed_element,
             data=processed_data[:, np.newaxis],  # convert to 2D
             column_headers=col_names,
         )
@@ -938,7 +1003,8 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
+        col_names: Sequence[str] | None = None,
+        processed_element: dict[str, str] | None = None,
     ) -> None:
         """Store timeseries.
 
@@ -952,12 +1018,17 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             The timeseries data to store.
         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, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         self._store_data(
-            kind="timeseries",
+            kind=StorageType.Timeseries,
             meta_md5=meta_md5,
             element=[element],  # convert to list
+            processed_element=processed_element,
             data=[data],  # convert to list
             column_headers=col_names,
             row_header_column_name="timepoint",
@@ -968,8 +1039,9 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         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[str, str] | None = None,
     ) -> None:
         """Store 2D timeseries.
 
@@ -985,6 +1057,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             The column labels (default None).
         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, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         store_timeseries_2d_checks(
@@ -993,9 +1069,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             col_names_len=len(col_names) if col_names is not None else 0,
         )
         self._store_data(
-            kind="timeseries_2d",
+            kind=StorageType.Timeseries2D,
             meta_md5=meta_md5,
             element=[element],  # convert to list
+            processed_element=processed_element,
             data=[data],  # convert to list
             column_headers=col_names,
             row_headers=row_names,
@@ -1006,9 +1083,10 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         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[str, str] | None = None,
     ) -> None:
         """Store table with scalar values.
 
@@ -1026,12 +1104,17 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             The row labels (default None).
         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, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         self._store_data(
-            kind="scalar_table",
+            kind=StorageType.ScalarTable,
             meta_md5=meta_md5,
             element=[element],  # convert to list
+            processed_element=processed_element,
             data=[data],  # convert to list
             column_headers=col_names,
             row_headers=row_names,
@@ -1065,8 +1148,12 @@ class HDF5FeatureStorage(BaseFeatureStorage):
         logger.info(
             f"Collecting metadata from {self.uri.parent}/*_{self.uri.name}"
         )
-        # Collect element files per feature MD5
+        # Collect element files per feature MD5, only the ones with elements
+        # of the feature that are not in other files: the data of an element
+        # can be in the files of several elements being processed (e.g., the
+        # data of a subject in the files of each of its tasks)
         elements_per_feature_md5 = defaultdict(list)
+        seen_elements = defaultdict(set)
         out_metadata = {}
         for file_ in tqdm(
             self.uri.parent.glob(f"*_{self.uri.name}"), desc="file-metadata"
@@ -1083,7 +1170,28 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             out_metadata.update(in_metadata)
 
             # Update element files for found MD5s
-            for feature_md5 in in_metadata.keys():
+            file_elements = _read_elements(str(file_), list(in_metadata))
+            for feature_md5, elements in file_elements.items():
+                # Only the files of single elements can be collected
+                if len(elements) != 1:
+                    raise_error(
+                        msg=(
+                            f"The file {file_} has {len(elements)} elements "
+                            f"for the feature {feature_md5}, but only the "
+                            "files with one element (as stored for each "
+                            "element with `single_output=False`) can be "
+                            "collected"
+                        ),
+                        klass=RuntimeError,
+                    )
+                (element,) = elements
+                if element in seen_elements[feature_md5]:
+                    logger.debug(
+                        f"Skipping {file_} for {feature_md5}: its element "
+                        "is in another file"
+                    )
+                    continue
+                seen_elements[feature_md5].add(element)
                 elements_per_feature_md5[feature_md5].append(file_)
 
         logger.info("Writing metadata to HDF5 file ...")
@@ -1096,7 +1204,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
 
         # Run loop to collect data per feature per file
         logger.info(
-            f"Collecting data from {self.uri.parent}/*_{self.uri.name}"  # type: ignore
+            f"Collecting data from {self.uri.parent}/*_{self.uri.name}"
         )
         logger.info(f"Will collect {len(elements_per_feature_md5)} features.")
 
@@ -1119,6 +1227,13 @@ class HDF5FeatureStorage(BaseFeatureStorage):
             elements = []
             static_data = None
             kind = None
+            # The data is written in chunks to a temporary key (h5io would
+            # delete the previous chunks), and moved to the feature at the end
+            fname = str(self.uri.resolve())
+            data_key = f"collecting_{feature_md5}"
+            with h5py.File(fname, mode="a") as fid:
+                if data_key in fid:
+                    del fid[data_key]
             for file_ in tqdm(element_files, desc="file-data"):
                 logger.debug(
                     f"Reading feature MD5: '{feature_md5}' "
@@ -1126,11 +1241,7 @@ class HDF5FeatureStorage(BaseFeatureStorage):
                 )
 
                 # Read the data
-                t_data = read_hdf5(
-                    fname=str(file_),
-                    title=feature_md5,
-                    slash="ignore",
-                )
+                t_data = _read_hdf5(fname=str(file_), title=feature_md5)
                 if i_file == 0:
                     # Store the "static" data
                     static_data = {
@@ -1150,35 +1261,32 @@ class HDF5FeatureStorage(BaseFeatureStorage):
                 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
-                    # Write chunked array
-                    to_write["data"] = _create_chunk(
+                    # elements, write the data in chunks to avoid memory
+                    # usage spikes
+                    _write_chunk(
+                        fname=fname,
+                        key=data_key,
                         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()),
-                        data=to_write,
-                        overwrite="update",
-                        compression=0,
-                        title=feature_md5,
-                        slash="error",
-                        use_json=False,
-                    )
 
                     # Increment counters and data
                     i_chunk += 1
                     chunk_data = []
+
+            # Write the feature without the data
+            write_hdf5(
+                fname=fname,
+                data={**static_data, "element": elements},
+                overwrite="update",
+                compression=0,
+                title=feature_md5,
+                slash="error",
+                use_json=False,
+            )
+            # Move the data to the feature (as written by h5io)
+            with h5py.File(fname, mode="a") as fid:
+                fid.move(data_key, f"{feature_md5}/key_data")