Evaluation Python API¶
Full evaluation is CLI-driven. dmri eval composes the Hydra configuration,
loads the checkpoint and input, chooses devices and batch sizes, runs the
pipeline, and dispatches exporters. There is no supported Python function that
accepts an eval config and replaces the CLI entry point; use subprocess or call
the dmri command for a complete run.
The package-level dmri.eval API exposes the metric configuration types and
metric runner. The additional functions below are selected module-level helpers
for code that already owns the required arrays, model, or configuration.
Configured Metrics¶
run_configured_metrics evaluates a metric set after inference. Callers must
construct a complete MetricContext, including spatial metadata used to embed
voxel values into NIfTI volumes. The CLI normally does this work.
Configurable evaluation metrics exported as NIfTI volumes.
This module replaces the previous ad-hoc metric helpers with a small registry that can be configured from Hydra. Each metric definition controls how the metric is evaluated, how the per-voxel samples are reduced, optional summary statistics, and where the resulting files are stored.
The public entry point run_configured_metrics is meant to be invoked by the
Hydra evaluation script once sampling has finished.
dataclass
¶
dataclass
¶
MetricSpec(key: str, type: str, output_filename: str, batch_size: int = 20000, enabled: bool = True, sample_reduction: str = 'mean', sample_axis: int | None = 1, requires_theta_samples: bool = True, write_summary: bool = True, summary_filename: str | None = None, aggregations: tuple[MetricAggregation, ...] = tuple(), preferred_device_kinds: tuple[str, ...] = ('gpu', 'tpu'), options: Mapping[str, Any] = dict())
Structured configuration for a metric.
class-attribute
instance-attribute
¶
aggregations: tuple[MetricAggregation, ...] = field(default_factory=tuple)
class-attribute
instance-attribute
¶
class-attribute
instance-attribute
¶
dataclass
¶
dataclass
¶
MetricContext(cfg: Any, sim_type: Any, acq: Any, full_data_flat_in_brain: ndarray, model_parameters_brain: ndarray | None, model_mask: ndarray | None, model_mask_samples: ndarray | None, brain_mask_flat: ndarray, brain_shape: tuple, orig_data: Any, out_path: str | None = None, true_model_parameters: ndarray | None = None, true_model_mask: ndarray | None = None)
class-attribute
instance-attribute
¶
run_configured_metrics(metrics_cfg: Mapping[str, Any] | Sequence[Any] | None, context: MetricContext, out_path: str, *, devices: Sequence[Device] | str | None = None, logger: Any | None = None) -> MutableMapping[str, MetricResult]
Evaluate and export all enabled metrics.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics_cfg
|
Mapping[str, Any] | Sequence[Any] | None
|
Mapping or list describing metric specifications. When
using Hydra, this typically comes from |
required |
context
|
MetricContext
|
Pre-computed tensors and metadata shared by all metrics. |
required |
out_path
|
str
|
Destination directory. |
required |
devices
|
Sequence[Device] | str | None
|
Optional device specification forwarded to batching helper. |
None
|
logger
|
Any | None
|
Optional logger for informational messages. |
None
|
Returns:
| Type | Description |
|---|---|
MutableMapping[str, MetricResult]
|
A mapping from metric key to :class: |
These names can also be imported directly from dmri.eval:
File Input And Preprocessing¶
Loading and normalisation of FSL/HCP-style diffusion data.
The brain mask is applied while reading, so only in-brain voxels are ever
materialised and everything downstream works on a flat (voxels, gradients)
array.
load_and_process_data(path, brain_mask='nodif_brain_mask.nii.gz', mri_data='data.nii.gz', bvals_data='bvals', bvecs_data='bvecs', round_bvals=True, mask_filter=None, clip_quantile=None)
Load a data folder and return the normalised in-brain signal.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mask_filter
|
optional callable applied to the brain mask before the signal is read, used to restrict processing to a sub-volume. Any voxel it drops is never loaded. |
None
|
|
clip_quantile
|
optional quantile in (0, 1]; signals above it are clipped. |
None
|
Returns:
| Type | Description |
|---|---|
|
|
|
|
has shape |
|
|
unread source image, kept for its affine and shape. |
Load a brain mask as a boolean array.
Read only the in-brain voxels of a 4-D image as (voxels, gradients).
order reorders the gradient axis (e.g. by b-value) without a second copy.
Sort b-values, b-vectors and the gradient axis of data together.
Divide the in-brain signal (voxels, gradients) by its mean b0 signal.
Voxels whose b0 signal is non-positive are set to zero rather than producing NaNs, and a missing b0 shell is an error instead of a silently empty volume.
Synthetic Input¶
Generate synthetic diffusion-MRI data from a simulator configuration.
Parameters¶
data_cfg:
Hydra/OmegaConf dict with num_voxels, acquisition_scheme or
bvals/bvecs, and optionally model_mask_hyperparameter.
sim_type:
A :class:MultiCompartment subclass used to sample masks and signals.
key:
JAX PRNG key.
Returns¶
tuple
(orig_data, data_norm, brain_mask, bvals, bvecs, model_masks, thetas, acq, key)
arranged as a pseudo-volume suitable for the evaluation pipeline.
Precision¶
Evaluation-time numeric precision selection.
The model config already carries dtype / param_dtype / precision /
preferred_element_type, which :func:dmri.train.build_model.build_model
forwards into every submodule. Because a checkpoint is rebuilt from its stored
config before its parameters are restored, choosing a precision at evaluation
time is just a matter of patching that config first — no layer needs to change.
Half precision applies to the network's matmul inputs only. Parameters stay
float32 (they are small; the memory goes into activations), and
preferred_element_type: float32 keeps the dot outputs in float32, so the
diffusion sampler, the correctors and the export never see half precision.
Validate a precision name, mapping None/empty to fp32.
Return a precision supported by the selected JAX backend.
CPU backends do not consistently implement fp16 dot products with fp32 accumulation, notably on macOS. Keep fp16 as a GPU optimization and fall back to fp32 elsewhere.
Patch cfg.model in place with the preset for precision.
Returns the normalized name. fp32 is a no-op.
Batching And Sampling Cost¶
eval_in_batches applies a prepared function over voxel batches. It does not
load a model or build an evaluation pipeline.
eval_in_batches(fn, key, *data, batch_size=10000, logger=None, min_batch_size=100, devices: Sequence[Device] | str | None = None, preferred_device_kinds: Sequence[str] | None = ('gpu', 'tpu'), cache_key: str | None = None, desc: str | None = None)
Evaluate a function on data batches, overlapping host prep with device compute.
Every batch is padded to batch_size so only a single shape is ever traced;
the padded rows are trimmed from the results.
Each row gets its key from a single up-front split over the whole input, so
the result depends only on key -- not on the batch size, the device count
or where the batch boundaries happen to fall.
Network evaluations the ODE solver performs for one posterior sample.
One to initialise, one per interval between the num_steps grid points,
and one more if a final Euler correction is taken.
Array And Export Helpers¶
Stack (theta, phi) angles into unit vectors of shape (..., 3).
Apply fn(theta, mask) over the (voxel, sample) axes in voxel chunks.
fn may return a pytree; the results are concatenated on the host as
float32 numpy arrays so nothing larger than one chunk stays on the device.
Embed a tensor in the full brain array