.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/06_torch/01_run_projection_layer.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_06_torch_01_run_projection_layer.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_06_torch_01_run_projection_layer.py:


pytorch parallelproj projection layer
=====================================

In this example, we show how to define a custom pytorch layer that can be used
to define a feed forward neural network that includes a parallelproj forward and back
backward projections (or any :class:`.LinearOperator`) that can be used with pytorch's
autograd engine.

.. GENERATED FROM PYTHON SOURCE LINES 12-31

.. code-block:: Python


    from __future__ import annotations

    import array_api_compat.torch as torch
    import matplotlib.pyplot as plt
    import parallelproj.operators
    import parallelproj.projectors
    import parallelproj.pet_scanners
    import parallelproj.pet_lors
    import parallelproj_core as ppc
    from array_api_compat import device

    # device variable (cpu or cuda) that determines whether calculations
    # are performed on the cpu or cuda gpu
    if torch.cuda.is_available() and ppc.cuda_enabled == 1:
        dev = "cuda"
    else:
        dev = "cpu"








.. GENERATED FROM PYTHON SOURCE LINES 32-38

Setup the forward projection layer
----------------------------------

We subclass :class:`torch.autograd.Function` to define a custom pytorch layer
that is compatible with pytorch's autograd engine.
see also: https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html

.. GENERATED FROM PYTHON SOURCE LINES 38-124

.. code-block:: Python



    class LinearSingleChannelOperator(torch.autograd.Function):
        """
        Function representing a linear operator acting on a mini batch of single channel images
        """

        @staticmethod
        def forward(
            ctx, x: torch.Tensor, operator: parallelproj.operators.LinearOperator
        ) -> torch.Tensor:
            """forward pass of the linear operator

            Parameters
            ----------
            ctx : context object
                that can be used to store information for the backward pass
            x : torch.Tensor
                mini batch of 3D images with dimension (batch_size, 1, num_voxels_x, num_voxels_y, num_voxels_z)
            operator : parallelproj.operators.LinearOperator
                linear operator that can act on a single 3D image

            Returns
            -------
            torch.Tensor
                mini batch of 3D images with dimension (batch_size, *operator.out_shape)
            """

            # https://pytorch.org/docs/stable/notes/extending.html#how-to-use
            ctx.set_materialize_grads(False)
            ctx.operator = operator

            batch_size = x.shape[0]
            y = torch.zeros(
                (batch_size,) + operator.out_shape, dtype=x.dtype, device=device(x)
            )

            # loop over all samples in the batch and apply linear operator
            # to the first channel
            for i in range(batch_size):
                y[i, ...] = operator(x[i, 0, ...].detach())

            return y

        @staticmethod
        def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
            """backward pass of the forward pass

            Parameters
            ----------
            ctx : context object
                that can be used to obtain information from the forward pass
            grad_output : torch.Tensor
                mini batch of dimension (batch_size, *operator.out_shape)

            Returns
            -------
            torch.Tensor, None
                mini batch of 3D images with dimension (batch_size, 1, *operator.in_shape)
            """

            # For details on how to implement the backward pass, see
            # https://pytorch.org/docs/stable/notes/extending.html#how-to-use

            # since forward takes two input arguments (x, operator)
            # we have to return two arguments (the latter is None)
            if grad_output is None:
                return None, None
            else:
                operator = ctx.operator

                batch_size = grad_output.shape[0]
                x = torch.zeros(
                    (batch_size, 1) + operator.in_shape,
                    dtype=grad_output.dtype,
                    device=device(grad_output),
                )

                # loop over all samples in the batch and apply linear operator
                # to the first channel
                for i in range(batch_size):
                    x[i, 0, ...] = operator.adjoint(grad_output[i, ...].detach())

                return x, None









.. GENERATED FROM PYTHON SOURCE LINES 125-131

Setup the back projection layer
-------------------------------

We subclass :class:`torch.autograd.Function` to define a custom pytorch layer
that is compatible with pytorch's autograd engine.
see also: https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html

