Tomography Model

Tomography Model#

The TomographyModel provides the basic interface for all specific geometries for tomographic projection and reconstruction.

Constructor#

class mbirjax.TomographyModel(sinogram_shape, **kwargs)[source]#

Bases: ParameterHandler

Represents a general model for tomographic reconstruction using MBIRJAX. This class encapsulates the parameters and methods for the forward and back projection processes required in tomographic imaging.

Note that this class is a template for specific subclasses. TomographyModel by itself does not implement projectors or recon. Use self.print_params() to print the parameters of the model after initialization.

Parameters:

sinogram_shape (tuple) – The shape of the sinogram array expected (num_views, num_det_rows, num_det_channels).
recon_shape (tuple) – The shape of the reconstruction array (num_rows, num_cols, num_slices).
**kwargs (dict) – Arbitrary keyword arguments for setting model parameters dynamically. See the full list of parameters and their descriptions at Parameter Documentation.

Sets up the reconstruction size and parameters.

Reconstruction and Projection#

TomographyModel.recon(sinogram, weights=None, init_recon=None, max_iterations=15, stop_threshold_change_pct=0.2, first_iteration=0, compute_prior_loss=False, logfile_path='~/.mbirjax/logs/recon.log', print_logs=True, output_sharded=False)[source]#

Perform MBIR reconstruction using the Multi-Granular Vector Coordinate Descent algorithm. This function takes care of generating its own partitions and partition sequence. TO restart a recon using the same partition sequence, set first_iteration to be the number of iterations completed so far, and set init_recon to be the output of the previous recon. This will continue using the same partition sequence from where the previous recon left off.

Reproducibility note: the pixel partitions are drawn from numpy’s global random number generator, so reconstructions vary slightly from run to run. For a reproducible result, call np.random.seed(seed) before calling this method.

Parameters:

sinogram (numpy or jax array) – 3D sinogram data with shape (num_views, num_det_rows, num_det_channels).
weights (numpy or jax array, optional) – 3D positive weights with same shape as error_sinogram. Defaults to None, in which case the weights are implicitly all 1.
init_recon (numpy or jax array, or None or 0, optional) – Initial reconstruction to use in reconstruction. If None, then direct_recon is called with default arguments. Defaults to None.
max_iterations (int, optional) – maximum number of iterations of the VCD algorithm to perform.
stop_threshold_change_pct (float, optional) – Stop reconstruction when 100 * ||delta_recon||_1 / ||recon||_1 change from one iteration to the next is below stop_threshold_change_pct. Defaults to 0.2. Set this to 0 to guarantee exactly max_iterations.
first_iteration (int, optional) – Set this to be the number of iterations previously completed when restarting a recon using init_recon. This defines the first index in the partition sequence. Defaults to 0.
compute_prior_loss (bool, optional) – Set true to calculate and return the prior model loss. This will lead to slower reconstructions and is meant only for small recons.
logfile_path (str, optional) – Path to the output log file (‘~’ expands to the user’s home directory). Defaults to ‘~/.mbirjax/logs/recon.log’. Set to None or ‘’ to skip file logging.
print_logs (bool, optional) – If true then print logs to console. Defaults to True.
output_sharded (bool, optional) – If False (default), return a numpy reconstruction array. If True, return the internal device form (slice-sharded across the model’s devices, no exit gather); on an unsharded model the output is the same either way.

Returns:

(recon, recon_dict) –

reconstruction array and a dict containing the recon parameters.

recon (numpy or jax array): the reconstruction volume
recon_dict (dict): A dict obtained from get_recon_dict() with entries
- ’recon_params’
- ’notes’
- ’recon_logs’
- ’model_params’

TomographyModel.direct_recon(sinogram, filter_name=None, output_sharded=False)[source]#

Do a direct (non-iterative) reconstruction, typically using a form of filtered backprojection. The implementation details are geometry specific, and direct_recon may not be available for all geometries.

Parameters:

sinogram (numpy or jax array) – 3D sinogram data with shape (num_views, num_det_rows, num_det_channels).
filter_name (string or None, optional) – The name of the filter to use, defaults to None, in which case the geometry specific method chooses a default, typically ‘ramp’.
output_sharded (bool, optional) – If False (default), return a numpy array. If True, return the internal device form (slice-sharded on a sharded model; on an unsharded model the output is the same single-device array either way).

