Changes to admm and svmbir wrapper to address issue #57

.gitignore

@@ -133,3 +133,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# macos files
+*.DS_Store

docs/source/team.rst

@@ -39,4 +39,5 @@ Contributors
 ============
 
 - `Oleg Korobkin <https://github.com/korobkin>`_ (BlockArray improvements)
-- `Yanpeng Yuan <https://github.com/yanpeng7>`_ (Improvements to ASTRA interface)
+- `Yanpeng Yuan <https://github.com/yanpeng7>`_ (ASTRA interface improvements)
+- `Soumendu Majee <https://github.com/smajee>`_ (SVMBIR interface improvements)

examples/scripts/ct_svmbir_ppp_bm3d_admm_cg.py

@@ -52,7 +52,7 @@
 """
 num_angles = int(N / 2)
 num_channels = N
-angles = snp.linspace(0, snp.pi, num_angles, dtype=snp.float32)
+angles = snp.linspace(0, snp.pi, num_angles, endpoint=False, dtype=snp.float32)
 A = ParallelBeamProjector(x_gt.shape, angles, num_channels)
 sino = A @ x_gt
 
@@ -88,7 +88,7 @@
 """
 y, x0, weights = jax.device_put([y, x_mrf, weights])
 
-ρ = 100  # ADMM penalty parameter
+ρ = 20  # ADMM penalty parameter
 σ = density * 0.2  # denoiser sigma
 
 f = SVMBIRWeightedSquaredL2Loss(y=y, A=A, W=Diagonal(weights), scale=0.5)

examples/scripts/ct_svmbir_ppp_bm3d_admm_prox.py

@@ -52,7 +52,7 @@
 """
 num_angles = int(N / 2)
 num_channels = N
-angles = snp.linspace(0, snp.pi, num_angles, dtype=snp.float32)
+angles = snp.linspace(0, snp.pi, num_angles, endpoint=False, dtype=snp.float32)
 A = ParallelBeamProjector(x_gt.shape, angles, num_channels)
 sino = A @ x_gt
 
@@ -91,7 +91,9 @@
 ρ = 10  # ADMM penalty parameter
 σ = density * 0.26  # denoiser sigma
 
-f = SVMBIRWeightedSquaredL2Loss(y=y, A=A, W=Diagonal(weights), scale=0.5)
+f = SVMBIRWeightedSquaredL2Loss(
+    y=y, A=A, W=Diagonal(weights), scale=0.5, prox_kwargs={"maxiter": 5}
+)
 g0 = σ * ρ * BM3D()
 g1 = NonNegativeIndicator()
 

examples/scripts/denoise_tv_iso_pgm.py

@@ -117,7 +117,7 @@ class IsoProjector(functional.Functional):
     def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         return 0.0
 
-    def prox(self, x: JaxArray, lam: float) -> JaxArray:
+    def prox(self, x: JaxArray, lam: float, **kwargs) -> JaxArray:
         norm_x_ptp = jnp.sqrt(jnp.sum(jnp.abs(x) ** 2, axis=0))
 
         x_out = x / jnp.maximum(jnp.ones(x.shape), norm_x_ptp)
@@ -172,7 +172,7 @@ class AnisoProjector(functional.Functional):
     def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         return 0.0
 
-    def prox(self, x: JaxArray, lam: float) -> JaxArray:
+    def prox(self, x: JaxArray, lam: float, **kwargs) -> JaxArray:
 
         return x / jnp.maximum(jnp.ones(x.shape), jnp.abs(x))
 

misc/conda/make_conda_env.sh

@@ -179,7 +179,7 @@ sort $ALLREQUIRE | uniq | $SED -E 's/(>|<|\|)/\\\1/g' \
     | $SED -E '/^-r.*|^jaxlib.*|^jax.*|^astra-toolbox.*/d' > $FLTREQUIRE
 # Remove requirements that cannot be installed via conda
 for nc in $NOCONDA; do
-    $SED -i "/^$nc\$/d" $FLTREQUIRE
+    $SED -i "/^$nc.*\$/d" $FLTREQUIRE
 done
 # Get list of requirements to be installed via conda
 CONDAREQ=$(cat $FLTREQUIRE | xargs)

requirements.txt

@@ -6,4 +6,4 @@ jax==0.2.19
 jaxlib==0.1.70
 flax
 bm3d
-svmbir
+svmbir>=0.2.6

scico/admm.py

@@ -208,7 +208,7 @@ def compute_rhs(self) -> Union[JaxArray, BlockArray]:
 
         if self.admm.f is not None:
             if isinstance(self.admm.f, WeightedSquaredL2Loss):
-                ATWy = self.admm.f.A.adj(self.admm.f.W.diagonal @ self.admm.f.y)
+                ATWy = self.admm.f.A.adj(self.admm.f.W.diagonal * self.admm.f.y)
                 rhs += 2.0 * self.admm.f.scale * ATWy
             else:
                 ATy = self.admm.f.A.adj(self.admm.f.y)
@@ -231,7 +231,7 @@ def solve(self, x0: Union[JaxArray, BlockArray]) -> Union[JaxArray, BlockArray]:
         """
         x0 = ensure_on_device(x0)
         rhs = self.compute_rhs()
