API Reference#

Use this page as a map of the public API. Narrative explanations are in the guide pages; this page lists the main entry points.

Data and Study Containers #

Data structures that keep arrays, units, and experiment context together.

DataSet bundles magnitude/k-space arrays plus optional noise level and derived results; shape and spatial_shape expose the spatial dimensions expected by models.
Quantity pairs any ndarray with an SI unit string so observables are explicitly typed.
Experiment declares dimension names and observables; validate checks dataset compatibility and observable types.
JointExperiment and join_experiments stack compatible experiments into a frame-major dataset for joint fitting or reconstruction; select_experiment_frames extracts sub-experiments by frame indices.
Dimension helpers in mrax.data_struct.dimensions infer non-spatial axis splits, reshape frame-major arrays, and expand observables to per-frame arrays.
Model ties a forward model, parametrization factory, and per-parameter units; it can emit regularizer terms via make_regularizer_terms.
Study keeps multiple experiments/datasets/models together and attaches shared study globals (e.g., field_strength_mhz) to each experiment/model.

Container for quantitative MRI data and optional raw k-space.

Non-Cartesian acquisitions store samples in kspace with the sample axis last and store the corresponding trajectory in trajectory as (num_samples, 2). Cartesian and EPI acquisitions can also provide sampling_mask for undersampled grids.

class mrax.data_struct.quantity.Quantity(values: Any, si_unit: str)#: Value paired with its SI unit string.

mrax.data_struct.quantity.attach_units(fields: dict[str, Array], units: dict[str, str]) → dict[str, Quantity]#: Pair raw parameter maps with their SI unit strings.

mrax.data_struct.quantity.unwrap_values(mapping: dict[str, Quantity]) → dict[str, Array]#: Extract raw arrays from a dict of quantities.

mrax.data_struct.quantity.require_quantity_dict(mapping: dict[str, Any], context: str) → dict[str, Quantity]#: Validate that all values are Quantity instances and return a typed dict.

class mrax.data_struct.experiment.Experiment(dims: ~typing.Dict[str, int], observables: ~typing.Dict[str, ~mrax.data_struct.quantity.Quantity] = <factory>, exp_type: str | None = None)#

Describes the acquisition setup (dims, observables such as TE/TR/offsets).

merged_observables(study_globals: Mapping[str, Quantity] | None = None) → dict[str, Quantity]#: Return observables merged with study-level globals (if available).

insert_observables(updates: Mapping[str, object] | None, *, units: Mapping[str, str] | None = None, overwrite: bool = True, in_place: bool = False) → Experiment#

Return an Experiment with additional observables merged in.

updates values may be Quantity, raw arrays (with units), or (values, unit) tuples.

sanity_check(model: Model, dataset: DataSet, *, study_globals: Mapping[str, Quantity] | None = None)#: Run a pre-fit sanity check for this experiment/model/dataset.

explain_sanity(model: Model, dataset: DataSet, *, study_globals: Mapping[str, Quantity] | None = None) → str#: Return a human-readable sanity check report.

class mrax.data_struct.joint.JointExperiment(experiment: Experiment, dataset: DataSet, sources: tuple[JointSource, ...], frame_map: Mapping[str, ndarray], source_observable: str | None = None)#: Joined experiment with frame indices for each source.

mrax.data_struct.joint.join_experiments(experiments: Sequence[Experiment], datasets: Sequence[DataSet], *, names: Sequence[str] | None = None, exp_type: str | None = None, joint_name: str = 'joint', frame_dim: str = 'frame', source_observable: str | None = 'joint_source', source_unit: str = '1', fill_value: float = nan, allow_missing_observables: bool = True, spatial_mode: str = 'require_equal', keep_kspace: bool = False) → JointExperiment#

Join multiple experiments into a frame-major stack.

Parameters#

spatial_mode{“require_equal”, “broadcast”}: How to handle spatial observables (e.g., B0/B1 maps) that differ across inputs. “require_equal” raises if multiple distinct maps are present. “broadcast” converts spatial maps into per-frame observables by repeating them per source.

mrax.data_struct.joint.select_experiment_frames(experiment: Experiment, dataset: DataSet | None = None, *, frame_indices: Sequence[int] | ndarray | None = None, frame_mask: Sequence[bool] | ndarray | None = None, frame_dim: str = 'frame', name: str | None = None) → tuple[Experiment, DataSet | None]#: Return a frame-subset experiment (and optional dataset).

class mrax.data_struct.dimensions.DimensionSpec(name: 'str', observable: 'str | None' = None, allow_singleton: 'bool' = True)#

mrax.data_struct.dimensions.build_experiment_dims(frame_count: int, spatial_shape: Iterable[int], observables: Mapping[str, Quantity] | None, *, exp_type: str | None = None, model_name: str | None = None, fallback_dim: str = 'frame') → dict[str, int]#

mrax.data_struct.dimensions.reshape_frames_to_dims(data: ndarray, dims: Mapping[str, int]) → ndarray#

mrax.data_struct.dimensions.reshape_dims_to_frames(data: ndarray, dims: Mapping[str, int], *, broadcast_spatial: bool = True) → ndarray#: Reshape a dim-ordered array into (frame, *spatial) layout.

mrax.data_struct.dimensions.reshape_observables_to_dims(observables: Mapping[str, Quantity], dims: Mapping[str, int], dim_to_observable: Mapping[str, str | None] | None = None, *, atol: float = 1e-06) → dict[str, Quantity]#

mrax.data_struct.dimensions.expand_observables_to_frames(observables: Mapping[str, Quantity], dims: Mapping[str, int], dim_to_observable: Mapping[str, str | None] | None = None) → dict[str, ndarray]#

mrax.data_struct.dimensions.pick_observable_for_frames(observables: Mapping[str, Quantity], frame_count: int | None, *, preferred_keys: Sequence[str] = ('TI', 'TE', 'TR', 'freq_offsets', 'offsets')) → tuple[ndarray | None, str | None]#: Pick an observable array matching the frame count (useful for spectra plots).

class mrax.data_struct.model.Model(name: str, model_apply: Callable[[dict, Array, object | None], Array], parametrization_factory: Callable[[tuple[int, int]], Tuple[dict, Callable[[dict], dict]]], regularizer_type: str | dict[str, str] = 'smoothness', regularizer_configs: dict[str, List[dict]] | None = None, reg_weights: dict[str, float] | None = None, parameter_units: Dict[str, str] | None = None, derived_map_units: Dict[str, str] | None = None, parameter_constant_dims: Dict[str, tuple[int | str, ...]] | None = None, observable_units: Dict[str, str] | None = None, required_observables: Sequence[str] | None = None, provided_observables: Sequence[str] | None = None)#: Model specification (forward, parametrization, regularization).

class mrax.data_struct.study.Study(experiments: Sequence[Experiment], datasets: Sequence[DataSet], models: Sequence[Model], names: Sequence[str] | None = None, study_globals: Dict[str, Quantity] | None = None)#

Collection of experiments/datasets/models that can be processed together.

iter_entries() → Iterator[tuple[str, Experiment, Model, DataSet]]#: Yield (name, experiment, model, dataset) tuples for the study.

execution_order(targets: Sequence[str] | None = None) → list[str]#: Return the study order (as provided), optionally filtered by targets.

describe(targets: Sequence[str] | None = None) → str#: Human-readable multi-line summary of the study configuration.

sanity_check(*, study_globals: Dict[str, Quantity] | None = None)#: Run per-experiment sanity checks for this study.

explain_sanity(*, study_globals: Dict[str, Quantity] | None = None) → str#: Return a combined sanity-check summary for all experiments.

Storage and Manifests #

Persist inputs/outputs for one or more experiments in an ISMRMRD-backed manifest. All files referenced in the manifest are created relative to the provided output_dir.

save_study_results(study, results, output_dir, ...) writes MR images, optional k-space, fitted parameter maps with units, plots when requested by the result payload, and optional NIfTI exports when save_nifti=True. Returns study_manifest.json.
load_study_results(manifest_path) reads the manifest plus all referenced arrays and reconstructs per-experiment dictionaries containing fitted maps, metadata, units, and optional plot paths.

mrax.data_struct.storage.save_study_results(study: Study, results: list[dict[str, Any]], output_dir: str | Path, study_name: str = 'study', targets: list[str] | None = None, file_extension: str = '.mrd', experiment_metadata: dict[str, Any] | None = None, save_nifti: bool = False)#: Persist study inputs/outputs to disk with a manifest that links files.

