From 3bab5472f71df0862aa7da61e58963aa49ebe543 Mon Sep 17 00:00:00 2001 From: Taylor Salo Date: Wed, 25 Mar 2026 13:43:38 -0400 Subject: [PATCH 1/2] Add a bunch of tests. --- test/test_cifti_cohort.py | 89 +++++++++++++++++++++++++++ test/test_h5_storage.py | 123 ++++++++++++++++++++++++++++++++++++++ test/test_parser_utils.py | 81 +++++++++++++++++++++++++ test/test_s3_utils.py | 23 +++++++ test/test_voxels_utils.py | 52 ++++++++++++++++ 5 files changed, 368 insertions(+) create mode 100644 test/test_cifti_cohort.py create mode 100644 test/test_h5_storage.py create mode 100644 test/test_parser_utils.py create mode 100644 test/test_s3_utils.py create mode 100644 test/test_voxels_utils.py diff --git a/test/test_cifti_cohort.py b/test/test_cifti_cohort.py new file mode 100644 index 0000000..8117387 --- /dev/null +++ b/test/test_cifti_cohort.py @@ -0,0 +1,89 @@ +"""Tests for CIFTI cohort normalization and greyordinate helpers.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from modelarrayio.utils.cifti import ( + _build_scalar_sources, + _cohort_to_long_dataframe, + brain_names_to_dataframe, +) + + +def test_cohort_long_format_preserves_rows() -> None: + df = pd.DataFrame( + { + 'scalar_name': ['THICK', 'THICK'], + 'source_file': ['sub-1.dscalar.nii', 'sub-2.dscalar.nii'], + 'extra_col': [1, 2], + } + ) + long_df = _cohort_to_long_dataframe(df) + assert len(long_df) == 2 + assert set(long_df.columns) == {'scalar_name', 'source_file'} + assert long_df.iloc[0]['scalar_name'] == 'THICK' + + +def test_cohort_long_format_strips_and_drops_empty() -> None: + df = pd.DataFrame( + { + 'scalar_name': [' THICK ', ''], + 'source_file': [' a.nii ', ' b.nii '], + } + ) + long_df = _cohort_to_long_dataframe(df) + assert len(long_df) == 1 + assert long_df.iloc[0]['scalar_name'] == 'THICK' + + +def test_cohort_wide_format_expands_columns() -> None: + df = pd.DataFrame( + { + 'id': [1, 2], + 'THICK': ['s1.nii', 's2.nii'], + 'FA': ['f1.nii', ''], + } + ) + long_df = _cohort_to_long_dataframe(df, scalar_columns=['THICK', 'FA']) + # Row 2 has empty FA — skipped + assert len(long_df) == 3 + scalars = set(long_df['scalar_name']) + assert scalars == {'THICK', 'FA'} + + +def test_cohort_wide_format_missing_scalar_column_raises() -> None: + df = pd.DataFrame({'THICK': ['a.nii']}) + with pytest.raises(ValueError, match='missing scalar columns'): + _cohort_to_long_dataframe(df, scalar_columns=['THICK', 'MISSING']) + + +def test_cohort_long_missing_required_raises() -> None: + df = pd.DataFrame({'only_this': [1]}) + with pytest.raises(ValueError, match='scalar_name'): + _cohort_to_long_dataframe(df) + + +def test_build_scalar_sources_ordering() -> None: + long_df = pd.DataFrame( + { + 'scalar_name': ['A', 'A', 'B'], + 'source_file': ['x1', 'x2', 'y1'], + } + ) + src = _build_scalar_sources(long_df) + assert list(src.keys()) == ['A', 'B'] + assert src['A'] == ['x1', 'x2'] + assert src['B'] == ['y1'] + + +def test_brain_names_to_dataframe() -> None: + names = np.array(['CORTEX_LEFT', 'CORTEX_LEFT', 'CORTEX_RIGHT']) + gdf, struct_strings = brain_names_to_dataframe(names) + assert len(gdf) == 3 + assert 'vertex_id' in gdf.columns + assert 'structure_id' in gdf.columns + assert gdf['vertex_id'].tolist() == [0, 1, 2] + assert len(struct_strings) == 2 # factorize unique structures diff --git a/test/test_h5_storage.py b/test/test_h5_storage.py new file mode 100644 index 0000000..4bc5add --- /dev/null +++ b/test/test_h5_storage.py @@ -0,0 +1,123 @@ +"""Unit tests for HDF5 storage helpers.""" + +from __future__ import annotations + +import h5py +import numpy as np +import pytest + +from modelarrayio.storage.h5_storage import ( + compute_chunk_shape_full_subjects, + create_scalar_matrix_dataset, + resolve_compression, + resolve_dtype, + write_rows_in_column_stripes, +) + + +def test_resolve_dtype() -> None: + assert resolve_dtype('float32') == np.float32 + assert resolve_dtype('float64') == np.float64 + assert resolve_dtype('FLOAT32') == np.float32 + assert resolve_dtype('unknown') == np.float32 + + +def test_resolve_compression_gzip_and_none() -> None: + comp, level, shuffle = resolve_compression('gzip', 9, True) + assert comp == 'gzip' + assert level == 9 + assert shuffle is True + + comp, level, shuffle = resolve_compression('none', 4, True) + assert comp is None + assert shuffle is False + + comp, level, shuffle = resolve_compression('lzf', 4, False) + assert comp == 'lzf' + assert shuffle is False + + +def test_resolve_compression_invalid_gzip_level_falls_back() -> None: + comp, level, _shuffle = resolve_compression('gzip', 'bad', True) + assert comp == 'gzip' + assert level == 4 + + +def test_compute_chunk_shape_full_subjects() -> None: + chunk = compute_chunk_shape_full_subjects( + num_subjects=3, + num_items=100, + item_chunk=10, + target_chunk_mb=2.0, + storage_np_dtype=np.float32, + ) + assert chunk == (3, 10) + + +def test_compute_chunk_shape_auto_item_chunk() -> None: + chunk = compute_chunk_shape_full_subjects( + num_subjects=2, + num_items=50, + item_chunk=0, + target_chunk_mb=1.0, + storage_np_dtype=np.float32, + ) + assert chunk[0] == 2 + assert 1 <= chunk[1] <= 50 + + +@pytest.mark.parametrize( + ('num_subjects', 'num_items'), + [(0, 10), (3, 0), (0, 0)], +) +def test_compute_chunk_shape_rejects_zero_dimension(num_subjects: int, num_items: int) -> None: + with pytest.raises(ValueError, match='zero-length'): + compute_chunk_shape_full_subjects( + num_subjects, + num_items, + item_chunk=0, + target_chunk_mb=2.0, + storage_np_dtype=np.float32, + ) + + +def test_create_scalar_matrix_dataset_writes_data_and_attrs(tmp_path) -> None: + path = tmp_path / 'out.h5' + data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) + with h5py.File(path, 'w') as h5: + create_scalar_matrix_dataset( + h5, + 'scalars/FA/values', + data, + ['sub-01', 'sub-02'], + compression='none', + shuffle=False, + ) + with h5py.File(path, 'r') as h5: + dset = h5['scalars/FA/values'] + np.testing.assert_array_equal(dset[...], data) + assert list(dset.attrs['column_names']) == ['sub-01', 'sub-02'] + + +def test_write_rows_in_column_stripes_matches_dense_write(tmp_path) -> None: + """Stripe writer should match assigning the full matrix.""" + path = tmp_path / 'stripe.h5' + num_subjects, num_elements = 3, 17 + full = np.arange(num_subjects * num_elements, dtype=np.float64).reshape(num_subjects, num_elements) + rows = [full[i].copy() for i in range(num_subjects)] + + with h5py.File(path, 'w') as h5: + dset = h5.create_dataset('m', shape=(num_subjects, num_elements), dtype='f8', chunks=(3, 5)) + write_rows_in_column_stripes(dset, rows) + + with h5py.File(path, 'r') as h5: + np.testing.assert_array_equal(h5['m'][...], full) + + +def test_write_rows_in_column_stripes_length_mismatch_raises(tmp_path) -> None: + with h5py.File(tmp_path / 'x.h5', 'w') as h5: + dset = h5.create_dataset('m', shape=(2, 4), dtype='f4') + with h5py.File(tmp_path / 'x.h5', 'a') as h5: + dset = h5['m'] + with pytest.raises(ValueError, match='rows length'): + write_rows_in_column_stripes(dset, [np.zeros(4)]) diff --git a/test/test_parser_utils.py b/test/test_parser_utils.py new file mode 100644 index 0000000..fc8ea1e --- /dev/null +++ b/test/test_parser_utils.py @@ -0,0 +1,81 @@ +"""Smoke tests for shared argparse helpers.""" + +from __future__ import annotations + +import argparse + +from modelarrayio.cli import parser_utils + + +def _parser_with_cohort() -> argparse.ArgumentParser: + p = argparse.ArgumentParser() + parser_utils.add_relative_root_arg(p) + parser_utils.add_output_hdf5_arg(p) + parser_utils.add_cohort_arg(p) + parser_utils.add_storage_args(p) + parser_utils.add_backend_arg(p) + parser_utils.add_s3_workers_arg(p) + return p + + +def test_minimal_invocation_defaults() -> None: + p = _parser_with_cohort() + args = p.parse_args(['--cohort-file', 'cohort.csv']) + assert args.cohort_file == 'cohort.csv' + assert args.storage_dtype == 'float32' + assert args.compression == 'gzip' + assert args.compression_level == 4 + assert args.shuffle is True + assert args.chunk_voxels == 0 + assert args.backend == 'hdf5' + assert args.s3_workers == 1 + assert args.log_level == 'INFO' + + +def test_storage_aliases_and_no_shuffle() -> None: + p = _parser_with_cohort() + args = p.parse_args( + [ + '--cohort-file', + 'c.csv', + '--dtype', + 'float64', + '--compression', + 'lzf', + '--no-shuffle', + '--chunk-voxels', + '128', + '--log-level', + 'WARNING', + ] + ) + assert args.storage_dtype == 'float64' + assert args.compression == 'lzf' + assert args.shuffle is False + assert args.chunk_voxels == 128 + assert args.log_level == 'WARNING' + + +def test_output_hdf5_default_name_override() -> None: + p = argparse.ArgumentParser() + parser_utils.add_output_hdf5_arg(p, default_name='custom.h5') + args = p.parse_args([]) + assert args.output_hdf5 == 'custom.h5' + + +def test_tiledb_args_group() -> None: + p = argparse.ArgumentParser() + parser_utils.add_output_tiledb_arg(p, default_name='arrays.tdb') + parser_utils.add_tiledb_storage_args(p) + args = p.parse_args([]) + assert args.output_tiledb == 'arrays.tdb' + assert args.tdb_compression == 'zstd' + assert args.tdb_tile_voxels == 0 + + +def test_scalar_columns_optional() -> None: + p = argparse.ArgumentParser() + parser_utils.add_cohort_arg(p) + parser_utils.add_scalar_columns_arg(p) + args = p.parse_args(['--cohort-file', 'c.csv', '--scalar-columns', 'c1', 'c2']) + assert args.scalar_columns == ['c1', 'c2'] diff --git a/test/test_s3_utils.py b/test/test_s3_utils.py new file mode 100644 index 0000000..ef03ff4 --- /dev/null +++ b/test/test_s3_utils.py @@ -0,0 +1,23 @@ +"""Tests for S3 path detection (no boto3 required).""" + +from __future__ import annotations + +import pytest + +from modelarrayio.utils.s3_utils import is_s3_path + + +@pytest.mark.parametrize( + ('path', 'expected'), + [ + ('s3://bucket/key.nii.gz', True), + ('s3://my-bucket/prefix/sub/file.nii', True), + ('S3://bucket/oops', False), + ('', False), + ('/local/path/file.nii.gz', False), + ('relative/path.nii', False), + ('https://example.com/x', False), + ], +) +def test_is_s3_path(path: str, expected: bool) -> None: + assert is_s3_path(path) is expected diff --git a/test/test_voxels_utils.py b/test/test_voxels_utils.py new file mode 100644 index 0000000..82d7c92 --- /dev/null +++ b/test/test_voxels_utils.py @@ -0,0 +1,52 @@ +"""Unit tests for voxel array helpers.""" + +from __future__ import annotations + +import nibabel as nb +import numpy as np +import pytest + +from modelarrayio.utils.voxels import flattened_image + + +def _eye_affine(): + return np.eye(4) + + +def test_flattened_image_extracts_group_masked_values() -> None: + shape = (3, 3, 3) + group_mask = np.zeros(shape, dtype=bool) + group_mask[0, 1, 2] = True + group_mask[2, 0, 1] = True + + scalar = np.zeros(shape, dtype=np.float32) + scalar[0, 1, 2] = 1.5 + scalar[2, 0, 1] = 2.5 + + indiv_mask = group_mask.copy() + scalar_img = nb.Nifti1Image(scalar, _eye_affine()) + mask_img = nb.Nifti1Image(indiv_mask.astype(np.float32), _eye_affine()) + + flat = flattened_image(scalar_img, mask_img, group_mask) + assert flat.shape == (2,) + assert np.allclose(flat, np.array([1.5, 2.5], dtype=np.float32)) + + +def test_flattened_image_nan_outside_individual_mask() -> None: + shape = (2, 2, 2) + group_mask = np.zeros(shape, dtype=bool) + group_mask[0, 0, 0] = True + group_mask[1, 1, 1] = True + + scalar = np.ones(shape, dtype=np.float32) * 3.0 + # Individual mask drops one group voxel + indiv = group_mask.copy() + indiv[1, 1, 1] = False + + scalar_img = nb.Nifti1Image(scalar, _eye_affine()) + mask_img = nb.Nifti1Image(indiv.astype(np.float32), _eye_affine()) + + flat = flattened_image(scalar_img, mask_img, group_mask) + assert flat.shape == (2,) + assert np.isnan(flat[1]) + assert flat[0] == pytest.approx(3.0) From 9e04c01f5e20fb244f61b9d7ec4482658fda3ed8 Mon Sep 17 00:00:00 2001 From: Taylor Salo Date: Wed, 25 Mar 2026 13:45:07 -0400 Subject: [PATCH 2/2] Run ruff. --- test/test_h5_storage.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test/test_h5_storage.py b/test/test_h5_storage.py index 4bc5add..fe67dab 100644 --- a/test/test_h5_storage.py +++ b/test/test_h5_storage.py @@ -103,11 +103,15 @@ def test_write_rows_in_column_stripes_matches_dense_write(tmp_path) -> None: """Stripe writer should match assigning the full matrix.""" path = tmp_path / 'stripe.h5' num_subjects, num_elements = 3, 17 - full = np.arange(num_subjects * num_elements, dtype=np.float64).reshape(num_subjects, num_elements) + full = np.arange(num_subjects * num_elements, dtype=np.float64).reshape( + num_subjects, num_elements + ) rows = [full[i].copy() for i in range(num_subjects)] with h5py.File(path, 'w') as h5: - dset = h5.create_dataset('m', shape=(num_subjects, num_elements), dtype='f8', chunks=(3, 5)) + dset = h5.create_dataset( + 'm', shape=(num_subjects, num_elements), dtype='f8', chunks=(3, 5) + ) write_rows_in_column_stripes(dset, rows) with h5py.File(path, 'r') as h5: