BSREM bb

BSREM.py

@@ -8,6 +8,7 @@
 # Copyright 2024 University College London
 
 import numpy
+import numpy as np 
 import sirf.STIR as STIR
 from sirf.Utilities import examples_data_path
 
@@ -59,6 +60,9 @@ def __init__(self, data, initial, initial_step_size, relaxation_eta,
 
         self.alpha = None 
 
+        self.x_prev = None 
+        self.x_update_prev = None 
+
     def subset_sensitivity(self, subset_num):
         raise NotImplementedError
 
@@ -74,12 +78,25 @@ def step_size(self):
     def update(self):
         g = self.subset_gradient(self.x, self.subset_order[self.subset])
 
-        self.alpha = self.step_size()
-        self.x_update = (self.x + self.eps) * g / self.average_sensitivity * self.alpha
+        self.x_update = (self.x + self.eps) * g / self.average_sensitivity 
+
+        if self.iteration == 0:
+            self.alpha = self.initial_step_size
+        else:
+            delta_x = self.x - self.x_prev
+            delta_g = self.x_update_prev - self.x_update 
+
+            self.alpha = delta_x.norm()**2 / (delta_x * delta_g).sum()
 
         if self.update_filter is not None:
             self.update_filter.apply(self.x_update)
-        self.x += self.x_update
+
+        self.x_prev = self.x.clone() 
+        self.x_update_prev = self.x_update.clone() 
+
+        #print("Use step size: ", self.alpha)
+        #print(self.x_update)
+        self.x += self.x_update * self.alpha
         # threshold to non-negative
         self.x.maximum(0, out=self.x)
         self.subset = (self.subset + 1) % self.num_subsets

main.py

@@ -106,7 +106,6 @@ def apply(self, algorithm, gradient, out=None):
 class Submission(BSREM):
     # note that `issubclass(BSREM1, Algorithm) == True`
     def __init__(self, data, 
-                       num_subsets: int = 7, 
                        update_objective_interval: int = 10,
                        preconditioner = "osem",
                        **kwargs):
@@ -116,7 +115,8 @@ def __init__(self, data,
         This is just an example. Try to modify and improve it!
         """
         mode = kwargs.get("mode", "sequential")
-        num_subsets = 1 
+        mode = kwargs.get("num_subsets", 1)
+
         data_sub, acq_models, obj_funs = partitioner.data_partition(data.acquired_data, data.additive_term,
                                                                     data.mult_factors, num_subsets,
                                                                     initial_image=data.OSEM_image,