mrax.data_struct.storage.load_study_results(manifest_path: str | Path) → dict[str, Any]#: Load study manifest plus arrays from ISMRMRD files.

Fitting Pipeline #

Build and run differentiable losses with Optax/Proximal variants. All functions accept per-parameter regularizer configs/weights that override model defaults.

build_loss(experiment, model, dataset, ...) → (loss_fn, init_params, unpack) where loss_fn(params) returns a scalar data+regularization loss.
run_fit(experiment, model, dataset, optimizer_name="adam", step_size=0.02, num_steps=200, ...) fits a single experiment; returns (fitted_fields_with_units, loss_history) or (fitted_fields_with_units, loss_history, model_fn) when return_model_callable=True (model_fn(fields, offsets=None, observables=None) applies study globals + observables). Supports Optax optimizers plus split_bregman/fista wrappers, box constraints, proximal regularizers, and explicit initial parameters.
make_model_callable(experiment, model, study_globals=None) builds a callable that merges study globals/observables before invoking model.model_apply. Use it in notebooks to evaluate predictions without manually attaching field_strength_mhz.
plan_study_fits(study, target_names=None, precomputed_fields=None) returns an ordered list of mrax.core.study_runner.PlannedFit entries for sequential fitting, optionally reusing cached fields.
execute_study_plan(plan, run_fit_kwargs=None, run_fit_overrides=None, ...) executes the plan via run_fit() and returns result entries plus the fitted field cache. Experiments missing required observables are skipped by default and marked with skipped=True in the results. Put shared run_fit options in run_fit_kwargs and per-experiment options in run_fit_overrides; dict-valued run_fit options such as regularizer_weights are passed through unchanged.
check_fit_sanity(experiment, model, dataset, study_globals=None) runs a fast pre-fit check and returns a mrax.core.fit_sanity.FitSanityReport.

mrax.core.loss_builder.build_loss(experiment: Experiment, model: Model, dataset: DataSet, regularizer_configs: dict[str, list[dict]] | None = None, regularizer_type: str | dict[str, str] | None = None, regularizer_weights: dict[str, float] | None = None, mask: object | None = None, mask_dims: tuple[str, ...] | None = None, mask_fill_value: float | int | Array | None = 0.0, study_globals: dict[str, Quantity] | None = None, frame_mask: object | None = None, auto_mask_nan_observables: bool = True) → Tuple[Callable[[dict], Array], dict, Callable[[dict], dict]]#: Build loss(params) combining data misfit and model regularization.

mrax.core.loss_builder.run_fit(experiment: Experiment, model: Model, dataset: DataSet, optimizer_name: str = 'adam', step_size: float = 0.02, num_steps: int = 200, regularizer_weight=0.0, regularizer_prox=None, optimizer_base_factory=None, box_constraints=None, initial_params: dict | None = None, regularizer_configs: dict[str, list[dict]] | None = None, regularizer_type: str | dict[str, str] | None = None, regularizer_weights: dict[str, float] | None = None, optimizer_factory=None, optimizer=None, mask: object | None = None, mask_dims: tuple[str, ...] | None = None, mask_fill_value: float | int | Array | None = 0.0, study_globals: dict[str, Quantity] | None = None, frame_mask: object | None = None, auto_mask_nan_observables: bool = True, return_model_callable: bool = False, compute_derived_maps: bool | None = None, progress_callback: Callable[[int, int, str], None] | None = None)#

Fit the model to data for a given experiment.

Supports Optax optimizers plus split_bregman and fista with optional proximal regularizers, box constraints, explicit parameter initializations, and per-parameter regularizer configs/weights (override model defaults). Custom Optax optimizers can be injected via optimizer or optimizer_factory. frame_mask and auto_mask_nan_observables let you ignore frames with missing per-frame observables (e.g., incomplete sat_time/b1 grids). mask_fill_value controls how masked output fields are filled (use None to skip masking). NaN/inf entries in data or masks are treated as masked during fitting. compute_derived_maps controls derived map evaluation. When None, derived maps are automatically computed for diffusion models that provide derived_maps.

mrax.core.loss_builder.make_model_callable(experiment: Experiment, model: Model, study_globals: dict[str, Quantity] | None = None, *, allow_raw_observables: bool = False)#: Return a callable(fields, offsets=None, observables=None) that merges study globals.

mrax.core.loss_builder.build_study_losses(study: Study)#: Build per-experiment losses for a collection bundled as a Study.

mrax.core.study_runner.plan_study_fits(study: Study, target_names: Sequence[str] | None = None, precomputed_fields: dict[str, dict] | None = None) → list[PlannedFit]#: Return a sequential plan for a Study (optionally filtered/cached).

mrax.core.study_runner.execute_study_plan(plan: Sequence[PlannedFit], run_fit_kwargs: dict | None = None, run_fit_overrides: dict[str, dict] | None = None, *, skip_missing_observables: bool = True, warn_on_skip: bool = True)#

Execute a planned Study fit in order, returning result entries and field cache.

When skip_missing_observables is True, experiments missing required observables are skipped with an explanatory warning and a skipped result entry.

run_fit_kwargs are passed to every fit. For backwards compatibility, a value can be a mapping keyed by plan entry name to provide per-experiment values, except for dict-valued run_fit options such as regularizer_weights and study_globals. Use run_fit_overrides for explicit per-experiment keyword overrides.

mrax.core.fit_sanity.check_fit_sanity(experiment: Experiment, model: Model, dataset: DataSet, *, study_globals: Mapping[str, Quantity] | None = None) → FitSanityReport#: Validate experiment/model/dataset consistency before running a fit.

class mrax.core.fit_sanity.FitSanityReport(experiment_name: str, model_name: str, dataset_shape: tuple[int, ...], errors: list[str] = <factory>, warnings: list[str] = <factory>, notes: list[str] = <factory>, missing_observables: list[str] = <factory>, missing_any_of: list[tuple[str, ...]] = <factory>)#: Structured report describing whether a fit is ready to run.

Batch Execution #

Utilities for looping over many studies on disk and exporting consistent manifests.

discover_study_dirs(root, selected=None) returns study folder paths (optionally filtered by name) with helpful error messages when nothing is found.
run_batch_studies(studies_root, study_loader, output_root, ...) iterates through discovered study folders, converts each to a Study via study_loader, runs analysis (default: study planner + executor), then calls save_study_results.

mrax.core.batch.discover_study_dirs(studies_root: str | Path, selected: Sequence[str] | None = None) → list[Path]#: Return study directories to process, optionally filtering by name.

mrax.core.batch.run_batch_studies(studies_root: str | Path, study_loader: Callable[[Path], Study | tuple[Study, dict[str, Any] | None]], output_root: str | Path, selected: Sequence[str] | None = None, analysis_fn: Callable[[...], list[dict[str, Any]]] | None = None, analysis_kwargs: dict[str, Any] | None = None, save_kwargs: dict[str, Any] | None = None) → list[dict[str, Any]]#

Load studies from a folder, run analysis/fitting, and persist results.

The batch runner assumes no interdependencies between studies; it simply loops through them. Provide a study_loader that converts a study folder into a Study (plus optional metadata), and an optional analysis_fn to override the default Study planner/executor flow. Results are saved under output_root/<study_name> using save_study_results().

Returns a list of dictionaries summarizing each run with study_name, manifest_path, results, metadata, and study_dir keys.

Runtime Helpers #

mrax.core.runtime.release_memory(*, gc_collect: bool = True, clear_jax_cache: bool = True) → None#: Release Python/JAX caches to reduce peak memory in long workflows.

Preprocessing #

Denoising and motion-correction helpers for image stacks.

denoise_filter applies Gaussian/median/uniform filtering on N-D arrays.
regularized_denoise runs TV, higher-order TV, and smoothness regularizers using project optimizers; see [Rudin1992] for the TV background.
denoise_bm3d and denoise_bm4d wrap the bundled external denoiser packages [Dabov2007] [Maggioni2013].
linear_drift_correct_framestack and linear_drift_correct_dataset fit one linear transform-parameter model over a frame stack for long-term drift.
TranslationTransform, Rigid2DTransform, Similarity2DTransform, and Affine2DTransform provide differentiable JAX transforms for that drift fit.
register_image and motion_correct_framestack provide SimpleITK-backed image registration.
motion_correct_dataset handles MRax experiment dimensions; use register_experiment_to_reference to align an experiment to a reference experiment image dataset.
resample_image applies a SimpleITK transform to another image on a chosen reference image grid.

