[Not ready for review] Reduce tomosipo exposure - More MPS compartibility for Mac users

.gitignore

@@ -15,4 +15,5 @@ dist
 example.py
 /slurm
 
-wandb
+wandb
+.DS_Store

LION/CTtools/ct_utils.py

@@ -15,7 +15,7 @@
 
 # AItomotools imports
 from LION.CTtools.ct_geometry import Geometry
-from LION.operators import CTProjectionOp
+from LION.operators.CTProjectionOp import CTProjectionOp
 
 
 def from_HU_to_normal(img):

LION/classical_algorithms/fista.py

@@ -5,7 +5,7 @@
 import torch
 from tqdm import tqdm
 
-from LION.operators import Operator
+from LION.operators.Operator import Operator
 from LION.utils.math import power_method
 
 

LION/losses/SUREpgImage.py

@@ -1,42 +1,55 @@
 import torch
 
-class SUREpgLoss():
 
-    def __init__(self, zeta: float, sigma2: float, eps1: float = 1e-3, eps2: float = 1e-3, kappa: float = 1.0):
+class SUREpgLoss:
+    def __init__(
+        self,
+        zeta: float,
+        sigma2: float,
+        eps1: float = 1e-3,
+        eps2: float = 1e-3,
+        kappa: float = 1.0,
+    ):
         self.zeta = zeta
         self.sigma2 = sigma2
         self.eps1 = eps1
         self.eps2 = eps2
         self.kappa = kappa
 
-        self.p = (1/2) * (1 + self.kappa / (self.kappa**2 + 4)**0.5)
+        self.p = (1 / 2) * (1 + self.kappa / (self.kappa**2 + 4) ** 0.5)
         self.q = 1 - self.p
-        self.a = (self.q / self.p)**0.5
-        self.b = (self.p / self.q)**0.5
+        self.a = (self.q / self.p) ** 0.5
+        self.b = (self.p / self.q) ** 0.5
 
     def __call__(self, model, y):
-        B=y.shape[0]
+        B = y.shape[0]
         N_per_img = y.shape[1] * y.shape[2] * y.shape[3]
 
-        fy=model(y)
-        loss = ((fy - y) ** 2).sum(dim=(1,2,3)) - self.zeta * y.sum(dim=(1,2,3)) - self.sigma2 * N_per_img
+        fy = model(y)
+        loss = (
+            ((fy - y) ** 2).sum(dim=(1, 2, 3))
+            - self.zeta * y.sum(dim=(1, 2, 3))
+            - self.sigma2 * N_per_img
+        )
 
-        #1st derivative MC
+        # 1st derivative MC
         delta1 = torch.randn_like(y)
         fy_perturbated = model(y + self.eps1 * delta1)
 
         u = self.zeta * y + self.sigma2
-        mc1 = (delta1 * u * (fy_perturbated - fy)).sum(dim=(1,2,3))
+        mc1 = (delta1 * u * (fy_perturbated - fy)).sum(dim=(1, 2, 3))
         loss += 2.0 * mc1 / self.eps1
 
-        #2nd derivative MC
+        # 2nd derivative MC
         u_rand = torch.rand_like(y)
-        delta2 = torch.where(u_rand < self.p,-self.a * torch.ones_like(y),+self.b * torch.ones_like(y))
+        delta2 = torch.where(
+            u_rand < self.p, -self.a * torch.ones_like(y), +self.b * torch.ones_like(y)
+        )
 
-        fy_plus  = model(y + self.eps2 * delta2)
+        fy_plus = model(y + self.eps2 * delta2)
         fy_minus = model(y - self.eps2 * delta2)
 
-        mc2 = (delta2 * (fy_plus - 2*fy + fy_minus)).sum(dim=(1,2,3))
+        mc2 = (delta2 * (fy_plus - 2 * fy + fy_minus)).sum(dim=(1, 2, 3))
         loss -= (2 * self.sigma2 * self.zeta / (self.eps2**2 * self.kappa)) * mc2
 
-        return loss.mean()/N_per_img
+        return loss.mean() / N_per_img

LION/optimizers/Noisier2Inverse.py

@@ -16,6 +16,7 @@
 
 import torchvision.transforms.functional as TF
 
