From c4e293ea9e46c81030ff4555ebbfbb9a231cbfde Mon Sep 17 00:00:00 2001 From: mcencini Date: Mon, 30 Sep 2024 23:40:36 -0400 Subject: [PATCH 01/10] add subspace projection model. --- docs/nufft.rst | 54 ++++++- examples/example_subspace.py | 231 +++++++++++++++++++++++++++++ src/mrinufft/extras/__init__.py | 12 +- src/mrinufft/extras/data.py | 155 ++++++++----------- src/mrinufft/extras/field_map.py | 106 +++++++++++++ src/mrinufft/extras/sim.py | 59 ++++++++ src/mrinufft/operators/__init__.py | 2 + src/mrinufft/operators/subspace.py | 151 +++++++++++++++++++ tests/case_trajectories.py | 10 +- tests/operators/test_subspace.py | 213 ++++++++++++++++++++++++++ 10 files changed, 894 insertions(+), 99 deletions(-) create mode 100644 examples/example_subspace.py create mode 100644 src/mrinufft/extras/field_map.py create mode 100644 src/mrinufft/extras/sim.py create mode 100644 src/mrinufft/operators/subspace.py create mode 100644 tests/operators/test_subspace.py diff --git a/docs/nufft.rst b/docs/nufft.rst index e8950681..8b5a44b5 100644 --- a/docs/nufft.rst +++ b/docs/nufft.rst @@ -166,6 +166,57 @@ Where :math:`E_mn = e^i\Delta\omega_0(u_n)t_m`. .. _nufft-algo: +Subspace Projection Model +~~~~~~~~~~~~~~~~~~~~~~~~~ +In several MRI applications, such as dynamic or quantitative MRI, a single acquisition provides a stack of images, each representing a single time frame (for dynamic MRI) or a single contrast (for quantitative MRI). +To achieve a clinically feasible scan time, each frame or contrast is acquired with a different aggressively undersampled k-space trajectory. In this context, the single-coil acquisition model becomes: + +.. math:: + + \tilde{\boldsymbol{y}} = \begin{bmatrix} + \mathcal{F}_\Omega_1 & 0 & \cdots & 0 \\ + 0 & \mathcal{F}_\Omega_2 & \cdots & 0 \\ + \vdots & \vdots & \ddots & \vdots \\ + 0 & 0 & \cdots & \mathcal{F}_\Omega_T + \end{bmatrix} + \boldsymbol{x} + \boldsymbol{n} + +where :math:`\mathcal{F}_\Omega_1, \dots, \mathcal{F}_\Omega_T` are the Fourier operators corresponding to each individual frame. Some applications (e.g., MR Fingerprinting [3]_) may consists of +thousands of total frames :math:`T`, leading to repeated Fourier Transform operations and high computational burden. However, the 1D signal series arising from similar voxels, e.g., with similar +relaxation properties, are typically highly correlated. For this reason, the image series can be represented as: + +.. math:: + + \boldsymbol{x} = \Phi\Phi^H \boldsymbol{x} + +where :math:`\Phi` is an orthonormal basis spanning a low dimensional subspace whose rank :math:`K \ll T` which can be obtained performing a Singular Value Decomposition of a low resolution fully sampled +training dataset or an ensemble of simulated Bloch responses. The signal model can be then written as: + +.. math:: + + \tilde{\boldsymbol{y}} = \begin{bmatrix} + \mathcal{F}_\Omega_1 & 0 & \cdots & 0 \\ + 0 & \mathcal{F}_\Omega_2 & \cdots & 0 \\ + \vdots & \vdots & \ddots & \vdots \\ + 0 & 0 & \cdots & \mathcal{F}_\Omega_T + \end{bmatrix} + \Phi \Phi^H \boldsymbol{x} + \boldsymbol{n} = + \begin{bmatrix} + \mathcal{F}_\Omega_1 & 0 & \cdots & 0 \\ + 0 & \mathcal{F}_\Omega_2 & \cdots & 0 \\ + \vdots & \vdots & \ddots & \vdots \\ + 0 & 0 & \cdots & \mathcal{F}_\Omega_T + \end{bmatrix} + \Phi \boldsymbol{\alpha} + \boldsymbol{n} + +where :math:`\boldsymbol{\alpha} = \Phi^H \boldsymbol{x}` are the spatial coefficients representing the image series. Since the elements of :math:`\Phi^H` do not depend on the specific k-space frequency points, +the projection operator :math:`\boldsymbol{\Phi}` commutes with the Fourier transform, and the signal equation finally becomes: + +.. math:: + + \tilde{\boldsymbol{y}} = \Phi \mathcal{F}_Omega(\boldsymbol{\alpha}) + \boldsymbol{n} + +that is, computation now involves :math:`K \ll T` Fourier Transform operations, each with the same sampling trajectory, which can be computed levaraging efficient NUFFT implementations for conventional static MRI. The Non Uniform Fast Fourier Transform ====================================== @@ -206,5 +257,6 @@ References ========== .. [1] https://en.m.wikipedia.org/wiki/Non-uniform_discrete_Fourier_transform -.. [2] Noll, D. C., Meyer, C. H., Pauly, J. M., Nishimura, D. G., Macovski, A., "A homogeneity correction method for magnetic resonance imaging with time-varying gradients", IEEE Transaction on Medical Imaging (1991), pp. 629-637 +.. [2] Noll, D. C., Meyer, C. H., Pauly, J. M., Nishimura, D. G., Macovski, A., "A homogeneity correction method for magnetic resonance imaging with time-varying gradients", IEEE Transaction on Medical Imaging (1991), pp. 629-637. .. [3] Fessler, J. A., Lee, S., Olafsson, V. T., Shi, H. R., Noll, D. C., "Toeplitz-based iterative image reconstruction for MRI with correction for magnetic field inhomogeneity", IEEE Transactions on Signal Processing 53.9 (2005), pp. 3393–3402. +.. [4] D. F. McGivney et al., "SVD Compression for Magnetic Resonance Fingerprinting in the Time Domain," IEEE Transactions on Medical Imaging (2014), pp. 2311-2322. \ No newline at end of file diff --git a/examples/example_subspace.py b/examples/example_subspace.py new file mode 100644 index 00000000..84149b01 --- /dev/null +++ b/examples/example_subspace.py @@ -0,0 +1,231 @@ +""" +====================== +Subspace NUFFT Operator +====================== + +An example to show how to setup a subspace NUFFT operator. + +This examples show how to use the Subspace NUFFT operator to acquire +and reconstruct data while projecting onto a low rank spatiotemporal subspace. +This can be useful e.g., for dynamic multi-contrast MRI. +Hereafter a 2D spiral trajectory is used for demonstration. + +""" + +import matplotlib.pyplot as plt +import numpy as np + +from mrinufft import display_2D_trajectory + +# %% +# Data preparation +# ================ +# +# Image loading +# ------------- +# +# For realistic a 2D image we will use the BrainWeb dataset, +# installable using ``pip install brainweb-dl``. + +from mrinufft.extras import get_brainweb_map + +M0, T1, T2 = get_brainweb_map(0) +M0 = np.flip(M0, axis=(0, 1, 2))[90] +T1 = np.flip(T1, axis=(0, 1, 2))[90] +T2 = np.flip(T2, axis=(0, 1, 2))[90] + +fig1, ax1 = plt.subplots(1, 3) + +im0 = ax1[0].imshow(M0, cmap="gray") +ax1[0].axis("off"), ax1[0].set_title("M0 [a.u.]") +fig1.colorbar(im0, ax=ax1[0], fraction=0.046, pad=0.04) + +im1 = ax1[1].imshow(T1, cmap="magma") +ax1[1].axis("off"), ax1[1].set_title("T1 [ms]") +fig1.colorbar(im1, ax=ax1[1], fraction=0.046, pad=0.04) + +im2 = ax1[2].imshow(T2, cmap="viridis") +ax1[2].axis("off"), ax1[2].set_title("T2 [ms]") +fig1.colorbar(im2, ax=ax1[2], fraction=0.046, pad=0.04) + +plt.tight_layout() +plt.show() + +# %% +# Sequence Parameters +# ------------------- +# +# As an example, we simulate a simple monoexponential Spin Echo acquisition. +# We assume that refocusing train is constant (flip angle set to 180 degrees) +# and k-space center is sampled at each TE (as for spiral or radial imaging). +# In this way, we obtain an image for each of the shots in the echo train. + +from mrinufft.extras import fse_simulation + +ETL = 48 # Echo Train Length +ESP = 6.0 # Echo Spacing [ms] +TE = np.arange(ETL, dtype=np.float32) * ESP # [ms] +TR = 3000.0 # [ms] + +# %% +# Subspace Generation +# ------------------- +# +# Here, we generate temporal subspace basis +# by Singular Value Decomposition of an ensemble +# of simulated signals in the same property range. +# our object. + + +def make_grid(T1range, T2range, natoms=100): + """Prepare parameter grid for basis estimation.""" + # Create linear grid + T1grid = np.linspace(T1range[0], T1range[1], num=natoms, dtype=np.float32) + T2grid = np.linspace(T2range[0], T2range[1], num=natoms, dtype=np.float32) + + # Create combined grid + T1grid, T2grid = np.meshgrid(T1grid, T2grid) + + return T1grid.ravel(), T2grid.ravel() + + +def estimate_subspace_basis(train_data, ncoeff=4): + """Estimate subspace data via SVD of simulated training data.""" + # Perform svd + _, _, basis = np.linalg.svd(train_data, full_matrices=False) + + # Calculate basis (retain only ncoeff coefficients) + basis = basis[:ncoeff] + + return basis + + +# Get range from tissues +T1range = (T1.min() + 1, T1.max()) # [ms] +T2range = (T2.min() + 1, T2.max()) # [ms] + +# Prepare tissue grid +T1grid, T2grid = make_grid(T1range, T2range) + +# Calculate training data +train_data = fse_simulation(1.0, T1grid, T2grid, TE, TR).astype(np.float32) + +# Calculate basis +basis = estimate_subspace_basis(train_data.T) + +fig2, ax2 = plt.subplots(1, 2) +ax2[0].plot(TE, train_data[:, ::100]), ax2[0].set( + xlabel="TE [ms]", ylabel="signal [a.u.]" +), ax2[0].set_title("training dataset") +ax2[1].plot(TE, basis.T), ax2[1].set(xlabel="TE [ms]", ylabel="signal [a.u.]"), ax2[ + 1 +].set_title("subspace basis") + +plt.show() + +# %% +# Here, we simulate Brainweb FSE data with the same +# sequence parameters used for the subspace estimation. + +mri_data = fse_simulation(M0, T1, T2, TE, TR).astype(np.float32) +mri_data = np.ascontiguousarray(mri_data) + +# Ground truth subspace coefficients +ground_truth = (mri_data.T @ basis.T).T +ground_truth = np.ascontiguousarray(ground_truth) +ground_truth_display = np.concatenate( + [abs(coeff) / abs(coeff).max() for coeff in ground_truth], axis=-1 +) +plt.figure() +plt.imshow(ground_truth_display, cmap="gray"), plt.axis("off"), plt.title( + "ground truth subspace coefficients" +) + +plt.show() + +# %% +# Trajectory generation +# --------------------- + +from mrinufft import initialize_2D_spiral +from mrinufft.density import voronoi + +samples = initialize_2D_spiral( + Nc=ETL * 16, Ns=1200, nb_revolutions=10, tilt="mri-golden" +) + +# assume trajectory is reordered as (ncontrasts, nshots_per_contrast, nsamples_per_shot, ndims) +samples = samples.reshape(16, ETL, *samples.shape[1:]) + +# flatten ncontrasts and nshots_per_contrast axes +samples = samples.reshape(-1, *samples.shape[2:]) + +# compute density compensation +density = voronoi(samples) + +display_2D_trajectory(samples) + +# %% +# Operator setup +# ============== + +from mrinufft import get_operator +from mrinufft.operators import MRISubspace + +# Generate standard NUFFT operator +nufft = get_operator("finufft")( + samples=samples, + shape=mri_data.shape[-2:], + density=density, +) + +# Generate subspace-projected NUFFT operator +multicontrast_nufft = MRISubspace(nufft, subspace_basis=np.eye(ETL)) +subspace_nufft = MRISubspace(nufft, subspace_basis=basis) + +# %% +# Generate K-Space +# ---------------- +# +# We generate the k-space data using a non-projected operator. +# This can be simply obtained by using an identity matrix +# with of shape (ETL, ETL) (= number of contrasts) as a subspace basis. + +kspace = multicontrast_nufft.op(mri_data) + +# %% +# Image reconstruction +# ----------------------- +# +# We now reconstruct both using the subspace expanded operator +# and the zero-filling operator followed by projection in image space + +# Reconstruct with projection in image space +zerofilled_data_adj = multicontrast_nufft.adj_op(kspace) +zerofilled_display = (zerofilled_data_adj.T @ basis.T).T +zerofilled_display = np.concatenate( + [abs(coeff) / abs(coeff).max() for coeff in zerofilled_display], axis=-1 +) + +# Reconstruct with projection in k-space +subspace_data_adj = subspace_nufft.adj_op(kspace) +subspace_display = np.concatenate( + [abs(coeff) / abs(coeff).max() for coeff in subspace_data_adj], axis=-1 +) + + +fig3, ax3 = plt.subplots(2, 1) +ax3[0].imshow(zerofilled_display, cmap="gray") +ax3[0].set_xticks([]) +ax3[0].set_yticks([]) +ax3[0].set(ylabel="kspace projection") +ax3[0].set_title("reconstructed subspace coefficients") +ax3[1].imshow(subspace_display, cmap="gray") +ax3[1].set_xticks([]) +ax3[1].set_yticks([]) +ax3[1].set(ylabel="zerofill + projection") + +plt.show() + +# %% +# The projected k-space is equivalent to the regular reconstruction followed by projection. diff --git a/src/mrinufft/extras/__init__.py b/src/mrinufft/extras/__init__.py index 3c08cb88..0a885d6b 100644 --- a/src/mrinufft/extras/__init__.py +++ b/src/mrinufft/extras/__init__.py @@ -1,13 +1,17 @@ -"""Sensitivity map estimation methods.""" +"""Additional supports routines.""" -from .data import make_b0map, make_t2smap +from .data import get_brainweb_map +from .field_map import make_b0map, make_t2smap +from .sim import fse_simulation from .smaps import low_frequency from .utils import get_smaps __all__ = [ + "fse_simulation", + "get_brainweb_map", + "get_smaps", + "low_frequency", "make_b0map", "make_t2smap", - "low_frequency", - "get_smaps", ] diff --git a/src/mrinufft/extras/data.py b/src/mrinufft/extras/data.py index 5dfcf281..92a797c6 100644 --- a/src/mrinufft/extras/data.py +++ b/src/mrinufft/extras/data.py @@ -1,106 +1,75 @@ -"""Field map generator module.""" +"""Data generator module.""" import numpy as np -def make_b0map(shape, b0range=(-300, 300), mask=None): - """ - Make radial B0 map. - - Parameters - ---------- - shape : tuple[int] - Matrix size. Only supports isotropic matrices. - b0range : tuple[float], optional - Frequency shift range in [Hz]. The default is (-300, 300). - mask : np.ndarray - Spatial support of the objec. If not provided, - build a radial mask with radius = 0.3 * shape - - Returns - ------- - np.ndarray - B0 map of shape (*shape) in [Hz], - with values included in (*b0range). - mask : np.ndarray, optional - Spatial support binary mask. - - """ - assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.") - ndim = len(shape) - if ndim == 2: - radial_mask, fieldmap = _make_disk(shape) - elif ndim == 3: - radial_mask, fieldmap = _make_sphere(shape) - if mask is None: - mask = radial_mask - - # build map - fieldmap *= mask - fieldmap = (b0range[1] - b0range[0]) * fieldmap / fieldmap.max() + b0range[0] # Hz - fieldmap *= mask - - # remove nan - fieldmap = np.nan_to_num(fieldmap, neginf=0.0, posinf=0.0) +def get_brainweb_map(sub_id: int) -> np.ndarray: + """Get MRI parametric maps from a brainweb crisp segmentation. - return fieldmap.astype(np.float32), mask - - -def make_t2smap(shape, t2svalue=15.0, mask=None): - """ - Make homogeneous T2* map. + Output maps have the same shape as the tissue segmentation. Parameters ---------- - shape : tuple[int] - Matrix size. - t2svalue : float, optional - Object T2* in [ms]. The default is 15.0. - mask : np.ndarray - Spatial support of the objec. If not provided, - build a radial mask with radius = 0.3 * shape + sub_id : int + Subject ID. + + Raises + ------ + ImportError + If brainweb-dl is not installed. Returns ------- - np.ndarray - T2* map of shape (*shape) in [ms]. - mask : np.ndarray, optional - Spatial support binary mask. + M0 : np.ndarray + Proton Density map. For sub_id > 0, + it is a binary mask. + T1 : np.ndarray + T1 map map in [ms]. + T2 : np.ndarray + T2 map map in [ms]. """ - assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.") - ndim = len(shape) - if ndim == 2: - radial_mask, fieldmap = _make_disk(shape) - elif ndim == 3: - radial_mask, fieldmap = _make_sphere(shape) - if mask is None: - mask = radial_mask - - # build map - fieldmap = t2svalue * mask # ms - - # remove nan - fieldmap = np.nan_to_num(fieldmap, neginf=0.0, posinf=0.0) - - return fieldmap.astype(np.float32), mask - - -def _make_disk(shape, frac_radius=0.3): - """Make circular binary mask.""" - ny, nx = shape - yy, xx = np.mgrid[:ny, :nx] - yy, xx = yy - ny // 2, xx - nx // 2 - yy, xx = yy / ny, xx / nx - rr = (xx**2 + yy**2) ** 0.5 - return rr < frac_radius, rr - - -def _make_sphere(shape, frac_radius=0.3): - """Make spherical binary mask.""" - nz, ny, nx = shape - zz, yy, xx = np.mgrid[:nz, :ny, :nx] - zz, yy, xx = zz - nz // 2, yy - ny // 2, xx - nx // 2 - zz, yy, xx = zz / nz, yy / ny, xx / nx - rr = (xx**2 + yy**2 + zz**2) ** 0.5 - return rr < frac_radius, rr + try: + import brainweb_dl + except ImportError as err: + raise ImportError( + "The brainweb-dl module is not available. Please install it using " + "the following command: pip install brainweb-dl" + ) from err + + # get segmentation + segmentation = brainweb_dl.get_mri(sub_id, "crisp") / 455 + segmentation = segmentation.astype(int) + + # get properties + model = brainweb_dl._brainweb.BrainWebTissueMap + if sub_id == 0: + properties = brainweb_dl._brainweb._load_tissue_map(model.v1) + else: + properties = brainweb_dl._brainweb._load_tissue_map(model.v2) + + # initialize maps + if sub_id == 0: + M0 = np.zeros_like(segmentation, dtype=np.float32) + T1 = np.zeros_like(segmentation, dtype=np.float32) + T2 = np.zeros_like(segmentation, dtype=np.float32) + + # fill maps + for tissue in properties: + idx = segmentation == int(tissue["Label"]) + if sub_id == 0: + M0 += float(tissue["PD (ms)"]) * idx + T1 += float(tissue["T1 (ms)"]) * idx + T2 += float(tissue["T2 (ms)"]) * idx + + if sub_id != 0: + M0 = (segmentation != 0).astype(np.float32) + + # pad to square + pad_width = segmentation.shape[1] - segmentation.shape[2] + pad = ((0, 0), (0, 0), (int(pad_width // 2), int(pad_width // 2))) + M0 = np.pad(M0, pad) + T1 = np.pad(T1, pad) + T2 = np.pad(T2, pad) + + return M0, T1, T2 diff --git a/src/mrinufft/extras/field_map.py b/src/mrinufft/extras/field_map.py new file mode 100644 index 00000000..5dfcf281 --- /dev/null +++ b/src/mrinufft/extras/field_map.py @@ -0,0 +1,106 @@ +"""Field map generator module.""" + +import numpy as np + + +def make_b0map(shape, b0range=(-300, 300), mask=None): + """ + Make radial B0 map. + + Parameters + ---------- + shape : tuple[int] + Matrix size. Only supports isotropic matrices. + b0range : tuple[float], optional + Frequency shift range in [Hz]. The default is (-300, 300). + mask : np.ndarray + Spatial support of the objec. If not provided, + build a radial mask with radius = 0.3 * shape + + Returns + ------- + np.ndarray + B0 map of shape (*shape) in [Hz], + with values included in (*b0range). + mask : np.ndarray, optional + Spatial support binary mask. + + """ + assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.") + ndim = len(shape) + if ndim == 2: + radial_mask, fieldmap = _make_disk(shape) + elif ndim == 3: + radial_mask, fieldmap = _make_sphere(shape) + if mask is None: + mask = radial_mask + + # build map + fieldmap *= mask + fieldmap = (b0range[1] - b0range[0]) * fieldmap / fieldmap.max() + b0range[0] # Hz + fieldmap *= mask + + # remove nan + fieldmap = np.nan_to_num(fieldmap, neginf=0.0, posinf=0.0) + + return fieldmap.astype(np.float32), mask + + +def make_t2smap(shape, t2svalue=15.0, mask=None): + """ + Make homogeneous T2* map. + + Parameters + ---------- + shape : tuple[int] + Matrix size. + t2svalue : float, optional + Object T2* in [ms]. The default is 15.0. + mask : np.ndarray + Spatial support of the objec. If not provided, + build a radial mask with radius = 0.3 * shape + + Returns + ------- + np.ndarray + T2* map of shape (*shape) in [ms]. + mask : np.ndarray, optional + Spatial support binary mask. + + """ + assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.") + ndim = len(shape) + if ndim == 2: + radial_mask, fieldmap = _make_disk(shape) + elif ndim == 3: + radial_mask, fieldmap = _make_sphere(shape) + if mask is None: + mask = radial_mask + + # build map + fieldmap = t2svalue * mask # ms + + # remove nan + fieldmap = np.nan_to_num(fieldmap, neginf=0.0, posinf=0.0) + + return fieldmap.astype(np.float32), mask + + +def _make_disk(shape, frac_radius=0.3): + """Make circular binary mask.""" + ny, nx = shape + yy, xx = np.mgrid[:ny, :nx] + yy, xx = yy - ny // 2, xx - nx // 2 + yy, xx = yy / ny, xx / nx + rr = (xx**2 + yy**2) ** 0.5 + return rr < frac_radius, rr + + +def _make_sphere(shape, frac_radius=0.3): + """Make spherical binary mask.""" + nz, ny, nx = shape + zz, yy, xx = np.mgrid[:nz, :ny, :nx] + zz, yy, xx = zz - nz // 2, yy - ny // 2, xx - nx // 2 + zz, yy, xx = zz / nz, yy / ny, xx / nx + rr = (xx**2 + yy**2 + zz**2) ** 0.5 + return rr < frac_radius, rr diff --git a/src/mrinufft/extras/sim.py b/src/mrinufft/extras/sim.py new file mode 100644 index 00000000..edf3532b --- /dev/null +++ b/src/mrinufft/extras/sim.py @@ -0,0 +1,59 @@ +"""Simple MR simulation module.""" + +import numpy as np + +from typing import Sequence + + +def fse_simulation( + M0: np.ndarray, + T1: np.ndarray, + T2: np.ndarray, + TE: float | Sequence[float], + TR: float | Sequence[float], +) -> np.ndarray: + """Perform simple analytical Fast Spin Echo simulation. + + Assume that refocusing angles are 180° and + k-space center is sampled for echo in the Echo Train + (e.g., spiral or radial trajectory). + + Parameters + ---------- + M0 : np.ndarray + Input equilibrium magnetization. + T1 : np.ndarray + Input T1 in [ms]. + T2 : np.ndarray + Input T2 in [ms]. + TE : float | Sequence[float] + Sequence Echo Time in [ms]. + TR : float | Sequence[float] + Sequence Repetition Time in [ms]. + + Returns + ------- + signal : np.ndarray + Simulated signal of shape (nTE*nTR, *M0). + + """ + # preprocess sequence parameters + TE, TR = np.broadcast_arrays(np.atleast_1d(TE), np.atleast_1d(TR)[:, None]) + TE, TR = TE.ravel().astype(np.float32), TR.ravel().astype(np.float32) + + # preprocess tissue parameters + M0, T1, T2 = np.atleast_1d(M0), np.atleast_1d(T1), np.atleast_1d(T2) + M0, T1, T2 = M0[..., None], T1[..., None], T2[..., None] + T1 += 1e-9 + T2 += 1e-9 + + # compute signal + signal = M0 * (1 - np.exp(-(TR - TE) / T1)) * np.exp(-TE / T2) + + # post process + signal = signal[None, ...].swapaxes(0, -1)[..., 0] + signal = signal.squeeze() + if signal.size == 1: + signal = signal.item() + + return signal diff --git a/src/mrinufft/operators/__init__.py b/src/mrinufft/operators/__init__.py index c642d47d..5fa657a1 100644 --- a/src/mrinufft/operators/__init__.py +++ b/src/mrinufft/operators/__init__.py @@ -12,6 +12,7 @@ ) from .off_resonance import MRIFourierCorrected, get_interpolators_from_fieldmap from .stacked import MRIStackedNUFFT +from .subspace import MRISubspace # # load all the interfaces modules @@ -25,6 +26,7 @@ "FourierOperatorBase", "MRIFourierCorrected", "MRIStackedNUFFT", + "MRISubspace", "check_backend", "get_operator", "list_backends", diff --git a/src/mrinufft/operators/subspace.py b/src/mrinufft/operators/subspace.py new file mode 100644 index 00000000..4b916d69 --- /dev/null +++ b/src/mrinufft/operators/subspace.py @@ -0,0 +1,151 @@ +"""Subspace NUFFT Operator wrapper.""" + +from mrinufft._array_compat import _get_device, _to_interface +from mrinufft._utils import ARRAY_LIBS, get_array_module + + +from .base import FourierOperatorBase + + +class MRISubspace(FourierOperatorBase): + """Fourier Operator with subspace projection. + + This is a wrapper around the Fourier Operator to project + data on a low-rank subspace (e.g., dynamic and multi-contrast MRI). + + Parameters + ---------- + fourier_op: object of class FourierBase + the fourier operator to wrap + subspace_basis : np.ndarray + Low rank subspace basis of shape (K, T), + where K is the rank of the subspace and T is the number + of time frames or contrasts in the original image series. + Also supports Cupy arrays and Torch tensors. + use_gpu : bool, optional + Whether to use the GPU. Default is False. + Ignored if the fourier operator internally use only GPU (e.g., cupy) + or CPU (e.g., numpy) + + Notes + ----- + This extension add an axis on the leftmost position for both + image and k-space data: + + * Image: ``(B, C, XYZ)`` -> ``(T, B, C, XYZ)`` + * K-Space: ``(B, C, K)`` -> ``(T, B, C, K)`` + + with ``T`` representing time domain or contrast space (for dynamic and m + multi-contrast MRI, respectively). + + Similarly, k-space trajectory is expected to have the following shape: + ``(, N_shots, N_samples, dim)``. The flatten + version is also accepted: ``(, N_shots * N_samples, dim)`` + + """ + + def __init__(self, fourier_op, subspace_basis, use_gpu=False): + self._fourier_op = fourier_op + self.backend = fourier_op.backend + + self.n_batchs = fourier_op.n_batchs + self.n_coils = fourier_op.n_coils + self.shape = fourier_op.shape + self.smaps = fourier_op.smaps + self.autograd_available = fourier_op.autograd_available + + self.subspace_basis = subspace_basis + self.n_coeffs, self.n_frames = self.subspace_basis.shape + + def op(self, data, *args): + """ + Compute Forward Operation time/contrast-domain backprojection. + + Parameters + ---------- + data: numpy.ndarray + N-D subspace-projected image. + Also supports Cupy arrays and Torch tensors. + + Returns + ------- + numpy.ndarray + Time/contrast-domain k-space data. + """ + xp = get_array_module(data) + device = _get_device(data) + subspace_basis = _to_interface(self.subspace_basis, xp, device) + + # perform computation + y = 0.0 + for idx in range(self.n_coeffs): + + # select basis element + basis_element = subspace_basis[idx] + + # actual transform + _y = self._fourier_op.op(data[idx], *args) + _y = _y.reshape(*_y.shape[:-1], self.n_frames, -1) + + # back-project on time domain + y += basis_element.conj() * _y.swapaxes(-2, -1) + + return y[None, ...].swapaxes(0, -1)[..., 0] # bring back time domain in front + + def adj_op(self, coeffs, *args): + """ + Compute Adjoint Operation with susbpace projection. + + Parameters + ---------- + coeffs: numpy.ndarray + Time/contrast-domain k-space data. + Also supports Cupy arrays and Torch tensors. + + Returns + ------- + numpy.ndarray + Inverse Fourier transform of the subspace-projected k-space. + """ + xp = get_array_module(coeffs) + device = _get_device(coeffs) + subspace_basis = _to_interface(self.subspace_basis, xp, device) + coeffs_d = coeffs[..., None].swapaxes(0, -1)[0, ...] + + # perform computation + y = [] + for idx in range(self.n_coeffs): + + # select basis element + basis_element = subspace_basis[idx] + + # project on current subspace basis element + _coeffs_d = basis_element * coeffs_d + _coeffs_d = _coeffs_d.swapaxes(-2, -1).reshape(*coeffs_d.shape[:-2], -1) + + # actual transform + y.append(self._fourier_op.adj_op(_coeffs_d, *args)) + + # stack coefficients + y = xp.stack(y, axis=0) + + return y + + +def _get_arraylib_from_operator( + fourier_op, use_gpu +): # maybe that is usefull for MRIFourierCorrected constructor? Shall we move it in array_compat? + LUT = { + "MRIBartNUFFT": ("numpy", "numpy"), + "MRICufiNUFFT": ("cupy", "cupy"), + "MRIfinufft": ("numpy", "numpy"), + "MRIGpuNUFFT": ("cupy", "cupy"), + "MRInfft": ("numpy", "numpy"), + "MRIPynufft": ("numpy", "numpy"), + "MRISigpyNUFFT": ("numpy", "cupy"), + "MRITensorflowNUFFT": ("tensorflow", "tensorflow"), + "MRITorchKbNufft": ("torch", "torch"), + "TorchKbNUFFTcpu": ("torch", "torch"), + "TorchKbNUFFTgpu": ("torch", "torch"), + } + return ARRAY_LIBS[LUT[fourier_op.__class__.__name__][use_gpu]][0] diff --git a/tests/case_trajectories.py b/tests/case_trajectories.py index 161022be..e57c1879 100644 --- a/tests/case_trajectories.py +++ b/tests/case_trajectories.py @@ -71,8 +71,16 @@ def case_grid3D(self, N=16): # 1D grid is only use once, so we don't want to include systematically -# fsin the cases collection. +# in the cases collection. def case_grid1D(N=256): """Create a 1D cartesian grid of frequencies locations.""" freq_1d = sp.fft.fftfreq(N) return freq_1d.reshape(-1, 1), (N,) + + +# multicontrast is only use once, so we don't want to include systematically +# in the cases collection. +def case_multicontrast2D(Nt=64, Nc=10, Ns=500, N=64): + """Create a 2D radial trajectory.""" + trajectory = initialize_2D_radial(Nc * Nt, Ns, tilt="mri-golden") + return trajectory.reshape(Nt, Nc, Ns, 2), (N, N) diff --git a/tests/operators/test_subspace.py b/tests/operators/test_subspace.py new file mode 100644 index 00000000..68f5bae6 --- /dev/null +++ b/tests/operators/test_subspace.py @@ -0,0 +1,213 @@ +"""Test NUFFT with simultaneous subspace projection.""" + +from itertools import product + +import numpy as np +import numpy.testing as npt + +from pytest_cases import parametrize_with_cases, parametrize, fixture +from helpers import ( + param_array_interface, + to_interface, + from_interface, +) +from case_trajectories import case_multicontrast2D + +from mrinufft import get_operator +from mrinufft.operators import MRISubspace + + +@fixture(scope="module") +@parametrize( + "n_coils, n_batchs, sense, use_gpu", + list(product([1, 2, 4], [1, 3], [False, True], [False, True])), +) +@parametrize_with_cases("kspace_locs, shape", cases=[case_multicontrast2D]) +@parametrize( + backend=["finufft", "cufinufft", "gpunufft", "torchkbnufft-cpu", "torchkbnufft-gpu"] +) +def operator( + request, + kspace_locs, + shape, + n_coils=1, + n_batchs=1, + sense=None, + use_gpu=None, + backend="finufft", +): + """Generate a subspace operator.""" + if sense: + n_coils = 2 + smaps = 1j * np.random.rand(n_coils, *shape) + smaps += np.random.rand(n_coils, *shape) + smaps = smaps.astype(np.complex64) + smaps /= np.linalg.norm(smaps, axis=0) + else: + smaps = None + kspace_locs = kspace_locs.astype(np.float32) + _op = get_operator(backend)( + kspace_locs, + shape, + n_coils=n_coils, + n_batchs=n_batchs, + smaps=smaps, + squeeze_dims=False, + ) + + # generate random orthonormal basis + H = 1j * np.random.rand(10000, 64) + H += np.random.rand(10000, 64) + u, s, vh = np.linalg.svd(H, full_matrices=False) + _basis = u[:3] @ vh + _basis = _basis.astype(_op.cpx_dtype) + + # get reference operator + ref_operator = [ + get_operator(backend)( + samples=kspace_locs[n], + shape=shape, + n_coils=n_coils, + n_batchs=n_batchs, + smaps=smaps, + squeeze_dims=False, + ) + for n in range(_basis.shape[-1]) + ] + + return MRISubspace(_op, _basis, use_gpu), ref_operator, _basis + + +@fixture(scope="module") +def image_data(operator): + """Generate a random image.""" + _operator = operator[0] + if _operator._fourier_op.uses_sense: + shape = (_operator.n_coeffs, _operator.n_batchs, *_operator.shape) + else: + shape = ( + _operator.n_coeffs, + _operator.n_batchs, + _operator.n_coils, + *_operator.shape, + ) + img = (1j * np.random.rand(*shape)).astype(_operator._fourier_op.cpx_dtype) + img += np.random.rand(*shape).astype(_operator._fourier_op.cpx_dtype) + return img + + +@fixture(scope="module") +def kspace_data(operator): + """Generate a random image.""" + _operator = operator[0] + n_samples = int(_operator._fourier_op.n_samples / _operator.n_frames) + shape = (_operator.n_frames, _operator.n_batchs, _operator.n_coils, n_samples) + kspace = (1j * np.random.rand(*shape)).astype(_operator._fourier_op.cpx_dtype) + kspace += np.random.rand(*shape).astype(_operator._fourier_op.cpx_dtype) + return kspace + + +@param_array_interface +def test_subspace_op(operator, array_interface, image_data): + subspace_op, ref_op, subspace_basis = operator + + # get reference + # 1. back-project to time/contrast-domain image space + tmp = image_data[..., None].swapaxes(0, -1)[0, ...] + tmp = tmp @ subspace_basis.conj() + tmp = tmp[None, ...].swapaxes(0, -1)[..., 0] + + # 2. apply NUFFT + kspace_ref = [ref_op[n].op(tmp[n]) for n in range(tmp.shape[0])] + kspace_ref = np.stack(kspace_ref, axis=0) + + # actual computation + image_data = to_interface(image_data, array_interface) + kspace = from_interface(subspace_op.op(image_data), array_interface) + + npt.assert_array_almost_equal(kspace, kspace_ref) + + +@param_array_interface +def test_subspace_op_adj(operator, array_interface, kspace_data): + subspace_op, ref_op, subspace_basis = operator + + # get reference + # 1. apply adjoint NUFFT + tmp = [ref_op[n].adj_op(kspace_data[n]) for n in range(kspace_data.shape[0])] + + # 2. project to subspace image space + tmp = np.stack(tmp, axis=-1) + tmp = tmp @ subspace_basis.T + image_ref = tmp[None, ...].swapaxes(0, -1)[..., 0] + + # actual computation + kspace_data = to_interface(kspace_data, array_interface) + image = from_interface(subspace_op.adj_op(kspace_data), array_interface) + + npt.assert_array_almost_equal(image, image_ref) + + +@param_array_interface +def test_data_consistency( + operator, + array_interface, + image_data, + kspace_data, +): + """Test the data consistency operation.""" + subspace_op, _, _ = operator + image_data = to_interface(image_data, array_interface) + kspace_data = to_interface(kspace_data, array_interface) + + res = subspace_op.data_consistency(image_data, kspace_data) + tmp = subspace_op.op(image_data) + res2 = subspace_op.adj_op(tmp - kspace_data) + + res = res.reshape(-1, *subspace_op.shape) + res2 = res2.reshape(-1, *subspace_op.shape) + + res = from_interface(res, array_interface) + res2 = from_interface(res2, array_interface) + atol = 1e-3 + rtol = 1e-3 + # FIXME 2D Sense is not very accurate... + if len(subspace_op.shape) == 2 and subspace_op._fourier_op.uses_sense: + print("Reduced accuracy for 2D Sense") + atol = 1e-1 + atol = 1e-1 + + npt.assert_allclose(res, res2, atol=atol, rtol=rtol) + + +def test_data_consistency_readonly(operator, image_data, kspace_data): + """Test that the data consistency does not modify the input parameters data.""" + subspace_op, _, _ = operator + kspace_tmp = kspace_data.copy() + image_tmp = image_data.copy() + kspace_tmp.setflags(write=False) + image_tmp.setflags(write=False) + subspace_op.data_consistency(image_data, kspace_tmp) + npt.assert_equal(kspace_tmp, kspace_data) + npt.assert_equal(image_tmp, image_data) + + +def test_gradient_lipschitz(operator, image_data, kspace_data): + """Test the gradient lipschitz constant.""" + subspace_op, _, _ = operator + C = 1 if subspace_op._fourier_op.uses_sense else subspace_op.n_coils + img = ( + image_data.copy() + .reshape(subspace_op.n_coeffs, subspace_op.n_batchs, C, *subspace_op.shape) + .squeeze() + ) + for _ in range(10): + grad = subspace_op.data_consistency(img, kspace_data) + norm = np.linalg.norm(grad) + grad /= norm + np.copyto(img, grad.squeeze()) + norm_prev = norm + + # TODO: check that the value is "not too far" from 1 + # TODO: to do the same with density compensation + assert (norm - norm_prev) / norm_prev < 1e-3 From 7ae833c8b4ae70c54a2d51efc4f53b0311f8cbd1 Mon Sep 17 00:00:00 2001 From: mcencini Date: Mon, 30 Sep 2024 23:49:12 -0400 Subject: [PATCH 02/10] fix: styling. --- src/mrinufft/operators/subspace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mrinufft/operators/subspace.py b/src/mrinufft/operators/subspace.py index 4b916d69..d71f538c 100644 --- a/src/mrinufft/operators/subspace.py +++ b/src/mrinufft/operators/subspace.py @@ -134,7 +134,7 @@ def adj_op(self, coeffs, *args): def _get_arraylib_from_operator( fourier_op, use_gpu -): # maybe that is usefull for MRIFourierCorrected constructor? Shall we move it in array_compat? +): # maybe that is usefull for MRIFourierCorrected constructor? LUT = { "MRIBartNUFFT": ("numpy", "numpy"), "MRICufiNUFFT": ("cupy", "cupy"), From f04a4bb2f0ed90f517eeeaf82c63bc6db8af91af Mon Sep 17 00:00:00 2001 From: mcencini Date: Tue, 1 Oct 2024 14:46:35 -0400 Subject: [PATCH 03/10] fix tests: switching from random to deterministic data --- examples/example_subspace.py | 8 ++--- src/mrinufft/operators/subspace.py | 2 +- tests/operators/test_subspace.py | 51 ++++++++++++++++++------------ 3 files changed, 35 insertions(+), 26 deletions(-) diff --git a/examples/example_subspace.py b/examples/example_subspace.py index 84149b01..543e5a5d 100644 --- a/examples/example_subspace.py +++ b/examples/example_subspace.py @@ -157,13 +157,11 @@ def estimate_subspace_basis(train_data, ncoeff=4): # assume trajectory is reordered as (ncontrasts, nshots_per_contrast, nsamples_per_shot, ndims) samples = samples.reshape(16, ETL, *samples.shape[1:]) -# flatten ncontrasts and nshots_per_contrast axes -samples = samples.reshape(-1, *samples.shape[2:]) - # compute density compensation density = voronoi(samples) +density = density.reshape(16, ETL, samples.shape[-2]) -display_2D_trajectory(samples) +display_2D_trajectory(samples.reshape(-1, *samples.shape[2:])) # %% # Operator setup @@ -176,7 +174,7 @@ def estimate_subspace_basis(train_data, ncoeff=4): nufft = get_operator("finufft")( samples=samples, shape=mri_data.shape[-2:], - density=density, + density=density.ravel(), ) # Generate subspace-projected NUFFT operator diff --git a/src/mrinufft/operators/subspace.py b/src/mrinufft/operators/subspace.py index d71f538c..c65ec1ac 100644 --- a/src/mrinufft/operators/subspace.py +++ b/src/mrinufft/operators/subspace.py @@ -94,7 +94,7 @@ def op(self, data, *args): def adj_op(self, coeffs, *args): """ - Compute Adjoint Operation with susbpace projection. + Compute Adjoint Operation with subspace projection. Parameters ---------- diff --git a/tests/operators/test_subspace.py b/tests/operators/test_subspace.py index 68f5bae6..0a091dff 100644 --- a/tests/operators/test_subspace.py +++ b/tests/operators/test_subspace.py @@ -14,6 +14,7 @@ from case_trajectories import case_multicontrast2D from mrinufft import get_operator +from mrinufft.density import voronoi from mrinufft.operators import MRISubspace @@ -45,21 +46,25 @@ def operator( smaps /= np.linalg.norm(smaps, axis=0) else: smaps = None + kspace_locs = kspace_locs.astype(np.float32) + density = voronoi(kspace_locs) + density = density.reshape(*kspace_locs.shape[:-1]) + _op = get_operator(backend)( kspace_locs, shape, n_coils=n_coils, n_batchs=n_batchs, smaps=smaps, + weights=density.ravel(), squeeze_dims=False, ) # generate random orthonormal basis - H = 1j * np.random.rand(10000, 64) - H += np.random.rand(10000, 64) - u, s, vh = np.linalg.svd(H, full_matrices=False) - _basis = u[:3] @ vh + train_data = np.exp(-4 * np.arange(64)[:, None] / np.arange(1, 1000, 1)) + _, _, _basis = np.linalg.svd(train_data.T, full_matrices=False) + _basis = _basis[:3] _basis = _basis.astype(_op.cpx_dtype) # get reference operator @@ -70,6 +75,7 @@ def operator( n_coils=n_coils, n_batchs=n_batchs, smaps=smaps, + weights=density[n].ravel(), squeeze_dims=False, ) for n in range(_basis.shape[-1]) @@ -82,18 +88,20 @@ def operator( def image_data(operator): """Generate a random image.""" _operator = operator[0] + _shape = _operator.shape if _operator._fourier_op.uses_sense: - shape = (_operator.n_coeffs, _operator.n_batchs, *_operator.shape) + shape = (_operator.n_batchs, *_shape) else: - shape = ( - _operator.n_coeffs, - _operator.n_batchs, - _operator.n_coils, - *_operator.shape, - ) - img = (1j * np.random.rand(*shape)).astype(_operator._fourier_op.cpx_dtype) - img += np.random.rand(*shape).astype(_operator._fourier_op.cpx_dtype) - return img + shape = (_operator.n_batchs, _operator.n_coils, *_shape) + + img = (1j * np.zeros(shape)).astype(_operator._fourier_op.cpx_dtype) + img += np.zeros(shape).astype(_operator._fourier_op.cpx_dtype) + img[..., 32, 32] = 1.0 + + sig = np.exp(-4 * np.arange(64) / 200.0) + weights = _operator.subspace_basis @ sig + + return np.stack([weights[n] * img for n in range(_operator.n_coeffs)], axis=0) @fixture(scope="module") @@ -101,10 +109,13 @@ def kspace_data(operator): """Generate a random image.""" _operator = operator[0] n_samples = int(_operator._fourier_op.n_samples / _operator.n_frames) - shape = (_operator.n_frames, _operator.n_batchs, _operator.n_coils, n_samples) - kspace = (1j * np.random.rand(*shape)).astype(_operator._fourier_op.cpx_dtype) - kspace += np.random.rand(*shape).astype(_operator._fourier_op.cpx_dtype) - return kspace + shape = (_operator.n_batchs, _operator.n_coils, n_samples) + kspace = (1j * np.zeros(shape)).astype(_operator._fourier_op.cpx_dtype) + kspace += np.ones(shape).astype(_operator._fourier_op.cpx_dtype) + + sig = np.exp(-4 * np.arange(64) / 200.0) + + return sig[:, None, None, None] * kspace[None, ...] @param_array_interface @@ -125,7 +136,7 @@ def test_subspace_op(operator, array_interface, image_data): image_data = to_interface(image_data, array_interface) kspace = from_interface(subspace_op.op(image_data), array_interface) - npt.assert_array_almost_equal(kspace, kspace_ref) + npt.assert_allclose(kspace, kspace_ref, rtol=1e-6) @param_array_interface @@ -145,7 +156,7 @@ def test_subspace_op_adj(operator, array_interface, kspace_data): kspace_data = to_interface(kspace_data, array_interface) image = from_interface(subspace_op.adj_op(kspace_data), array_interface) - npt.assert_array_almost_equal(image, image_ref) + npt.assert_allclose(image, image_ref, rtol=5e-4) @param_array_interface From 2ee9d29ce30968a544dd1c20add8abb6c3d5c327 Mon Sep 17 00:00:00 2001 From: mcencini Date: Tue, 1 Oct 2024 14:51:01 -0400 Subject: [PATCH 04/10] fix typo --- examples/example_subspace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/example_subspace.py b/examples/example_subspace.py index 543e5a5d..15c8268f 100644 --- a/examples/example_subspace.py +++ b/examples/example_subspace.py @@ -5,7 +5,7 @@ An example to show how to setup a subspace NUFFT operator. -This examples show how to use the Subspace NUFFT operator to acquire +This examples shows how to use the Subspace NUFFT operator to acquire and reconstruct data while projecting onto a low rank spatiotemporal subspace. This can be useful e.g., for dynamic multi-contrast MRI. Hereafter a 2D spiral trajectory is used for demonstration. From 081eb8f996b7e947c748f30d1d3403f19e34806e Mon Sep 17 00:00:00 2001 From: mcencini Date: Tue, 1 Oct 2024 18:40:24 -0400 Subject: [PATCH 05/10] update test --- tests/operators/test_subspace.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/tests/operators/test_subspace.py b/tests/operators/test_subspace.py index 0a091dff..c276e4b7 100644 --- a/tests/operators/test_subspace.py +++ b/tests/operators/test_subspace.py @@ -14,7 +14,6 @@ from case_trajectories import case_multicontrast2D from mrinufft import get_operator -from mrinufft.density import voronoi from mrinufft.operators import MRISubspace @@ -48,8 +47,6 @@ def operator( smaps = None kspace_locs = kspace_locs.astype(np.float32) - density = voronoi(kspace_locs) - density = density.reshape(*kspace_locs.shape[:-1]) _op = get_operator(backend)( kspace_locs, @@ -57,7 +54,6 @@ def operator( n_coils=n_coils, n_batchs=n_batchs, smaps=smaps, - weights=density.ravel(), squeeze_dims=False, ) @@ -75,7 +71,6 @@ def operator( n_coils=n_coils, n_batchs=n_batchs, smaps=smaps, - weights=density[n].ravel(), squeeze_dims=False, ) for n in range(_basis.shape[-1]) @@ -136,7 +131,7 @@ def test_subspace_op(operator, array_interface, image_data): image_data = to_interface(image_data, array_interface) kspace = from_interface(subspace_op.op(image_data), array_interface) - npt.assert_allclose(kspace, kspace_ref, rtol=1e-6) + npt.assert_allclose(kspace, kspace_ref, rtol=5e-6) @param_array_interface From 20b46329f3228ff0a0fa8b87611dcc17439807b4 Mon Sep 17 00:00:00 2001 From: mcencini Date: Tue, 1 Oct 2024 18:41:12 -0400 Subject: [PATCH 06/10] fix typo --- examples/example_subspace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/example_subspace.py b/examples/example_subspace.py index 15c8268f..de6402d9 100644 --- a/examples/example_subspace.py +++ b/examples/example_subspace.py @@ -5,7 +5,7 @@ An example to show how to setup a subspace NUFFT operator. -This examples shows how to use the Subspace NUFFT operator to acquire +This example shows how to use the Subspace NUFFT operator to acquire and reconstruct data while projecting onto a low rank spatiotemporal subspace. This can be useful e.g., for dynamic multi-contrast MRI. Hereafter a 2D spiral trajectory is used for demonstration. From 50c86fd3423788410b000b850418dbc2df1eef22 Mon Sep 17 00:00:00 2001 From: mcencini Date: Tue, 1 Oct 2024 20:17:25 -0400 Subject: [PATCH 07/10] increase test tolerance --- tests/operators/test_subspace.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/operators/test_subspace.py b/tests/operators/test_subspace.py index c276e4b7..3aff81df 100644 --- a/tests/operators/test_subspace.py +++ b/tests/operators/test_subspace.py @@ -131,7 +131,7 @@ def test_subspace_op(operator, array_interface, image_data): image_data = to_interface(image_data, array_interface) kspace = from_interface(subspace_op.op(image_data), array_interface) - npt.assert_allclose(kspace, kspace_ref, rtol=5e-6) + npt.assert_allclose(kspace, kspace_ref, rtol=1e-3, atol=1e-3) @param_array_interface @@ -151,7 +151,7 @@ def test_subspace_op_adj(operator, array_interface, kspace_data): kspace_data = to_interface(kspace_data, array_interface) image = from_interface(subspace_op.adj_op(kspace_data), array_interface) - npt.assert_allclose(image, image_ref, rtol=5e-4) + npt.assert_allclose(image, image_ref, rtol=1e-3, atol=1e-3) @param_array_interface From d8c7278626eb4f38defd3a37fa3234655cdbb044 Mon Sep 17 00:00:00 2001 From: mcencini Date: Wed, 2 Oct 2024 22:08:43 -0400 Subject: [PATCH 08/10] address PR review \!docs_build --- docs/nufft.rst | 4 +- examples/example_subspace.py | 43 ++++++++++++-------- src/mrinufft/extras/data.py | 3 +- src/mrinufft/extras/field_map.py | 12 ++++-- src/mrinufft/extras/sim.py | 4 +- src/mrinufft/operators/subspace.py | 64 ++++++++++++++++++++++++------ tests/case_trajectories.py | 2 +- tests/operators/test_subspace.py | 46 +++++++++++---------- 8 files changed, 117 insertions(+), 61 deletions(-) diff --git a/docs/nufft.rst b/docs/nufft.rst index 8b5a44b5..7051bf75 100644 --- a/docs/nufft.rst +++ b/docs/nufft.rst @@ -168,7 +168,7 @@ Where :math:`E_mn = e^i\Delta\omega_0(u_n)t_m`. Subspace Projection Model ~~~~~~~~~~~~~~~~~~~~~~~~~ -In several MRI applications, such as dynamic or quantitative MRI, a single acquisition provides a stack of images, each representing a single time frame (for dynamic MRI) or a single contrast (for quantitative MRI). +In several MRI applications, such as dynamic or quantitative MRI, a single acquisition provides a stack of two- or three-dimensional images, each representing a single time frame (for dynamic MRI) or a single contrast (for quantitative MRI). To achieve a clinically feasible scan time, each frame or contrast is acquired with a different aggressively undersampled k-space trajectory. In this context, the single-coil acquisition model becomes: .. math:: @@ -216,7 +216,7 @@ the projection operator :math:`\boldsymbol{\Phi}` commutes with the Fourier tran \tilde{\boldsymbol{y}} = \Phi \mathcal{F}_Omega(\boldsymbol{\alpha}) + \boldsymbol{n} -that is, computation now involves :math:`K \ll T` Fourier Transform operations, each with the same sampling trajectory, which can be computed levaraging efficient NUFFT implementations for conventional static MRI. +that is, computation now involves :math:`K \ll T` Fourier Transform operations, each with the same sampling trajectory, which can be computed by levaraging efficient NUFFT implementations for conventional static MRI. The Non Uniform Fast Fourier Transform ====================================== diff --git a/examples/example_subspace.py b/examples/example_subspace.py index de6402d9..1d24e863 100644 --- a/examples/example_subspace.py +++ b/examples/example_subspace.py @@ -5,8 +5,9 @@ An example to show how to setup a subspace NUFFT operator. -This example shows how to use the Subspace NUFFT operator to acquire -and reconstruct data while projecting onto a low rank spatiotemporal subspace. +This example shows how to use the Subspace NUFFT operator to acquire data +and reconstruct an image (or a volume for 3D imaging) +while projecting onto a low rank spatiotemporal subspace. This can be useful e.g., for dynamic multi-contrast MRI. Hereafter a 2D spiral trajectory is used for demonstration. @@ -56,8 +57,8 @@ # ------------------- # # As an example, we simulate a simple monoexponential Spin Echo acquisition. -# We assume that refocusing train is constant (flip angle set to 180 degrees) -# and k-space center is sampled at each TE (as for spiral or radial imaging). +# We assume that k-space center is sampled at each TE (as in spiral or radial imaging) +# and a constant flip angle train is constant set to 180 degrees. # In this way, we obtain an image for each of the shots in the echo train. from mrinufft.extras import fse_simulation @@ -72,9 +73,9 @@ # ------------------- # # Here, we generate temporal subspace basis -# by Singular Value Decomposition of an ensemble -# of simulated signals in the same property range. -# our object. +# by performing the Singular Value Decomposition of an ensemble +# of simulated signals corresponding to the same T1-T2 range +# of our object. def make_grid(T1range, T2range, natoms=100): @@ -83,7 +84,7 @@ def make_grid(T1range, T2range, natoms=100): T1grid = np.linspace(T1range[0], T1range[1], num=natoms, dtype=np.float32) T2grid = np.linspace(T2range[0], T2range[1], num=natoms, dtype=np.float32) - # Create combined grid + # Cartesian product T1grid, T2grid = np.meshgrid(T1grid, T2grid) return T1grid.ravel(), T2grid.ravel() @@ -125,7 +126,7 @@ def estimate_subspace_basis(train_data, ncoeff=4): # %% # Here, we simulate Brainweb FSE data with the same -# sequence parameters used for the subspace estimation. +# sequence parameters as those used for the subspace estimation. mri_data = fse_simulation(M0, T1, T2, TE, TR).astype(np.float32) mri_data = np.ascontiguousarray(mri_data) @@ -155,11 +156,12 @@ def estimate_subspace_basis(train_data, ncoeff=4): ) # assume trajectory is reordered as (ncontrasts, nshots_per_contrast, nsamples_per_shot, ndims) -samples = samples.reshape(16, ETL, *samples.shape[1:]) +# with contrast axis sorted by ascending TEs +samples = samples.reshape(ETL, 16, *samples.shape[1:]) # compute density compensation density = voronoi(samples) -density = density.reshape(16, ETL, samples.shape[-2]) +density = density.reshape(ETL, 16, samples.shape[-2]) display_2D_trajectory(samples.reshape(-1, *samples.shape[2:])) @@ -178,7 +180,6 @@ def estimate_subspace_basis(train_data, ncoeff=4): ) # Generate subspace-projected NUFFT operator -multicontrast_nufft = MRISubspace(nufft, subspace_basis=np.eye(ETL)) subspace_nufft = MRISubspace(nufft, subspace_basis=basis) # %% @@ -189,17 +190,27 @@ def estimate_subspace_basis(train_data, ncoeff=4): # This can be simply obtained by using an identity matrix # with of shape (ETL, ETL) (= number of contrasts) as a subspace basis. -kspace = multicontrast_nufft.op(mri_data) +multicontrast_nufft = [ + get_operator("finufft")( + samples=samples[n], + shape=mri_data.shape[-2:], + density=density[n].ravel(), + ) + for n in range(basis.shape[-1]) +] +kspace = np.stack([multicontrast_nufft[n].op(mri_data[n]) for n in range(ETL)], axis=0) # %% # Image reconstruction # ----------------------- # -# We now reconstruct both using the subspace expanded operator -# and the zero-filling operator followed by projection in image space +# We now reconstruct both using the subspace expanded adjoint operator +# and the zero-filling adjoint operator followed by projection in image space # Reconstruct with projection in image space -zerofilled_data_adj = multicontrast_nufft.adj_op(kspace) +zerofilled_data_adj = np.stack( + [multicontrast_nufft[n].adj_op(kspace[n]) for n in range(ETL)], axis=0 +) zerofilled_display = (zerofilled_data_adj.T @ basis.T).T zerofilled_display = np.concatenate( [abs(coeff) / abs(coeff).max() for coeff in zerofilled_display], axis=-1 diff --git a/src/mrinufft/extras/data.py b/src/mrinufft/extras/data.py index 92a797c6..0f1ada64 100644 --- a/src/mrinufft/extras/data.py +++ b/src/mrinufft/extras/data.py @@ -4,7 +4,8 @@ def get_brainweb_map(sub_id: int) -> np.ndarray: - """Get MRI parametric maps from a brainweb crisp segmentation. + """ + Get M0, T1 and T2 parametric maps from a brainweb crisp segmentation. Output maps have the same shape as the tissue segmentation. diff --git a/src/mrinufft/extras/field_map.py b/src/mrinufft/extras/field_map.py index 5dfcf281..0bcde4de 100644 --- a/src/mrinufft/extras/field_map.py +++ b/src/mrinufft/extras/field_map.py @@ -14,7 +14,7 @@ def make_b0map(shape, b0range=(-300, 300), mask=None): b0range : tuple[float], optional Frequency shift range in [Hz]. The default is (-300, 300). mask : np.ndarray - Spatial support of the objec. If not provided, + Spatial support of the object. If not provided, build a radial mask with radius = 0.3 * shape Returns @@ -26,7 +26,9 @@ def make_b0map(shape, b0range=(-300, 300), mask=None): Spatial support binary mask. """ - assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.") + if np.unique(shape).size != 1: + raise ValueError("Only isotropic matriex are supported.") + ndim = len(shape) if ndim == 2: radial_mask, fieldmap = _make_disk(shape) @@ -57,7 +59,7 @@ def make_t2smap(shape, t2svalue=15.0, mask=None): t2svalue : float, optional Object T2* in [ms]. The default is 15.0. mask : np.ndarray - Spatial support of the objec. If not provided, + Spatial support of the object. If not provided, build a radial mask with radius = 0.3 * shape Returns @@ -68,7 +70,9 @@ def make_t2smap(shape, t2svalue=15.0, mask=None): Spatial support binary mask. """ - assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.") + if np.unique(shape).size != 1: + raise ValueError("Only isotropic matriex are supported.") + ndim = len(shape) if ndim == 2: radial_mask, fieldmap = _make_disk(shape) diff --git a/src/mrinufft/extras/sim.py b/src/mrinufft/extras/sim.py index edf3532b..7307838a 100644 --- a/src/mrinufft/extras/sim.py +++ b/src/mrinufft/extras/sim.py @@ -15,8 +15,8 @@ def fse_simulation( """Perform simple analytical Fast Spin Echo simulation. Assume that refocusing angles are 180° and - k-space center is sampled for echo in the Echo Train - (e.g., spiral or radial trajectory). + k-space center is sampled for each echo in the Echo Train + (e.g., as in spiral or radial imaging). Parameters ---------- diff --git a/src/mrinufft/operators/subspace.py b/src/mrinufft/operators/subspace.py index c65ec1ac..3bdcce59 100644 --- a/src/mrinufft/operators/subspace.py +++ b/src/mrinufft/operators/subspace.py @@ -11,32 +11,31 @@ class MRISubspace(FourierOperatorBase): """Fourier Operator with subspace projection. This is a wrapper around the Fourier Operator to project - data on a low-rank subspace (e.g., dynamic and multi-contrast MRI). + data onto a low-rank subspace (e.g., dynamic and multi-contrast MRI). Parameters ---------- fourier_op: object of class FourierBase the fourier operator to wrap subspace_basis : np.ndarray - Low rank subspace basis of shape (K, T), + Low rank subspace basis of shape ``(K, T)``, where K is the rank of the subspace and T is the number of time frames or contrasts in the original image series. Also supports Cupy arrays and Torch tensors. use_gpu : bool, optional Whether to use the GPU. Default is False. - Ignored if the fourier operator internally use only GPU (e.g., cupy) - or CPU (e.g., numpy) + Ignored if the Fourier operator internally use only GPU (e.g., Cupy) + or CPU (e.g., Numpy) Notes ----- - This extension add an axis on the leftmost position for both - image and k-space data: + This extension adds a new axis for both image and k-space data: - * Image: ``(B, C, XYZ)`` -> ``(T, B, C, XYZ)`` - * K-Space: ``(B, C, K)`` -> ``(T, B, C, K)`` + * Image: ``(B, C, XYZ)`` -> ``(B, S, C, XYZ)`` + * K-Space: ``(B, C, K)`` -> ``(B, T, C, K)`` - with ``T`` representing time domain or contrast space (for dynamic and m - multi-contrast MRI, respectively). + with ``S`` representing the subspace index and ``T`` representing time + domain or contrast space (for dynamic and multi-contrast MR, respectively). Similarly, k-space trajectory is expected to have the following shape: ``(, N_shots, N_samples, dim)``. The flatten @@ -76,6 +75,18 @@ def op(self, data, *args): device = _get_device(data) subspace_basis = _to_interface(self.subspace_basis, xp, device) + # if required, move subspace index axis to leftmost position + if self.n_batchs != 1 or data.shape[0] == 1: # non-squeezed data + data_d = data.swapaxes(0, 1) + else: + data_d = data + + # enforce data contiguity + if xp.__name__ == "torch": + data_d = data_d.contiguous() + else: + data_d = xp.ascontiguousarray(data_d) + # perform computation y = 0.0 for idx in range(self.n_coeffs): @@ -84,13 +95,19 @@ def op(self, data, *args): basis_element = subspace_basis[idx] # actual transform - _y = self._fourier_op.op(data[idx], *args) + _y = self._fourier_op.op(data_d[idx], *args) _y = _y.reshape(*_y.shape[:-1], self.n_frames, -1) # back-project on time domain y += basis_element.conj() * _y.swapaxes(-2, -1) - return y[None, ...].swapaxes(0, -1)[..., 0] # bring back time domain in front + y = y[None, ...].swapaxes(0, -1)[..., 0] + + # bring back time/contrast axis to original position (B, T, ...) + if self.n_batchs != 1 or data.shape[0] == 1: # non-squeezed data + y = y.swapaxes(0, 1) + + return y def adj_op(self, coeffs, *args): """ @@ -110,7 +127,14 @@ def adj_op(self, coeffs, *args): xp = get_array_module(coeffs) device = _get_device(coeffs) subspace_basis = _to_interface(self.subspace_basis, xp, device) - coeffs_d = coeffs[..., None].swapaxes(0, -1)[0, ...] + + # if required, move time/contrast axis to leftmost position + if self.n_batchs != 1 or coeffs.shape[0] == 1: # non-squeezed data + coeffs_d = coeffs.swapaxes(0, 1) + else: + coeffs_d = coeffs + + coeffs_d = coeffs_d[..., None].swapaxes(0, -1)[0, ...] # perform computation y = [] @@ -123,14 +147,28 @@ def adj_op(self, coeffs, *args): _coeffs_d = basis_element * coeffs_d _coeffs_d = _coeffs_d.swapaxes(-2, -1).reshape(*coeffs_d.shape[:-2], -1) + # enforce data contiguity + if xp.__name__ == "torch": + _coeffs_d = _coeffs_d.contiguous() + else: + _coeffs_d = xp.ascontiguousarray(_coeffs_d) + # actual transform y.append(self._fourier_op.adj_op(_coeffs_d, *args)) # stack coefficients y = xp.stack(y, axis=0) + # bring back subspace index to original position (B, S, ...) + if self.n_batchs != 1 or coeffs.shape[0] == 1: + y = y.swapaxes(0, 1) + return y + @property + def n_samples(self): + return self._fourier_op.n_samples + def _get_arraylib_from_operator( fourier_op, use_gpu diff --git a/tests/case_trajectories.py b/tests/case_trajectories.py index e57c1879..4fc04b8c 100644 --- a/tests/case_trajectories.py +++ b/tests/case_trajectories.py @@ -80,7 +80,7 @@ def case_grid1D(N=256): # multicontrast is only use once, so we don't want to include systematically # in the cases collection. -def case_multicontrast2D(Nt=64, Nc=10, Ns=500, N=64): +def case_multicontrast2D(Nt=48, Nc=10, Ns=500, N=64): """Create a 2D radial trajectory.""" trajectory = initialize_2D_radial(Nc * Nt, Ns, tilt="mri-golden") return trajectory.reshape(Nt, Nc, Ns, 2), (N, N) diff --git a/tests/operators/test_subspace.py b/tests/operators/test_subspace.py index 3aff81df..a166c8d8 100644 --- a/tests/operators/test_subspace.py +++ b/tests/operators/test_subspace.py @@ -58,9 +58,9 @@ def operator( ) # generate random orthonormal basis - train_data = np.exp(-4 * np.arange(64)[:, None] / np.arange(1, 1000, 1)) + train_data = np.exp(-4 * np.arange(48)[:, None] / np.arange(1, 1000, 1)) _, _, _basis = np.linalg.svd(train_data.T, full_matrices=False) - _basis = _basis[:3] + _basis = _basis[:5] _basis = _basis.astype(_op.cpx_dtype) # get reference operator @@ -93,24 +93,24 @@ def image_data(operator): img += np.zeros(shape).astype(_operator._fourier_op.cpx_dtype) img[..., 32, 32] = 1.0 - sig = np.exp(-4 * np.arange(64) / 200.0) + sig = np.exp(-4 * np.arange(48) / 200.0) weights = _operator.subspace_basis @ sig - - return np.stack([weights[n] * img for n in range(_operator.n_coeffs)], axis=0) + image = np.stack([weights[n] * img for n in range(_operator.n_coeffs)], axis=0) + return np.ascontiguousarray(image.swapaxes(0, 1)) @fixture(scope="module") def kspace_data(operator): """Generate a random image.""" _operator = operator[0] - n_samples = int(_operator._fourier_op.n_samples / _operator.n_frames) + n_samples = int(_operator.n_samples / _operator.n_frames) shape = (_operator.n_batchs, _operator.n_coils, n_samples) kspace = (1j * np.zeros(shape)).astype(_operator._fourier_op.cpx_dtype) kspace += np.ones(shape).astype(_operator._fourier_op.cpx_dtype) - sig = np.exp(-4 * np.arange(64) / 200.0) - - return sig[:, None, None, None] * kspace[None, ...] + sig = np.exp(-4 * np.arange(48) / 200.0) + kspace = sig[:, None, None, None] * kspace[None, ...] + return np.ascontiguousarray(kspace.swapaxes(0, 1)) @param_array_interface @@ -119,13 +119,15 @@ def test_subspace_op(operator, array_interface, image_data): # get reference # 1. back-project to time/contrast-domain image space - tmp = image_data[..., None].swapaxes(0, -1)[0, ...] + tmp = image_data.swapaxes(0, 1) + tmp = tmp[..., None].swapaxes(0, -1)[0, ...] tmp = tmp @ subspace_basis.conj() tmp = tmp[None, ...].swapaxes(0, -1)[..., 0] # 2. apply NUFFT - kspace_ref = [ref_op[n].op(tmp[n]) for n in range(tmp.shape[0])] - kspace_ref = np.stack(kspace_ref, axis=0) + tmp = [ref_op[n].op(tmp[n]) for n in range(tmp.shape[0])] + tmp = np.stack(tmp, axis=0) + kspace_ref = tmp.swapaxes(0, 1) # actual computation image_data = to_interface(image_data, array_interface) @@ -140,12 +142,14 @@ def test_subspace_op_adj(operator, array_interface, kspace_data): # get reference # 1. apply adjoint NUFFT - tmp = [ref_op[n].adj_op(kspace_data[n]) for n in range(kspace_data.shape[0])] + tmp = kspace_data.swapaxes(0, 1) + tmp = [ref_op[n].adj_op(tmp[n]) for n in range(tmp.shape[0])] # 2. project to subspace image space tmp = np.stack(tmp, axis=-1) tmp = tmp @ subspace_basis.T - image_ref = tmp[None, ...].swapaxes(0, -1)[..., 0] + tmp = tmp[None, ...].swapaxes(0, -1)[..., 0] + image_ref = tmp.swapaxes(0, 1) # actual computation kspace_data = to_interface(kspace_data, array_interface) @@ -201,17 +205,15 @@ def test_data_consistency_readonly(operator, image_data, kspace_data): def test_gradient_lipschitz(operator, image_data, kspace_data): """Test the gradient lipschitz constant.""" subspace_op, _, _ = operator - C = 1 if subspace_op._fourier_op.uses_sense else subspace_op.n_coils - img = ( - image_data.copy() - .reshape(subspace_op.n_coeffs, subspace_op.n_batchs, C, *subspace_op.shape) - .squeeze() - ) + + img = image_data.copy() + kspace = kspace_data.copy() + for _ in range(10): - grad = subspace_op.data_consistency(img, kspace_data) + grad = subspace_op.data_consistency(img, kspace) norm = np.linalg.norm(grad) grad /= norm - np.copyto(img, grad.squeeze()) + np.copyto(img, grad.reshape(*img.shape)) norm_prev = norm # TODO: check that the value is "not too far" from 1 From f7c9c296dbd21ab777d3768a10eaed505f767726 Mon Sep 17 00:00:00 2001 From: mcencini Date: Wed, 2 Oct 2024 22:10:56 -0400 Subject: [PATCH 09/10] lint \!docs_build --- src/mrinufft/operators/subspace.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mrinufft/operators/subspace.py b/src/mrinufft/operators/subspace.py index 3bdcce59..bc2151c9 100644 --- a/src/mrinufft/operators/subspace.py +++ b/src/mrinufft/operators/subspace.py @@ -167,6 +167,7 @@ def adj_op(self, coeffs, *args): @property def n_samples(self): + """Return the number of samples used by the operator.""" return self._fourier_op.n_samples From 1849bcb16b98c016b9d0100c05bb5bae50ae9bcf Mon Sep 17 00:00:00 2001 From: mcencini Date: Wed, 2 Oct 2024 23:03:13 -0400 Subject: [PATCH 10/10] fieldmap actually worked for rectangular matrices too \!docs_build --- examples/example_subspace.py | 4 ++-- src/mrinufft/extras/field_map.py | 8 +------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/examples/example_subspace.py b/examples/example_subspace.py index 1d24e863..4d04307e 100644 --- a/examples/example_subspace.py +++ b/examples/example_subspace.py @@ -1,7 +1,7 @@ """ -====================== +======================= Subspace NUFFT Operator -====================== +======================= An example to show how to setup a subspace NUFFT operator. diff --git a/src/mrinufft/extras/field_map.py b/src/mrinufft/extras/field_map.py index 0bcde4de..d51d53d4 100644 --- a/src/mrinufft/extras/field_map.py +++ b/src/mrinufft/extras/field_map.py @@ -10,7 +10,7 @@ def make_b0map(shape, b0range=(-300, 300), mask=None): Parameters ---------- shape : tuple[int] - Matrix size. Only supports isotropic matrices. + Matrix size. b0range : tuple[float], optional Frequency shift range in [Hz]. The default is (-300, 300). mask : np.ndarray @@ -26,9 +26,6 @@ def make_b0map(shape, b0range=(-300, 300), mask=None): Spatial support binary mask. """ - if np.unique(shape).size != 1: - raise ValueError("Only isotropic matriex are supported.") - ndim = len(shape) if ndim == 2: radial_mask, fieldmap = _make_disk(shape) @@ -70,9 +67,6 @@ def make_t2smap(shape, t2svalue=15.0, mask=None): Spatial support binary mask. """ - if np.unique(shape).size != 1: - raise ValueError("Only isotropic matriex are supported.") - ndim = len(shape) if ndim == 2: radial_mask, fieldmap = _make_disk(shape)