BASS-Python

Python/NumPy and optional PyTorch implementation of Bias Accelerated Subset Selection for learning parallel MRI sampling masks.

Rahul implemented the BASS sampling loop around SENSE-CG reconstruction: constrained mask initialization, epsilon and r heuristic maps, add/remove proposals, objective-based accept/reject, train/validation metrics, saved artifacts, visualization utilities, and synthetic tests that can run without private MRI data.

Core capabilities:

• Learn 2D or 2D+time sampling masks from fully sampled k-space and coil maps.
• Preserve a fully sampled center region and support conjugate-symmetry style selection constraints.
• Reconstruct with SENSE and conjugate gradient using NumPy, with optional PyTorch CPU/CUDA execution.
• Track k-space NRMSE, image NRMSE, SSIM, proposal history, masks, reconstructions, and ground-truth artifacts when provided.
• Run multi-slice train/validation experiments from .mat inputs.
• Verify core behavior with pytest fixtures built from synthetic data.

flowchart TD
    inputs["Input .mat files<br/>k-space + coil maps<br/>optional ground truth"] --> init["Initial sampling mask<br/>center lock + constraints"]
    init --> recon["SENSE-CG reconstruction<br/>NumPy or optional PyTorch"]
    recon --> residual["Residual and metric computation<br/>k-space NRMSE, image NRMSE, SSIM"]
    residual --> maps["Heuristic maps<br/>epsilon for add candidates<br/>r for remove candidates"]
    maps --> proposal["Add/remove proposal<br/>selection probabilities + constraints"]
    proposal --> accept["Accept if objective does not worsen<br/>otherwise shrink proposal budget"]
    accept --> recon
    accept --> artifacts["Saved artifacts<br/>mask, recon, history, ground truth"]

Loading

Verify It Quickly

Fresh clones do not include MRI datasets. The fastest truthful verification path uses synthetic fixtures from the test suite:

conda env create -f environment.yml
conda activate bass-dev

python -m pytest tests/test_bass_algorithm.py tests/test_solvers.py -q
python run_bass.py --help
python visualize_bass.py --help

What this verifies:

• BASS converges toward the target sample budget on synthetic data.
• Center locking is respected.
• Multi-frame masks run consistently.
• Validation image metrics are recorded when ground truth is available.
• The optional PyTorch backend is exercised when the local environment supports it.
• SENSE-CG solver behavior is tested on known fixtures.

Important files to inspect:

• bass/algorithm.py: core BASS loop, maps, proposals, acceptance, metrics, artifact state.
• bass/torch_ops.py: optional PyTorch encoding operator and CG solver.
• run_bass.py: experiment CLI and artifact writing.
• visualize_bass.py: mask, heuristic, metric, reconstruction, and VD-comparison views.
• matlab_parity/: parity-oriented Python structure matching the legacy MATLAB-style implementation.
• tests/: synthetic tests for algorithm behavior, operators, solvers, masks, edge cases, and visualization loading.

What This Repository Implements

This repository implements a research-software version of BASS around SENSE-CG:

Area	Implemented here
Sampling optimization	Add/remove proposal loop, proposal budget shrink on rejected candidates, fixed target sample budget.
Heuristics	Epsilon maps for unsampled add candidates and r maps for sampled remove candidates.
Constraints	Fully sampled center block, optional center unlock, 2D and 2D+time mask modes, frame-aware center locking.
Reconstruction	SENSE forward/adjoint operators and conjugate-gradient reconstruction.
Backends	NumPy default path and optional PyTorch backend via --use-torch.
Metrics	K-space residual/NRMSE, image NRMSE, SSIM when ground truth is provided, train/validation summaries.
Artifacts	.npy mask/reconstruction/ground-truth files and .npz history with BassIteration objects.
Visualization	Heuristic maps, mask evolution, metric curves, reconstructions, and optional variable-density comparison.
Validation	Pytest fixtures using synthetic data, plus MATLAB-parity utilities for reference-style experiments.

Paper, Repository, And Measured Evidence

The README deliberately separates paper results from repository evidence.

Category	What is claimed
Original BASS paper	Zibetti, Herman, and Regatte report faster convergence and lower NRMSE than selected baselines for their experimental setup. Those are paper-reported results, not automatically reproduced by this repository.
This Python repository	Implements the BASS+SENSE-CG mechanics, parity-oriented utilities, visualization, CLI artifact generation, optional PyTorch execution, and tests.
Checked-in evidence	Synthetic pytest coverage and source-level parity utilities are included. No public MRI datasets or checked-in result images are committed in this repository.

Do not treat the paper's "50x faster" or "45% lower NRMSE" findings as independently reproduced here unless you add public artifacts that prove those exact comparisons.

Why This Project Matters

This is a research-to-engineering project. The interesting part is not only translating equations into Python, but turning an MRI sampling method into inspectable software:

• deterministic CLI entrypoints
• explicit data contracts for k-space, coil maps, and optional ground truth
• testable numerical components
• per-iteration metrics and history
• backend choice between NumPy and PyTorch
• artifact names that make experiments easier to audit

For a technical reviewer, this is evidence of scientific-computing and ML-data infrastructure work: handling complex arrays, numerical solvers, validation metrics, reproducible artifacts, and careful limits around private data.