Returns:

recon (numpy or jax array) – The reconstructed volume after direct reconstruction.

TomographyModel.prox_map(prox_input, sinogram, sigma_prox=None, weights=None, init_recon=None, do_initialization=True, stop_threshold_change_pct=0.2, max_iterations=3, first_iteration=0, logfile_path='~/.mbirjax/logs/prox.log', print_logs=True, output_sharded=False)[source]#

Proximal Map function for use in Plug-and-Play applications. This function is similar to recon, but it essentially uses a prior with a mean of prox_input and a standard deviation of sigma_prox.

Reproducibility note: the pixel partitions are drawn from numpy’s global random number generator, so results vary slightly from run to run. For a reproducible result, call np.random.seed(seed) before calling this method.

Parameters:

prox_input (numpy or jax array) – proximal map input with same shape as reconstruction.
sinogram (numpy or jax array) – 3D sinogram data with shape (num_views, num_det_rows, num_det_channels).
sigma_prox (None or float, optional) – The standard deviation of the proximal map prior term. If None, then set automatically from the sinogram. Defaults to None.
weights (numpy or jax array, optional) – 3D positive weights with same shape as sinogram. Defaults to None, in which case the weights are implicitly all 1s.
init_recon (numpy or jax array, optional) – optional reconstruction to be used for initialization. Defaults to None, in which case the initial recon is determined by vcd_recon.
do_initialization (bool, optional) – If True, then initialize parameters and place arrays on appropriate devices. Defaults to True. Set to False if initialization (partitions and regularization parameters) has already been performed on this sinogram by a previous prox_map call on this model.
stop_threshold_change_pct (float, optional) – Stop reconstruction when NMAE percent change from one iteration to the next is below stop_threshold_change_pct. Defaults to 0.2.
max_iterations (int, optional) – maximum number of iterations of the VCD algorithm to perform.
first_iteration (int, optional) – Set this to be the number of iterations previously completed when restarting a recon using init_recon. This defines the first index in the partition sequence. Defaults to 0.
logfile_path (str, optional) – Path to the output log file (‘~’ expands to the user’s home directory). Defaults to ‘~/.mbirjax/logs/prox.log’. Set to None or ‘’ to skip file logging.
print_logs (bool, optional) – If true then print logs to console. Defaults to True.
output_sharded (bool, optional) – If False (default), return a numpy reconstruction array. If True, return the internal device form (no exit gather); on an unsharded model the output is the same either way.

Returns:

(recon, recon_dict) –

reconstruction array and a dict containing the recon parameters.

recon (numpy or jax array): the reconstruction volume
recon_dict (dict): A dict obtained from get_recon_dict() with entries
- ’recon_params’
- ’notes’
- ’recon_logs’
- ’model_params’

TomographyModel.forward_project(recon, output_sharded=False)[source]#

Perform a full forward projection at all voxels in the field-of-view.

The projection automatically uses whatever devices the model is configured for: on one device it runs there, and on several it is spread across them. recon may be an ordinary array or one already distributed across the model’s devices (e.g. the output of an earlier on-device step) – either is accepted, and the returned form is set by output_sharded regardless of which you pass.

Note

This method should generally not be used directly for iterative reconstruction. For iterative reconstruction, use recon().

Parameters:

recon (numpy or jax array) – The 3D reconstruction array, either ordinary or already distributed across the model’s devices.
output_sharded (bool, optional) – Choose the form of the returned sinogram. If False (default), return an ordinary host NumPy array (assembled on the host – safe at any size). If True, leave it distributed (view-sharded) across the model’s devices, so a following on-device step can use it without gathering it back.

Returns:

numpy or jax array – The 3D sinogram – an ordinary host NumPy array by default, or view-sharded across the devices if output_sharded=True.

TomographyModel.back_project(sinogram, output_sharded=False)[source]#

Perform a full back projection at all voxels in the field-of-view.

The back projection automatically uses whatever devices the model is configured for: on one device it runs there, and on several it is spread across them. sinogram may be an ordinary array or one already distributed across the model’s devices (e.g. a sinogram from prepare_sino_for_devices()) – either is accepted, and the returned form is set by output_sharded regardless of which you pass.

Note

This method should generally not be used directly for iterative reconstruction. For iterative reconstruction, use recon().

Parameters:

sinogram (numpy or jax array) – The 3D sinogram, either ordinary or already distributed across the model’s devices.
output_sharded (bool, optional) – Choose the form of the returned recon. If False (default), return an ordinary host NumPy array (assembled on the host – safe at any size). If True, leave it distributed (slice-sharded) across the model’s devices, so a following on-device step (such as a sharded initialization for recon()) can use it without gathering it back; on a single device the two forms are identical.

Returns:

numpy or jax array – The reconstructed 3D volume – an ordinary host NumPy array by default, or slice-sharded across the devices if output_sharded=True.

Parameter Handling#

TomographyModel.set_params(no_warning=False, no_compile=False, **kwargs)[source]#

Update parameters using keyword arguments.

This method updates internal model parameters. If any key geometry-related parameters are modified, it triggers recompilation of the projector system unless suppressed via the no_compile flag.

Parameters:

no_warning (bool, optional) – If True, disables validity checking and warning messages. Defaults to False.
no_compile (bool, optional) – If True, suppresses projector recompilation after updates. Defaults to False.
**kwargs – Arbitrary keyword arguments specifying parameter names and values to update.

Example

>>> import mbirjax as mj
>>> ct_model = mj.ParallelBeamModel(sinogram_shape, angles)
>>> ct_model.set_params(recon_shape=(128, 128, 128), sharpness=0.7)

Note

use_gpu is DEPRECATED: device selection is an execution-environment choice, not a model parameter (a persisted value silently follows a saved model to a different machine). Use configure_devices() instead – configure_devices('cpu') for CPU-only, configure_devices(None) for automatic. During the deprecation window the request is HONORED by forwarding to configure_devices, so it now also takes effect over a previous device pin (it was formerly ignored, silently, once configure_devices had been called).

ParameterHandler.get_params(parameter_names: Literal['geometry_type', 'file_format', 'sinogram_shape', 'delta_det_channel', 'delta_det_row', 'det_row_offset', 'det_channel_offset', 'sigma_y', 'alu_unit', 'alu_value', 'recon_shape', 'delta_voxel', 'voxel_row_aspect', 'voxel_slice_aspect', 'sigma_x', 'sigma_prox', 'p', 'q', 'T', 'qggmrf_nbr_wts', 'auto_regularize_flag', 'positivity_flag', 'snr_db', 'sharpness', 'granularity', 'partition_sequence', 'verbose', 'use_gpu', 'max_overrelaxation', 'use_ror_mask'] | list[Literal['geometry_type', 'file_format', 'sinogram_shape', 'delta_det_channel', 'delta_det_row', 'det_row_offset', 'det_channel_offset', 'sigma_y', 'alu_unit', 'alu_value', 'recon_shape', 'delta_voxel', 'voxel_row_aspect', 'voxel_slice_aspect', 'sigma_x', 'sigma_prox', 'p', 'q', 'T', 'qggmrf_nbr_wts', 'auto_regularize_flag', 'positivity_flag', 'snr_db', 'sharpness', 'granularity', 'partition_sequence', 'verbose', 'use_gpu', 'max_overrelaxation', 'use_ror_mask']]) → Any[source]#

Get the values of the listed parameter names from the internal parameter dictionary.

This method retrieves the current values of one or more parameters managed by the model.

Parameters:: parameter_names (str or list of str) – Name of a parameter, or a list of parameter names.
Returns:: Any or list – Single parameter value if a string is passed, or a list of values if a list is passed.
Raises:: NameError – If any of the provided parameter names are not recognized.

Example

>>> sharpness = model.get_params('sharpness')
>>> recon_shape, sharpness = model.get_params(['recon_shape', 'sharpness'])

ParameterHandler.print_params()[source]#

Print the current parameter values in the model.

This method prints all parameters stored in the model’s internal dictionary. If the model’s verbosity level is less than 3, it prints only the parameter names and values. If verbosity is 3 or higher, it also includes the recompile_flag status for each parameter.

Example

>>> ct_model = mj.ParallelBeamModel(sinogram_shape, angles)
>>> ct_model.set_params(sharpness=0.7, recon_shape=(128, 128, 128))
>>> ct_model.print_params()

TomographyModel.get_recon_dict(recon_params=None, notes=None, save_log=True, save_model=True, str_format=False)[source]#

