use a small number of subsets for TOF data

SyneRBI · Sep 30, 2024 · 65e128d · 65e128d
1 parent b785402
commit 65e128d
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 9 deletions.
diff --git a/bsrem_bb.py b/bsrem_bb.py
@@ -102,7 +102,7 @@ def __init__(self, data, initial,
 
         self.x_update = initial.get_uniform_copy(0)
 
-        self.rdp_hessian_freq = 4
+        self.rdp_hessian_freq = 3
 
     def subset_sensitivity(self, subset_num):
         raise NotImplementedError
@@ -117,8 +117,16 @@ def step_size(self):
         return self.initial_step_size / (1 + self.relaxation_eta * self.epoch())
 
     def update(self):
-
-        g = self.subset_gradient(self.x, self.subset_order[self.subset])
+        if self.iteration == 0:
+            g = self.x.get_uniform_copy(0)
+            for i in range(self.num_subsets):
+                gm = self.subset_gradient(self.x, self.subset_order[i]) 
+                g.add(gm, out=g)
+
+            g /= self.num_subsets
+        else:
+            g = self.subset_gradient(self.x, self.subset_order[self.subset])
+
         if self.iteration == 0:
             prior_grad = self.dataset.prior.gradient(self.x)
             if prior_grad.norm()/g.norm() > 0.5:

diff --git a/main.py b/main.py
@@ -17,7 +17,7 @@
 
 assert issubclass(BSREM, Algorithm)
 
-
+from utils.number_of_subsets import divisorGenerator
 
 
 import setup_postprocessing 
@@ -38,7 +38,18 @@ def __init__(self, data,
                        update_objective_interval: int = 2,
                        **kwargs):
 
-        num_subsets = 1
+        tof = data.acquired_data.shape[0]
+        views = data.acquired_data.shape[2]
+        if tof == 1:
+            num_subsets = 1
+        else:
+            # use a small number of subsets for TOF data
+            num_divisors = list(divisorGenerator(views))
+
+            if len(num_divisors) <= 2:
+                num_subsets = num_divisors[-1]
+            else:
+                num_subsets = num_divisors[2]
 
         data_sub, _, obj_funs = partitioner.data_partition(data.acquired_data, data.additive_term,
                                                                     data.mult_factors, num_subsets,

diff --git a/utils/number_of_subsets.py b/utils/number_of_subsets.py
@@ -121,11 +121,15 @@ def divisorGenerator(n):
 
 if __name__ == "__main__":
 
-    for num in [50, 252, 128]:
-        print("Test for ", num, " -> ", compute_number_of_subsets(num, tof=False), " subsets")
+    #for num in [50, 252, 128]:
+    #    print("Test for ", num, " -> ", compute_number_of_subsets(num, tof=False), " subsets")
 
-    for num in [50, 252, 128]:
-        print("Test TOF for ", num, " -> ", compute_number_of_subsets(num, tof=True), " subsets")
+    #for num in [50, 252, 128]:
+    #    print("Test TOF for ", num, " -> ", compute_number_of_subsets(num, tof=True), " subsets")
+    views = 50
+    num_divisors = list(divisorGenerator(views))
+
+    print(num_divisors)
 
     #for num in range(40, 400):
     #    print(num, compute_number_of_subsets(num, tof=True))