mrax.preprocessing.denoise.denoise_filter(array: Any, *, method: str = 'gaussian', sigma: float | tuple[float, ...] = 1.0, size: int | tuple[int, ...] = 3, mode: str = 'reflect', truncate: float = 4.0, preserve_nan: bool = False, complex_mode: str = 'real_imag') → Array#: Denoise with simple multidimensional filters (JAX implementation).

mrax.preprocessing.denoise.regularized_denoise(data: Any, *, regularizer: str | Callable | None = 'tv', weight: float = 0.1, num_steps: int = 50, step_size: float = 0.1, optimizer_name: str = 'adam', axes: tuple[int, ...] | None = None, mask: Any | None = None, init: Any | None = None, complex_mode: str = 'real_imag', return_history: bool = False) → Array | tuple[Array, list[float] | dict[str, list[float]]]#: Denoise using a regularizer optimized with the project optimizers.

mrax.preprocessing.denoise.denoise_bm3d(array: Any, *args, sigma: float | None = None, stage: str = 'full', **kwargs) → Array#

BM3D denoising via the external bm3d package.

Pass package-specific options through kwargs.

mrax.preprocessing.denoise.denoise_bm4d(array: Any, *args, sigma: float | None = None, stage: str = 'full', **kwargs) → Array#

BM4D denoising via the external bm4d package.

Pass package-specific options through kwargs.

mrax.preprocessing.drift.linear_drift_correct_framestack(frames: ~numpy.ndarray, *, transform: ~mrax.preprocessing.drift.ImageTransform | ~typing.Callable[[~jax.Array, ~jax.Array], ~jax.Array] | None = None, metric: ~typing.Callable[[~jax.Array, ~jax.Array], ~jax.Array | float] = <function normalized_cross_correlation>, reference: str | int | ~numpy.ndarray = 'first', times: ~typing.Iterable[float] | None = None, anchor_frame: int | None = None, identity_parameters: ~typing.Iterable[float] | None = None, fit_intercept: bool = False, frame_weights: ~typing.Iterable[float] | None = None, mask: ~numpy.ndarray | ~jax.Array | None = None, max_parameter_change: float | ~typing.Iterable[float] | None = None, optimizer_name: str = 'adam', step_size: float = 0.05, num_steps: int = 200, optimizer_factory=None, optimizer=None) → FrameStackLinearDriftResult#

Fit and apply a global linear transform-parameter drift to a frame stack.

The JAX optimizer does not fit one transform per frame. It optimizes the coefficients of identity + design(time) @ coefficients and evaluates the chosen image transform and similarity metric for each frame under that shared model. By default the line is anchored at identity for the reference frame or the first frame; set fit_intercept=True to fit a free line. mask restricts every objective and similarity evaluation to selected reference-grid pixels. Custom metrics and transforms used in this objective must be JAX-compatible; when mask is passed the metric must accept a mask= keyword.

mrax.preprocessing.drift.linear_drift_correct_dataset(experiment: Experiment, dataset: DataSet, *, dimensions: Iterable[str] | None = None, reference: str | int | ndarray = 'first', times: Iterable[float] | None = None, **drift_kwargs) → LinearDriftCorrectionResult#: Fit grouped linear drift corrections and restore MRax dataset dimensions.

mrax.preprocessing.drift.apply_image_transform_stack(frames: ndarray, transform: ImageTransform | Callable[[Array, Array], Array], parameters: ndarray) → ndarray#: Apply one per-frame image-transform parameter vector to a frame stack.

class mrax.preprocessing.drift.TranslationTransform(order: int | str = 'linear', mode: str = 'constant', cval: float = 0.0)#

N-D JAX subpixel translation in MRax array-axis order.

Parameters are shifts in array-axis order. For a 2D (x, y) MRax image, a parameter vector (dx, dy) moves the sampled image along those axes. order accepts "linear"/1 or "nearest"/0.

class mrax.preprocessing.drift.Rigid2DTransform(order: int | str = 'linear', mode: str = 'constant', cval: float = 0.0)#

2D translation plus rotation about the image center.

Parameters are (dx, dy, angle_radians) in MRax array-axis order. order accepts "linear"/1 or "nearest"/0.

class mrax.preprocessing.drift.Similarity2DTransform(order: int | str = 'linear', mode: str = 'constant', cval: float = 0.0)#

2D translation, rotation, and isotropic log-scale.

Parameters are (dx, dy, angle_radians, log_scale). order accepts "linear"/1 or "nearest"/0.

class mrax.preprocessing.drift.Affine2DTransform(order: int | str = 'linear', mode: str = 'constant', cval: float = 0.0)#

2D translation plus a small affine update about the image center.

Parameters are (dx, dy, a00, a01, a10, a11) and the affine matrix is eye(2) + [[a00, a01], [a10, a11]]. order accepts "linear"/1 or "nearest"/0.

mrax.preprocessing.motion.register_image(fixed: ndarray, moving: ndarray, *, registration_method: Any | Callable[[Any, Any], Any] | None = None, initial_transform: Any | Callable[[Any, Any], Any] | None = None, transform_in_place: bool = False, interpolator: str | Any = 'linear', default_value: float = 0.0, output_pixel_type: Any | None = None) → ImageRegistrationResult#

Without registration_method this uses a rigid Mattes mutual-information registration. Pass a configured SimpleITK.ImageRegistrationMethod or a factory accepting (fixed_image, moving_image) to use any SimpleITK registration configuration. initial_transform can likewise be a transform or a factory when it should be set by the wrapper.

mrax.preprocessing.motion.motion_correct_framestack(frames: ndarray, *, reference: str | int | ndarray = 'first', **registration_kwargs) → FrameStackMotionResult#: Register a (frame, *spatial) stack to one reference image.

mrax.preprocessing.motion.motion_correct_dataset(experiment: Experiment, dataset: DataSet, *, dimensions: Iterable[str] | None = None, reference: str | int | ndarray = 'first', **registration_kwargs) → MotionCorrectionResult#

Motion-correct MRax dataset data while preserving experiment dimensions.

With dimensions=None all non-spatial experiment dimensions are collapsed into one frame stack in experiment order. When dimensions names a subset such as ("t",), each combination of the remaining dimensions receives a separate registration stack before the corrected data is expanded back.

mrax.preprocessing.motion.register_experiment_to_reference(reference_experiment: Experiment, reference_dataset: DataSet, moving_experiment: Experiment, moving_dataset: DataSet, *, reference_frame: str | int | ndarray = 'first', moving_frame: str | int | ndarray = 'first', interpolator: str | Any = 'linear', default_value: float = 0.0, output_pixel_type: Any | None = None, **registration_kwargs) → ExperimentRegistrationResult#

The transform is estimated from moving_frame to reference_frame and then applied to every moving frame. The returned experiment keeps the moving non-spatial dimensions and adopts the reference spatial grid.

mrax.preprocessing.motion.resample_image(moving: ndarray, reference: ndarray, transform, *, interpolator: str | Any = 'linear', default_value: float = 0.0, output_pixel_type: Any | None = None) → ndarray#: Resample one image onto a reference image grid with a SimpleITK transform.

Post-fit Analysis #

Utilities for extracting ROI statistics, ROI timecourses, and performing group-level comparisons on batch outputs.

ROI defines an ROI mask with optional dimension names for broadcasting.
collect_roi_stats produces a long-form table of ROI summaries (mean/variance/etc.).
extract_timecourse reduces an image stack to an ROI timecourse.
analyze_groups builds subject-level metrics, group summaries, and pairwise statistics.

class mrax.analysis.roi.ROI(name: str, mask: ndarray, dims: tuple[str, ...] | None=None, metadata: dict[str, ~typing.Any]=<factory>)#: Region-of-interest definition with an optional dimension label list.

mrax.analysis.roi.collect_roi_stats(maps: Mapping[str, Any], rois: Sequence[ROI], *, stats: Iterable[str] = ('mean', 'variance'), ddof: int = 1, target_dims: Sequence[str] | None = None, metadata: Mapping[str, Any] | None = None) → list[dict[str, Any]]#

Collect ROI statistics across multiple parameter maps.

Returns a long-form list of rows with keys: map, roi, stat, value, unit (optional).

mrax.analysis.roi.match_spatial_mask(mask: Any | None, spatial_shape: Sequence[int], *, allow_transpose: bool = True) → ndarray | None#: Validate or transpose a 2D/3D mask to match a spatial shape.

