Adding PyTorch functionality to SIRF

src/torch/SIRF_torch.py

@@ -8,29 +8,31 @@
 # https://github.com/educating-dip/pet_deep_image_prior/blob/main/src/deep_image_prior/torch_wrapper.py
 
 
-def sirf_to_torch(sirf_src, torch_dest, requires_grad=False):
+def sirf_to_torch(sirf_src, device, requires_grad=False):
     if requires_grad:
-        return torch.tensor(sirf_src.as_array(), requires_grad=True).to(torch_dest.device)
+        # use torch.tensor to infer data type
+        return torch.tensor(sirf_src.as_array(), requires_grad=True).to(device)
     else:
-        return torch.tensor(sirf_src.as_array()).to(torch_dest.device)
+        return torch.tensor(sirf_src.as_array()).to(device)
 
 def torch_to_sirf(torch_src, sirf_dest): 
     return sirf_dest.fill(torch_src.detach().cpu().numpy())
 
 
-class _ObjectiveFunctionModule(torch.autograd.Function):
+class _ObjectiveFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx,
             torch_image,
             sirf_image_template, 
             sirf_obj_func
             ):
 
+        device = torch_image.device
         sirf_image_template = torch_to_sirf(torch_image, sirf_image_template)
         value_np = sirf_obj_func.get_value(sirf_image_template).as_array()
         if torch_image.requires_grad:
-            ctx.save_for_backward(torch_image, sirf_image_template, sirf_obj_func)
-            return torch.tensor(value_np, requires_grad=True).to(torch_image.device)
+            ctx.save_for_backward(device, sirf_image_template, sirf_obj_func)
+            return torch.tensor(value_np, requires_grad=True).to(device)
         else:
             return torch.tensor(value_np).to(torch_image.device)
 
@@ -39,70 +41,75 @@ def backward(ctx,
             grad_output
             ):
 
-        torch_image, sirf_image_template, sirf_obj_func = ctx.saved_tensors
+        device, sirf_image_template, sirf_obj_func = ctx.saved_tensors
         tmp_grad = sirf_obj.get_gradient(sirf_image_template)
-        grad = sirf_to_torch(tmp_grad, torch_image, requires_grad=True)
+        grad = sirf_to_torch(tmp_grad, device, requires_grad=True)
         return grad_output*grad, None, None, None
 
 class _AcquisitionModelForward(torch.autograd.Function):
     @staticmethod
     def forward(ctx,
             torch_image,
-            torch_measurements_template,
             sirf_image_template,
             sirf_acq_mdl
             ):
 
+        device = torch_image.device
         sirf_image_template = torch_to_sirf(torch_image, sirf_image_template)
         sirf_forward_projected = sirf_acq_mdl.forward(sirf_image_template)
         if torch_image.requires_grad:
-            ctx.torch_image = torch_image
+            ctx.device = device
             ctx.sirf_forward_projected = sirf_forward_projected
             ctx.sirf_acq_mdl = sirf_acq_mdl
-            return sirf_to_torch(sirf_forward_projected, torch_measurements_template, requires_grad=True)
+            return sirf_to_torch(sirf_forward_projected, device, requires_grad=True)
         else:
-            return sirf_to_torch(sirf_forward_projected, torch_measurements_template)
+            return sirf_to_torch(sirf_forward_projected, device)
 
     @staticmethod
     def backward(ctx,
             grad_output
             ):
         sirf_image = ctx.sirf_acq_mdl.backward(torch_to_sirf(grad_output, ctx.sirf_forward_projected))
-        grad = sirf_to_torch(sirf_image, ctx.torch_image, requires_grad=True)
-        return grad, None, None, None, None
+        grad = sirf_to_torch(sirf_image, ctx.device, requires_grad=True)
+        return grad, None, None, None
 
 
 
 class _AcquisitionModelBackward(torch.autograd.Function):
     @staticmethod
     def forward(ctx,
             torch_measurements,
-            torch_image_template,
             sirf_measurements_template,
             sirf_acq_mdl
             ):
 
