.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/03-regularization/02-operators-and-solvers.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. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_03-regularization_02-operators-and-solvers.py: ============================= 6. Custom encoding and solver ============================= In lesson 5 you handed ``pics`` the data, the maps and a penalty, and it did the rest. In this lesson you assemble the same reconstruction from its parts. You build the SENSE encoding operator, check its adjoint, run the solver on it, and obtain exactly what ``pics`` returns, bit for bit. Then you extend the operator with one written in Python, and differentiate through it. The mental model: :func:`bartorch.apps.pics` builds three objects and runs BART's iteration on them: the encoding operator :math:`A = PFS` (sensitivities, Fourier transform, sampling), the regularization terms of :mod:`bartorch.priors`, and an iterative algorithm of :mod:`bartorch.optim`. You need the parts whenever BART has no application for what you want to do: an encoding with an extra factor, such as an off-resonance or motion phase; a solver called from your own outer loop; an operator written in Python; or a gradient with respect to the data, to train a network. **Learning objectives** - Build :class:`bartorch.linop.CartesianSense` and read the plan it was lowered into. - Check an operator's adjoint with the dot-product test. - Prepare the data as ``pics`` does and reproduce :func:`bartorch.apps.pics` bit for bit with :class:`bartorch.optim.FISTA`. - Compose operators, include one defined in Python, and differentiate through an application. Previous: :doc:`01-regularized-reconstruction`. Next: :doc:`../04-non-cartesian/01-trajectories-and-transforms`, where you replace the FFT with a NUFFT and leave the Cartesian grid. .. GENERATED FROM PYTHON SOURCE LINES 38-159 .. GENERATED FROM PYTHON SOURCE LINES 160-179 The encoding operator --------------------- Your scan is the one of lesson 2 [#brainweb]_, with a tenth of its noise: eight channels, ``kspace`` of shape ``(coils, 1, y, x)`` with one phase-encoding line in three acquired around a 24-line ACS block. ESPIRiT gives you the maps, as before. :class:`bartorch.linop.CartesianSense` is :math:`A = PFS` as one operator. You give it the sensitivities, the shape of the image, and the sampling pattern; the shape of the k-space it maps to follows from those. :func:`bartorch.tools.pattern` reads the pattern off the data, as you would for a prospectively undersampled scan. ``modulated=True`` selects the uncentred sample convention that BART's applications iterate in; the default is the centred convention of :func:`bartorch.fft`. The two differ by a modulation of the samples and give the same image. The choice matters only when you apply the operator to data already in one of them, as you do below. .. GENERATED FROM PYTHON SOURCE LINES 179-199 .. code-block:: Python import torch import bartorch import bartorch.tools as bt from bartorch import apps, linop, optim, priors maps = bt.ecalib(kspace, maps=1, calib_size=CALIBRATION, crop=0.8) pattern = bt.pattern(kspace) A = linop.CartesianSense( maps.squeeze(1), (SIZE, SIZE), pattern.squeeze(), modulated=True, ) print(f"{A.ishape} -> {A.oshape}") print(A.plan) print(f"fused: {A.plan.fused}") .. rst-class:: sphx-glr-script-out .. code-block:: none (192, 192) -> (8, 192, 192) Plan(transform=fft, image=sensitivities, normal=transform, coil_batch=1, streamed=coils, executor=slab) fused: True .. GENERATED FROM PYTHON SOURCE LINES 200-214 The operator maps an image of ``(y, x)`` to k-space of ``(coils, y, x)``. ``A.plan`` reports how it is computed: which transform, what multiplies the image and the samples, and how the normal operator :math:`A^H A` is applied. It is read back from the operator that was built, not predicted. ``plan.fused`` is false where the composition could not be expressed as one encoding and fell back to a chain, which gives the same numbers more slowly. The adjoint ----------- A solver needs :math:`A^H`, not :math:`A^T`, and one built on the wrong one converges to the wrong image. You can check it with the definition, :math:`\langle Ax, y\rangle = \langle x, A^H y\rangle`, which holds for any pair of vectors, random ones included: .. GENERATED FROM PYTHON SOURCE LINES 214-222 .. code-block:: Python generator = torch.Generator().manual_seed(0) probe = torch.randn(A.ishape, dtype=torch.complex64, generator=generator) samples = torch.randn(A.oshape, dtype=torch.complex64, generator=generator) forward = (A(probe).conj() * samples).sum() adjoint = (probe.conj() * A.H(samples)).sum() print(f"relative difference {abs(forward - adjoint) / abs(forward):.2e}") .. rst-class:: sphx-glr-script-out .. code-block:: none relative difference 1.20e-07 .. GENERATED FROM PYTHON SOURCE LINES 223-233 The two inner products agree to single-precision round-off. Solving ------- A solver is called as ``solver(y, A)``. To get the answer ``pics`` gets, you give it the data ``pics`` iterates on rather than the array the scanner wrote: the sampling pattern applied, the modulation into the uncentred convention, and the data divided by the scale :func:`bartorch.optim.data_scaling` estimates. .. GENERATED FROM PYTHON SOURCE LINES 233-240 .. code-block:: Python measured = bartorch.fftmod(kspace * pattern, axes=(-1, -2, -3), inverse=True) scale = optim.data_scaling(measured) data = (measured / scale).squeeze(1) term = priors.Wavelet(axes=(-1, -2), weight=0.002) assembled = optim.FISTA(term, maxiter=100)(data, A) * scale .. GENERATED FROM PYTHON SOURCE LINES 241-245 The solution is multiplied back by the scale, which puts the image into the units of the k-space, as ``pics`` does unless given ``rescale=False``. Now run ``pics`` on the raw ``kspace`` with the same term and solver, and compare: .. GENERATED FROM PYTHON SOURCE LINES 245-275 .. code-block:: Python tool = apps.pics(kspace, maps, regularizers=term, solver="fista", maxiter=100) same = torch.equal(assembled.squeeze(), tool.squeeze()) print(f"identical to pics: {same}") def nrmse(estimate): return bt.nrmse(image.abs(), estimate.abs(), scaled=True) print(f"NRMSE, adjoint {nrmse(A.H(data)):.3f}") print(f"NRMSE, FISTA {nrmse(assembled):.3f}") .. image-sg:: /auto_examples/03-regularization/images/sphx_glr_02-operators-and-solvers_001.png :alt: the object, adjoint, $A^H y$, FISTA, wavelet, |error|, FISTA :srcset: /auto_examples/03-regularization/images/sphx_glr_02-operators-and-solvers_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none identical to pics: True NRMSE, adjoint 0.119 NRMSE, FISTA 0.032 .. GENERATED FROM PYTHON SOURCE LINES 276-293 The two are the same iteration over the same operator, so they return the same bits, not merely close ones. In the figure, the adjoint :math:`A^H y` is not the inverse: it is the sensitivity-weighted coil combination of the zero-filled k-space, and it carries the aliasing of the undersampling and the shading of :math:`\sum_c |S_c|^2`. The solve removes both, as the NRMSE printed above shows, and its remaining error sits at the tissue boundaries. Operator algebra ---------------- ``@`` composes, ``+`` adds, ``A.H`` is the adjoint and ``A.gram()`` the normal operator :math:`A^H A`. A composition builds one BART operator rather than a Python chain, so a solver iterating on it does not return to Python between applications. :func:`bartorch.optim.maxeigen` runs the power iteration on an operator; the largest eigenvalue of :math:`A^H A` is the Lipschitz constant of the least-squares gradient, which bounds the gradient step size. .. GENERATED FROM PYTHON SOURCE LINES 293-295 .. code-block:: Python print(f"largest eigenvalue of A^H A: {optim.maxeigen(A.gram()):.3f}") .. rst-class:: sphx-glr-script-out .. code-block:: none largest eigenvalue of A^H A: 0.999 .. GENERATED FROM PYTHON SOURCE LINES 296-301 An operator of your own is composed with BART's operators through :meth:`~bartorch.linop.LinearOperator.from_callbacks`: you give it the input and output shapes, a forward and an adjoint, and BART calls them back. Here it is a spatially varying phase, as an off-resonance or eddy-current phase would be, placed between the image and the encoding: .. GENERATED FROM PYTHON SOURCE LINES 301-311 .. code-block:: Python field = torch.exp(1j * 0.4 * torch.pi * grid_x).to(torch.complex64) phase = linop.LinearOperator.from_callbacks( (SIZE, SIZE), (SIZE, SIZE), lambda u: field * u, lambda u: field.conj() * u, ) composed = A @ phase print(f"{composed.ishape} -> {composed.oshape}, fused: {composed.plan.fused}") .. rst-class:: sphx-glr-script-out .. code-block:: none (192, 192) -> (8, 192, 192), fused: True .. GENERATED FROM PYTHON SOURCE LINES 312-324 ``composed`` has the same interface as ``A`` (forward, adjoint, normal operator) and is passed to a solver in the same way. Differentiation --------------- Apply an operator to a tensor that requires a gradient, and autograd records it. For :math:`y = Ax` the gradient torch propagates back is :math:`A^H g`, not :math:`A^T g`, which is torch's conjugate Wirtinger convention for complex tensors. A real :math:`A` has :math:`A^H = A^T`, so only a complex check tells the two apart. Here the gradient of :math:`\|Ax - y\|^2` at :math:`x = 0` should be :math:`-2A^H y`: .. GENERATED FROM PYTHON SOURCE LINES 324-332 .. code-block:: Python variable = data.new_zeros(A.ishape).requires_grad_(True) residual = A(variable) - data (residual.abs() ** 2).sum().backward() expected = 2 * A.H(-data) difference = (variable.grad - expected).abs().max() / expected.abs().max() print(f"relative difference from 2 A^H (Ax - y): {float(difference):.2e}") .. rst-class:: sphx-glr-script-out .. code-block:: none relative difference from 2 A^H (Ax - y): 2.06e-07 .. GENERATED FROM PYTHON SOURCE LINES 333-343 :doc:`../../explanation/learned-reconstruction` describes the backward passes of the solvers, and :doc:`../../explanation/solvers` which algorithm applies to which problem. Exercise: FISTA convergence --------------------------- ``optim.FISTA`` is a loop over :class:`bartorch.optim.FISTABlock`, one iteration per call. Run the loop yourself for 200 iterations and record the error of the image after each one: .. GENERATED FROM PYTHON SOURCE LINES 343-361 .. code-block:: Python block = optim.FISTABlock(term) state = block.start(data, A) trace = [] for _ in range(200): state = block(state, A) trace.append(nrmse(block.output(state, A))) .. image-sg:: /auto_examples/03-regularization/images/sphx_glr_02-operators-and-solvers_002.png :alt: 02 operators and solvers :srcset: /auto_examples/03-regularization/images/sphx_glr_02-operators-and-solvers_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 362-385 The error falls steeply over the first few tens of iterations, reaches its minimum before iteration 50, and then stays close to that level with small fluctuations up to iteration 200. For this data and weight, the 100 iterations used above are past the point where more iterations reduce the error. Try ``step=0.5`` in :class:`~bartorch.optim.FISTABlock`, or a larger weight, and compare the curves. As a spec --------- What this lesson built, stated the way you would ask an agent for it: .. code-block:: text With bartorch, build linop.CartesianSense from ESPIRiT maps, the image shape and the sampling pattern of a 2D Cartesian 8-channel k-space (modulated=True), print its plan, and check its adjoint with a random dot-product test. Prepare the data as pics does (pattern, fftmod inverse, divided by optim.data_scaling), solve with optim.FISTA and priors.Wavelet over the last two axes (weight 0.002, 100 iterations), and check torch.equal against apps.pics with the same term. Compose the operator with a Python phase operator from from_callbacks, and check that autograd's gradient of ||Ax - y||^2 at zero is -2 A^H y. .. GENERATED FROM PYTHON SOURCE LINES 387-394 References ---------- .. [#brainweb] Collins DL, Zijdenbos AP, Kollokian V, Sled JG, Kabani NJ, Holmes CJ, Evans AC. Design and construction of a realistic digital brain phantom. *IEEE Trans Med Imaging* 17(3):463-468 (1998). https://doi.org/10.1109/42.712135 .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 6.414 seconds) .. _sphx_glr_download_auto_examples_03-regularization_02-operators-and-solvers.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 02-operators-and-solvers.ipynb <02-operators-and-solvers.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 02-operators-and-solvers.py <02-operators-and-solvers.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 02-operators-and-solvers.zip <02-operators-and-solvers.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_