Add more tests

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

test/test_h5_storage.py

@@ -0,0 +1,127 @@
+"""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)])

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']

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

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)