Consolidate TileDB- and HDF5-related parameters (#28)

Author: tsalo

docs/examples/plot_fixel_workflow.py

@@ -102,7 +102,7 @@
 #         --index-file      /home/username/myProject/data/FD/index.mif \
 #         --directions-file /home/username/myProject/data/FD/directions.mif \
 #         --cohort-file     /home/username/myProject/data/cohort_FD.csv \
-#         --output-hdf5     /home/username/myProject/data/FD.h5
+#         --output          /home/username/myProject/data/FD.h5
 #
 # This produces ``FD.h5`` in ``/home/username/myProject/data``.  You can then use
 # `ModelArray <https://pennlinc.github.io/ModelArray/>`_ to run statistical analyses on it.

docs/examples/plot_voxel_workflow.py

@@ -107,7 +107,7 @@
 #     modelarrayio nifti-to-h5 \
 #         --group-mask-file /home/username/myProject/data/group_mask.nii.gz \
 #         --cohort-file     /home/username/myProject/data/cohort_FA.csv \
-#         --output-hdf5     /home/username/myProject/data/FA.h5
+#         --output          /home/username/myProject/data/FA.h5
 #
 # This produces ``FA.h5`` in ``/home/username/myProject/data``.  You can then use
 # `ModelArray <https://pennlinc.github.io/ModelArray/>`_ to run statistical analyses on it.

src/modelarrayio/cli/cifti_to_h5.py

@@ -12,12 +12,10 @@
 from modelarrayio.cli.parser_utils import (
     add_backend_arg,
     add_cohort_arg,
-    add_output_hdf5_arg,
-    add_output_tiledb_arg,
+    add_output_arg,
     add_s3_workers_arg,
     add_scalar_columns_arg,
     add_storage_args,
-    add_tiledb_storage_args,
 )
 from modelarrayio.storage import h5_storage, tiledb_storage
 from modelarrayio.utils.cifti import (
@@ -34,68 +32,53 @@
 def cifti_to_h5(
     cohort_file,
     backend='hdf5',
-    output_hdf5='fixeldb.h5',
-    output_tiledb='arraydb.tdb',
+    output='fixelarray.h5',
     storage_dtype='float32',
     compression='gzip',
     compression_level=4,
     shuffle=True,
     chunk_voxels=0,
     target_chunk_mb=2.0,
-    tdb_compression='zstd',
-    tdb_compression_level=5,
-    tdb_shuffle=True,
-    tdb_tile_voxels=0,
-    tdb_target_tile_mb=2.0,
-    tdb_workers=None,
+    workers=None,
     scalar_columns=None,
     s3_workers=1,
 ):
-    """Load all CIFTI data and write to an HDF5 file with configurable storage.
+    """Load all CIFTI data and write to an HDF5 or TileDB file.
 
     Parameters
     ----------
     cohort_file : :obj:`str`
         Path to a csv with demographic info and paths to data
     backend : :obj:`str`
-        Backend to use for storage
-    output_hdf5 : :obj:`str`
-        Path to a new .h5 file to be written
-    output_tiledb : :obj:`str`
-        Path to a new .tdb file to be written
+        Backend to use for storage (``'hdf5'`` or ``'tiledb'``)
+    output : :obj:`str`
+        Output path. For the hdf5 backend, path to an .h5 file;
+        for the tiledb backend, path to a .tdb directory.
     storage_dtype : :obj:`str`
         Floating type to store values
     compression : :obj:`str`
-        HDF5 compression filter
+        Compression filter. ``gzip`` works for both backends;
+        ``lzf`` is HDF5-only; ``zstd`` is TileDB-only.
     compression_level : :obj:`int`
-        Gzip compression level (0-9)
+        Compression level (codec-dependent)
     shuffle : :obj:`bool`
-        Enable HDF5 shuffle filter
+        Enable shuffle filter
     chunk_voxels : :obj:`int`
-        Chunk size along the voxel axis
+        Chunk/tile size along the greyordinate axis (0 = auto)
     target_chunk_mb : :obj:`float`
-        Target chunk size in MiB when auto-computing chunk_voxels
-    tdb_compression : :obj:`str`
-        TileDB compression filter
-    tdb_compression_level : :obj:`int`
-        TileDB compression level
-    tdb_shuffle : :obj:`bool`
-        Enable TileDB shuffle filter
-    tdb_tile_voxels : :obj:`int`
-        Tile size along the voxel axis
-    tdb_target_tile_mb : :obj:`float`
-        Target tile size in MiB when auto-computing tdb_tile_voxels
-    tdb_workers : :obj:`int`
-        Number of workers to use for parallel loading
+        Target chunk/tile size in MiB when auto-computing the spatial axis length
+    workers : :obj:`int`
+        Maximum number of parallel TileDB write workers (``None`` = auto).
+        Has no effect when ``backend='hdf5'``.
     scalar_columns : :obj:`list`
         List of scalar columns to use
     s3_workers : :obj:`int`
-        Number of workers to use for parallel loading
+        Number of workers for parallel S3 downloads
 
     Returns
     -------
     status : :obj:`int`
-        Status of the operation. 0 if successful, 1 if failed.
+        0 if successful, 1 if failed.
     """
     cohort_df = pd.read_csv(cohort_file)
     cohort_long = _cohort_to_long_dataframe(cohort_df, scalar_columns=scalar_columns)
@@ -108,7 +91,7 @@ def cifti_to_h5(
     if backend == 'hdf5':
         scalars, last_brain_names = _load_cohort_cifti(cohort_long, s3_workers)
 
-        f = h5py.File(output_hdf5, 'w')
+        f = h5py.File(output, 'w')
 
         greyordinate_table, structure_names = brain_names_to_dataframe(last_brain_names)
         greyordinatesh5 = f.create_dataset(
@@ -136,9 +119,9 @@ def cifti_to_h5(
 
             h5_storage.write_rows_in_column_stripes(dset, scalars[scalar_name])
         f.close()
-        return int(not os.path.exists(output_hdf5))
+        return int(not os.path.exists(output))
     else:
-        os.makedirs(output_tiledb, exist_ok=True)
+        os.makedirs(output, exist_ok=True)
         if not scalar_sources:
             return 0
 
@@ -161,26 +144,26 @@ def _process_scalar_job(scalar_name, source_files):
                 return scalar_name
             num_items = rows[0].shape[0]
             tiledb_storage.create_empty_scalar_matrix_array(
-                output_tiledb,
+                output,
                 dataset_path,
                 num_subjects,
                 num_items,
                 storage_dtype=storage_dtype,
-                compression=tdb_compression,
-                compression_level=tdb_compression_level,
-                shuffle=tdb_shuffle,
-                tile_voxels=tdb_tile_voxels,
-                target_tile_mb=tdb_target_tile_mb,
+                compression=compression,
+                compression_level=compression_level,
+                shuffle=shuffle,
+                tile_voxels=chunk_voxels,
+                target_tile_mb=target_chunk_mb,
                 sources_list=source_files,
             )
             # write column names array for ModelArray compatibility
-            tiledb_storage.write_column_names(output_tiledb, scalar_name, source_files)
-            uri = os.path.join(output_tiledb, dataset_path)
+            tiledb_storage.write_column_names(output, scalar_name, source_files)
+            uri = os.path.join(output, dataset_path)
             tiledb_storage.write_rows_in_column_stripes(uri, rows)
             return scalar_name
 
         scalar_names = list(scalar_sources.keys())
-        worker_count = tdb_workers if isinstance(tdb_workers, int) and tdb_workers > 0 else None
+        worker_count = workers if isinstance(workers, int) and workers > 0 else None
         if worker_count is None:
             cpu_count = os.cpu_count() or 1
             worker_count = min(len(scalar_names), max(1, cpu_count))
@@ -207,20 +190,14 @@ def _process_scalar_job(scalar_name, source_files):
 def cifti_to_h5_main(
     cohort_file,
     backend='hdf5',
-    output_hdf5='fixelarray.h5',
-    output_tiledb='arraydb.tdb',
+    output='fixelarray.h5',
     storage_dtype='float32',
     compression='gzip',
     compression_level=4,
     shuffle=True,
     chunk_voxels=0,
     target_chunk_mb=2.0,
-    tdb_compression='zstd',
-    tdb_compression_level=5,
-    tdb_shuffle=True,
-    tdb_tile_voxels=0,
-    tdb_target_tile_mb=2.0,
-    tdb_workers=None,
+    workers=None,
     scalar_columns=None,
     s3_workers=1,
     log_level='INFO',
@@ -233,44 +210,37 @@ def cifti_to_h5_main(
     return cifti_to_h5(
         cohort_file=cohort_file,
         backend=backend,
-        output_hdf5=output_hdf5,
-        output_tiledb=output_tiledb,
+        output=output,
         storage_dtype=storage_dtype,
         compression=compression,
         compression_level=compression_level,
         shuffle=shuffle,
         chunk_voxels=chunk_voxels,
         target_chunk_mb=target_chunk_mb,
-        tdb_compression=tdb_compression,
-        tdb_compression_level=tdb_compression_level,
-        tdb_shuffle=tdb_shuffle,
-        tdb_tile_voxels=tdb_tile_voxels,
-        tdb_target_tile_mb=tdb_target_tile_mb,
-        tdb_workers=tdb_workers,
+        workers=workers,
         scalar_columns=scalar_columns,
         s3_workers=s3_workers,
     )
 
 
 def _parse_cifti_to_h5():
     parser = argparse.ArgumentParser(
-        description='Create a hdf5 file of CIDTI2 dscalar data',
+        description='Create a hdf5 file of CIFTI2 dscalar data',
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     add_cohort_arg(parser)
     add_scalar_columns_arg(parser)
-    add_output_hdf5_arg(parser, default_name='fixelarray.h5')
-    add_output_tiledb_arg(parser, default_name='arraydb.tdb')
+    add_output_arg(parser, default_name='fixelarray.h5')
     add_backend_arg(parser)
     add_storage_args(parser)
-    add_tiledb_storage_args(parser)
     parser.add_argument(
-        '--tdb-workers',
-        '--tdb_workers',
+        '--workers',
         type=int,
         help=(
-            'Maximum number of TileDB write workers. Default 0 (auto, uses CPU count). '
-            'Set to 1 to disable parallel writes.'
+            'Maximum number of parallel TileDB write workers. '
+            'Default 0 (auto, uses CPU count). '
+            'Set to 1 to disable parallel writes. '
+            'Has no effect when --backend=hdf5.'
         ),
         default=0,
     )