Artifact Naming

When --save-mask, --save-recon, --save-history, or --save-gt are passed, outputs are written to --output-dir using this pattern:

bass_mask_<recon-method>_<dataname>_<slice-tag>_<accel>x_<iterations>i_<seed>s.npy
bass_recon_<recon-method>_<dataname>_<slice-tag>_<accel>x_<iterations>i_<seed>s.npy
bass_history_<recon-method>_<dataname>_<slice-tag>_<accel>x_<iterations>i_<seed>s.npz
bass_ground_truth_<recon-method>_<dataname>_<slice-tag>_<accel>x_<iterations>i_<seed>s.npy

The history archive contains:

• history: BassIteration objects for each iteration
• initial_mask: the initial mask used for the run

Data Availability

No MRI data is committed to this repository.

The CLI expects a directory like:

data/<Dataset>/
  kdata/kdata<slice>.mat
  coilmap/coilmap<slice>.mat
  recon_*/images<slice>.mat        # optional ground truth
  sampling_pattern/*.mat           # optional reference masks

Expected array roles:

• kdata: fully sampled complex k-space.
• coilmap: complex coil sensitivity maps.
• recon_*: optional reference images for image NRMSE/SSIM.

Public data can be used if you convert it to the expected .mat layout. For example, public fastMRI brain data may be appropriate subject to its license and access terms. Restricted NYU Langone or CAI2R data must not be committed or redistributed.

Running With Your Own Data

Example command for a supplied dataset:

python run_bass.py \
  --data-root data \
  --dataname T1rho \
  --train-slices 31-35 \
  --val-slices 101-105 \
  --accel 4 \
  --iterations 10 \
  --points 200 \
  --cg-iters 10 \
  --center 0.25 \
  --seed 0 \
  --ground-truth data/T1rho/recon_DIC \
  --ground-truth-slices 31-35,101-105 \
  --ground-truth-key images \
  --save-mask --save-recon --save-history --save-gt \
  --output-dir outputs/t1rho_demo

This command will fail from a fresh clone unless the referenced data files exist.

See docs/USAGE.md for detailed CLI options, visualization commands, dataset layout notes, and MATLAB-parity notes.

Visualization

After a run writes artifacts, use visualize_bass.py:

python visualize_bass.py \
  --history outputs/t1rho_demo/bass_history_sense-cg_t1rho_s31-35_n5_4x_10i_0s.npz \
  --final-mask outputs/t1rho_demo/bass_mask_sense-cg_t1rho_s31-35_n5_4x_10i_0s.npy \
  --reconstruction outputs/t1rho_demo/bass_recon_sense-cg_t1rho_s31-35_n5_4x_10i_0s.npy \
  --ground-truth-vol outputs/t1rho_demo/bass_ground_truth_sense-cg_t1rho_s31-35_n5_4x_10i_0s.npy \
  --mode summary \
  --compare-vd \
  --dataname T1rho \
  --slice 31 \
  --vd-samples 2048

No screenshot is checked in because the repository has no public MRI inputs. A valid screenshot should only be added after it can be reproduced from public or redistributable data.

Repository Structure

Path	Purpose
bass/algorithm.py	Core BASS loop, reconstruction orchestration, heuristics, metrics, and saved state.
bass/config.py	Configuration and dataset dataclasses.
bass/result.py	Result and per-iteration history structures.
bass/torch_ops.py	Optional PyTorch SENSE operator and CG solver.
operators.py	NumPy SENSE forward/adjoint operator.
solvers.py	NumPy conjugate-gradient solver.
mask_generation.py	Variable-density mask utility for baselines/comparison.
run_bass.py	Main experiment CLI.
visualize_bass.py	Visualization CLI for saved artifacts.
matlab_parity/	MATLAB-style parity-oriented implementation and helpers.
tests/	Pytest suite using synthetic fixtures.

MATLAB Parity Utilities

matlab_parity/ mirrors the legacy file/module organization:

• DDLSP_optimize_SP_BASS_SENSE.py
• DDLSP_optimize_SP_BASS_SENSE_1D.py
• select_add*.py and select_remove*.py
• Emat_xyt/*
• L2_CG.py

These files are useful for comparing logic against MATLAB-style reference code, but they still require compatible local data to run full experiments.

Tests

Run all tests:

python -m pytest -q

Focused checks:

python -m pytest tests/test_bass_algorithm.py -q
python -m pytest tests/test_solvers.py -q
python -m pytest tests/test_operators.py -q

Some tests skip optional PyTorch execution when the local PyTorch build is unavailable.

Limitations

• No MRI datasets are committed.
• No public result montage is committed.
• The primary implemented reconstruction path is SENSE-CG, not every reconstruction family from the paper.
• Paper-reported benchmark results are cited as paper results only.
• The code is research software, not a clinical tool.
• No license file is included, so reuse rights are not specified.

Reference

Marcelo V. W. Zibetti, Gabor T. Herman, and Ravinder R. Regatte. "Fast data-driven learning of parallel MRI sampling patterns for large scale problems." Scientific Reports 11, 19312, 2021. https://doi.org/10.1038/s41598-021-97995-w