.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated/autoexamples/04-model-based-imaging/02-nonlinear-inversion.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code or to run this example in your browser via Binder. .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_autoexamples_04-model-based-imaging_02-nonlinear-inversion.py: ============================================== Nonlinear inversion from k-space ============================================== The scope of this notebook is to reconstruct T2 maps straight from k-space, with the signal model inside the forward operator, and to say where the time and the memory go. Physics-based reconstruction removes the intermediate images. The forward operator is a chain .. math:: F = P \, \mathcal{F} \, C \, M -- sampling, Fourier encoding, coil sensitivities, and the **signal model** -- and the parameter maps are solved for directly against the k-space that was measured. Only the last factor changes with the sequence, and it is the only one TorchSim supplies: :class:`~torchsim.recon.ModelOperator` turns any simulator into it, and the encoding comes from mri-nufft. Unlike a subspace this stays nonlinear, so it needs a starting guess and a loop around it, and it pays for that with a model of any number of parameters where a basis would have to span their product. The comparison here is against gridding, by iteratively regularized Gauss-Newton. Wang X, Tan Z, Scholand N, Roeloffs V, Uecker M. *Physics-based reconstruction methods for magnetic resonance imaging.* Phil Trans R Soc A 379:20200196 (2021). .. GENERATED FROM PYTHON SOURCE LINES 34-38 .. colab-link:: :needs_gpu: 1 !pip install torchsim brainweb-dl cmap mri-nufft[finufft,cufinufft] deepinv .. GENERATED FROM PYTHON SOURCE LINES 40-45 The phantom is BrainWeb's, reached through ``brainweb-dl``: ``get_mri`` fetches the fuzzy tissue memberships, and the package ships the table of relaxation times that goes with them -- which is what the two standard library imports read. .. GENERATED FROM PYTHON SOURCE LINES 46-138 .. code-block:: Python import csv from pathlib import Path import brainweb_dl from brainweb_dl import get_mri .. GENERATED FROM PYTHON SOURCE LINES 139-146 The Fourier encoding is not TorchSim's and never will be. ``mri-nufft`` supplies the radial trajectory and the non-uniform transform that plays it; ``deepinv`` supplies the :class:`~deepinv.physics.LinearPhysics` base class the encoding operator is written against, and the linear solver a Gauss-Newton step hands its linearized problem to. Anything exposing ``A`` and ``A_adjoint`` composes with what TorchSim supplies. .. GENERATED FROM PYTHON SOURCE LINES 147-151 .. code-block:: Python import mrinufft from deepinv.physics import LinearPhysics from mrinufft.trajectories import initialize_2D_radial .. GENERATED FROM PYTHON SOURCE LINES 152-156 From TorchSim: the sequence, the estimator the contrast-then-fit routes need, and :class:`~torchsim.recon.ModelOperator`, which is the signal model as a factor of the forward operator. .. GENERATED FROM PYTHON SOURCE LINES 157-167 .. code-block:: Python import time import numpy as np import torch from torchsim.estimators import DictionaryMatcher from torchsim.recon import GaussNewton, ModelOperator, Schedule, iterative from torchsim.simulators import MultiEchoSimulator .. GENERATED FROM PYTHON SOURCE LINES 168-171 What the experiment is: a 96 matrix read as 16 radial spokes per echo, eight echoes, and the rank the baseline's estimator compresses to. .. GENERATED FROM PYTHON SOURCE LINES 172-184 .. code-block:: Python SIZE = 96 ECHOES = 8 SPOKES = 16 SAMPLES = 192 RANK = 3 # The GPU transform is used when it is both installed and usable; the # simulation follows it, so the images and the operator meet on one device. on_gpu = torch.cuda.is_available() and mrinufft.check_backend("cufinufft") device = "cuda" if on_gpu else "cpu" backend = "cufinufft" if on_gpu else "finufft" .. GENERATED FROM PYTHON SOURCE LINES 185-193 Phantom ------- BrainWeb subject 0, slice 90, resampled to the matrix reconstructed here. BrainWeb publishes fuzzy memberships rather than labels, so weighting the tabulated relaxation times by them gives a T2 map whose mixed voxels sit between the pure ones, known everywhere. .. GENERATED FROM PYTHON SOURCE LINES 194-231 .. GENERATED FROM PYTHON SOURCE LINES 232-240 Sequence and sampling --------------------- A multi-echo spin echo on a golden-angle radial trajectory that rotates between echoes. Sixteen spokes per echo across a 96-sample matrix is roughly ninefold undersampled. The protocol stays on the host; :class:`~torchsim.recon.ModelOperator` takes it wherever the maps are. .. GENERATED FROM PYTHON SOURCE LINES 241-296 .. code-block:: Python TE = torch.linspace(10.0, 150.0, ECHOES) simulator = MultiEchoSimulator(TE=TE) images = ( torch.as_tensor(simulator.to(device).simulate(T2=T2_true)).to(torch.complex64) * M0_true.to(torch.complex64)[..., None] ) trajectory = ( initialize_2D_radial(SPOKES * ECHOES, SAMPLES, tilt="golden") .astype(np.float32) .reshape(ECHOES, SPOKES * SAMPLES, 2) ) build = mrinufft.get_operator(backend) per_echo = [ build(trajectory[echo], (SIZE, SIZE), n_coils=1, squeeze_dims=False, density=True) for echo in range(ECHOES) ] class RadialEncoding(LinearPhysics): """``(batch, echoes, x, y)`` images to k-space, one trajectory per echo. This is the whole of ``P F C`` for this experiment, and none of it is TorchSim's: it wraps mri-nufft, which is what a real pipeline would do with its own trajectory, its own density compensation and its own coils. """ def A(self, x, **kwargs): return torch.stack( [per_echo[e].op(x[:, e][:, None])[:, 0] for e in range(ECHOES)], 1 ) def A_adjoint(self, y, **kwargs): return torch.stack( [per_echo[e].adj_op(y[:, e][:, None])[:, 0] for e in range(ECHOES)], 1 ) encoding = RadialEncoding() kspace = encoding.A(images.movedim(-1, 0)[None]) # The k-space is scaled so the adjoint image peaks at one. Every damping # weight below is then a number about the model rather than about the # receiver gain, which is what makes one choice of it transferable. gridded = encoding.A_adjoint(kspace)[0].movedim(0, -1) scale = float(gridded.abs().max()) kspace, gridded = kspace / scale, gridded / scale .. rst-class:: sphx-glr-script-out .. code-block:: none 16 spokes per echo: 9x undersampled .. GENERATED FROM PYTHON SOURCE LINES 297-301 The maps on the left are what every route recovers; the spokes on the right are all that is measured of them, one echo's worth, rotated by the golden angle from the echo before. .. GENERATED FROM PYTHON SOURCE LINES 302-337 .. image-sg:: /generated/autoexamples/04-model-based-imaging/images/sphx_glr_02-nonlinear-inversion_001.png :alt: ground truth, proton density, 16 spokes per echo, 3 of 8 shown :srcset: /generated/autoexamples/04-model-based-imaging/images/sphx_glr_02-nonlinear-inversion_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 338-346 Estimator for the baseline -------------------------- The baseline reconstructs images and then fits them, so it needs an estimator, stated over a compressed basis: three directions hold essentially all of an eight-echo exponential. The nonlinear route has no such step -- its answer is the maps. .. GENERATED FROM PYTHON SOURCE LINES 347-379 .. code-block:: Python grid = torch.linspace(20.0, 400.0, 500) mapping = DictionaryMatcher(simulator).fit(T2=grid, M0=1.0, rank=RANK, seed=0) .. rst-class:: sphx-glr-script-out .. code-block:: none rank 3 of 8 contrasts keeps 0.999992 .. GENERATED FROM PYTHON SOURCE LINES 380-389 Baseline reconstruction ----------------------- Gridding is the adjoint with a density weighting -- one pass, smooth, biased -- and the estimator above turns its eight images into a T2 map. Sixteen spokes of 192 samples is 3072 measurements against 9216 unknowns, so each echo alone is underdetermined and iterating has nothing to converge to. Accuracy comes from a constraint across the echoes, which is the model. .. GENERATED FROM PYTHON SOURCE LINES 390-400 .. code-block:: Python adjoint = mapping(gridded)["T2"] .. rst-class:: sphx-glr-script-out .. code-block:: none adjoint per echo 0.0s T2 error 30.6 ms (20.6%) tensor([[ 52.7455, 61.8838, 58.0762, ..., 96.9138, 74.0681, 63.4068], [ 47.4148, 49.6994, 50.4609, ..., 91.5832, 66.4529, 77.8758], [ 52.7455, 60.3607, 74.0681, ..., 76.3527, 71.0220, 78.6373], ..., [ 71.0220, 80.1603, 106.0521, ..., 74.8297, 51.9840, 64.1683], [103.0060, 123.5671, 128.1363, ..., 93.8677, 64.9299, 56.5531], [122.0441, 109.0982, 111.3828, ..., 82.4449, 81.6834, 69.4990]]) .. GENERATED FROM PYTHON SOURCE LINES 401-416 Nonlinear model --------------- The signal model stays inside the forward operator and the maps are solved for against k-space directly. Two things are declared: * **what is unknown** -- ``T2``, plus the complex amplitude the operator carries for it, which is proton density and receive phase together; * **what T2 may be** -- a box bound, kept by solving for a transformed variable so no iterate leaves it. That matters more here than in a fit: the model is evaluated at every voxel to predict every k-space sample, so one unphysical voxel corrupts the whole residual. An equality constraint would be written into the model instead. .. GENERATED FROM PYTHON SOURCE LINES 417-425 .. code-block:: Python operator = ModelOperator(simulator, "T2", bounds={"T2": (20.0, 400.0)}) # The amplitude starts from the first gridded echo, which is nearly free and # is most of what makes the first Newton step sensible. initial = operator.initial((1, SIZE, SIZE), T2=100.0).to(device) initial[0, ..., 1] = gridded[..., 0].real initial[0, ..., 2] = gridded[..., 0].imag .. GENERATED FROM PYTHON SOURCE LINES 426-432 An iteratively regularized Gauss-Newton: linearize, solve, step, lower the damping. TorchSim supplies the loop and the derivative but not the linear solver -- :func:`~torchsim.recon.iterative` hands the linearized problem to the same deepinv routine the baseline called. A proximal solver under a wavelet prior is a change to that one argument. .. GENERATED FROM PYTHON SOURCE LINES 433-454 .. code-block:: Python found = GaussNewton( Schedule(initial=1e-3, factor=0.5, minimum=1e-7), solve=iterative(max_iter=20), max_iterations=8, ).minimize(operator, kspace, initial, encoding=encoding) # No fit afterwards: the maps are what was solved for. modelled = operator.split(found.x)["T2"][0] .. rst-class:: sphx-glr-script-out .. code-block:: none model-based 2.3s T2 error 17.1 ms (12.6%) residual 5.529e+03 -> 5.577e-01, damping 1e-03 -> 8e-06 .. GENERATED FROM PYTHON SOURCE LINES 455-468 Timing ------ Each conjugate-gradient step costs one product with the Jacobian and one with its adjoint, and each is the encoding operator once and the model once. Timing the four says which half a faster reconstruction would come from; here they are comparable. Neither product builds the Jacobian. That is a memory argument: the blocks are ``voxels x channels x contrasts`` where a signal is ``voxels x contrasts``, so what is not held is the channel count times the signal, every iteration. .. GENERATED FROM PYTHON SOURCE LINES 469-506 .. code-block:: Python tangent = torch.randn_like(initial) predicted = operator.A_jvp(initial, tangent) adjoint_image = encoding.A_adjoint(kspace).movedim(1, -1) .. rst-class:: sphx-glr-script-out .. code-block:: none per conjugate-gradient step, 3 channels solved for: model J v 2.4 ms model J^H v 1.0 ms encoding A 4.3 ms encoding A^H 4.0 ms the Jacobian this avoids holding: 1.7 MiB, against 0.6 MiB for a signal .. GENERATED FROM PYTHON SOURCE LINES 507-510 Maps ---- .. GENERATED FROM PYTHON SOURCE LINES 511-535 .. image-sg:: /generated/autoexamples/04-model-based-imaging/images/sphx_glr_02-nonlinear-inversion_002.png :alt: truth, adjoint, model-based :srcset: /generated/autoexamples/04-model-based-imaging/images/sphx_glr_02-nonlinear-inversion_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 536-545 Writing a different model ------------------------- The model is the only thing above that names a relaxation time, and it is an ordinary :class:`~torchsim.model.Simulator` -- the same object the fitting and sequence-design notebooks use. Water-fat separation, T2* with a field map, a Look-Locker inversion recovery: each is a different ``evaluate``, and the operator, the loop and the encoding are unchanged. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 3.123 seconds) .. _sphx_glr_download_generated_autoexamples_04-model-based-imaging_02-nonlinear-inversion.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: binder-badge .. image:: images/binder_badge_logo.svg :target: https://mybinder.org/v2/gh/firmlab-pisa/torchsim/gh-pages?urlpath=lab/tree/v0.0.8/examples/generated/autoexamples/04-model-based-imaging/02-nonlinear-inversion.ipynb :alt: Launch binder :width: 150 px .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 02-nonlinear-inversion.ipynb <02-nonlinear-inversion.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 02-nonlinear-inversion.py <02-nonlinear-inversion.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 02-nonlinear-inversion.zip <02-nonlinear-inversion.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_