[BUG]: HDF5Storage fails to collect #198

Merged
fraimondo merged 9 commits from fix/hdf5_collect into main 2023-03-24 09:51:24 +00:00
3 changed files with 357 additions and 209 deletions

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

View file

@ -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
synchon commented 2023-03-22 09:41:54 +00:00 (Migrated from github.com)

Missing docstring here.

Missing docstring here.
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 data.dtype == np.dtype("float64") and self.force_float32:
data = data.astype(dtype=np.dtype("float32"), casting="same_kind")
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"
)
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
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_)
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",
)
# 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]
logger.debug(
f"Reading feature MD5: '{feature_md5}' "
f"from HDF5 file: {file_} ..."
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 the data
t_data = read_hdf5(
fname=str(file_),
title=feature_md5,
slash="ignore",
)
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"]
}
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
# Write chunked array
to_write["data"] = _create_chunk(
chunk_data=chunk_data,
kind=kind,
element_count=element_count,
chunk_size=t_chunk_size,
i_chunk=i_chunk,
)
# Read from HDF5 and collect data
stored_data_for_chunk.append(
read_hdf5(
fname=str(file_),
title=feature_md5,
slash="ignore",
)
if i_file == element_count:
to_write["element"] = elements
# Write to HDF5
write_hdf5(
fname=str(self.uri.resolve()), # type: ignore
data=to_write,
overwrite="update", # type: ignore
compression=0,
title=feature_md5,
slash="error",
use_json=False,
)
# 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")
}
# 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],
[],
)
# 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,
)
# Write to HDF5
write_hdf5(
fname=str(self.uri.resolve()), # type: ignore
data=data_to_write,
overwrite=self.overwrite, # type: ignore
compression=0,
title=feature_md5,
slash="error",
use_json=False,
)
# Increment counters and data
i_chunk += 1
chunk_data = []

View file

@ -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)
synchon commented 2023-03-22 09:51:42 +00:00 (Migrated from github.com)

Missing type annotations.

Missing type annotations.
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)
# 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
# Store metadata for tables
storage.store_metadata(
meta_md5=hash_1,
element=element_to_store_1,
meta=meta_to_store_1,
)
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,
storage = HDF5FeatureStorage(
uri=uri, single_output=False, chunk_size=chunk_size
)
# Store tables
storage.store_vector(
meta_md5=hash_1,
element=element_to_store_1,
data=data_1,
col_names=col_headers,
)
storage.store_vector(
meta_md5=hash_2,
element=element_to_store_2,
data=data_2,
col_names=col_headers,
)
storage.store_vector(
meta_md5=hash_3,
element=element_to_store_3,
data=data_3,
col_names=col_headers,
)
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=meta_md5,
element=t_data["element"],
meta=t_data["meta"],
)
if kind == "vector":
# Store tables
storage.store_vector(
meta_md5=meta_md5,
element=t_data["element"],
**t_data["data"],
)
elif kind == "matrix":
# Store tables
storage.store_matrix(
meta_md5=meta_md5,
element=t_data["element"],
**t_data["data"],
)
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()
# 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
for t_data in all_data:
# Convert element to preifx
prefix = element_to_prefix(t_data["element"])
# URIs for data storage
elem_uri = uri.parent / f"{prefix}{uri.name}"
# Check URIs for data storage exist
assert elem_uri.exists()
# 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}"
# Check URIs for data storage exist
assert uri_1.exists()
assert uri_2.exists()
assert uri_3.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"])
# Check if metadata are equal
assert read_meta_1 == read_meta_2
assert read_meta_2 == read_meta_3
# Read stored metadata from different files using element
read_meta = storage._read_metadata(element=t_data["element"])
# Check if metadata are equal
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: