Adding PyTorch functionality to SIRF

CHANGES.md

@@ -5,6 +5,17 @@
   - `ScatterEstimation` has extra methods that allow setting masks for the tail-fitting
   - `ImageData` has extra method to zoom image using information from a template image, `zoom_image_as_template`.
   - Error raised in `AcquisitionSensitivityModel.[un]normalise` methods applied to a read-only object.
+  - Error raised if `AcquisitionModel.adjoint` ran when the model is not linear.
+* SIRF-torch
+  - `torch/torch.py` has wrappers for pytorch objective functions, objective function gradient and operators
+  - `torch/tests/gradchecks.py` has gradchecks for the wrappers 2d/3d PET and 2d MRI.
+  - `torch/tests/use_cases.py` has use cases for 2d PET using all the wrappers.
+  - `torch/README.md` includes user directions for the wrappers.
+  - `torch/CMakeList.txt` installation of sirf.torch
+  - `src/CMakeList.txt` installation of sirf.torch
+* SIRF
+  - `cmake/sirf.__init__.py.in` import sirf.SIRF content into the `sirf` namespace for convenience
+  - `common/SIRF.py` adding adjoint operator
 
 
 ## v3.8.1

cmake/sirf.__init__.py.in

@@ -1,4 +1,7 @@
 __version_major__ = '@VERSION_MAJOR@'
 __version_minor__ = '@VERSION_MINOR@'
 __version_patch__ = '@VERSION_PATCH@'
-__version__ = '@VERSION_MAJOR@.@VERSION_MINOR@.@VERSION_PATCH@'
+__version__ = '@VERSION_MAJOR@.@VERSION_MINOR@.@VERSION_PATCH@'
+
+# import sirf.SIRF content into the `sirf` namespace for convenience
+from sirf.SIRF import *

src/CMakeLists.txt

@@ -160,3 +160,4 @@ else()
 endif()
 
 ADD_SUBDIRECTORY(common)
+ADD_SUBDIRECTORY(torch)

src/common/SIRF.py

@@ -1,7 +1,4 @@
-''' 
-Object-Oriented wrap for the cSIRF-to-Python interface pysirf.py
-'''
-
+'''Object-Oriented wrap for the cSIRF-to-Python interface pysirf.py'''
 ## SyneRBI Synergistic Image Reconstruction Framework (SIRF)
 ## Copyright 2015 - 2020 Rutherford Appleton Laboratory STFC
 ## Copyright 2015 - 2020 University College London
@@ -43,6 +40,10 @@
 else:
     ABC = abc.ABCMeta('ABC', (), {})
 
+# In future, would be good to explicitly list all objects to import when doing `from sirf.SIRF import *`.
+# However, we will keep this for later to avoid mistakes in updating this variable.
+# __all__ = ['DataContainer', 'ImageData', 'GeometricalInfo', 'AdjointOperator']
+
 
 class DataContainer(ABC):
     '''
@@ -63,7 +64,7 @@ def __add__(self, other):
         '''
         Overloads + for data containers.
 
-        Returns the sum of the container data with another container 
+        Returns the sum of the container data with another container
         data viewed as vectors.
         other: DataContainer
         '''
@@ -73,7 +74,7 @@ def __sub__(self, other):
         '''
         Overloads - for data containers.
 
-        Returns the difference of the container data with another container 
+        Returns the difference of the container data with another container
         data viewed as vectors.
         other: DataContainer
         '''
@@ -251,7 +252,7 @@ def squared_norm(self):
 
     def dot(self, other):
         '''
-        Returns the dot product of the container data with another container 
+        Returns the dot product of the container data with another container
         data viewed as vectors.
         other: DataContainer
         '''
@@ -417,10 +418,10 @@ def axpby(self, a, b, y, out=None, **kwargs):
         '''
         Linear combination for data containers.
 
-        Returns the linear combination of the self data with another container 
+        Returns the linear combination of the self data with another container
         data y viewed as vectors.
         a: multiplier to self, can be a number or a DataContainer
-        b: multiplier to y, can be a number or a DataContainer 
+        b: multiplier to y, can be a number or a DataContainer
         y: DataContainer
         out:   DataContainer to store the result to.
         '''
@@ -430,10 +431,10 @@ def sapyb(self, a, y, b, out=None, **kwargs):
         '''
         Linear combination for data containers: new interface.
 
-        Returns the linear combination of the self data with another container 
+        Returns the linear combination of the self data with another container
         data y viewed as vectors.
         a: multiplier to self, can be a number or a DataContainer
-        b: multiplier to y, can be a number or a DataContainer 
+        b: multiplier to y, can be a number or a DataContainer
         y: DataContainer
         out:   DataContainer to store the result to, can be self or y.
         '''
@@ -599,7 +600,7 @@ def dtype(self):
         else:
             dt = 'float%s' % bits
         return numpy.dtype(dt)
-    
+
 
 class ImageData(DataContainer):
     '''
@@ -716,7 +717,7 @@ def get_spacing(self):
         arr = numpy.ndarray((3,), dtype = numpy.float32)
         try_calling (pysirf.cSIRF_GeomInfo_get_spacing(self.handle, arr.ctypes.data))
         return tuple(arr)
-    
+
     def get_size(self):
         """Size is the number of voxels in each dimension."""
         arr = numpy.ndarray((3,), dtype = cpp_int_dtype())
@@ -734,3 +735,20 @@ def get_index_to_physical_point_matrix(self):
         arr = numpy.ndarray((4,4), dtype = numpy.float32)
         try_calling (pysirf.cSIRF_GeomInfo_get_index_to_physical_point_matrix(self.handle, arr.ctypes.data))
         return arr
+
+class AdjointOperator(object):
+    """
+    Creates the adjoint operator of a linear operator `lin_op`.
+    """
+    def __init__(self, operator):
+        self.operator = operator
+
+    def forward(self, x):
+        """Calls the adjoint method of the original linear operator"""
+        # Note: calling `adjoint` will raise an error in SIRF if the operator is not linear.
+        return self.operator.adjoint(x)
+
+    def backward(self, x):
+        """Calls the `direct` method of the original linear operator"""
+        # Note: calling `direct` will raise an error in SIRF if the operator is not linear.
+        return self.operator.direct(x)

src/torch/CMakeLists.txt

@@ -0,0 +1,22 @@
+#========================================================================
+# Author: Evgueni Ovtchinnikov
+# Copyright 2020 University College London
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0.txt
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+#=========================================================================
+
+
+if (BUILD_PYTHON)
+  INSTALL(FILES torch.py DESTINATION "${PYTHON_DEST}/sirf")
+endif()

src/torch/README.md