mrax.analysis.timecourse.extract_timecourse(array: Any, roi: ROI | Any, *, stat: str = 'mean', spatial_axes: Sequence[int] | None = None, mask_dims: Sequence[str] | None = None, target_dims: Sequence[str] | None = None) → Any#: Extract a timecourse by reducing over ROI spatial axes.

mrax.analysis.timecourse.extract_timecourse_from_experiment(data: Any, experiment: Experiment, roi: ROI, *, stat: str = 'mean') → dict[str, Any]#: Extract a timecourse with dimension metadata from an Experiment.

mrax.analysis.group.analyze_groups(groups: Mapping[str, Sequence[str | Path]], *, experiment_name: str, rois: Sequence[ROI], map_names: Sequence[str] | None = None, stats: Sequence[str] = ('mean', 'variance'), ddof: int = 1, output_root: str | Path | None = None, paired: bool = False, equal_var: bool = False, nonparametric: bool = True, p_adjust: str | None = None) → dict[str, Any]#: End-to-end group analysis: subject stats, group summaries, and comparisons.

mrax.analysis.group.collect_subject_roi_stats(groups: Mapping[str, Sequence[str | Path]], *, experiment_name: str, rois: Sequence[ROI], map_names: Sequence[str] | None = None, stats: Sequence[str] = ('mean', 'variance'), ddof: int = 1, output_root: str | Path | None = None) → list[dict[str, Any]]#: Collect per-subject ROI statistics across groups.

mrax.analysis.group.summarize_group_stats(subject_rows: Sequence[Mapping[str, Any]], *, ddof: int = 1) → list[dict[str, Any]]#: Summarize subject-level ROI stats within each group.

mrax.analysis.group.compare_group_stats(subject_rows: Sequence[Mapping[str, Any]], *, paired: bool = False, equal_var: bool = False, nonparametric: bool = True, p_adjust: str | None = None) → list[dict[str, Any]]#: Compute pairwise group comparisons for each map/ROI/stat.

Synthetic Examples #

Synthetic datasets that use the same containers as imported studies.

make_example_suite(noise_level=0.05, seed=0, model_names=("t1_ir","t2_se")) → list of mrax.core.synthetic.Example objects.
make_synthetic_study(...) → (Study, metadata); the metadata includes ground truth fields for benchmarking.

class mrax.core.synthetic.Example(name: 'str', dataset: 'DataSet', experiment: 'Experiment', model: 'Model', truth_fields: 'dict | None' = None)#

mrax.core.synthetic.make_example_suite(noise_level: float = 0.05, seed: int = 0, model_names: Sequence[str] | None = None) → List[Example]#

mrax.core.synthetic.make_synthetic_study(noise_level: float = 0.05, seed: int = 0, model_names: Sequence[str] | None = None, study_globals: dict[str, Quantity] | None = None) → Tuple[Study, dict]#: Build a Study from synthetic examples to mirror importer outputs.

Optimization and Regularization #

Core primitives for solvers and penalties. All regularizer functions operate on ndarrays and return scalar costs.

make_optimizer(optimizer_name, step_size, regularizer_weight, regularizer_prox, box_constraints) resolves Optax optimizers plus proximal split_bregman/fista variants.
spatial_smoothness_penalty and total_variation_penalty operate along configurable axes. higher_order_total_variation_penalty adds finite-difference TV of configurable order.
Low-level ops like soft_threshold, isotropic_shrinkage, gradient/divergence helpers, and conjugate_gradient are available for custom solvers.

mrax.opt.optimizers.make_optimizer(optimizer_name: str, step_size: float, regularizer_weight=0.0, regularizer_prox: Callable[[Array, float], Array] | None = None, base_optimizer_factory: Callable[[float], GradientTransformation] | None = None, box_constraints: Tuple[object, object] | None = None) → GradientTransformation#

Single entry point to build optimizers, with optional proximal variants.

split_bregman / fista support arbitrary proximal regularizers via regularizer_prox with per-parameter weights (scalars, arrays, or pytrees) and optional box constraints.
All other optimizers are resolved from Optax and only support L2 weight decay plus optional box constraints.

mrax.opt.regularizers.spatial_smoothness_penalty(field: Array, axes: tuple[int, ...] | None = None, mask: Array | None = None) → Array#: Encourage smooth parameter maps via finite-difference gradients.

mrax.opt.regularizers.total_variation_penalty(field: Array, axes: tuple[int, ...] | None = None, mask: Array | None = None) → Array#: Isotropic TV penalty across selected axes.

mrax.opt.regularizers.higher_order_total_variation_penalty(field: Array, order: int = 2, axes: tuple[int, ...] | None = None, mask: Array | None = None) → Array#: Higher-order TV using repeated forward differences.

mrax.opt.regularizers.make_regularizer_fn(reg_type: str, axes: tuple[int, ...] | None = None)#

mrax.opt.regularizers.estimate_field_scales(fields: Mapping[str, object], *, mask: ndarray | None = None, method: str = 'median', min_scale: float = 1e-06) → dict[str, float]#: Estimate per-field scales for regularizer normalization.

mrax.opt.regularizers.estimate_param_scales(model, field_shape: tuple[int, ...], *, mask: ndarray | None = None, method: str = 'median', min_scale: float = 1e-06, param_specs: Mapping[str, Mapping[str, object]] | None = None) → dict[str, float]#: Estimate parameter scales from a model’s initialization or param specs.

mrax.opt.regularizers.autoscale_regularizer_configs(regularizer_configs: Mapping[str, list[dict]] | None, param_scales: Mapping[str, float], *, type_exponents: Mapping[str, float] | None = None, min_scale: float = 1e-06) → dict[str, list[dict]] | None#: Scale regularizer weights by parameter magnitude.

mrax.opt.ops.soft_threshold(x: Array, threshold: float) → Array#: Element-wise soft thresholding (supports complex inputs).

mrax.opt.ops.isotropic_shrinkage(v: Array, threshold: float) → Array#: Vector soft thresholding used for isotropic total variation.

mrax.opt.ops.forward_gradient2d(x: Array) → Array#: Forward finite differences with zero Neumann boundary.

mrax.opt.ops.divergence2d(v: Array) → Array#: Adjoint of forward_gradient2d (negative divergence).

mrax.opt.ops.conjugate_gradient(matvec: Callable[[Array], Array], b: Array, x0: Array | None = None, max_iters: int = 25, tol: float = 1e-06) → Array#: JAX-compatible conjugate gradient solver for SPD systems.

Reconstruction #

Vendor-agnostic reconstruction utilities for both Cartesian and non-Cartesian k-space.

RecoPipeline (plus stages like FFTStage/POCSStage) provides a compact, NumPy-first pipeline (accepting JAX arrays) used by importers when turning raw k-space into images.
reconstruct_ifft provides zero-filled reconstruction for Cartesian k-space.
reconstruct_regrid provides a non-Cartesian gridding preview and optional vendor-style post-processing controls for oversampled grids, crop offsets, deapodization, support masks, RSS coil combination, and magnitude output.
mrax.reco.processing exposes importer-independent k-space preprocessing blocks: receiver scale/phase factors, per-sample complex corrections, FID/readout transforms, readout crops, radial density estimates, and fractional readout phase shifts.
reconstruct_variational solves variational reconstruction problems with Optax and proximal variants, for example TV regularization via FISTA plus a TV proximal operator [BeckTeboulle2009].
reconstruct_model_based reconstructs quantitative parameter maps directly from k-space by chaining a mrax.data_struct.model.Model forward model with the MRI encoding operator and, when needed, a reference-image term.
jax-finufft from the optional mrax[nufft] extra is used when you supply non-Cartesian trajectories.

class mrax.reco.pipeline.RecoPipeline(stages: Iterable[RecoStage] | None = None)#: Chain of reconstruction stages.

class mrax.reco.pipeline.RecoState(kspace: ndarray, image: ndarray | None = None, trajectory: ndarray | None = None, density: ndarray | None = None, mask: ndarray | None = None, metadata: dict[str, ~typing.Any]=<factory>)#: State flowing through the reconstruction pipeline.

mrax.reco.pipeline.cartesian_fft(array: ndarray, axes: Sequence[int] = (-2, -1)) → ndarray#: Centered FFT with orthonormal scaling.

mrax.reco.pipeline.cartesian_ifft(array: ndarray, axes: Sequence[int] = (-2, -1)) → ndarray#: Centered inverse FFT with orthonormal scaling.

