feat: Introduce Fast_OS_SART and Adaptive-Weighted TV to ART algorithms

Python/tigre/algorithms/art_family_algorithms.py

@@ -1,5 +1,5 @@
 import copy
-
+import numpy as np
 from tigre.algorithms.iterative_recon_alg import IterativeReconAlg
 from tigre.algorithms.iterative_recon_alg import decorator
 from tigre.utilities.im_3d_denoise import im3ddenoise
@@ -76,77 +76,191 @@ def __init__(self, proj, geo, angles, niter, **kwargs):
         if "blocksize" in kwargs and kwargs['blocksize']>1:
             print('Warning: blocksize is set to 1, please use an OS version of the algorithm for blocksize > 1')
         kwargs.update(dict(blocksize=1))
-        self.tvlambda = 50 if 'tvlambda' not in kwargs else kwargs['tvlambda']
-        self.tviter = 50 if 'tviter' not in kwargs else kwargs['tviter']
-        # these two settings work well for nVoxel=[254,254,254]
-
-        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
-
-    # Override
-    def run_main_iter(self):
-        """
-        Goes through the main iteration for the given configuration.
-        :return: None
-        """
-        Quameasopts = self.Quameasopts
-
-        for i in range(self.niter):
-
-            res_prev = None
-            if Quameasopts is not None:
-                res_prev = copy.deepcopy(self.res)
-            if self.verbose:
-                self._estimate_time_until_completion(i)
-
-            getattr(self, self.dataminimizing)()
-            # print("run_main_iter: gpuids = {}", self.gpuids)
-            self.res = im3ddenoise(self.res, self.tviter, self.tvlambda, self.gpuids)
-            if Quameasopts is not None:
-                self.error_measurement(res_prev, i)
-
-
-sart_tv = decorator(SART_TV, name="sart_tv")
-
-
-class OSSART_TV(IterativeReconAlg):  
-    __doc__ = (
-        "OSSART_TV solves Cone Beam CT image reconstruction using Oriented Subsets\n"
-        "Simultaneous Algebraic Reconstruction Technique with TV regularization algorithm\n"
-        "OSSART_TV(PROJ,GEO,ALPHA,NITER,BLOCKSIZE=20,TVLAMBDA=50,TVITER=50) \n"
-        "solves the reconstruction problem using the projection data PROJ taken\n"
-        "over ALPHA angles, corresponding to the geometry described in GEO,\n"
-        "using NITER iterations.\n"
-    ) + IterativeReconAlg.__doc__
-
-    def __init__(self, proj, geo, angles, niter, **kwargs):
-
-        self.blocksize = 20 if 'blocksize' not in kwargs else kwargs['blocksize']
-        self.tvlambda = 50 if 'tvlambda' not in kwargs else kwargs['tvlambda']
-        self.tviter = 50 if 'tviter' not in kwargs else kwargs['tviter']
-        # these two settings work well for nVoxel=[254,254,254]
-
-        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
-
-   # Override
-    def run_main_iter(self):
-        """
-        Goes through the main iteration for the given configuration.
-        :return: None
-        """
-        Quameasopts = self.Quameasopts
-
-        for i in range(self.niter):
-
-            res_prev = None
-            if Quameasopts is not None:
-                res_prev = copy.deepcopy(self.res)
-            if self.verbose:
-                self._estimate_time_until_completion(i)
-
-            getattr(self, self.dataminimizing)()
-            # print("run_main_iter: gpuids = {}", self.gpuids)
-            self.res = im3ddenoise(self.res, self.tviter, self.tvlambda, self.gpuids)
-            if Quameasopts is not None:
-                self.error_measurement(res_prev, i)
-
-ossart_tv = decorator(OSSART_TV, name="ossart_tv")
+import copy
+import numpy as np
+from tigre.algorithms.iterative_recon_alg import IterativeReconAlg
+from tigre.algorithms.iterative_recon_alg import decorator
+from tigre.utilities.im_3d_denoise import im3ddenoise
+
+
+
+class SART(IterativeReconAlg):  
+    __doc__ = (
+        "SART solves Cone Beam CT image reconstruction using \n"
+        "Simultaneous Algebraic Reconstruction Technique algorithm\n"
+        "SART(PROJ,GEO,ALPHA,NITER) solves the reconstruction problem\n"
+        "using the projection data PROJ taken over ALPHA angles, corresponding\n"
+        "to the geometry described in GEO, using NITER iterations. \n"
+    ) + IterativeReconAlg.__doc__
+
+    def __init__(self, proj, geo, angles, niter, **kwargs):
+        if "blocksize" in kwargs and kwargs['blocksize']>1:
+            print('Warning: blocksize is set to 1, please use an OS version of the algorithm for blocksize > 1')
+        kwargs.update(dict(blocksize=1))
+        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
+
+
+sart = decorator(SART, name="sart")
+
+
+class SIRT(IterativeReconAlg):  
+    __doc__ = (
+        "SIRT solves Cone Beam CT image reconstruction using \n"
+        "Simultaneous Iterative Reconstructive Technique algorithm\n"
+        "SIRT(PROJ,GEO,ALPHA,NITER) solves the reconstruction problem\n"
+        "using the projection data PROJ taken over ALPHA angles, corresponding\n"
+        "to the geometry described in GEO, using NITER iterations.\n"
+    ) + IterativeReconAlg.__doc__
+
+    def __init__(self, proj, geo, angles, niter, **kwargs):
+        if "blocksize" in kwargs and kwargs['blocksize']>1:
+            print('Warning: blocksize is set to {}, please do not specify blocksize for this algorithm'.format(angles.shape[0]))
+        kwargs.update(dict(blocksize=angles.shape[0]))
+        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
+
+
+sirt = decorator(SIRT, name="sirt")
+
+
+class OS_SART(IterativeReconAlg):  
+    __doc__ = (
+        "OS_SART solves Cone Beam CT image reconstruction using Oriented Subsets\n"
+        "Simultaneous Algebraic Reconstruction Technique algorithm\n"
+        "OS_SART(PROJ,GEO,ALPHA,NITER,BLOCKSIZE=20) solves the reconstruction problem\n"
+        "using the projection data PROJ taken over ALPHA angles, corresponding\n"
+        "to the geometry described in GEO, using NITER iterations.\n"
+    ) + IterativeReconAlg.__doc__
+
+    def __init__(self, proj, geo, angles, niter, **kwargs):
+
+        self.blocksize = 20 if 'blocksize' not in kwargs else kwargs["blocksize"]       
+        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
+
+
+ossart = decorator(OS_SART, name="ossart")
+
+
+class SART_TV(IterativeReconAlg):  
+    __doc__ = (
+        "SART_TV solves Cone Beam CT image reconstruction using Simultaneous \n"
+        "Algebraic Reconstruction Technique with TV regularization algorithm\n"
+        "SART_TV(PROJ,GEO,ALPHA,NITER,TVLAMBDA=50,TVITER=50) solves the reconstruction\n"
+        "problem using the projection data PROJ taken over ALPHA angles\n"
+        "corresponding to the geometry described in GEO, using NITER iterations. \n"
+    ) + IterativeReconAlg.__doc__
+
+    def __init__(self, proj, geo, angles, niter, **kwargs):
+
+        if "blocksize" in kwargs and kwargs['blocksize']>1:
+            print('Warning: blocksize is set to 1, please use an OS version of the algorithm for blocksize > 1')
+        kwargs.update(dict(blocksize=1))
+        self.tvlambda = 50 if 'tvlambda' not in kwargs else kwargs['tvlambda']
+        self.tviter = 50 if 'tviter' not in kwargs else kwargs['tviter']
+        # these two settings work well for nVoxel=[254,254,254]
+
+        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
+
+    # Override
+    def run_main_iter(self):
+        """
+        Goes through the main iteration for the given configuration.
+        :return: None
+        """
+        Quameasopts = self.Quameasopts
+
+        nesterov = False
+        if isinstance(self.lmbda, str) and self.lmbda.lower() == "nesterov":
+            nesterov = True
+            self.lmbda = 1.0
+            t = 1.0
+            y_rec = copy.deepcopy(self.res)
+
+        for i in range(self.niter):
+
+            res_prev = None
+            if Quameasopts is not None:
+                res_prev = copy.deepcopy(self.res)
+            if self.verbose:
+                self._estimate_time_until_completion(i)
+
+            if nesterov:
+                x_rec_old = copy.deepcopy(self.res)
+                self.res = copy.deepcopy(y_rec)
+
+            getattr(self, self.dataminimizing)()
+            # print("run_main_iter: gpuids = {}", self.gpuids)
+            self.res = im3ddenoise(self.res, self.tviter, self.tvlambda, self.gpuids)
+
+            if nesterov:
+                t_old = t
+                t = (1.0 + np.sqrt(1.0 + 4.0 * t ** 2)) / 2.0
+                gamma = np.float32((t_old - 1.0) / t)
+                y_rec = self.res + gamma * (self.res - x_rec_old)
+
+            if Quameasopts is not None:
+                self.error_measurement(res_prev, i)
+
+
+sart_tv = decorator(SART_TV, name="sart_tv")
+
+
+class OSSART_TV(IterativeReconAlg):  
+    __doc__ = (
+        "OSSART_TV solves Cone Beam CT image reconstruction using Oriented Subsets\n"
+        "Simultaneous Algebraic Reconstruction Technique with TV regularization algorithm\n"
+        "OSSART_TV(PROJ,GEO,ALPHA,NITER,BLOCKSIZE=20,TVLAMBDA=50,TVITER=50) \n"
+        "solves the reconstruction problem using the projection data PROJ taken\n"
+        "over ALPHA angles, corresponding to the geometry described in GEO,\n"
+        "using NITER iterations.\n"
+    ) + IterativeReconAlg.__doc__
+
+    def __init__(self, proj, geo, angles, niter, **kwargs):
+
+        self.blocksize = 20 if 'blocksize' not in kwargs else kwargs['blocksize']
+        self.tvlambda = 50 if 'tvlambda' not in kwargs else kwargs['tvlambda']
+        self.tviter = 50 if 'tviter' not in kwargs else kwargs['tviter']
+        # these two settings work well for nVoxel=[254,254,254]
+
+        IterativeReconAlg.__init__(self, proj, geo, angles, niter, **kwargs)
+
+   # Override
+    def run_main_iter(self):
+        """
+        Goes through the main iteration for the given configuration.
+        :return: None
+        """
+        Quameasopts = self.Quameasopts
+
+        nesterov = False
+        if isinstance(self.lmbda, str) and self.lmbda.lower() == "nesterov":
+            nesterov = True
+            self.lmbda = 1.0
+            t = 1.0
+            y_rec = copy.deepcopy(self.res)
+
+        for i in range(self.niter):
+
+            res_prev = None
+            if Quameasopts is not None:
+                res_prev = copy.deepcopy(self.res)
+            if self.verbose:
+                self._estimate_time_until_completion(i)
+
+            if nesterov:
+                x_rec_old = copy.deepcopy(self.res)
+                self.res = copy.deepcopy(y_rec)
+
+            getattr(self, self.dataminimizing)()
+            # print("run_main_iter: gpuids = {}", self.gpuids)
+            self.res = im3ddenoise(self.res, self.tviter, self.tvlambda, self.gpuids)
+
+            if nesterov:
+                t_old = t
+                t = (1.0 + np.sqrt(1.0 + 4.0 * t ** 2)) / 2.0
+                gamma = np.float32((t_old - 1.0) / t)
+                y_rec = self.res + gamma * (self.res - x_rec_old)
+
+            if Quameasopts is not None:
+                self.error_measurement(res_prev, i)
+
+ossart_tv = decorator(OSSART_TV, name="ossart_tv")