Encapsulate the recon parameters, logs, notes, and optionally all model parameters to a text-based dict with entries ‘recon_params’, ‘recon_log’, ‘notes’, and optionally ‘model_params’. This dict can be used with mbirjax.viewer.slice_viewer() and TomographyModel.save_recon_hdf5().

This dict from this function is returned by TomographyModel.recon().

Parameters:

recon_params (dict, optional) – dict of reconstruction parameters. Defaults to None.
notes (str, optional) – User-supplied notes to attach to the dataset. Defaults to None.
save_log (bool, optional) – If True, saves the internal log buffer (if available). Defaults to True.
save_model (bool, optional) – If True, saves the model parameters as a YAML string. Defaults to True.
str_format (bool, optional) – If True, then each top level entry is a string, which is a yaml string when the entries could be saved as a dict.

Returns:

dict –

A dict with entries

’recon_params’
’notes’
’recon_log’
’model_params’.

Example

>>> recon, recon_dict = ct_model.recon(sinogram)
>>> print(recon_dict['recon_log'])

Recon Shape and Voxel Spacing#

TomographyModel.auto_set_recon_geometry(no_compile=False, no_warning=False)[source]#

Set the automatic value of the recon shape and voxel pitch using the geometry parameters and sinogram shape.

Note: This function should be run after changing geometry parameters such as delta_det_channel. It will set reconstruction parameters such as recon_shape and delta_voxel parameters to reasonable values.

Parameters:

no_compile (bool, optional) – If True, do not recompile the JAX projector functions. Defaults to False.
no_warning (bool, optional) – If True, do not issue warnings. Defaults to False.

Example

>>> import mbirjax as mj
>>> import jax.numpy as jnp
>>> sinogram = jnp.zeros(shape=(100, 100, 100))
>>> angles = jnp.linspace(0, jnp.pi, 100)
>>> ct_model = mj.ParallelBeamModel(sinogram.shape, angles)
>>> ct_model.set_params(delta_det_channel=100.0)
>>>
>>> # Required reset of recon shape and voxel spacing parameters
>>> ct_model.auto_set_recon_geometry()

TomographyModel.scale_recon_shape(row_scale=1.0, col_scale=1.0, slice_scale=1.0)[source]#

Scale the reconstruction shape by the given scale factors.

This can be used before starting a reconstruction to improve results when part of the object projects outside the detector. The method updates the internal recon_shape parameter.

Parameters:

row_scale (float) – Scale factor for the number of rows in the reconstruction.
col_scale (float) – Scale factor for the number of columns in the reconstruction.
slice_scale (float) – Scale factor for the number of slices in the reconstruction.

Returns:

tuple[int, int, int] – A 3-tuple representing the number of pixels added to the (rows, columns, slices) dimensions due to scaling.

Example

>>> old_shape = model.get_params('recon_shape')
>>> added_padding = model.scale_recon_shape(row_scale=1.2, col_scale=1.1)
>>> new_shape = model.get_params('recon_shape')
>>> print(f"Shape increased by: {added_padding}")

TomographyModel.get_magnification()#

Compute the scale factor from a voxel at iso (at the origin on the center of rotation) to its projection on the detector. For parallel beam, this is 1, but it may be parameter-dependent for other geometries.

Returns:: (float) – magnification

Device Configuration#

On a machine with multiple GPUs, MBIRJAX automatically divides a reconstruction across them to increase the available memory and reduce reconstruction time – with no change to your script, and for every geometry. The methods below give explicit control over which devices are used and report what was chosen. See Multi-GPU Reconstruction for a full discussion.

TomographyModel.configure_devices(devices=None)[source]#

Configure which devices the reconstruction runs on (the user-facing control surface).

Resolves devices to a concrete device list and PINS it: a later set_params (which re-runs set_devices) keeps these devices rather than re-auto-selecting.

None – automatic (the default construction behavior): shard across all available GPUs (or, with no GPU, all available CPU devices), trimmed by _auto_device_count so the last shard is never entirely padding. A single available device gives a trivial 1-device layout.

'cpu' / 'gpu' – all devices of that platform, trimmed as for automatic. configure_devices('cpu') forces CPU-only operation (this replaces the deprecated set_params(use_gpu='none')); 'gpu' raises when no GPU backend is available.

int n – use the first n devices of the default platform (GPUs if any, else CPUs); configure_devices(1) is the way to force single-device operation.

sequence of ints – use those indices into the default device list (jax.devices()).

sequence of jax devices – use exactly those devices.