mrax.reco.utils.estimate_sensitivities_cartesian(kspace: ndarray | Array, *, frame: int = 0, eps: float = 1e-08) → ndarray | None#: Estimate coil sensitivity maps from Cartesian k-space via RSS normalization.

mrax.reco.utils.pick_sensitivity_frame(kspace: ndarray | Array, *, strategy: Literal['first', 'max_energy'] | str = 'max_energy') → int#: Choose a reference frame for sensitivity estimation.

mrax.reco.utils.pocs_complete_kspace(kspace: ndarray | Array, *, iterations: int = 20, enforce_real: bool = True, mask: ndarray | None = None) → ndarray#: Fill missing Cartesian k-space via POCS (useful for partial Fourier).

mrax.reco.processing.apply_channel_scale_phase(data: ndarray, *, scale: Sequence[float] | ndarray | None = None, phase_degrees: Sequence[float] | ndarray | None = None, coil_axis: int = 1) → ndarray#

Apply per-coil scale factors and phase offsets to k-space or coil images.

Parameters#

data:: Array with one coil/channel dimension. Typical shapes are (frame, coil, sample) for non-Cartesian streams or (frame, coil, *spatial) for Cartesian grids and coil images.
scale:: Optional real multiplicative factors, one per coil.
phase_degrees:: Optional phase offsets in degrees, one per coil. Positive values multiply the corresponding coil by exp(1j * phase).
coil_axis:: Axis containing the coil/channel dimension.

Returns#

np.ndarray: Corrected array. Inputs are not modified in-place.

Notes#

Use this for receiver normalization from any vendor. Bruker maps RecoScaleChan and RecoPhaseChan to these arguments; other importers can map equivalent receiver gains and phase references the same way.

mrax.reco.processing.apply_samplewise_correction(data: ndarray, correction: Sequence[complex] | ndarray, *, sample_axis: int = -1) → ndarray#

Multiply a k-space stream by a per-sample complex correction.

The correction can be a one-dimensional vector matching sample_axis or an already-broadcastable array. This is useful for B0/demodulation phase corrections, trajectory-specific scaling, or calibration vectors derived outside MRax.

mrax.reco.processing.estimate_radial_density(trajectory: ndarray, *, spatial_shape: tuple[int, ...] | None = None, power: float | None = None) → ndarray#

Estimate radial density compensation weights from a 2D or 3D trajectory.

Parameters#

trajectory:: (sample, D) array in normalized k-space coordinates, usually cycles/pixel.
spatial_shape:: Optional output image shape. When supplied, weights are normalized to sum to the number of spatial pixels/voxels. Otherwise they are normalized to mean 1.
power:: Radial power. Defaults to 1 for 2D trajectories and 2 for 3D trajectories, matching common rho and rho^2 density compensation.

Notes#

This is a simple analytic estimate. Prefer sequence-provided or iteratively optimized density weights when they are available and validated.

mrax.reco.processing.fid_profiles_to_kspace(profiles: ndarray, *, axis: int = -1, transform: Literal['none', 'fft', 'ifft', 'fft_ifftshift', 'ifft_ifftshift']='fft_ifftshift', output_points: int | None = None, crop_offset: int | None = None, readout_shift: float | None = None, dtype: dtype | type | None = <class 'numpy.complex64'>) → ndarray#

Transform time-domain readout profiles into frequency-domain k-space profiles.

Parameters#

profiles:: Input readout profiles. The selected axis is interpreted as the readout time/frequency axis.
transform:: Readout transform to apply before cropping. "fft_ifftshift" matches the common pattern where a centered FID-like readout is shifted to FFT order and transformed along the readout.
output_points, crop_offset:: Optional crop on the transformed readout axis. crop_offset=None means centered crop.
readout_shift:: Optional fractional sample shift applied after the crop as a phase ramp.
dtype:: Optional output dtype. Use None to keep NumPy’s transform dtype.

Returns#

np.ndarray: Frequency-domain readout profiles.

Notes#

This is not Bruker-specific. It is useful whenever a sequence stores raw FIDs or oversampled readouts that must be Fourier-transformed, cropped, and phase shifted before samples are paired with a trajectory.

mrax.reco.processing.apply_readout_phase_shift(data: ndarray, shift_samples: float, *, axis: int = -1) → ndarray#

Apply a fractional readout shift as a linear phase ramp.

shift_samples is expressed in samples over the selected axis. Positive values use exp(+2j*pi*shift*n/N). Pass a negative shift when your convention requires the inverse ramp.

mrax.reco.processing.crop_along_axis(data: ndarray, output_size: int, *, axis: int = -1, offset: int | None = None) → ndarray#: Crop one axis of an array, centered by default.

mrax.reco.processing.readout_crop_bounds(input_size: int, output_size: int, offset: int | None = None) → tuple[int, int]#

Return (start, stop) bounds for a readout crop.

offset=None performs a centered crop. Explicit offsets are useful when a vendor reconstruction records an asymmetric cutoff.

class mrax.reco.variational.MRIEncodingOperator(spatial_shape: tuple[int, int], sensitivities: Array | None = None, mask: Array | None = None, trajectory: Array | None = None, density: Array | None = None, axes: tuple[int, int] = (-2, -1), nufft_eps: float = 1e-06)#

Encoding operator A for multi-coil MRI: y = M F S x (Cartesian) or NUFFT (non-Cartesian).

forward(image: Array) → Array#: Apply A(image) -> k-space (shape: batch..., coils, *spatial or batch..., coils, num_samples).

adjoint(kspace: Array) → Array#: Apply A^H(kspace) -> combined image (shape: batch..., *spatial).

class mrax.reco.variational.RegridApodizationCorrection(kernel_width: float, shape_parameter: float, full_shape: tuple[int, ...] | None = None, crop_offset: tuple[int, ...] | None = None, strength: float = 1.0)#

Kaiser-Bessel deapodization settings for gridding-style reconstructions.

The correction is applied in the image domain after optional cropping. full_shape and crop_offset describe the grid coordinate system used to evaluate the separable correction when the reconstruction is cropped from an oversampled grid.

mrax.reco.variational.build_operator(*, spatial_shape: tuple[int, int] | None = None, kspace: Array | ndarray | None = None, sensitivities: Array | ndarray | None = None, mask: Array | ndarray | None = None, trajectory: Array | ndarray | None = None, density: Array | ndarray | None = None, axes: Sequence[int] = (-2, -1), nufft_eps: float = 1e-06) → MRIEncodingOperator#: Build a vendor-agnostic MRI encoding operator from generic inputs.

mrax.reco.variational.center_crop_spatial(image: Array | ndarray, output_shape: tuple[int, ...], crop_offset: tuple[int, ...] | None = None) → Array#

Crop the last len(output_shape) axes of an image or coil-image stack.

Use this after reconstructing on an oversampled grid. crop_offset=None performs a centered crop; explicit offsets reproduce vendor cutoff windows or asymmetric grids.

mrax.reco.variational.checkerboard_phase(shape: tuple[int, ...]) → Array#

Return a +1/-1 checkerboard phase array for spatial image axes.

This is useful after FFT-style reconstructions when a scanner convention or gridding path leaves the image on the opposite centered-grid parity. It is the public version of the phase step optionally used by reconstruct_regrid().

mrax.reco.variational.apply_regrid_apodization_correction(image: Array | ndarray, correction: RegridApodizationCorrection, spatial_shape: tuple[int, ...]) → Array#

Apply separable Kaiser-Bessel deapodization to a gridded reconstruction.

Apply this only when the gridding kernel matches the correction. For example, a vendor reconstruction using a Kaiser-Bessel convolution kernel can require this image-domain correction; applying it to nearest-neighbor gridding is usually only a diagnostic experiment.

mrax.reco.variational.elliptical_support_mask(shape: tuple[int, ...], radius: float = 1.0) → Array#: Return a circular/elliptical support mask for 2D or 3D image shapes.

mrax.reco.variational.apply_image_support_mask(image: Array | ndarray, image_mask: str | Array | ndarray | None, spatial_shape: tuple[int, ...]) → Array#

Apply an image-domain support mask to the last spatial axes.

image_mask may be "ellipse"/"circle" or an explicit mask matching spatial_shape. Passing None returns the input unchanged.

mrax.reco.variational.combine_coil_images(coil_images: Array | ndarray, *, method: Literal['sum', 'rss', 'sense'] = 'sense', sensitivities: Array | ndarray | None = None, coil_axis: int | None = None, eps: float = 1e-08) → Array#