-        x, self.info = self.cg(self.lhs_op, rhs, x0=x0, **self.cg_kwargs)
+        x, self.info = self.cg(self.lhs_op, rhs, x0, **self.cg_kwargs)
         return x
 
 
@@ -604,7 +604,7 @@ def z_and_u_step(self, u_list, z_list):
             zip(self.rho_list, self.g_list, self.C_list, z_list, u_list)
         ):
             Cix = Ci(self.x)
-            zi = gi.prox(Cix + ui, 1 / rhoi)
+            zi = gi.prox(Cix + ui, 1 / rhoi, v0=zi)
             ui = ui + Cix - zi
             z_list[i] = zi
             u_list[i] = ui

scico/functional/_denoiser.py

@@ -54,7 +54,7 @@ def __init__(self, is_rgb: Optional[bool] = False):
 
         super().__init__()
 
-    def prox(self, x: JaxArray, lam: float = 1) -> JaxArray:
+    def prox(self, x: JaxArray, lam: float = 1.0, **kwargs) -> JaxArray:
         r"""Apply BM3D denoiser with noise level ``lam``.
 
         Args:

scico/functional/_flax.py

@@ -38,7 +38,7 @@ def __init__(self, model: Callable[..., nn.Module], variables: PyTree):
         self.variables = variables
         super().__init__()
 
-    def prox(self, x: JaxArray, lam: float = 1) -> JaxArray:
+    def prox(self, x: JaxArray, lam: float = 1.0, **kwargs) -> JaxArray:
         r"""Apply trained flax model.
 
         *Warning*: The ``lam`` parameter is ignored, and has no effect on

scico/functional/_functional.py

@@ -75,7 +75,9 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
                 f"Functional {type(self)} cannot be evaluated; has_eval={self.has_eval}"
             )
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Scaled proximal operator of functional.
 
         Evaluate scaled proximal operator of this functional, with
@@ -90,15 +92,18 @@ def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray
 
         Args:
             x : Point at which to evaluate prox function.
-            lam : Proximal parameter :math:`\lambda`
+            lam : Proximal parameter :math:`\lambda`.
+            kwargs : Additional arguments that may be used by derived
+                classes. These include ``v0``, an initial guess for the minimizer.
+
         """
         if not self.has_prox:
             raise NotImplementedError(
                 f"Functional {type(self)} does not have a prox; has_prox={self.has_prox}"
             )
 
     def conj_prox(
-        self, x: Union[JaxArray, BlockArray], lam: float = 1
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
     ) -> Union[JaxArray, BlockArray]:
         r"""Scaled proximal operator of convex conjugate of functional.
 