+        device = torch_measurements.device
         sirf_measurements_template = torch_to_sirf(torch_measurements, sirf_measurements_template)
         sirf_backward_projected = sirf_acq_mdl.backward(sirf_measurements_template)
-        if torch_image_template.requires_grad:
-            ctx.torch_measurements = torch_measurements
+        if torch_measurements.requires_grad:
+            ctx.device = device
             ctx.sirf_backward_projected = sirf_backward_projected
             ctx.sirf_acq_mdl = sirf_acq_mdl
-            return sirf_to_torch(sirf_backward_projected, torch_image_template, requires_grad=True)
+            return sirf_to_torch(sirf_backward_projected, device, requires_grad=True)
         else:
-            return sirf_to_torch(sirf_backward_projected, torch_image_template)
+            return sirf_to_torch(sirf_backward_projected, device)
 
     @staticmethod
     def backward(ctx,
             grad_output
             ):
 
         sirf_measurements = ctx.sirf_acq_mdl.forward(torch_to_sirf(grad_output, ctx.sirf_backward_projected))
-        grad = sirf_to_torch(sirf_measurements, ctx.torch_measurements, requires_grad=True)
+        grad = sirf_to_torch(sirf_measurements, ctx.device, requires_grad=True)
         return grad, None, None, None, None
 
 class AcquisitionModelForward(torch.nn.Module):
-    def __init__(self, acq_mdl, sirf_image_template, sirf_measurements_template, device = "cpu", ):
+    def __init__(self,
+            acq_mdl, 
+            sirf_image_template, 
+            sirf_measurements_template, 
+            device = "cpu", 
+            ):
         super(AcquisitionModelForward, self).__init__()
         # get the shape of image and measurements
         self.acq_mdl = acq_mdl