Combine coil images with sum, root-sum-of-squares, or SENSE weighting.

Parameters#

coil_images:: Array containing a coil axis and trailing spatial axes.
method:: "sum" for direct complex summation, "rss" for magnitude RSS, or "sense" for sensitivity-weighted combination. "sense" falls back to sum when sensitivities are not supplied, matching reconstruct_regrid().
sensitivities:: Optional (coil, *spatial) sensitivity map for SENSE combination.
coil_axis:: Axis containing coils. If omitted, MRax assumes (frame, coil, *spatial) for non-SENSE data and infers the axis from sensitivities for SENSE data.

mrax.reco.variational.reconstruct_ifft(kspace: Array | ndarray, *, sensitivities: Array | ndarray | None = None, mask: Array | ndarray | None = None, axes: Sequence[int] = (-2, -1)) → Array#: Direct (zero-filled) reconstruction for Cartesian k-space via \(A^H y\).

Quick non-Cartesian reconstruction via regridding + IFFT (not iterative).

If density is omitted, a simple radial density compensation (\(|k|\)) normalized to sum to the number of spatial pixels is applied. Pass density=np.ones(num_samples) to disable.

Optional arguments expose common vendor-style gridding post-processing without changing the default MRax path:

grid_shape can reconstruct on an oversampled grid before cropping to spatial_shape.
phase_correction="checkerboard" applies a \(\pi\) phase alternation on all spatial image axes.
apodization_correction applies a separable Kaiser-Bessel deapodization.
image_mask="ellipse" masks the final image support.
coil_combine="rss" uses root-sum-of-squares coil combination.

mrax.reco.variational.reconstruct_variational(kspace: Array | ndarray, *, spatial_shape: tuple[int, int] | None = None, sensitivities: Array | ndarray | None = None, mask: Array | ndarray | None = None, trajectory: Array | ndarray | None = None, density: Array | ndarray | None = None, frame_shifts: Array | ndarray | None = None, frame_axis: int = 0, shift_order: int = 1, shift_mode: str = 'constant', shift_cval: float = 0.0, optimizer_name: str = 'fista', step_size: float = 0.5, num_steps: int = 50, regularizer: Literal['none', 'l2', 'l1', 'tv'] = 'tv', regularizer_weight: float = 0.01, tv_iterations: int = 30, nufft_eps: float = 1e-06, initial_image: Array | ndarray | None = None, return_history: bool = False, grad_clip_norm: float | None = None) → tuple[Array, list[float]] | Array#

Solve a variational reconstruction with Optax/proximal variants.

Minimizes: \(\tfrac12 \|\|A x - y\|\|_2^2 + R(x)\)

Cartesian operators use centered FFT and optional sampling masks.
Non-Cartesian operators use FINUFFT (jax_finufft) when trajectory is provided.
TV regularization uses a proper proximal operator (Chambolle) when optimizer_name is "fista" or "split_bregman"; otherwise TV is treated as a differentiable penalty.
If density is omitted for non-Cartesian data, a simple radial density compensation (\(|k|\)) normalized so the weights sum to the number of spatial pixels is applied to the data term. Pass density=np.ones(num_samples) to disable.
If frame_shifts is provided, the data term is built from per-frame shifted images: y_i = A(shift(x, shift_i)). Shifts are specified in index units (row, col). In this mode, kspace must include a frame axis.
For proximal solvers, the effective shrinkage threshold is regularizer_weight * step_size.
Optional gradient clipping helps stabilize long iterative runs.

mrax.reco.variational.reconstruct_model_based(kspace: Array | ndarray, *, experiment: Experiment, model: Model, spatial_shape: tuple[int, int] | None = None, sensitivities: Array | ndarray | None = None, mask: Array | ndarray | None = None, completion_kspace: Array | ndarray | None = None, completion_weight: float = 0.0, trajectory: Array | ndarray | None = None, density: Array | ndarray | None = None, optimizer_name: str = 'adam', step_size: float = 0.02, num_steps: int = 200, regularizer_weight=0.0, regularizer_prox=None, optimizer_base_factory=None, box_constraints=None, initial_params: dict | None = None, regularizer_configs: dict[str, list[dict]] | None = None, regularizer_type: str | dict[str, str] | None = None, regularizer_weights: dict[str, float] | None = None, roi_mask: object | None = None, roi_mask_dims: tuple[str, ...] | None = None, study_globals: dict[str, Quantity] | None = None, frame_mask: object | None = None, auto_mask_nan_observables: bool = True, nufft_eps: float = 1e-06, reference_image: Array | ndarray | Quantity | None = None, reference_kspace: Array | ndarray | None = None, reference_mask: Array | ndarray | None = None, reference_weight: float = 1.0, initial_reference_image: Array | ndarray | None = None, frame_phase_correction: str | None = 'none', return_history: bool = False, return_model_callable: bool = False) → tuple[dict[str, Quantity], list[float], Callable[[dict, Array | None, dict[str, Any] | None], Array]] | tuple[dict[str, Quantity], list[float]] | dict[str, Quantity]#

Directly reconstruct parameter maps from k-space via a chained model+encoding operator.

Minimizes: \(\tfrac12 \|\|A(M(\theta)) - y\|\|_2^2 + R(\theta)\)

Notes#

The signal model comes from mrax.data_struct.model.Model and is treated as a (generally non-linear) operator mapping parameter fields -> image series.
The encoding operator \(A\) is the same as in reconstruct_variational() (FFT/NUFFT, optional sampling mask, optional coil sensitivities).
roi_mask weights regularizers and masks output fields; when provided, the model prediction is also zeroed outside the mask to prevent NaN observables from propagating through the k-space data term.
frame_mask (or auto_mask_nan_observables=True) can drop frames with missing per-frame observables; masked frames do not contribute to the data term.
reference_image / reference_kspace can be used to model acquisitions where the signal model predicts a relative image series (e.g., Z-spectra). In that case the encoded prediction becomes \(A(reference\_image \cdot M(\theta))\). When reference_kspace is provided, the reference image is reconstructed jointly by adding a second data term \(reference\_weight \|\|A_{ref}(reference\_image) - reference\_kspace\|\|^2\).
completion_kspace / completion_weight can add an optional prior on the missing k-space samples (i.e., outside mask) by penalizing the difference to a filled k-space estimate (useful for partial Fourier / POCS-style completion). To match the weighting of a full-data mean-squared-error objective, set completion_weight=(1-f)/f where \(f=\mathrm{mean}(mask)\).
frame_phase_correction can apply a global phase alignment per frame (recommended for magnitude-only models when raw k-space contains frame-dependent phase flips).

Import Helpers #

Import definitions #

Normalize vendor metadata, infer experiment types, extract observables, and materialize vendor-neutral experiments/datasets from importer payloads.

Normalized acquisition metadata used for experiment typing and observables.

Vendor-specific importers should translate scanner metadata into this schema before creating MRax Experiment and DataSet objects. The fields are limited to concepts the internal factory understands; raw scanner dictionaries should remain importer metadata and not leak into downstream model code.

class mrax.io.import_definitions.ImportDefinition(name: str, data: ~numpy.ndarray, frame_count: int, spatial_shape: tuple[int, ...], exp_type: str, observables: ~typing.Mapping[str, ~mrax.data_struct.quantity.Quantity] = <factory>, processed: ~typing.Mapping[str, object] | None = None, kspace: ~numpy.ndarray | ~mrax.io.import_definitions.KSpaceImportDefinition | None = None, trajectory: ~numpy.ndarray | None = None, density: ~numpy.ndarray | None = None, sampling_mask: ~numpy.ndarray | None = None, kspace_kind: str = 'cartesian', sequence_type: str = 'standard', kspace_metadata: ~typing.Mapping[str, object] = <factory>)#: Vendor-neutral payload used to materialize MRax experiment and dataset objects.

class mrax.io.import_definitions.KSpaceImportDefinition(samples: ndarray, kind: str = 'cartesian', sequence_type: str = 'standard', trajectory: ndarray | None = None, density: ndarray | None = None, sampling_mask: ndarray | None = None, spatial_shape: tuple[int, ...] | None=None, sample_axis: int = -1, metadata: Mapping[str, object]=<factory>)#

Vendor-neutral raw k-space payload.

kind is one of "cartesian", "epi", or "non_cartesian". Cartesian and EPI data use grid-shaped k-space with spatial axes last. Non-Cartesian data uses samples with the sample axis last and a matching trajectory array with columns (kx, ky) in cycles/pixel. sequence_type describes the sequence- specific acquisition ordering used to interpret the samples (for example "rare" for Turbo Spin Echo / RARE style line ordering).