@@ -116,9 +121,10 @@ def conj_prox(
 
         Args:
             x : Point at which to evaluate prox function.
-            lam : Proximal parameter :math:`\lambda`
+            lam : Proximal parameter :math:`\lambda`.
+            kwargs : additional keyword args, passed directly to ``self.prox``.
         """
-        return x - lam * self.prox(x / lam, 1.0 / lam)
+        return x - lam * self.prox(x / lam, 1.0 / lam, **kwargs)
 
     def grad(self, x: Union[JaxArray, BlockArray]):
         r"""Evaluates the gradient of this functional at point :math:`\mb{x}`.
@@ -174,7 +180,9 @@ def __init__(self, functional: Functional, scale: float):
     def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         return self.scale * self.functional(x)
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         return self.functional.prox(x, lam * self.scale)
 
 
@@ -216,7 +224,7 @@ def __call__(self, x: BlockArray) -> float:
                 f"Number of blocks in x, {len(x.shape)}, and length of functional_list, {len(self.functional_list)}, do not match"
             )
 
-    def prox(self, x: BlockArray, lam: float = 1) -> BlockArray:
+    def prox(self, x: BlockArray, lam: float = 1.0, **kwargs) -> BlockArray:
         r"""Evaluate proximal operator of the separable functional.
 
         Evaluate proximal operator of the separable functional (see Theorem 6.6 of :cite:`beck-2017-first`).
@@ -252,5 +260,7 @@ class ZeroFunctional(Functional):
     def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         return 0.0
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         return x

scico/functional/_indicator.py

@@ -46,7 +46,9 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         # snp.inf if snp.any(x < 0) else 0.0
         return jax.lax.cond(snp.any(x < 0), lambda x: snp.inf, lambda x: 0.0, None)
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Evaluate proximal operator of indicator over non-negative orthant:
 
         .. math::
@@ -95,7 +97,9 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         # snp.inf if norm(x) > self.radius else 0.0
         return jax.lax.cond(norm(x) > self.radius, lambda x: snp.inf, lambda x: 0.0, None)
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Evaluate proximal operator of indicator over :math:`\ell_2` ball:
 
         .. math::

scico/functional/_norm.py

@@ -36,7 +36,9 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
 
     @staticmethod
     @jit
-    def prox(x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Evaluate proximal operator of :math:`\ell_0` norm
 
 
@@ -71,7 +73,7 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         return snp.abs(x).sum()
 
     @staticmethod
-    def prox(x: Union[JaxArray, BlockArray], lam: float = 1) -> JaxArray:
+    def prox(x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs) -> JaxArray:
         r"""Evaluate proximal operator of :math:`\ell_1` norm
 
         .. math::
@@ -115,7 +117,9 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         # behavior of snp.norm(x) at 0.
         return (snp.abs(x) ** 2).sum()
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Evaluate proximal operator of squared :math:`\ell_2` norm:
 
         .. math::
@@ -143,7 +147,9 @@ class L2Norm(Functional):
     def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         return norm(x)
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Evaluate proximal operator of :math:`\ell_2` norm:
 
         .. math::
@@ -199,7 +205,9 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
         l2 = norm(x, axis=self.l2_axis)
         return snp.abs(l2).sum()
 
-    def prox(self, x: Union[JaxArray, BlockArray], lam: float = 1) -> Union[JaxArray, BlockArray]:
+    def prox(
+        self, x: Union[JaxArray, BlockArray], lam: float = 1.0, **kwargs
+    ) -> Union[JaxArray, BlockArray]:
         r"""Evaluate proximal operator of the :math:`\ell_{2,1}` norm.
 
         In two dimensions,

scico/linop/radon_astra.py

@@ -22,10 +22,7 @@
 try:
     import astra
 except ImportError:
-    raise ImportError(
-        "Could not import astra, please refer to INSTALL.rst "
-        "for instructions on how to install the ASTRA toolbox."
-    )
+    raise ImportError("Could not import astra; please install the ASTRA toolbox.")
 
 
 from jaxlib.xla_extension import GpuDevice
@@ -38,7 +35,11 @@
 
 
 class ParallelBeamProjector(LinearOperator):
-    r"""Parallel beam projector based on ASTRA."""
+    r"""Parallel beam Radon transform based on the ASTRA toolbox.
+
+    Perform tomographic projection of an image at specified angles,
+    using the `ASTRA toolbox <https://github.com/astra-toolbox/astra-toolbox>`_.
+    """
 
     def __init__(
         self,