.. GENERATED FROM PYTHON SOURCE LINES 131-216

.. code-block:: Python



    class AdjointLinearSingleChannelOperator(torch.autograd.Function):
        """
        Function representing the adjoint of a linear operator acting on a mini batch of single channel images
        """

        @staticmethod
        def forward(
            ctx, x: torch.Tensor, operator: parallelproj.operators.LinearOperator
        ) -> torch.Tensor:
            """forward pass of the adjoint of the linear operator

            Parameters
            ----------
            ctx : context object
                that can be used to store information for the backward pass
            x : torch.Tensor
                mini batch of 3D images with dimension (batch_size, *operator.out_shape)
            operator : parallelproj.operators.LinearOperator
                linear operator that can act on a single 3D image

            Returns
            -------
            torch.Tensor
                mini batch of 3D images with dimension (batch_size, 1, *operator.in_shape)
            """

            ctx.set_materialize_grads(False)
            ctx.operator = operator

            batch_size = x.shape[0]
            y = torch.zeros(
                (batch_size, 1) + operator.in_shape, dtype=x.dtype, device=device(x)
            )

            # loop over all samples in the batch and apply linear operator
            # to the first channel
            for i in range(batch_size):
                y[i, 0, ...] = operator.adjoint(x[i, ...].detach())

            return y

        @staticmethod
        def backward(ctx, grad_output):
            """backward pass of the forward pass

            Parameters
            ----------
            ctx : context object
                that can be used to obtain information from the forward pass
            grad_output : torch.Tensor
                mini batch of dimension (batch_size, 1, *operator.in_shape)

            Returns
            -------
            torch.Tensor, None
                mini batch of 3D images with dimension (batch_size, *operator.out_shape)
            """

            # For details on how to implement the backward pass, see
            # https://pytorch.org/docs/stable/notes/extending.html#how-to-use

            # since forward takes two input arguments (x, operator)
            # we have to return two arguments (the latter is None)
            if grad_output is None:
                return None, None
            else:
                operator = ctx.operator

                batch_size = grad_output.shape[0]
                x = torch.zeros(
                    (batch_size,) + operator.out_shape,
                    dtype=grad_output.dtype,
                    device=device(grad_output),
                )

                # loop over all samples in the batch and apply linear operator
                # to the first channel
                for i in range(batch_size):
                    x[i, ...] = operator(grad_output[i, 0, ...].detach())

                return x, None









.. GENERATED FROM PYTHON SOURCE LINES 217-222

Setup a minimal non-TOF PET projector
-------------------------------------

We setup a minimal non-TOF PET projector of small scanner with
three rings.

.. GENERATED FROM PYTHON SOURCE LINES 222-247

.. code-block:: Python


    num_rings = 3
    scanner = parallelproj.pet_scanners.RegularPolygonPETScannerGeometry(
        torch,
        dev,
        radius=35.0,
        num_sides=12,
        num_lor_endpoints_per_side=6,
        lor_spacing=3.0,
        ring_positions=torch.linspace(-4, 4, num_rings, device=dev),
        symmetry_axis=2,
    )

    # setup the LOR descriptor that defines the sinogram
    lor_desc = parallelproj.pet_lors.RegularPolygonPETLORDescriptor(
        scanner,
        parallelproj.pet_lors.Michelogram(scanner.num_rings, max_ring_difference=1, span=1),
        radial_trim=10,
        sinogram_order=parallelproj.pet_lors.SinogramSpatialAxisOrder.RVP,
    )

    proj = parallelproj.projectors.RegularPolygonPETProjector(
        lor_desc, img_shape=(20, 20, 5), voxel_size=(2.0, 2.0, 2.0)
    )








.. GENERATED FROM PYTHON SOURCE LINES 248-250

Define a mini batch of input and output tensors
-----------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 250-267