+
 class Noisier2Inverse(LIONsolver):
     def __init__(
         self,
@@ -39,20 +40,19 @@ def __init__(
         )
 
         self.operator = ct_utils.make_operator(self.geometry)
-        self.model.geometry = self.geometry  
+        self.model.geometry = self.geometry
         self.model.operator = self.operator
         self.projector = to_autograd(self.operator, num_extra_dims=1)
         self.recon_fn = self.solver_params.recon_fn
 
-
     @staticmethod
     def default_parameters() -> LIONParameter:
         params = LIONParameter()
-        params.sigma=3
-        params.delta=1
+        params.sigma = 3
+        params.delta = 1
         params.recon_fn = fdk
         return params
-    
+
     def mini_batch_step(self, sinos, targets):
         sigma = self.solver_params.sigma
         delta = self.solver_params.delta
@@ -62,21 +62,20 @@ def mini_batch_step(self, sinos, targets):
         N = TF.gaussian_blur(N, kernel_size=[ks, ks], sigma=[sigma, sigma])
         z = sinos + N
 
-        input_recon = self.recon_fn(z,self.model.operator)
+        input_recon = self.recon_fn(z, self.model.operator)
         output_recon = self.model(input_recon)
         output_sino = self.projector(output_recon)
         target_sino = sinos - N
 
-        #Sobolev Loss
-        #res = output_sino - target_sino
-        #grad_x = res[:, :, :, 1:] - res[:, :, :, :-1]
-        #grad_y = res[:, :, 1:, :] - res[:, :, :-1, :]
+        # Sobolev Loss
+        # res = output_sino - target_sino
+        # grad_x = res[:, :, :, 1:] - res[:, :, :, :-1]
+        # grad_y = res[:, :, 1:, :] - res[:, :, :-1, :]
 
-        #batch_loss = ((output_sino - target_sino)**2).mean() + (grad_x**2).mean() + (grad_y**2).mean()
-        batch_loss= ((output_sino - target_sino)**2).mean()
+        # batch_loss = ((output_sino - target_sino)**2).mean() + (grad_x**2).mean() + (grad_y**2).mean()
+        batch_loss = ((output_sino - target_sino) ** 2).mean()
         return batch_loss
 
-
     # No validation in Noisier2Inverse
     def validate(self):
         return 0

LION/optimizers/Proj2ProjSolver.py

@@ -14,6 +14,7 @@
 import LION.CTtools.ct_utils as ct_utils
 from tomosipo.torch_support import to_autograd
 
+
 class Proj2ProjSolver(LIONsolver):
     def __init__(
         self,
@@ -37,57 +38,63 @@ def __init__(
         )
 
         self.operator = ct_utils.make_operator(self.geometry)
-        self.model.geometry = self.geometry  
+        self.model.geometry = self.geometry
         self.model.operator = self.operator
         self.projector = to_autograd(self.operator, num_extra_dims=1)
         self.recon_fn = self.solver_params.recon_fn
-        self.global_step=0
+        self.global_step = 0
 
     def get_mask(self, shape, step):
         # shape: (B, C, H, W)
         mask = torch.ones(shape, device=self.device)
-        grid = self.solver_params.grid_size 
-        
+        grid = self.solver_params.grid_size
+
         for b in range(shape[0]):
-            idx = (step+b) % (grid * grid)
+            idx = (step + b) % (grid * grid)
             r = idx // grid
             c = idx % grid
 
             mask[b, :, r::grid, c::grid] = 0
         return mask
-    
+
     def fill_mean(self, sinos, mask):
-        kernel = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=torch.float32, device=self.device) / 4.0
-
+        kernel = (
+            torch.tensor(
+                [[0, 1, 0], [1, 0, 1], [0, 1, 0]],
+                dtype=torch.float32,
+                device=self.device,
+            )
+            / 4.0
+        )
+
         kernel = kernel.view(1, 1, 3, 3)
         local_mean_sino = F.conv2d(sinos, kernel, padding=1)
-        
+
         filled_sinos = (sinos * mask) + (local_mean_sino * (1 - mask))
         return filled_sinos
 
     @staticmethod
     def default_parameters() -> LIONParameter:
         params = LIONParameter()
-        params.grid_size = 4 
+        params.grid_size = 4
         params.recon_fn = fdk
         return params
-    
+
     def mini_batch_step(self, sinos, targets):
-        mask=self.get_mask(sinos.shape, self.global_step)
+        mask = self.get_mask(sinos.shape, self.global_step)
         self.global_step += sinos.shape[0]
-        input_sino=self.fill_mean(sinos,mask)
-        
-        input_recon=self.recon_fn(input_sino,self.model.operator)
+        input_sino = self.fill_mean(sinos, mask)
+
+        input_recon = self.recon_fn(input_sino, self.model.operator)
         output_recon = self.model(input_recon)
-        output_sino=self.projector(output_recon)
+        output_sino = self.projector(output_recon)
 
-        output_sino_mask=output_sino*(1-mask)
-        target_sino=sinos*(1-mask)
+        output_sino_mask = output_sino * (1 - mask)
+        target_sino = sinos * (1 - mask)
 
         batch_loss = ((output_sino_mask - target_sino) ** 2).mean()
         return batch_loss
 
-
     # No validation in Proj2Proj
     def validate(self):
         return 0