@@ -0,0 +1,101 @@
+# SIRF-PyTorch Wrapper
+This wrapper provides a bridge between the [SIRF](https://github.com/SyneRBI/SIRF) (Synergistic Image Reconstruction Framework) library and [PyTorch](https://github.com/pytorch/pytorch), enabling the use of SIRF's image reconstruction operators and objective functions within PyTorch's automatic differentiation (autodiff) framework.
+
+## Usage and Use Cases
+
+The `sirf.torch.Operator`, `sirf.torch.ObjectiveFunction`, and `sirf.torch.ObjectiveFunctionGradient` classes are designed to be used as standard PyTorch `nn.Module`s.  You would initialise them with the appropriate SIRF objects and then use them in your forward pass like any other PyTorch layer.
+
+`tests/use_cases.py` demonstrates `sirf.torch` integration in PyTorch with minimal 2D PET examples:
+
+*   **Learned Primal-Dual:** Implements a learned primal-dual network for PET image reconstruction, showcasing the use of `sirf.torch.Operator` for handling the forward and adjoint projection operations.
+*   **PET Variational Network (PETVarNet):**  Demonstrates a variational network approach, combining convolutional blocks with gradient information from a SIRF objective function using `sirf.torch.ObjectiveFunctionGradient`.
+*   **ADAM Gradient Descent Comparison:**  Compares two gradient descent implementations: one leveraging the `sirf.torch.Operator` for the acquisition model within the loss calculation, and another directly utilising the `sirf.torch.ObjectiveFunction` for a more traditional optimisation approach.  This highlights the flexibility of the wrapper in different optimisation strategies.
+
+## Dimensions used by the wrapper
+
+The wrappers prioritise SIRF's data formats, meaning that the torch arrays must have the shape:
+* [batch, [channel,] *SIRF.DataContainer.shape], where the channel dimension is optional.
+
+This requires the **user** to ensure the dimensionality to match between layers.
+
+### Example dimension manipulation
+For example, a sinogram in SIRF has shape [tof bins, sinograms, views, tang pos]. For a single non-tof sinogram this is [1, 1, views, tang pos]. The expected torch tensor shape for this wrapper is [batch, [channel,] 1, 1, views, tang pos]. On the otherhand a 2D convolution requires [batch, [channel,] height, width].
+
+```python
+conv_1 = torch.nn.Conv2D()
+conv_2 = torch.nn.Conv2D()
+adjoint_operator = sirf.SIRF_torch.Operator(sirf.AdjointOperator(acquisition_model))
+y # sinogram of dimension [batch, [channel,] views, tang pos]
+
+y_filtered = conv_1(y) # filtered sinogram of dimension [batch, [channel,] views, tang pos]
+y_filtered = y_filtered.unsqueeze(-3).unsqueeze(-3) # filtered sinogram of dimension [batch, [channel,] 1, 1, views, tang pos]
+x_bp = adjoint_operator(y_filtered) # back-projected image of dimension [batch, [channel,] 1, height, width]
+x_bp = x_bp.squeeze(-3)  # back-projected image of dimension [batch, [channel,] height, width]
+x_bp_filtered = conv_2(x_bp)  # filtered back-projected image of dimension [batch, [channel,] height, width]
+```
+
+## Forward and backward clarification
+
+The use of the terms forward and backward have different meaning given the context:
+* Automatic differentiation: Forward (tangent) mode autodiff computes the Jacobian-Vector-Product (JVP). This propagates derivatives forward along with the function evaluation. Backward (or reverse/adjoint) mode autodiff is the Vector-Jacobian-Product (VJP) that propagates derivative information in the reverse direction of the function's evaluation. 
+* Backward autodiff cont.: Forward pass evaluates the function saving intermediate values. Backward pass uses the chain rule and intermediate values computing the derivatives in the reverse direction with the VJP.
+* `torch.autograd.Function`: the `forward` method (forward pass) is the function evaluation. The `backward` method (backward pass) computes the VJP. More specifically, the `backward(*grad_output)` method multiplies the `grad_output` which represents the gradient(s) of a subsequent function/operator (evaluated at the output of `forward`), via chain-rule, by the adjoint of the Jacobian of the `forward` method. 
+
+This SIRF-PyTorch wrapper is **only** for reverse-mode automatic differentiation via subclassing `torch.autograd.Function`.
+
+## Wrapper Design
+
+The wrapper provides three main classes:
+
+1.  `sirf.torch.Operator`: Wraps a SIRF `Operator` (e.g., a projection operator). Applies the operator forward pass, and applies the adjoint of the Jacobian in backward pass.
+2.  `sirf.torch.ObjectiveFunction`: Wraps a SIRF `ObjectiveFunction` for computing its value in the forward pass, and multiplying with the objective function gradient in the backward pass.
+3.  `sirf.torch.ObjectiveFunctionGradient`: Wraps a SIRF `ObjectiveFunction` that computes the objective function gradients in the forward pass and the Hessian-vector product in the backward pass. In the backward the Hessian is evaluated at the point which the objective function's gradient was evaluated.
+
+These classes use custom `torch.autograd.Function` implementations (`_Operator`, `_ObjectiveFunction`, and `_ObjectiveFunctionGradient`) to define the forward and backward passes, handling the conversions between PyTorch tensors and SIRF objects.
+
+### `_Operator` (Forward and Backward Passes)
+
+*   **Forward Pass:**
+    1.  Converts the input PyTorch tensor to a SIRF object.
+    2.  Applies the SIRF `Operator.forward()` method.
+    3.  Converts the result back to a PyTorch tensor.
+    4.  If the input tensor requires gradients, it saves relevant information (the output SIRF object and the operator) in the context (`ctx`) for use in the backward pass.
+
+*   **Backward Pass (VJP):**
+    1.  Receives the "upstream gradient" (`grad_output`).
+    2.  Converts `grad_output` to a SIRF object.
+    3.  Applies the SIRF `Operator.backward()` method. This will apply the **Jacobian adjoint** of the operator to upstream gradient (the vector).
+    4.  Converts the resulting SIRF object back to a PyTorch tensor and returns it.
+
+### `_ObjectiveFunction` (Forward and Backward Passes)
+
+*   **Forward Pass:**
+    1.  Converts the input PyTorch tensor (representing an image for instance) to a SIRF object.
+    2.  Calls the SIRF `ObjectiveFunction.__call__()` method.
+    3.  Returns the of the objective function value as a PyTorch tensor.
+    4.  Saves relevant information to the `ctx` if gradients are needed.
+
+*   **Backward Pass (VJP):**
+    1.  Receives the upstream gradient (`grad_output`), in this case it is always a scalar.
+    2.  Gets the gradient of the objective function using `sirf_obj_func.gradient()`, which computed at the input and multiplied by the upstream gradient.
+    3.  Converts the SIRF gradient to a PyTorch tensor.
+    4.  Returns the gradient multiplied by `grad_output`.
+
+
+### `_ObjectiveFunctionGradient` (Forward and Backward Passes)
+
+*   **Forward Pass:**
+    1.  Converts the input PyTorch tensor to a SIRF object.
+    2.  Computes the *gradient* of the SIRF objective function using `sirf_obj_func.gradient()`, which is computed on the input.
+    3.  Returns the gradient as a PyTorch tensor.
+
+*   **Backward Pass (VJP):**
+    1.  Receives the upstream gradient (`grad_output`), which now represents a *vector* (not a scalar) of the same shape as the output of `forward`.
+    2.  Converts `grad_output` to a SIRF object.
+    3.  Multiples the Hessian evaluated at the input of `forward` with the "upstream gradient" using `sirf_obj_func.multiply_with_Hessian()`.
+    4.  Returns the Hessian multiplied with a vector as a tensor.
+
+# TODO
+
+* Extend to subsets in the wrapper
+* Extend objective functions that vary between batch items