mrax.io.import_definitions.build_imported_experiment_dataset(definition: ImportDefinition) → tuple[Experiment, DataSet]#: Build MRax containers from a normalized import definition.

mrax.io.import_definitions.build_cartesian_kspace_definition(kspace: ndarray, *, sequence_type: str = 'standard', sampling_mask: ndarray | None = None, spatial_shape: tuple[int, ...] | None = None, metadata: Mapping[str, object] | None = None) → KSpaceImportDefinition#: Create a Cartesian k-space definition.

mrax.io.import_definitions.build_epi_kspace_definition(kspace: ndarray, *, samples_per_line: int | None = None, line_count: int | None = None, reverse_odd_lines: bool = True, sequence_type: str = 'standard', sampling_mask: ndarray | None = None, spatial_shape: tuple[int, ...] | None = None, metadata: Mapping[str, object] | None = None) → KSpaceImportDefinition#: Create an EPI k-space definition, optionally reshaping serial ADC samples.

mrax.io.import_definitions.build_non_cartesian_kspace_definition(samples: ndarray, *, trajectory: ndarray, density: ndarray | None = None, spatial_shape: tuple[int, ...] | None = None, sample_axis: int = -1, sequence_type: str = 'standard', metadata: Mapping[str, object] | None = None) → KSpaceImportDefinition#: Create a non-Cartesian k-space definition with trajectory and density.

mrax.io.import_definitions.detect_experiment_type(meta: NormalizedMetadata) → str#: Infer experiment type from normalized acquisition metadata.

mrax.io.import_definitions.extract_observables_from_metadata(exp_type: str, meta: NormalizedMetadata, frame_count: int) → dict[str, Quantity]#: Collect common observables (TE/TR/TI/flip/offsets) from normalized metadata.

mrax.io.import_definitions.reshape_epi_samples(samples: ndarray, *, samples_per_line: int, line_count: int | None = None, reverse_odd_lines: bool = True) → ndarray#

Reshape serial EPI ADC samples into (..., ky, kx) order.

Pulseq and vendor raw streams commonly store EPI data in acquisition order: line 0 left-to-right, line 1 right-to-left, and so on. Setting reverse_odd_lines=True reverses the readout direction on odd ky lines so the returned grid has a consistent kx ordering.

mrax.io.import_definitions.reshape_raw_kspace(samples: ndarray, *, frame_count: int, spatial_shape: tuple[int, ...], coil_count: int | None = None, is_multi_echo: bool = False, enc_steps: Sequence[int] | None = None, target_ky: int | None = None) → ndarray#: Reshape serial complex samples into k-space, applying optional ordering hints.

Pulseq Importers #

Read Pulseq .seq files, calculate ADC trajectories when pypulseq is installed, and combine sequence definitions with supplied k-space or image arrays. This importer is intended for synthetic demos and importer-development experiments; it has not been validated on measured scanner data.

class mrax.io.pulseq_import.PulseqSequenceInfo(path: Path, definitions: Mapping[str, object]=<factory>, trajectory: ndarray | None = None, trajectory_physical: ndarray | None = None, adc_times: ndarray | None = None, adc_counts: tuple[int, ...]=(), sequence_kind: str = 'unknown', sequence_type: str = 'standard', spatial_shape: tuple[int, int] | None=None, normalized_metadata: NormalizedMetadata = <factory>)#: Parsed Pulseq sequence metadata and ADC trajectory.

mrax.io.pulseq_import.read_pulseq(seq_path: str | Path, *, normalize_trajectory: bool = True) → PulseqSequenceInfo#: Read a Pulseq .seq file and return definitions plus ADC trajectory.

mrax.io.pulseq_import.read_pulseq_sidecar_metadata(seq_path: str | Path) → dict[str, object]#: Read optional Pulseq sidecar metadata from JSON files next to the sequence.

Build MRax containers from a Pulseq sequence and optional acquired data.

Pulseq files define the sequence and ADC trajectory, but not measured signal samples. Supply either kspace or image_data; if only kspace is provided, a quick preview image is reconstructed for DataSet.data.

mrax.io.pulseq_import.detect_pulseq_kspace_kind(trajectory: ndarray | None, adc_counts: Sequence[int] = ()) → str#: Infer whether a Pulseq ADC trajectory is Cartesian, EPI, or non-Cartesian.

Bruker Importers and Models #

Parse ParaVision folders, infer experiment types, extract observables, and map them to models. The Bruker importer has been tested on ParaVision 360 3.6 data. Other ParaVision versions and sequence families have not been validated in MRax.

load_bruker_experiment(exp_dir, override_type=None, model_map_by_type=None, study_globals=None, pdata_index="1", recipe=None) returns an (Experiment, DataSet, Model, metadata) tuple with parsed JCAMP dictionaries.
load_bruker_study(study_root, ..., pdata_index="1", recipe=None) walks a subject folder, loads valid scans, and returns (Study, metadata); scans are kept as a flat study.
prepare_bruker_cartesian_kspace(kspace, reco, ...) applies Bruker reco metadata that act directly in k-space and returns a reconstruction-ready Cartesian grid plus a summary of applied steps.
prepare_bruker_non_cartesian_kspace(samples, reco) applies Bruker receiver scale and phase corrections to non-Cartesian (frame, coil, sample) streams, while read_bruker_regrid_trajectory(exp_dir, ...) reads ParaVision RecoRegridN trajectory files and read_bruker_regrid_density(exp_dir, ...) derives the matching ParaVision density compensation weights for 2D spiral/radial and 3D UTE/ZTE raw-data imports. bruker_vendor_regrid_options(reco, ...) converts ParaVision image-side regridding steps into optional MRax reconstruct_regrid arguments for vendor-comparison paths.
reconstruct_bruker_cartesian_kspace(kspace, reco, ...) reconstructs a vendor-like magnitude image from Bruker Cartesian k-space, and format_bruker_cartesian_image(image, reco, ...) applies the corresponding post-FFT formatting steps to an already reconstructed complex image.

pdata_index accepts a single index (default "1") or a dict mapping scan numbers to indices; when a dict is used, unspecified scans fall back to "1" (or the value under "default"). The optional recipe callback can override the detected experiment type, inject observables, and choose a default model based on raw Bruker metadata.

Model factories in mrax.library.default_models cover the built-in model families listed in Models. Each returns a mrax.data_struct.model.Model instance. Importer-detected experiment types can be remapped with model_map_by_type or recipe.

mrax.io.bruker_import.load_bruker_experiment(exp_dir: Path, override_type: str | None = None, model_map_by_type: Dict[str, Model | Callable[[], Model]] | None = None, model_map_by_name: Dict[str, Model | Callable[[], Model]] | None = None, study_globals: Dict[str, Quantity] | None = None, pdata_index: str | int | Dict[str | int, str | int] = '1', recipe: Callable[[dict, dict, int], str | Tuple[str | None, dict | None, Model | Callable[[], Model] | None]] | None = None, load_raw_kspace: bool = False, reco_pipeline: RecoPipeline | Callable[[dict], RecoPipeline] | None = None, prefer_raw_image: bool = False, rotate_bvecs_to_image: bool = True, slice_crop_mode: str | None = None) → Tuple[Experiment, DataSet, Model, dict]#

mrax.io.bruker_import.load_bruker_study(study_root: str | Path, type_overrides: Dict[str, str] | None = None, model_map_by_type: Dict[str, Model | Callable[[], Model]] | None = None, model_map_by_name: Dict[str, Model | Callable[[], Model]] | None = None, study_globals: Dict[str, Quantity] | None = None, field_strength_mhz: float | Quantity | None = None, pdata_index: str | int | Dict[str | int, str | int] = '1', recipe: Callable[[dict, dict, int], str | Tuple[str | None, dict | None, Model | Callable[[], Model] | None]] | None = None, load_raw_kspace: bool = False, reco_pipeline: RecoPipeline | Callable[[dict], RecoPipeline] | None = None, prefer_raw_image: bool = False, rotate_bvecs_to_image: bool = True, slice_crop_mode: str | None = None, progress_callback: Callable[[int, int, str], None] | None = None)#

mrax.io.bruker_import.prepare_bruker_cartesian_kspace(kspace: ndarray, reco: dict, *, target_kx: int | None = None) → tuple[ndarray, dict]#

