.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/03_torch_projection.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_torch_projection.py: Torch Projection Layer ====================== This minimal example demonstrates how to use the `parallelproj_core` library with torch tensors and how to backward propagate gradients through the forward projection operation. .. GENERATED FROM PYTHON SOURCE LINES 8-14 .. code-block:: Python import parallelproj_core import array_api_compat.torch as torch from array_api_compat import device .. GENERATED FROM PYTHON SOURCE LINES 15-19 Setup a custom torch autograd function to wraps the forward projection. We also implement the backward path (Jacobian-vector product) which in case of a linear operator is equivalent to its adjoint (back projection). Here we operator on a singel 3D tensor. In practice, this would probably be a batch of 3D tensors, with a signle channel (5D tensor). .. GENERATED FROM PYTHON SOURCE LINES 19-50 .. code-block:: Python class FwdProjLayer(torch.autograd.Function): @staticmethod def forward(ctx, x, lor_start, lor_end, img_origin, voxel_size): dev = device(x) img_fwd = torch.zeros(lor_start.shape[0], dtype=torch.float32, device=dev) parallelproj_core.joseph3d_fwd( lor_start, lor_end, x, img_origin, voxel_size, img_fwd ) # save context variables for backward pass image_shape = torch.asarray(x.shape) ctx.save_for_backward(lor_start, lor_end, img_origin, voxel_size, image_shape) return img_fwd @staticmethod def backward(ctx, grad_output): dev = device(grad_output) lor_start, lor_end, img_origin, voxel_size, image_shape = ctx.saved_tensors grad_input = torch.zeros(tuple(image_shape), dtype=torch.float32, device=dev) parallelproj_core.joseph3d_back( lor_start, lor_end, grad_input, img_origin, voxel_size, grad_output, ) return grad_input, None, None, None, None .. GENERATED FROM PYTHON SOURCE LINES 51-52 select device .. GENERATED FROM PYTHON SOURCE LINES 52-58 .. code-block:: Python if torch.cuda.is_available() and parallelproj_core.cuda_enabled == 1: dev = "cuda" else: dev = "cpu" .. GENERATED FROM PYTHON SOURCE LINES 59-60 Print backend and device info. .. GENERATED FROM PYTHON SOURCE LINES 60-64 .. code-block:: Python print(f"parallelproj_core version: {parallelproj_core.__version__}") print(f"parallelproj_core cuda enabled: {parallelproj_core.cuda_enabled}") print(f"using array API compatible library: {torch.__name__} on device {dev}") .. rst-class:: sphx-glr-script-out .. code-block:: none parallelproj_core version: v2.0.5 parallelproj_core cuda enabled: 0 using array API compatible library: array_api_compat.torch on device cpu .. GENERATED FROM PYTHON SOURCE LINES 65-66 Define a mini sparse demo image. .. GENERATED FROM PYTHON SOURCE LINES 66-74 .. code-block:: Python image = torch.zeros((3, 3, 3), dtype=torch.float32, device=dev) image[0, 0, 2] = 0.25 image[2, 0, 0] = 0.25 image[2, 2, 2] = 0.5 voxel_size = torch.asarray([2.0, 2.0, 2.0], device=dev, dtype=torch.float32) img_origin = torch.asarray([-1.0, -1.0, -1.0], device=dev, dtype=torch.float32) .. GENERATED FROM PYTHON SOURCE LINES 75-76 Define LOR start and end points. .. GENERATED FROM PYTHON SOURCE LINES 76-87 .. code-block:: Python lor_start = torch.asarray( [[7.0, 2.0, 2.0], [-1, -1, -5], [3, -3, -3]], device=dev, dtype=torch.float32, ) lor_end = torch.asarray( [[-5.0, 2.0, 2.0], [-1, -1, 7], [3, 6, 6]], device=dev, dtype=torch.float32, ) .. GENERATED FROM PYTHON SOURCE LINES 88-92 .. code-block:: Python fwd_proj_layer = FwdProjLayer.apply img_fwd = fwd_proj_layer(image, lor_start, lor_end, img_origin, voxel_size) print("Forward projection result:", img_fwd) .. rst-class:: sphx-glr-script-out .. code-block:: none Forward projection result: tensor([0.2500, 0.5000, 2.1213]) .. GENERATED FROM PYTHON SOURCE LINES 93-96 Run a torch gradient check to verify that the backward projection correctly computes gradients. Since all parallelproj_core functions use float32, we set the eps, atol, and rtol parameters to higher values to account for numerical precision issues. .. GENERATED FROM PYTHON SOURCE LINES 96-110 .. code-block:: Python x = torch.rand(image.shape, dtype=torch.float32, device=dev, requires_grad=True) print("Running forward projection layer gradient test") grad_test_fwd = torch.autograd.gradcheck( fwd_proj_layer, (x, lor_start, lor_end, img_origin, voxel_size), eps=1e-2, atol=1e-5, rtol=1e-5, ) assert grad_test_fwd, "Gradient check failed!" .. rst-class:: sphx-glr-script-out .. code-block:: none Running forward projection layer gradient test /home/docs/checkouts/readthedocs.org/user_builds/libparallelproj/conda/v2.0.5/lib/python3.14/site-packages/torch/autograd/gradcheck.py:2106: UserWarning: Input #0 requires gradient and is not a double precision floating point or complex. This check will likely fail if all the inputs are not of double precision floating point or complex. _check_inputs(tupled_inputs) .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 1.388 seconds) .. _sphx_glr_download_auto_examples_03_torch_projection.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 03_torch_projection.ipynb <03_torch_projection.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 03_torch_projection.py <03_torch_projection.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 03_torch_projection.zip <03_torch_projection.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_