.. code-block:: Python


    batch_size = 2

    x = torch.rand(
        (batch_size, 1) + proj.in_shape,
        device=dev,
        dtype=torch.float32,
        requires_grad=True,
    )

    y = torch.rand(
        (batch_size,) + proj.out_shape,
        device=dev,
        dtype=torch.float32,
        requires_grad=True,
    )








.. GENERATED FROM PYTHON SOURCE LINES 268-270

Define the forward and backward projection layers
-------------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 270-284

.. code-block:: Python


    fwd_op_layer = LinearSingleChannelOperator.apply
    adjoint_op_layer = AdjointLinearSingleChannelOperator.apply

    f1 = fwd_op_layer(x, proj)
    print("forward projection (Ax) .:", f1.shape, type(f1), device(f1))

    b1 = adjoint_op_layer(y, proj)
    print("back projection (A^T y) .:", b1.shape, type(b1), device(b1))

    fb1 = adjoint_op_layer(fwd_op_layer(x, proj), proj)
    print("back + forward projection (A^TAx) .:", fb1.shape, type(fb1), device(fb1))






.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    forward projection (Ax) .: torch.Size([2, 51, 36, 7]) <class 'torch.Tensor'> cpu
    back projection (A^T y) .: torch.Size([2, 1, 20, 20, 5]) <class 'torch.Tensor'> cpu
    back + forward projection (A^TAx) .: torch.Size([2, 1, 20, 20, 5]) <class 'torch.Tensor'> cpu




.. GENERATED FROM PYTHON SOURCE LINES 285-287

Define a dummy loss function and trigger the backpropagation
------------------------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 287-296

.. code-block:: Python


    # define a dummy loss function
    dummy_loss = (fb1**2).sum()
    # trigger the backpropagation
    dummy_loss.backward()

    print(f" backpropagted gradient {x.grad}")