Sharding scheme (uniform across geometries): sinogram-like arrays are sharded by view (axis 0), recon-like arrays by slice. NEITHER axis constrains the device count: a non-dividing view or slice axis is zero-padded to the next multiple of the count and the padding is kept exactly inert (it cannot affect the results), so the result is independent of the device count. The pad metadata is re-derived from the current shapes on every recompile, so device selection is order-independent with respect to shape changes.

Parameters:: devices (None, str, int, or sequence of ints / jax devices) – see above.
Returns:: Nothing; builds recon_placement / sino_placement and the empirical device-to-device safety flag (see _set_device_layout).

TomographyModel.prepare_sino_for_devices(sinogram, weights=None)[source]#

Place a sinogram (and optionally weights) in the model’s device form, once.

The device form is the view-sharded layout the reconstruction methods use internally: the sinogram is distributed across the configured devices, and when the view count does not divide the device count it is zero-padded to the next multiple (the padding is exactly inert – it cannot affect the results). The transfer streams shard-by-shard from the host, so no padded host copy is ever created; the only transient is one shard on one device.

Calling this is OPTIONAL: every reconstruction method applies the same placement automatically to a plain input. Use it to pay the host-to- device transfer once when running several reconstructions on the same large sinogram – a prepared array passes through the entry placement untouched. If the device configuration changes afterwards (e.g. a new configure_devices), the prepared array no longer matches and the entry placement raises with instructions to re-run this method.

Parameters:

sinogram (numpy or jax array) – sinogram in the model’s sinogram_shape.
weights (numpy or jax array, optional) – weights of the same shape; the zero-filled padding makes padded views weightless as well.

Returns:

The prepared sinogram, or a (sinogram, weights) tuple when weights were given.

property TomographyModel.device_summary#

A read-only summary of the devices the reconstruction will actually use.

A property – read it as an attribute (model.device_summary), not a call.

This is the resolved OUTCOME of the device configuration – the use_gpu parameter and configure_devices are the request; this reports what was chosen. E.g. '4 x GPU (sharded)', with notes appended when the view axis is padded or when automatic selection left GPUs idle. The same line is logged at the start of every reconstruction.

Type:: str

Saving and Loading#

TomographyModel.save_recon_hdf5(filepath, recon, recon_dict=None)[source]#

Save the reconstruction array and optionally the recon_dict from recon().

This method creates a file that contains a single dataset named ‘recon’, with the entries in recon_dict serialized to strings and saved as hdf5 dataset attributes.

The resulting file can be loaded with load_recon_hdf5() or mbirjax.viewer.slice_viewer().

Parameters:

filepath (str or Path) – Path to the output HDF5 file. Should typically end with a .h5 extension.
recon (array-like) – The reconstruction volume as a NumPy or JAX array.
recon_dict (dict or None, optional) – The dictionary of recon attributes from get_recon_dict()

Raises:

Exception – If saving the file or directory creation fails.

Example

>>> recon, recon_dict = ct_model.recon(sinogram)
>>> recon_dict['notes'] += 'Test scan'
>>> ct_model.save_recon_hdf5("output/my_recon.h5", recon, recon_dict=recon_dict)

static TomographyModel.load_recon_hdf5(filepath, recreate_model=False)[source]#

This function loads a numpy array stored in an HDF5 file created by save_recon_hdf5(). It also loads any associated attribute dict and can use the model parameters in that dict to create a new model.

Parameters:

filepath (str) – Path to the HDF5 file containing the reconstructed volume.
recreate_model (bool, optional) – Deprecated. Will raise a ValueError if set to True.

Returns:

(recon, recon_dict)

recon (ndarray): The tensor saved by save_data_hdf5()
recon_dict (dict): A dict with the attributes for the data array as in get_recon_dict()

Raises:

FileNotFoundError – If the file does not exist.
ValueError – If more than one dataset is not found in the file or if recreate_model is set to True.

Example

>>> recon, recon_dict = ct_model.load_recon_hdf5("output/recon_volume.h5")
>>> recon.shape
(64, 256, 256)

Parameter Documentation#

See the Primary Parameters page.

Tomography Model

Contents

Tomography Model#

Constructor#

Reconstruction and Projection#

Parameter Handling#

Recon Shape and Voxel Spacing#

Device Configuration#

Saving and Loading#

Parameter Documentation#