Apply Bruker Cartesian raw-data corrections and return reconstruction-ready k-space.

The returned array stays in MRax’s internal Cartesian ordering (frame, coil?, ky, kx) and is suitable for direct FFT-based reconstruction. The helper applies the Bruker metadata that act in k-space or as per-channel global factors:

RECO_qopts quadrature options
read-direction zero-filling/placement implied by RECO_ft_size
RECO_rotate as a phase ramp
RecoScaleChan / RecoPhaseChan receive-channel normalization

mrax.io.bruker_import.prepare_bruker_non_cartesian_kspace(samples: ndarray, reco: dict) → tuple[ndarray, dict]#

Apply Bruker receiver corrections to non-Cartesian sample streams.

Non-Cartesian samples stay in MRax order (frame, coil?, sample). ParaVision’s regridding pipeline applies the receive-channel scale and phase factors before channel combination, so those factors belong on DataSet.kspace just like the Cartesian path.

mrax.io.bruker_import.bruker_vendor_regrid_options(reco: dict, *, spatial_shape: tuple[int, ...]) → tuple[dict, dict]#

Build optional MRax regridding options that mirror ParaVision image steps.

The returned options can be passed to mrax.reco.variational.reconstruct_regrid(). They implement the reproducible parts of Bruker’s RecoStageNodes chain: oversampled regridding, phase checkerboarding, center cutoff, approximate Kaiser-Bessel deapodization, image support masking, RSS coil combination, and magnitude output.

mrax.io.bruker_import.read_bruker_regrid_density(exp_dir: Path, trajectory: ndarray, *, spatial_shape: tuple[int, ...], method: dict | None = None, reco: dict | None = None) → tuple[ndarray, dict]#

Read or derive ParaVision RecoRegridN density compensation weights.

Bruker records the density compensation operation in RecoStageNodes. Radial trajectories use type="rho" and spiral trajectories commonly use type="jacobian" with a companion trajDC trajectory. The returned weights are shaped (sample,) and normalized like MRax’s generic density estimates.

mrax.io.bruker_import.read_bruker_regrid_trajectory(exp_dir: Path, *, method: dict | None = None, reco: dict | None = None) → tuple[ndarray, dict]#

Read a Bruker ParaVision RecoRegridN trajectory file.

The returned trajectory is shaped (num_samples, dim) and uses the values written by ParaVision, which are already normalized around the MRax cycles/pixel convention for validated ParaVision 360 3.6 2D regridding scans.

mrax.io.bruker_import.reconstruct_bruker_cartesian_kspace(kspace: ndarray, reco: dict, *, frame_transposition: Sequence[int] | None = None, assume_corrections_applied: bool = False, target_kx: int | None = None) → ndarray#: Reconstruct a Bruker Cartesian k-space stack using the standard vendor steps.

mrax.io.bruker_import.format_bruker_cartesian_image(image: ndarray, reco: dict, *, frame_transposition: Sequence[int] | None = None) → ndarray#

Apply Bruker post-FFT image-domain formatting steps to a Cartesian image stack.

This helper expects the image produced from vendor-corrected Cartesian k-space and applies the image-domain steps that define the standard Bruker magnitude output:

FFT normalization adjustment from MRax’s orthonormal FFT to Bruker’s FFT
RECO_size / RECO_offset cutoff
sum-of-squares coil combination
RECO_transposition / VisuCoreTransposition

mrax.io.study_selection.auto_select_scans(study: Any, meta: Mapping[str, Any] | None, *, fair_exp_types: Sequence[str] = ('fair_asl',), t2_exp_types: Sequence[str] = ('mse',), sat_trans_type: str = 'sat_trans', wasabi_keywords: Iterable[str] = ('wasabi', 'wobble', 'b0', 'b1'), min_frames: int = 1) → tuple[str | None, str | None, str | None, list[str]]#: Heuristically select FAIR, T2, WASABI, and sat-trans scans from a study.

mrax.io.study_selection.pick_scan_by_type(study: Any, exp_type: str, *, min_frames: int = 1) → str | None#: Pick the first scan name matching an experiment type and minimum frame count.

mrax.library.default_models.make_t1_ir_model(scale: float = 1.0, param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#

Inversion recovery T1 model.

scale controls the softplus magnitude for amplitude/bias maps to match data range.

mrax.library.default_models.make_t1_vtr_model(scale: float = 1.0, param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#: Variable-TR T1 model.

mrax.library.default_models.make_t1_vfa_model(scale: float = 1.0, param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#: Variable flip-angle T1 model using spoiled GRE steady state.

mrax.library.default_models.make_exp_decay_model(scale: float = 1.0, param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None, *, name: str = 'exp_decay', decay_parameter: str = 'decay_time', decay_bounds: tuple[float, float] | None = None)#

Mono-exponential decay model S = amplitude * exp(-TE / tau) + bias.

make_t2_se_model and make_t2star_se_model use this factory with only the decay parameter renamed to t2 or t2star.

mrax.library.default_models.make_t2_se_model(scale: float = 1.0, param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#: T2 spin-echo model backed by the shared exponential-decay model.

mrax.library.default_models.make_t2star_se_model(scale: float = 1.0, param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#: T2* gradient-echo decay model backed by the shared exponential-decay model.

mrax.library.default_models.make_wasabi_model(name: str = 'wasabi', param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#

mrax.library.default_models.make_multipool_lorentzian_model(pools: Sequence[str] | None = None, pool_overrides: dict | None = None, name: str = 'multipool_lorentzian', baseline_range: tuple[float, float] = (0.3, 1.0), param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None)#

Multi-pool Lorentzian Z-spectra model for CEST/CESL scans.

Parameters#

poolsSequence[str] | None: Ordered pool names to include. Defaults to [“water”, “mt”].
pool_overridesdict | None: Optional mapping pool -> {amp, gamma, dw, train_dw} or full param-spec dicts to tweak ranges or add new pools.
namestr: Model name to attach to the returned Model.
baseline_rangetuple[float, float]: Lower/upper bounds for the baseline term.
param_keystr: Parameter-set key to pull from the parameter library.
param_overridesdict | None: Param-spec overrides (per-parameter) merged on top of the library.

mrax.library.default_models.make_AB_MT_n_model(pools: Sequence[str] | None = None, pool_overrides: dict | None = None, name: str = 'ab_mt_n', param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None, enable_water_offset: bool | None = None, water_offset_bounds: tuple[float, float] = (-1.0, 1.0), enable_water_offset_linear_drift: bool = False, water_offset_drift_bounds: tuple[float, float] = (-0.05, 0.05))#

AB-MT+n Z-spectra model for CEST/CESL scans.

Parameters#

poolsSequence[str] | None: Ordered CEST pool names to include. Defaults to [“mt”].
pool_overridesdict | None: Optional mapping pool -> {...} to tweak pool parameter specs or metadata such as enable_mt_offset and line_shape.
namestr: Model name to attach to the returned Model.
param_keystr: Parameter-set key to pull from the parameter library.
param_overridesdict | None: Param-spec overrides (per-parameter) merged on top of the library.
enable_water_offsetbool | None: Add a trainable residual water-offset parameter dw_water in ppm. If None, this is enabled when param_overrides defines dw_water.
water_offset_boundstuple[float, float]: Default bounded range for dw_water when no explicit override is provided.
enable_water_offset_linear_driftbool: Add a trainable dw_water_drift slope in ppm/min. The measured b0_map observable remains the fixed intercept, so the modeled water shift is b0_map + dw_water_drift * water_offset_time.
water_offset_drift_boundstuple[float, float]: Default bounded range in ppm/min for dw_water_drift.

mrax.library.default_models.make_zaiss_full_model(pools: Sequence[str] | None = None, pool_overrides: dict | None = None, name: str = 'zaiss_full', param_key: str = 'default', param_overrides: dict[str, dict[str, object]] | None = None, enable_water_offset: bool | None = None, water_offset_bounds: tuple[float, float] = (-1.0, 1.0))#

Full analytical cw saturation-transfer model following Zaiss et al.

This model evaluates the approximate Bloch-McConnell Z-spectrum solution used as the starting point for AB-MT+n: R1rho = Reff + sum(Rex_i) with each exchange contribution described by the Zaiss/Trott-Palmer analytical rate term. Frequency offsets and pool chemical shifts are converted to the same angular-ppm convention used by the other MRax saturation-transfer models: delta_omega = 2*pi*freq_offsets_ppm and omega1 = gamma*sat_b1_uT/field_strength_mhz.