.. rst-class:: sphx-glr-script-out

 .. code-block:: none

     backpropagted gradient tensor([[[[[17495450.0000,  6518800.5000, 28074182.0000,  6646377.0000,
                17822030.0000],
               [19793590.0000,  8793182.0000, 31344058.0000,  8945131.0000,
                20153762.0000],
               [21380784.0000, 11194555.0000, 33405930.0000, 11441207.0000,
                21832030.0000],
               ...,
               [20844806.0000, 10831412.0000, 32681274.0000, 11206055.0000,
                21609734.0000],
               [19244670.0000,  8546403.0000, 30617602.0000,  8826238.0000,
                19936428.0000],
               [17100338.0000,  6352868.0000, 27551406.0000,  6556106.0000,
                17697986.0000]],

              [[19757946.0000,  8831423.0000, 31384560.0000,  8987548.0000,
                20153092.0000],
               [22802976.0000, 12202413.0000, 35755264.0000, 12454075.0000,
                23259120.0000],
               [23755596.0000, 14729286.0000, 36777880.0000, 15114716.0000,
                24278262.0000],
               ...,
               [23157664.0000, 14308314.0000, 35898624.0000, 14828040.0000,
                24001510.0000],
               [22219154.0000, 11832925.0000, 34974544.0000, 12234699.0000,
                23020054.0000],
               [19328408.0000,  8548561.0000, 30769048.0000,  8822910.0000,
                19989850.0000]],

              [[21426788.0000, 11202656.0000, 33631320.0000, 11437637.0000,
                21869290.0000],
               [23780670.0000, 14791346.0000, 36912276.0000, 15173322.0000,
                24290600.0000],
               [24513332.0000, 17687122.0000, 37594552.0000, 18178342.0000,
                25093628.0000],
               ...,
               [23966778.0000, 17158000.0000, 36745600.0000, 17841090.0000,
                24900944.0000],
               [23237582.0000, 14319027.0000, 36112852.0000, 14851176.0000,
                24098096.0000],
               [20956654.0000, 10921059.0000, 32913918.0000, 11282532.0000,
                21682244.0000]],

              ...,

              [[20933906.0000, 10972925.0000, 33181398.0000, 11319848.0000,
                21422960.0000],
               [23258556.0000, 14457134.0000, 36525192.0000, 14888950.0000,
                23810750.0000],
               [24026788.0000, 17352714.0000, 37293584.0000, 17895870.0000,
                24576776.0000],
               ...,
               [24249558.0000, 17229752.0000, 36820196.0000, 17793656.0000,
                24602620.0000],
               [23430030.0000, 14348408.0000, 35968452.0000, 14785201.0000,
                23785596.0000],
               [21120774.0000, 10856907.0000, 32809800.0000, 11167794.0000,
                21433376.0000]],

              [[19314104.0000,  8582776.0000, 31071800.0000,  8878143.0000,
                19803672.0000],
               [22224844.0000, 11921524.0000, 35439436.0000, 12311323.0000,
                22803068.0000],
               [23210396.0000, 14477292.0000, 36511704.0000, 14918756.0000,
                23767770.0000],
               ...,
               [23424928.0000, 14330553.0000, 36113116.0000, 14796163.0000,
                23858628.0000],
               [22403776.0000, 11813002.0000, 34915080.0000, 12167533.0000,
                22817628.0000],
               [19453328.0000,  8538034.0000, 30629586.0000,  8787860.0000,
                19774366.0000]],

              [[17099338.0000,  6384269.0000, 27851626.0000,  6600580.0000,
                17555274.0000],
               [19285316.0000,  8611253.0000, 31086938.0000,  8913234.0000,
                19794846.0000],
               [20918878.0000, 10937782.0000, 33252762.0000, 11304432.0000,
                21443482.0000],
               ...,
               [21082566.0000, 10891806.0000, 32918388.0000, 11239042.0000,
                21500214.0000],
               [19431540.0000,  8510878.0000, 30680086.0000,  8777471.0000,
                19854250.0000],
               [17231868.0000,  6332819.0000, 27462982.0000,  6535246.5000,
                17562404.0000]]]],



            [[[[17379904.0000,  6633539.0000, 28294268.0000,  6508132.5000,
                17283006.0000],
               [19589548.0000,  8894056.0000, 31482552.0000,  8751283.0000,
                19518568.0000],
               [21228892.0000, 11351533.0000, 33589452.0000, 11185089.0000,
                21187638.0000],
               ...,
               [21178478.0000, 11431438.0000, 34026676.0000, 11216227.0000,
                21180158.0000],
               [19543158.0000,  9008566.0000, 31884204.0000,  8859779.0000,
                19534504.0000],
               [17360798.0000,  6683149.5000, 28621252.0000,  6554375.5000,
                17309226.0000]],

              [[19632808.0000,  8949840.0000, 31540784.0000,  8798499.0000,
                19540552.0000],
               [22591198.0000, 12344454.0000, 35902876.0000, 12167669.0000,
                22509434.0000],
               [23537322.0000, 14899740.0000, 36913672.0000, 14713468.0000,
                23515378.0000],
               ...,
               [23516198.0000, 15143856.0000, 37377492.0000, 14778929.0000,
                23529394.0000],
               [22568612.0000, 12485469.0000, 36370560.0000, 12238899.0000,
                22551352.0000],
               [19639034.0000,  8997627.0000, 31926366.0000,  8841146.0000,
                19592778.0000]],

              [[21354198.0000, 11376594.0000, 33849392.0000, 11194712.0000,
                21211358.0000],
               [23623384.0000, 14978292.0000, 37065804.0000, 14783460.0000,
                23533470.0000],
               [24344816.0000, 17904950.0000, 37753476.0000, 17674478.0000,
                24294738.0000],
               ...,
               [24328996.0000, 18145674.0000, 38121744.0000, 17662996.0000,
                24352628.0000],
               [23601578.0000, 15134516.0000, 37388124.0000, 14760272.0000,
                23622982.0000],
               [21283740.0000, 11483916.0000, 33989468.0000, 11255213.0000,
                21265298.0000]],

              ...,

              [[21322036.0000, 11424591.0000, 33814292.0000, 11205672.0000,
                21234622.0000],
               [23697596.0000, 15001081.0000, 37156244.0000, 14725089.0000,
                23596826.0000],
               [24487336.0000, 18018104.0000, 37799808.0000, 17624504.0000,
                24345568.0000],
               ...,
               [24127466.0000, 17827506.0000, 37486780.0000, 17681082.0000,
                24622126.0000],
               [23313046.0000, 14850903.0000, 36801276.0000, 14779467.0000,
                23821426.0000],
               [21033666.0000, 11239511.0000, 33594428.0000, 11175231.0000,
                21440972.0000]],

              [[19638384.0000,  8954660.0000, 31653720.0000,  8797732.0000,
                19644124.0000],
               [22635162.0000, 12412142.0000, 36070396.0000, 12185245.0000,
                22633674.0000],
               [23662684.0000, 15026525.0000, 37032696.0000, 14731470.0000,
                23579848.0000],
               ...,
               [23324734.0000, 14794837.0000, 36714840.0000, 14713341.0000,
                23807764.0000],
               [22309272.0000, 12207168.0000, 35660124.0000, 12154627.0000,
                22793386.0000],
               [19380196.0000,  8828988.0000, 31278718.0000,  8765756.0000,
                19714808.0000]],

              [[17385516.0000,  6649961.0000, 28393588.0000,  6552749.0000,
                17433072.0000],
               [19644596.0000,  8989104.0000, 31657850.0000,  8839217.0000,
                19669424.0000],
               [21325934.0000, 11390315.0000, 33778992.0000, 11179188.0000,
                21271052.0000],
               ...,
               [21004864.0000, 11241240.0000, 33451454.0000, 11167654.0000,
                21382100.0000],
               [19371466.0000,  8795409.0000, 31262978.0000,  8738416.0000,
                19746092.0000],
               [17188670.0000,  6555337.5000, 28065080.0000,  6495116.5000,
                17470212.0000]]]]])




