Add proximal average implementation (#469)

* Work in progress on proximal average

* Resolve jit issues due to Python conditionals

* Add __repr__ method

* Fix Functional __repr__ method

* Improve weight handling

* Add API docs

* Minor edits

* Add test

* Improve functional __repr__ formatting

* Clean up bibtex

* Fix dtype issues

* Add example

* Rename examples

* Edit example docs

* Update submodule

* Edit example docs

* Improve example script

* Add option for excluding indicator function from sum of components in evaluation

* Rename some example scripts

* Update submodule

* Add tests

* Rename some example scripts

docs/source/examples.rst

@@ -53,7 +53,7 @@ Deconvolution
    examples/deconv_microscopy_tv_admm
    examples/deconv_microscopy_allchn_tv_admm
    examples/deconv_ppp_bm3d_admm
-   examples/deconv_ppp_bm3d_pgm
+   examples/deconv_ppp_bm3d_apgm
    examples/deconv_ppp_dncnn_admm
    examples/deconv_ppp_dncnn_padmm
    examples/deconv_ppp_bm4d_admm
@@ -67,11 +67,12 @@ Sparse Coding
 .. toctree::
    :maxdepth: 1
 
-   examples/sparsecode_admm
+   examples/sparsecode_nn_admm
+   examples/sparsecode_nn_apgm
    examples/sparsecode_conv_admm
    examples/sparsecode_conv_md_admm
-   examples/sparsecode_pgm
-   examples/sparsecode_poisson_pgm
+   examples/sparsecode_apgm
+   examples/sparsecode_poisson_apgm
 
 
 Miscellaneous
@@ -84,7 +85,7 @@ Miscellaneous
    examples/superres_ppp_dncnn_admm
    examples/denoise_l1tv_admm
    examples/denoise_tv_admm
-   examples/denoise_tv_pgm
+   examples/denoise_tv_apgm
    examples/denoise_tv_multi
    examples/denoise_cplx_tv_nlpadmm
    examples/denoise_cplx_tv_pdhg
@@ -113,7 +114,7 @@ Plug and Play Priors
    examples/ct_svmbir_ppp_bm3d_admm_prox
    examples/ct_fan_svmbir_ppp_bm3d_admm_prox
    examples/deconv_ppp_bm3d_admm
-   examples/deconv_ppp_bm3d_pgm
+   examples/deconv_ppp_bm3d_apgm
    examples/deconv_ppp_dncnn_admm
    examples/deconv_ppp_dncnn_padmm
    examples/deconv_ppp_bm4d_admm
@@ -142,7 +143,7 @@ Total Variation
    examples/deconv_microscopy_allchn_tv_admm
    examples/denoise_l1tv_admm
    examples/denoise_tv_admm
-   examples/denoise_tv_pgm
+   examples/denoise_tv_apgm
    examples/denoise_tv_multi
    examples/denoise_cplx_tv_nlpadmm
    examples/denoise_cplx_tv_pdhg
@@ -157,11 +158,12 @@ Sparsity
    :maxdepth: 1
 
    examples/diffusercam_tv_admm
-   examples/sparsecode_admm
+   examples/sparsecode_nn_admm
+   examples/sparsecode_nn_apgm
    examples/sparsecode_conv_admm
    examples/sparsecode_conv_md_admm
-   examples/sparsecode_pgm
-   examples/sparsecode_poisson_pgm
+   examples/sparsecode_apgm
+   examples/sparsecode_poisson_apgm
    examples/video_rpca_admm
 
 
@@ -215,7 +217,7 @@ ADMM
    examples/deconv_ppp_dncnn_admm
    examples/deconv_ppp_bm4d_admm
    examples/diffusercam_tv_admm
-   examples/sparsecode_admm
+   examples/sparsecode_nn_admm
    examples/sparsecode_conv_admm
    examples/sparsecode_conv_md_admm
    examples/demosaic_ppp_bm3d_admm
@@ -274,10 +276,11 @@ PGM
 .. toctree::
    :maxdepth: 1
 
-   examples/deconv_ppp_bm3d_pgm
-   examples/sparsecode_pgm
-   examples/sparsecode_poisson_pgm
-   examples/denoise_tv_pgm
+   examples/deconv_ppp_bm3d_apgm
+   examples/sparsecode_apgm
+   examples/sparsecode_nn_apgm
+   examples/sparsecode_poisson_apgm
+   examples/denoise_tv_apgm
 
 
 PCG

docs/source/references.bib

@@ -110,7 +110,7 @@ @InCollection {beck-2010-gradient
   pages =	 {42--88},
   publisher =	 {Cambridge University Press},
   year =	 2010,
-  doi =	         {10.1017/CBO9780511804458.003},
+  doi =		 {10.1017/CBO9780511804458.003},
   url =          {http://www.math.tau.ac.il/~teboulle/papers/gradient_chapter.pdf}
 }
 
@@ -124,7 +124,7 @@ @Software {bradbury-2018-jax
                   {P}ython+{N}um{P}y programs},
   url =		 {http://github.com/google/jax},
   version =	 {0.2.5},
-  year =	 {2018}
+  year =	 2018
 }
 
 @Book {beck-2017-first,
@@ -420,11 +420,11 @@ @Article {kamilov-2023-plugandplay
                   Learned Models in Computational Imaging},
   journal =	 {IEEE Signal Processing Magazine},
   year =	 2023,
-  month =        Jan,
-  volume =       40,
-  number =       1,
-  pages =        {85--97},
-  doi =          {10.1109/MSP.2022.3199595}
+  month =	 Jan,
+  volume =	 40,
+  number =	 1,
+  pages =	 {85--97},
+  doi =		 {10.1109/MSP.2022.3199595}
 }
 
 @Article {liu-2018-first,
@@ -722,6 +722,18 @@ @Article {yang-2012-linearized
   pages =	 {301--329}
 }
 
+@InProceedings {yu-2013-better,
+  author =	 {Yu, Yao-Liang},
+  booktitle =	 {Advances in Neural Information Processing Systems},
+  editor =	 {C.J. Burges and L. Bottou and M. Welling and
+                  Z. Ghahramani and K.Q. Weinberger},
+  title =	 {Better Approximation and Faster Algorithm Using the
+                  Proximal Average},
+  url =          {https://proceedings.neurips.cc/paper_files/paper/2013/file/49182f81e6a13cf5eaa496d51fea6406-Paper.pdf},
+  volume =	 26,
+  year =	 2013
+}
+
 @Article {zhang-2017-dncnn,
   author =	 {Kai Zhang and Wangmeng Zuo and Yunjin Chen and Deyu
                   Meng and Lei Zhang},

examples/scripts/README.rst

@@ -61,7 +61,7 @@ Deconvolution
       Deconvolution Microscopy (All Channels)
    `deconv_ppp_bm3d_admm.py <deconv_ppp_bm3d_admm.py>`_
       PPP (with BM3D) Image Deconvolution (ADMM Solver)
-   `deconv_ppp_bm3d_pgm.py <deconv_ppp_bm3d_pgm.py>`_
+   `deconv_ppp_bm3d_apgm.py <deconv_ppp_bm3d_apgm.py>`_
       PPP (with BM3D) Image Deconvolution (APGM Solver)
    `deconv_ppp_dncnn_admm.py <deconv_ppp_dncnn_admm.py>`_
       PPP (with DnCNN) Image Deconvolution (ADMM Solver)
@@ -78,15 +78,17 @@ Deconvolution
 Sparse Coding
 ^^^^^^^^^^^^^
 
-   `sparsecode_admm.py <sparsecode_admm.py>`_
+   `sparsecode_nn_admm.py <sparsecode_nn_admm.py>`_
       Non-Negative Basis Pursuit DeNoising (ADMM)
+   `sparsecode_nn_apgm.py <sparsecode_nn_apgm.py>`_
+      Non-Negative Basis Pursuit DeNoising (APGM)
    `sparsecode_conv_admm.py <sparsecode_conv_admm.py>`_
       Convolutional Sparse Coding (ADMM)
    `sparsecode_conv_md_admm.py <sparsecode_conv_md_admm.py>`_
       Convolutional Sparse Coding with Mask Decoupling (ADMM)
-   `sparsecode_pgm.py <sparsecode_pgm.py>`_
+   `sparsecode_apgm.py <sparsecode_apgm.py>`_
       Basis Pursuit DeNoising (APGM)
-   `sparsecode_poisson_pgm.py <sparsecode_poisson_pgm.py>`_
+   `sparsecode_poisson_apgm.py <sparsecode_poisson_apgm.py>`_
       Non-negative Poisson Loss Reconstruction (APGM)
 
 
@@ -101,7 +103,7 @@ Miscellaneous
       ℓ1 Total Variation Denoising
    `denoise_tv_admm.py <denoise_tv_admm.py>`_
       Total Variation Denoising (ADMM)
-   `denoise_tv_pgm.py <denoise_tv_pgm.py>`_
+   `denoise_tv_apgm.py <denoise_tv_apgm.py>`_
       Total Variation Denoising with Constraint (APGM)
    `denoise_tv_multi.py <denoise_tv_multi.py>`_
       Comparison of Optimization Algorithms for Total Variation Denoising
@@ -139,7 +141,7 @@ Plug and Play Priors
       PPP (with BM3D) Fan-Beam CT Reconstruction
    `deconv_ppp_bm3d_admm.py <deconv_ppp_bm3d_admm.py>`_
       PPP (with BM3D) Image Deconvolution (ADMM Solver)
-   `deconv_ppp_bm3d_pgm.py <deconv_ppp_bm3d_pgm.py>`_
+   `deconv_ppp_bm3d_apgm.py <deconv_ppp_bm3d_apgm.py>`_
       PPP (with BM3D) Image Deconvolution (APGM Solver)
    `deconv_ppp_dncnn_admm.py <deconv_ppp_dncnn_admm.py>`_
       PPP (with DnCNN) Image Deconvolution (ADMM Solver)
@@ -186,7 +188,7 @@ Total Variation
       ℓ1 Total Variation Denoising
    `denoise_tv_admm.py <denoise_tv_admm.py>`_
       Total Variation Denoising (ADMM)
-   `denoise_tv_pgm.py <denoise_tv_pgm.py>`_
+   `denoise_tv_apgm.py <denoise_tv_apgm.py>`_
       Total Variation Denoising with Constraint (APGM)
    `denoise_tv_multi.py <denoise_tv_multi.py>`_
       Comparison of Optimization Algorithms for Total Variation Denoising
@@ -204,15 +206,17 @@ Sparsity
 
    `diffusercam_tv_admm.py <diffusercam_tv_admm.py>`_
       TV-Regularized 3D DiffuserCam Reconstruction
-   `sparsecode_admm.py <sparsecode_admm.py>`_
+   `sparsecode_nn_admm.py <sparsecode_nn_admm.py>`_
       Non-Negative Basis Pursuit DeNoising (ADMM)
+   `sparsecode_nn_apgm.py <sparsecode_nn_apgm.py>`_
+      Non-Negative Basis Pursuit DeNoising (APGM)
    `sparsecode_conv_admm.py <sparsecode_conv_admm.py>`_
       Convolutional Sparse Coding (ADMM)
    `sparsecode_conv_md_admm.py <sparsecode_conv_md_admm.py>`_
       Convolutional Sparse Coding with Mask Decoupling (ADMM)
-   `sparsecode_pgm.py <sparsecode_pgm.py>`_
+   `sparsecode_apgm.py <sparsecode_apgm.py>`_
       Basis Pursuit DeNoising (APGM)
-   `sparsecode_poisson_pgm.py <sparsecode_poisson_pgm.py>`_
+   `sparsecode_poisson_apgm.py <sparsecode_poisson_apgm.py>`_
       Non-negative Poisson Loss Reconstruction (APGM)
    `video_rpca_admm.py <video_rpca_admm.py>`_
       Video Decomposition via Robust PCA
@@ -289,7 +293,7 @@ ADMM
       PPP (with BM4D) Volume Deconvolution
    `diffusercam_tv_admm.py <diffusercam_tv_admm.py>`_
       TV-Regularized 3D DiffuserCam Reconstruction
-   `sparsecode_admm.py <sparsecode_admm.py>`_
+   `sparsecode_nn_admm.py <sparsecode_nn_admm.py>`_
       Non-Negative Basis Pursuit DeNoising (ADMM)
    `sparsecode_conv_admm.py <sparsecode_conv_admm.py>`_
       Convolutional Sparse Coding (ADMM)
@@ -352,13 +356,15 @@ PDHG
 PGM
 ^^^
 
-   `deconv_ppp_bm3d_pgm.py <deconv_ppp_bm3d_pgm.py>`_
+   `deconv_ppp_bm3d_apgm.py <deconv_ppp_bm3d_apgm.py>`_
       PPP (with BM3D) Image Deconvolution (APGM Solver)
-   `sparsecode_pgm.py <sparsecode_pgm.py>`_
+   `sparsecode_apgm.py <sparsecode_apgm.py>`_
       Basis Pursuit DeNoising (APGM)
-   `sparsecode_poisson_pgm.py <sparsecode_poisson_pgm.py>`_
+   `sparsecode_nn_apgm.py <sparsecode_nn_apgm.py>`_
+      Non-Negative Basis Pursuit DeNoising (APGM)
+   `sparsecode_poisson_apgm.py <sparsecode_poisson_apgm.py>`_
       Non-negative Poisson Loss Reconstruction (APGM)
-   `denoise_tv_pgm.py <denoise_tv_pgm.py>`_
+   `denoise_tv_apgm.py <denoise_tv_apgm.py>`_
       Total Variation Denoising with Constraint (APGM)
 
 

examples/scripts/index.rst

@@ -37,7 +37,7 @@ Deconvolution
    - deconv_microscopy_tv_admm.py
    - deconv_microscopy_allchn_tv_admm.py
    - deconv_ppp_bm3d_admm.py
-   - deconv_ppp_bm3d_pgm.py
+   - deconv_ppp_bm3d_apgm.py
    - deconv_ppp_dncnn_admm.py
    - deconv_ppp_dncnn_padmm.py
    - deconv_ppp_bm4d_admm.py
@@ -48,11 +48,12 @@ Deconvolution
 Sparse Coding
 ^^^^^^^^^^^^^
 
-   - sparsecode_admm.py
+   - sparsecode_nn_admm.py
+   - sparsecode_nn_apgm.py
    - sparsecode_conv_admm.py
    - sparsecode_conv_md_admm.py
-   - sparsecode_pgm.py
-   - sparsecode_poisson_pgm.py
+   - sparsecode_apgm.py
+   - sparsecode_poisson_apgm.py
 
 
 Miscellaneous
@@ -62,7 +63,7 @@ Miscellaneous
    - superres_ppp_dncnn_admm.py
    - denoise_l1tv_admm.py
    - denoise_tv_admm.py
-   - denoise_tv_pgm.py
+   - denoise_tv_apgm.py
    - denoise_tv_multi.py
    - denoise_cplx_tv_nlpadmm.py
    - denoise_cplx_tv_pdhg.py
@@ -85,7 +86,7 @@ Plug and Play Priors
    - ct_svmbir_ppp_bm3d_admm_prox.py
    - ct_fan_svmbir_ppp_bm3d_admm_prox.py
    - deconv_ppp_bm3d_admm.py
-   - deconv_ppp_bm3d_pgm.py
+   - deconv_ppp_bm3d_apgm.py
    - deconv_ppp_dncnn_admm.py
    - deconv_ppp_dncnn_padmm.py
    - deconv_ppp_bm4d_admm.py
@@ -111,7 +112,7 @@ Total Variation
    - deconv_microscopy_allchn_tv_admm.py
    - denoise_l1tv_admm.py
    - denoise_tv_admm.py
-   - denoise_tv_pgm.py
+   - denoise_tv_apgm.py
    - denoise_tv_multi.py
    - denoise_cplx_tv_nlpadmm.py
    - denoise_cplx_tv_pdhg.py
@@ -123,11 +124,12 @@ Sparsity
 ^^^^^^^^
 
    - diffusercam_tv_admm.py
-   - sparsecode_admm.py
+   - sparsecode_nn_admm.py
+   - sparsecode_nn_apgm.py
    - sparsecode_conv_admm.py
    - sparsecode_conv_md_admm.py
-   - sparsecode_pgm.py
-   - sparsecode_poisson_pgm.py
+   - sparsecode_apgm.py
+   - sparsecode_poisson_apgm.py
    - video_rpca_admm.py
 
 
@@ -172,7 +174,7 @@ ADMM
    - deconv_ppp_dncnn_admm.py
    - deconv_ppp_bm4d_admm.py
    - diffusercam_tv_admm.py
-   - sparsecode_admm.py
+   - sparsecode_nn_admm.py
    - sparsecode_conv_admm.py
    - sparsecode_conv_md_admm.py
    - demosaic_ppp_bm3d_admm.py
@@ -216,10 +218,11 @@ PDHG
 PGM
 ^^^
 
-   - deconv_ppp_bm3d_pgm.py
-   - sparsecode_pgm.py
-   - sparsecode_poisson_pgm.py
-   - denoise_tv_pgm.py
+   - deconv_ppp_bm3d_apgm.py
+   - sparsecode_apgm.py
+   - sparsecode_nn_apgm.py
+   - sparsecode_poisson_apgm.py
+   - denoise_tv_apgm.py
 
 
 PCG

examples/scripts/sparsecode_pgm.py → examples/scripts/sparsecode_apgm.py

@@ -35,10 +35,10 @@
 σ = 0.5  # Noise level
 
 np.random.seed(12345)
-D = np.random.randn(m, n)
+D = np.random.randn(m, n).astype(np.float32)
 L0 = np.linalg.norm(D, 2) ** 2
 
-x_gt = np.zeros(n)  # true signal
+x_gt = np.zeros(n, dtype=np.float32)  # true signal
 idx = np.random.permutation(list(range(0, n - 1)))
 x_gt[idx[0:s]] = np.random.randn(s)
 y = D @ x_gt + σ * np.random.randn(m)  # synthetic signal

examples/scripts/sparsecode_admm.py → examples/scripts/sparsecode_nn_admm.py

@@ -17,6 +17,9 @@
 where $D$ the dictionary, $\mathbf{y}$ the signal to be represented,
 $\mathbf{x}$ is the sparse representation, and $I(\mathbf{x} \geq 0)$
 is the non-negative indicator.
+
+In this example the problem is solved via ADMM, while Accelerated PGM is
+used in a [companion example](sparsecode_nn_apgm.rst).
 """
 
 import numpy as np
@@ -36,10 +39,10 @@
 s = 10  # sparsity level
 
 np.random.seed(1)
-D = np.random.randn(m, n)
+D = np.random.randn(m, n).astype(np.float32)
 D = D / np.linalg.norm(D, axis=0, keepdims=True)  # normalize dictionary
 
-xt = np.zeros(n)  # true signal
+xt = np.zeros(n, dtype=np.float32)  # true signal
 idx = np.random.randint(low=0, high=n, size=s)  # support of xt
 xt[idx] = np.random.rand(s)
 y = D @ xt + 5e-2 * np.random.randn(m)  # synthetic signal