Modify cg_kwargs mechanism in admm.LinearSubproblemSolver

examples/scripts/ct_astra_tv_admm.py

@@ -57,7 +57,8 @@
 λ = 2e-0  # L1 norm regularization parameter
 ρ = 5e-0  # ADMM penalty parameter
 maxiter = 25  # number of ADMM iterations
-num_inner_iter = 20  # number of CG iterations per ADMM iteration
+cg_tol = 1e-4  # CG relative tolerance
+cg_maxiter = 25  # maximum CG iterations per ADMM iteration
 
 g = λ * functional.L1Norm()  # regularization functionals gi
 C = linop.FiniteDifference(input_shape=x_gt.shape)  # analysis operators Ci
@@ -73,7 +74,7 @@
     rho_list=[ρ],
     x0=x0,
     maxiter=maxiter,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": num_inner_iter}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": cg_tol, "maxiter": cg_maxiter}),
     itstat_options={"display": True, "period": 5},
 )
 

examples/scripts/ct_astra_weighted_tv_admm.py

@@ -106,7 +106,8 @@ def postprocess(x):
 lambda_unweighted = 2.56e2  # regularization strength
 
 maxiter = 50  # number of ADMM iterations
-max_inner_iter = 10  # number of CG iterations per ADMM iteration
+cg_tol = 1e-5  # CG relative tolerance
+cg_maxiter = 10  # maximum CG iterations per ADMM iteration
 
 f = loss.SquaredL2Loss(y=y, A=A)
 
@@ -117,7 +118,7 @@ def postprocess(x):
     rho_list=[ρ],
     x0=x0,
     maxiter=maxiter,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": max_inner_iter}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": cg_tol, "maxiter": cg_maxiter}),
     itstat_options={"display": True, "period": 10},
 )
 print(f"Solving on {device_info()}\n")
@@ -147,7 +148,7 @@ def postprocess(x):
     rho_list=[ρ],
     maxiter=maxiter,
     x0=x0,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": max_inner_iter}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": cg_tol, "maxiter": cg_maxiter}),
     itstat_options={"display": True, "period": 10},
 )
 admm_weighted.solve()

examples/scripts/ct_svmbir_ppp_bm3d_admm_cg.py

@@ -103,7 +103,7 @@
     rho_list=[ρ, ρ],
     x0=x0,
     maxiter=20,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 100}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-4, "maxiter": 100}),
     itstat_options={"display": True, "period": 1},
 )
 

examples/scripts/ct_svmbir_ppp_bm3d_admm_prox.py

@@ -105,7 +105,7 @@
     rho_list=[ρ, ρ, ρ],
     x0=x0,
     maxiter=20,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 100}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 100}),
     itstat_options={"display": True},
 )
 

examples/scripts/ct_svmbir_tv_multi.py

@@ -102,7 +102,7 @@
     rho_list=[2e1],
     x0=x0,
     maxiter=50,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 10}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-4, "maxiter": 10}),
     itstat_options={"display": True, "period": 10},
 )
 print(f"Solving on {device_info()}\n")

examples/scripts/deconv_ppp_bm3d_admm.py

@@ -68,7 +68,7 @@
     rho_list=[ρ],
     x0=A.T @ y,
     maxiter=maxiter,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 100}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 100}),
     itstat_options={"display": True},
 )
 

examples/scripts/deconv_ppp_dncnn_admm.py

@@ -73,7 +73,7 @@
     rho_list=[ρ],
     x0=A.T @ y,
     maxiter=maxiter,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 30}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 30}),
     itstat_options={"display": True},
 )
 

examples/scripts/demosaic_ppp_bm3d_admm.py

@@ -114,7 +114,7 @@ def demosaic(cfaimg):
     rho_list=[ρ],
     x0=imgb,
     maxiter=maxiter,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 100}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 100}),
     itstat_options={"display": True},
 )
 

examples/scripts/denoise_tv_iso_admm.py

@@ -69,7 +69,7 @@
     rho_list=[1e1],
     x0=y,
     maxiter=100,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 20}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 20}),
     itstat_options={"display": True, "period": 10},
 )
 
@@ -92,7 +92,7 @@
     rho_list=[1e1],
     x0=y,
     maxiter=100,
-    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"maxiter": 20}),
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 20}),
     itstat_options={"display": True, "period": 10},
 )
 

scico/optimize/admm.py

@@ -158,25 +158,29 @@ class LinearSubproblemSolver(SubproblemSolver):
             :math:`\mb{x}` update step.
     """
 
-    def __init__(self, cg_kwargs: dict = {"maxiter": 100}, cg_function: str = "scico"):
+    def __init__(self, cg_kwargs: Optional[dict] = None, cg_function: str = "scico"):
         """Initialize a :class:`LinearSubproblemSolver` object.
 
         Args:
             cg_kwargs: Dictionary of arguments for CG solver. See
-                :func:`scico.solver.cg` or
-                :func:`jax.scipy.sparse.linalg.cg`, documentation,
+                documentation for :func:`scico.solver.cg` or
+                :func:`jax.scipy.sparse.linalg.cg`,
                 including how to specify a preconditioner.
             cg_function: String indicating which CG implementation to
                 use. One of "jax" or "scico"; default "scico".  If
                 "scico", uses :func:`scico.solver.cg`. If "jax", uses
-                :func:`jax.scipy.sparse.linalg.cg`.  The "jax" option is
+                :func:`jax.scipy.sparse.linalg.cg`. The "jax" option is
                 slower on small-scale problems or problems involving
                 external functions, but can be differentiated through.
                 The "scico" option is faster on small-scale problems, but
                 slower on large-scale problems where the forward
                 operator is written entirely in jax.
         """
-        self.cg_kwargs = cg_kwargs
+
+        default_cg_kwargs = {"tol": 1e-4, "maxiter": 100}
+        if cg_kwargs:
+            default_cg_kwargs.update(cg_kwargs)
+        self.cg_kwargs = default_cg_kwargs
         if cg_function == "scico":
             self.cg = scico_cg
         elif cg_function == "jax":