@@ -111,21 +118,49 @@ def __init__(self, acq_mdl, sirf_image_template, sirf_measurements_template, dev
         self.sirf_image_template = sirf_image_template
         self.sirf_measurements_template = sirf_measurements_template
 
-        self.torch_measurements_template = torch.tensor(sirf_measurements_template.as_array(), requires_grad=False).to(device)*0
-
-    def forward(self, torch_image):
-        # view as torch sometimes doesn't like singleton dimensions
-        torch_image = torch_image.view(self.sirf_image_shape)
-        return _AcquisitionModelForward.apply(torch_image, self.torch_measurements_template, self.sirf_image_template, self.acq_mdl).squeeze()
-
-
+    def forward(self, image, **kwargs):
+        # PyTorch image (2D) is size [batch, channel, height, width] or [batch, height, width]
+        # PyTorch volume (3D) is size [batch, channel, depth, height, width] or [batch, depth, height, width]
+        # check keys of kwargs
+        if not set(kwargs.keys()).issubset({'attenuation', 'background', 'sensitivity','norm','csm'}):
+            raise ValueError("Invalid keyword arguments, only 'attenuation', 'background', 'sensitivity','norm','csm' are allowed. Assumes lists of same length.")
+
+        n_batch = image.shape[0]
+        if self.sirf_image_shape[0] == 1:
+            # if 2D
+            if len(image.shape) == 3:
+                # if 2D and no channel then add singleton
+                # add singleton for SIRF
+                image = image.unsqueeze(1)
+        else:
+            # if 3D
+            if len(image.shape) == 4:
+                # if 3D and no channel then add singleton
+                image = image.unsqueeze(1)
+        n_channel = image.shape[1]
+
+        # This looks horrible, but PyTorch should be able to trace.
+        batch_images = []
+        for batch in n_batch:
+            channel_images = []
+            for channel in n_channel:
+                torch_image = image[batch, channel]
+                torch_image = torch_image.view(self.sirf_image_shape)
+                # if there are kwargs then raise error
+                if kwargs:
+                    raise NotImplementedError("Keyword arguments are not implemented yet.")
+                channel_images.append(_AcquisitionModelForward.apply(torch_image, self.sirf_image_template, self.sirf_measurements_template, self.acq_mdl))
+            batch_images.append(torch.stack(channel_images, dim=0))
+        # [batch, channel, *forward_projected.shape]
+        return torch.stack(batch_images, dim=0) 
+
 if __name__ == '__main__':
     import os
     import numpy
     # Import the PET reconstruction engine
     import sirf.STIR as pet
     # Set the verbosity
-    #pet.set_verbosity(1)
+    pet.set_verbosity(1)
     # Store temporary sinograms in RAM
     pet.AcquisitionData.set_storage_scheme("memory")
     import sirf
@@ -155,7 +190,7 @@ def forward(self, torch_image):
 
 
     print("Comparing the forward projected")
-    torch_measurements = _AcquisitionModelForward.apply(torch_image, torch_measurements_template, sirf_image_template, sirf_acq_mdl)
+    torch_measurements = _AcquisitionModelForward.apply(torch_image, sirf_image_template, sirf_acq_mdl)
     print("Sum of torch: ", torch_measurements.detach().cpu().numpy().sum(), "Sum of sirf: ", sirf_measurements.sum(), "Sum of Differences: ", numpy.abs(torch_measurements.detach().cpu().numpy() - sirf_measurements.as_array()).sum())
     print("Comparing the backward of forward projected")
     # TO TEST THAT WE ARE BACKWARDING CORRECTLY RETAIN GRAD AND SUM THEN BACKWARD
@@ -173,7 +208,7 @@ def forward(self, torch_image):
     bp_sirf_measurements = acq_mdl.backward(sirf_measurements)
 
     print("Comparing the backward projected")
-    torch_image = _AcquisitionModelBackward.apply(torch_measurements, torch_image_template, sirf_measurements_template, sirf_acq_mdl)
+    torch_image = _AcquisitionModelBackward.apply(torch_measurements, sirf_measurements_template, sirf_acq_mdl)
     print("Sum of torch: ", torch_image.detach().cpu().numpy().sum(), "Sum of sirf: ", bp_sirf_measurements.sum(), \
         "Sum of Differences: ", numpy.abs(torch_image.detach().cpu().numpy() - bp_sirf_measurements.as_array()).sum())
     torch_measurements.retain_grad()
@@ -184,11 +219,14 @@ def forward(self, torch_image):
         numpy.abs(torch_measurements.grad.detach().cpu().numpy() - comparison.as_array()).sum())
 
 
-    """ # form objective function
+    # form objective function
     print("Objective Function Test")
     obj_fun = pet.make_Poisson_loglikelihood(sirf_measurements)
-    obj_fun.set_acquisition_model(acq_mdl)
+    print("Made poisson loglikelihood")
+    obj_fun.set_acquisition_model(sirf_acq_mdl)
+    print("Set acquisition model")
     obj_fun.set_up(sirf_image_template)
+    print("Set up")
 
     torch_image = torch.tensor(image.as_array(), requires_grad=True).to(device)
     sirf_image_template = image.get_uniform_copy(0)
@@ -200,8 +238,8 @@ def forward(self, torch_image):
     # grad sum(f(x)) = f'(x)
     comparison = obj_fun.gradient(sirf_image_template)
     print("Sum of torch: ", torch_image.grad.sum(), "Sum of sirf: ", comparison.sum(), "Sum of Differences: ", (torch_image.grad.detach().cpu().numpy() - comparison.as_array()).sum())
-    """
-    from sirf.Gadgetron import *
+
+    """ from sirf.Gadgetron import *
     print("2D MRI TEST")
     data_path = examples_data_path('MR')
     AcquisitionData.set_storage_scheme('memory')
@@ -227,7 +265,7 @@ def forward(self, torch_image):
 
     am = AcquisitionModel(processed_data, complex_images)
     am.set_coil_sensitivity_maps(csms)
-    print(am)
+    print(am) """
 
 # Objective function
 # Acquistion model