.. GENERATED FROM PYTHON SOURCE LINES 297-306

Check whether the gradients are calculated correctly
----------------------------------------------------

We use pytorch's gradcheck function to check whether the implementation
of the backward pass, needed to calculate the gradients, is correct.

This test can be slow which is why we only execute it on the gpu.
Note that parallelproj's projectors use single precision precision
which is why we have to use a larger atol and rtol.

.. GENERATED FROM PYTHON SOURCE LINES 306-320

.. code-block:: Python


    if dev == "cpu":
        print("skipping (slow) gradient checks on cpu")
    else:
        print("Running forward projection layer gradient test")
        grad_test_fwd = torch.autograd.gradcheck(
            fwd_op_layer, (x, proj), eps=1e-1, atol=1e-3, rtol=1e-3
        )

        print("Running adjoint projection layer gradient test")
        grad_test_fwd = torch.autograd.gradcheck(
            adjoint_op_layer, (y, proj), eps=1e-1, atol=1e-3, rtol=1e-3
        )





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    skipping (slow) gradient checks on cpu




.. GENERATED FROM PYTHON SOURCE LINES 321-323

Visualize the scanner geometry and image FOV
--------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 323-330

.. code-block:: Python


    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection="3d")
    ax.view_init(elev=-30, azim=160, roll=180, vertical_axis="y")
    proj.show_geometry(ax)
    fig.tight_layout()
    fig.show()



.. image-sg:: /auto_examples/06_torch/images/sphx_glr_01_run_projection_layer_001.png
   :alt: 01 run projection layer
   :srcset: /auto_examples/06_torch/images/sphx_glr_01_run_projection_layer_001.png
   :class: sphx-glr-single-img






.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.169 seconds)


.. _sphx_glr_download_auto_examples_06_torch_01_run_projection_layer.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: 01_run_projection_layer.ipynb <01_run_projection_layer.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 01_run_projection_layer.py <01_run_projection_layer.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 01_run_projection_layer.zip <01_run_projection_layer.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_