Python/tigre/algorithms/iterative_recon_alg.py

@@ -312,6 +312,13 @@ def run_main_iter(self):
         """
         Quameasopts = self.Quameasopts
 
+        nesterov = False
+        if isinstance(self.lmbda, str) and self.lmbda.lower() == "nesterov":
+            nesterov = True
+            self.lmbda = 1.0
+            t = 1.0
+            y_rec = copy.deepcopy(self.res)
+
         for i in range(self.niter):
 
             res_prev = None
@@ -320,7 +327,18 @@ def run_main_iter(self):
             if self.verbose:
                 self._estimate_time_until_completion(i)
 
+            if nesterov:
+                x_rec_old = copy.deepcopy(self.res)
+                self.res = copy.deepcopy(y_rec)
+
             getattr(self, self.dataminimizing)()
+
+            if nesterov:
+                t_old = t
+                t = (1.0 + np.sqrt(1.0 + 4.0 * t ** 2)) / 2.0
+                gamma = np.float32((t_old - 1.0) / t)
+                y_rec = self.res + gamma * (self.res - x_rec_old)
+
             self.error_measurement(res_prev, i)
 
     def art_data_minimizing(self):

generate_benchmarks.py

@@ -0,0 +1,70 @@
+import sys
+import time
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Attempt to import TIGRE
+try:
+    import tigre
+    import tigre.algorithms as algs
+    from tigre.utilities.sample_loader import load_head_phantom
+    from tigre.utilities.Measure_Quality import Measure_Quality
+except ImportError:
+    print("ERROR: TIGRE is not properly installed or compiled.")
+    print("Please run this script in an environment with TIGRE's C++/CUDA backend compiled.")
+    sys.exit(1)
+
+def run_benchmarks():
+    print("--- Setting up TIGRE Geometry & Phantom ---")
+    # 1. Setup geometry and phantom
+    geo = tigre.geometry_default(high_resolution=False)
+    geo.nVoxel = np.array([64, 64, 64]) # Use small voxel size for fast benchmarking
+
+    # Generate angles
+    angles = np.linspace(0, 2 * np.pi, 100)
+
+    # Load phantom
+    head = load_head_phantom(geo.nVoxel)
+
+    # Generate projection data
+    print("Generating forward projections...")
+    proj = tigre.Ax(head, geo, angles)
+
+    niter = 30
+    blocksize = 20
+
+    print("\n--- Benchmark 1: Convergence Speed & Time (OS_SART vs OS_SART with Nesterov) ---")
+
+    # Standard OS_SART
+    print("Running standard OS_SART...")
+    start_time = time.time()
+    res_os_sart, err_os_sart = algs.ossart(proj, geo, angles, niter=niter, blocksize=blocksize, computel2=True)
+    time_os_sart = time.time() - start_time
+
+    # Fast OS_SART
+    print("Running OS_SART with Nesterov acceleration...")
+    start_time = time.time()
+    res_fast, err_fast = algs.ossart(proj, geo, angles, niter=niter, blocksize=blocksize, lmbda='nesterov', computel2=True)
+    time_fast = time.time() - start_time
+
+    print(f"OS_SART Total Time: {time_os_sart:.2f}s ({time_os_sart/niter:.3f}s per iteration)")
+    print(f"OS_SART (Nesterov) Total Time: {time_fast:.2f}s ({time_fast/niter:.3f}s per iteration)")
+    print(f"Final L2 Error -> OS_SART: {err_os_sart[0][-1]:.4f} | OS_SART (Nesterov): {err_fast[0][-1]:.4f}")
+
+    # Plot convergence
+    plt.figure(figsize=(8, 5))
+    plt.plot(err_os_sart[0], label="OS_SART", linewidth=2)
+    plt.plot(err_fast[0], label="OS_SART (Nesterov)", linewidth=2)
+    plt.title("Convergence Speed: OS_SART vs OS_SART (Nesterov)")
+    plt.xlabel("Iteration")
+    plt.ylabel("L2 Error")
+    plt.legend()
+    plt.grid(True)
+    plt.savefig("convergence_benchmark.png")
+    print("Saved convergence plot to convergence_benchmark.png")
+
+    print("\n--- Benchmarks Complete! ---")
+    print("You can copy these results into your GitHub PR.")
+
+if __name__ == '__main__':
+    run_benchmarks()