CUQIpy 1.2.0.post0.dev90__py3-none-any.whl → 1.2.0.post0.dev245__py3-none-any.whl

{CUQIpy-1.2.0.post0.dev90.dist-info → CUQIpy-1.2.0.post0.dev245.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: CUQIpy
-Version: 1.2.0.post0.dev90
+Version: 1.2.0.post0.dev245
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License: Apache License

{CUQIpy-1.2.0.post0.dev90.dist-info → CUQIpy-1.2.0.post0.dev245.dist-info}/RECORD

@@ -1,6 +1,6 @@
 cuqi/__init__.py,sha256=LsGilhl-hBLEn6Glt8S_l0OJzAA1sKit_rui8h-D-p0,488
 cuqi/_messages.py,sha256=fzEBrZT2kbmfecBBPm7spVu7yHdxGARQB4QzXhJbCJ0,415
-cuqi/_version.py,sha256=fadCQ-al0LVIaJsUncww5HgsVEcSF53R-lWs5uar-ow,509
+cuqi/_version.py,sha256=Td4M9WCq7hYHaBteaGyYBJsI-t7L2iZcgoErbmahT4I,510
 cuqi/config.py,sha256=wcYvz19wkeKW2EKCGIKJiTpWt5kdaxyt4imyRkvtTRA,526
 cuqi/diagnostics.py,sha256=5OrbJeqpynqRXOe5MtOKKhe7EAVdOEpHIqHnlMW9G_c,3029
 cuqi/array/__init__.py,sha256=-EeiaiWGNsE3twRS4dD814BIlfxEsNkTCZUc5gjOXb0,30
@@ -35,26 +35,27 @@ cuqi/distribution/_smoothed_laplace.py,sha256=p-1Y23mYA9omwiHGkEuv3T2mwcPAAoNlCr
 cuqi/distribution/_truncated_normal.py,sha256=sZkLYgnkGOyS_3ZxY7iw6L62t-Jh6shzsweRsRepN2k,4240
 cuqi/distribution/_uniform.py,sha256=KA8yQ6ZS3nQGS4PYJ4hpDg6Eq8EQKQvPsIpYfR8fj2w,1967
 cuqi/experimental/__init__.py,sha256=vhZvyMX6rl8Y0haqCzGLPz6PSUKyu75XMQbeDHqTTrw,83
-cuqi/experimental/mcmc/__init__.py,sha256=1sn0U6Ep0x5zv2602og2DkV3Bs8hNFOiq7C3VcMimVw,4472
+cuqi/experimental/mcmc/__init__.py,sha256=zSqLZmxOqQ-F94C9-gPv7g89TX1XxlrlNm071Eb167I,4487
 cuqi/experimental/mcmc/_conjugate.py,sha256=VNPQkGity0mposcqxrx4UIeXm35EvJvZED4p2stffvA,9924
 cuqi/experimental/mcmc/_conjugate_approx.py,sha256=uEnY2ea9su5ivcNagyRAwpQP2gBY98sXU7N0y5hTADo,3653
 cuqi/experimental/mcmc/_cwmh.py,sha256=50v3uZaWhlVnfrEB5-lB_7pn8QoUVBe-xWxKGKbmNHg,7234
 cuqi/experimental/mcmc/_direct.py,sha256=9pQS_2Qk2-ybt6m8WTfPoKetcxQ00WaTRN85-Z0FrBY,777
 cuqi/experimental/mcmc/_gibbs.py,sha256=evgxf2tLFLlKB3hN0qz9a9NcZQSES8wdacnn3uNWocQ,12005
 cuqi/experimental/mcmc/_hmc.py,sha256=8p4QxZBRpFLzwamH-DWHSdZE0aXX3FqonBzczz_XkDw,19340
-cuqi/experimental/mcmc/_langevin_algorithm.py,sha256=yNO7ABxmkixzcLG-lv57GOTyeTr7HwFs2DrrhuZW9OI,8398
+cuqi/experimental/mcmc/_langevin_algorithm.py,sha256=LtPdC1IAeF_gS3T93FDLFutXy7THw-JqZeywZExpefo,14527
 cuqi/experimental/mcmc/_laplace_approximation.py,sha256=rdiE3cMQFq6FLQcOQwPpuGIxrTAp3aoGPxMDSdeopV0,5688
 cuqi/experimental/mcmc/_mh.py,sha256=MXo0ahXP4KGFkaY4HtvcBE-TMQzsMlTmLKzSvpz7drU,2941
 cuqi/experimental/mcmc/_pcn.py,sha256=wqJBZLuRFSwxihaI53tumAg6AWVuceLMOmXssTetd1A,3374
-cuqi/experimental/mcmc/_rto.py,sha256=OtzgiYCxDoTdXp7y4mkLa2upj74qadesoqHYpr11ZCg,10061
+cuqi/experimental/mcmc/_rto.py,sha256=Ub5rDe_yfkzxqcnimEArXWVb3twuGUJmvxEQNPKQWfU,10061
 cuqi/experimental/mcmc/_sampler.py,sha256=xtoT70T8xe3Ye7yYdIFQD_kivjXlqUImyV3bMt406nk,20106
 cuqi/experimental/mcmc/_utilities.py,sha256=kUzHbhIS3HYZRbneNBK41IogUYX5dS_bJxqEGm7TQBI,525
 cuqi/geometry/__init__.py,sha256=Tz1WGzZBY-QGH3c0GiyKm9XHN8MGGcnU6TUHLZkzB3o,842
 cuqi/geometry/_geometry.py,sha256=SDRZdiN2CIuS591lXxqgFoPWPIpwY-MHk75116QvdYY,46901
-cuqi/implicitprior/__init__.py,sha256=CaDQGYtmeFzN37vf3QUmKhcN9-H5lO66ZbK035k4qUw,246
+cuqi/implicitprior/__init__.py,sha256=6z3lvw-tWDyjZSpB3pYzvijSMK9Zlf1IYqOVTtMD2h4,309
 cuqi/implicitprior/_regularizedGMRF.py,sha256=IR9tKzNMoz-b0RKu6ahVgMx_lDNB3jZHVWFMQm6QqZk,6259
 cuqi/implicitprior/_regularizedGaussian.py,sha256=cQtrgzyJU2pwoK4ORGl1erKLE9VY5NqwZTiqiViDswA,12371
 cuqi/implicitprior/_regularizedUnboundedUniform.py,sha256=H2fTOSqYTlDiLxQ7Ya6wnpCUIkpO4qKrkTOsOPnBBeU,3483
+cuqi/implicitprior/_restorator.py,sha256=ixnH8RGcLpqlaIUdR5Dwjx72sO9f3BeotNFRC7Z7qZo,9198
 cuqi/likelihood/__init__.py,sha256=QXif382iwZ5bT3ZUqmMs_n70JVbbjxbqMrlQYbMn4Zo,1776
 cuqi/likelihood/_likelihood.py,sha256=z3AXAbIrv_DjOYh4jy3iDHemuIFUUJu6wdvJ5e2dgW0,6913
 cuqi/model/__init__.py,sha256=IcN4aZCnyp9o-8TNIoZ8vew99QQgi0EmZvnsIuR6qYI,49
@@ -75,19 +76,19 @@ cuqi/sampler/_langevin_algorithm.py,sha256=o5EyvaR6QGAD7LKwXVRC3WwAP5IYJf5GoMVWl
 cuqi/sampler/_laplace_approximation.py,sha256=u018Z5eqlcq_cIwD9yNOaA15dLQE_vUWaee5Xp8bcjg,6454
 cuqi/sampler/_mh.py,sha256=V5tIdn-KdfWo4J_Nbf-AH6XwKWblWUyc4BeuSikUHsE,7062
 cuqi/sampler/_pcn.py,sha256=F0h9-nUFtkqn-o-1s8BCsmr8V7u6R7ycoCOeeV1uhj0,8601
-cuqi/sampler/_rto.py,sha256=-AtMiYq4fh7pF9zVqfYjYtQbIIEGayrWyRGTj8KecfE,11518
+cuqi/sampler/_rto.py,sha256=eJe7_gN_1NpHHc_okKmFtLcOrvoe6cBoVLdf9ULuB_w,11518
 cuqi/sampler/_sampler.py,sha256=TkZ_WAS-5Q43oICa-Elc2gftsRTBd7PEDUMDZ9tTGmU,5712
 cuqi/samples/__init__.py,sha256=vCs6lVk-pi8RBqa6cIN5wyn6u-K9oEf1Na4k1ZMrYv8,44
 cuqi/samples/_samples.py,sha256=hUc8OnCF9CTCuDTrGHwwzv3wp8mG_6vsJAFvuQ-x0uA,35832
-cuqi/solver/__init__.py,sha256=DGl8IdUnochRXHNDEy_13o_VT0vLFY6FjMmmSH6YUkY,169
-cuqi/solver/_solver.py,sha256=eRmpBkHv_RXFdZTWhYqebH-toNbQcPgEgklNd5zOyOw,22803
+cuqi/solver/__init__.py,sha256=3eoTTgBHe3M6ygrbgUVG3GlqaZVe5lGajNV9rolXZJ8,179
+cuqi/solver/_solver.py,sha256=4LdfxLaU-fUHltZw7Sq-Xohyxd_6RvKy03xxtIMW6Zs,29488
 cuqi/testproblem/__init__.py,sha256=DWTOcyuNHMbhEuuWlY5CkYkNDSAqhvsKmJXBLivyblU,202
 cuqi/testproblem/_testproblem.py,sha256=x769LwwRdJdzIiZkcQUGb_5-vynNTNALXWKato7sS0Q,52540
 cuqi/utilities/__init__.py,sha256=H7xpJe2UinjZftKvE2JuXtTi4DqtkR6uIezStAXwfGg,428
 cuqi/utilities/_get_python_variable_name.py,sha256=QwlBVj2koJRA8s8pWd554p7-ElcI7HUwY32HknaR92E,1827
 cuqi/utilities/_utilities.py,sha256=Jc4knn80vLoA7kgw9FzXwKVFGaNBOXiA9kgvltZU3Ao,11777
-CUQIpy-1.2.0.post0.dev90.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
-CUQIpy-1.2.0.post0.dev90.dist-info/METADATA,sha256=KBSZdCAb8ZYWIzYvHOZ4iqrog8QGiBynjOw0gbo_sis,18495
-CUQIpy-1.2.0.post0.dev90.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-CUQIpy-1.2.0.post0.dev90.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
-CUQIpy-1.2.0.post0.dev90.dist-info/RECORD,,
+CUQIpy-1.2.0.post0.dev245.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
+CUQIpy-1.2.0.post0.dev245.dist-info/METADATA,sha256=ibU2b50SIsnhPjnUdOJpuwftMbE4nU_jRugFzTbhOj4,18496
+CUQIpy-1.2.0.post0.dev245.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+CUQIpy-1.2.0.post0.dev245.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
+CUQIpy-1.2.0.post0.dev245.dist-info/RECORD,,

cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-11-03T22:18:33+0100",
+ "date": "2024-11-08T12:37:05+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "8f8b00804a857370d46fd7bdf26cb9542a6b8f34",
- "version": "1.2.0.post0.dev90"
+ "full-revisionid": "113dd1dc30ade5f182e79d003153bcce9aee1894",
+ "version": "1.2.0.post0.dev245"
 }
 '''  # END VERSION_JSON
 

cuqi/experimental/mcmc/__init__.py

@@ -109,7 +109,7 @@ Main changes for users
 
 
 from ._sampler import Sampler, ProposalBasedSampler
-from ._langevin_algorithm import ULA, MALA
+from ._langevin_algorithm import ULA, MALA, MYULA, PnPULA
 from ._mh import MH
 from ._pcn import PCN
 from ._rto import LinearRTO, RegularizedLinearRTO

cuqi/experimental/mcmc/_langevin_algorithm.py

@@ -1,14 +1,18 @@
 import numpy as np
 import cuqi
 from cuqi.experimental.mcmc import Sampler
+from cuqi.implicitprior import RestorationPrior, MoreauYoshidaPrior
 from cuqi.array import CUQIarray
+from copy import deepcopy
 
 class ULA(Sampler): # Refactor to Proposal-based sampler?
     """Unadjusted Langevin algorithm (ULA) (Roberts and Tweedie, 1996)
 
-    Samples a distribution given its logpdf and gradient (up to a constant) based on
-    Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt,  where L_t is 
-    the Langevin diffusion and W_t is the `dim`-dimensional standard Brownian motion.
+    It approximately samples a distribution given its logpdf gradient based on
+    the Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt, where
+    W_t is the `dim`-dimensional standard Brownian motion.
+    ULA results from the Euler-Maruyama discretization of this Langevin stochastic
+    differential equation (SDE). 
 
     For more details see: Roberts, G. O., & Tweedie, R. L. (1996). Exponential convergence
     of Langevin distributions and their discrete approximations. Bernoulli, 341-363.
@@ -23,9 +27,10 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
     initial_point : ndarray
         Initial parameters. *Optional*
 
-    scale : int
-        The Langevin diffusion discretization time step (In practice, a scale of 1/dim**2 is
-        recommended but not guaranteed to be the optimal choice).
+    scale : float
+        The Langevin diffusion discretization time step (In practice, scale must
+        be smaller than 1/L, where L is the Lipschitz of the gradient of the log
+        target density, logd).
 
     callback : callable, *Optional*
         If set this function will be called after every sample.
@@ -61,26 +66,30 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
     # TODO: update demo once sampler merged
     """
 
-    _STATE_KEYS = Sampler._STATE_KEYS.union({'current_target_logd', 'scale', 'current_target_grad'})
+    _STATE_KEYS = Sampler._STATE_KEYS.union({'scale', 'current_target_grad'})
 
     def __init__(self, target=None, scale=1.0, **kwargs):
 
         super().__init__(target, **kwargs)
-
         self.initial_scale = scale
 
     def _initialize(self):
         self.scale = self.initial_scale
-        self.current_target_logd = self.target.logd(self.current_point)
-        self.current_target_grad = self.target.gradient(self.current_point)
+        self.current_target_grad = self._eval_target_grad(self.current_point)
 
     def validate_target(self):
         try:
-            self.target.gradient(np.ones(self.dim))
+            self._eval_target_grad(np.ones(self.dim))
             pass
         except (NotImplementedError, AttributeError):
             raise ValueError("The target needs to have a gradient method")
 
+    def _eval_target_logd(self, x):
+        return None
+
+    def _eval_target_grad(self, x):
+        return self.target.gradient(x)
+
     def _accept_or_reject(self, x_star, target_eval_star, target_grad_star):
         """
         Accepts the proposed state and updates the sampler's state accordingly, i.e.,
@@ -102,14 +111,11 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
         scalar
             1 (accepted)
         """
-        acc = 0
-        if (not np.isnan(target_eval_star)) and \
-           (not np.isinf(target_eval_star)):
-            self.current_point = x_star
-            self.current_target_logd = target_eval_star
-            self.current_target_grad = target_grad_star
-            acc = 1
-            
+
+        self.current_point = x_star
+        self.current_target_grad = target_grad_star
+        acc = 1
+
         return acc
 
     def step(self):
@@ -118,7 +124,8 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
         x_star = self.current_point + 0.5*self.scale*self.current_target_grad + xi
 
         # evaluate target
-        target_eval_star, target_grad_star = self.target.logd(x_star), self.target.gradient(x_star)
+        target_eval_star = self._eval_target_logd(x_star)
+        target_grad_star = self._eval_target_grad(x_star)
 
         # accept or reject proposal
         acc = self._accept_or_reject(x_star, target_eval_star, target_grad_star)
@@ -133,9 +140,11 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
     """  Metropolis-adjusted Langevin algorithm (MALA) (Roberts and Tweedie, 1996)
 
     Samples a distribution given its logd and gradient (up to a constant) based on
-    Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt,  where L_t is 
-    the Langevin diffusion and W_t is the `dim`-dimensional standard Brownian motion. 
-    The sample is then accepted or rejected according to Metropolis–Hastings algorithm.
+    Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt,
+    W_t is the `dim`-dimensional standard Brownian motion.
+    A sample is firstly proposed by ULA and is then accepted or rejected according
+    to a Metropolis–Hastings step.
+    This accept-reject step allows us to remove the asymptotic bias of ULA.
 
     For more details see: Roberts, G. O., & Tweedie, R. L. (1996). Exponential convergence
     of Langevin distributions and their discrete approximations. Bernoulli, 341-363.
@@ -150,8 +159,10 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
     initial_point : ndarray
         Initial parameters. *Optional*
 
-    scale : int
-        The Langevin diffusion discretization time step.
+    scale : float
+        The Langevin diffusion discretization time step (In practice, scale must
+        be smaller than 1/L, where L is the Lipschitz of the gradient of the log
+        target density, logd).
 
     callback : callable, *Optional*
         If set this function will be called after every sample.
@@ -187,9 +198,20 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
     # TODO: update demo once sampler merged
     """
 
+    _STATE_KEYS = ULA._STATE_KEYS.union({'current_target_logd'})
+
+    def _initialize(self):
+        super()._initialize()
+        self.current_target_logd = self.target.logd(self.current_point)
+
+    def _eval_target_logd(self, x):
+        return self.target.logd(x)
+
     def _accept_or_reject(self, x_star, target_eval_star, target_grad_star):
         """
-        Accepts the proposed state according to a Metropolis step and updates the sampler's state accordingly, i.e., current_point, current_target_eval, and current_target_grad_eval.
+        Accepts the proposed state according to a Metropolis step and updates
+        the sampler's state accordingly, i.e., current_point, current_target_eval,
+        and current_target_grad_eval.
 
         Parameters
         ----------
@@ -231,3 +253,137 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
         mu = theta_k + ((self.scale)/2)*g_logpi_k
         misfit = theta_star - mu
         return -0.5*((1/(self.scale))*(misfit.T @ misfit))
+
+
+class MYULA(ULA):
+    """Moreau-Yoshida Unadjusted Langevin algorithm (MYUULA) (Durmus et al., 2018)
+
+    Samples a smoothed target distribution given its smoothed logpdf gradient.
+    It is based on the Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt, 
+    where W_t is a `dim`-dimensional standard Brownian motion.
+    It targets a differentiable density (partially) smoothed by the Moreau-Yoshida
+    envelope. The smoothed target density can be made arbitrarily closed to the
+    true unsmoothed target density.
+
+    For more details see: Durmus, Alain, Eric Moulines, and Marcelo Pereyra.
+    "Efficient Bayesian
+    computation by proximal Markov chain Monte Carlo: when Langevin meets Moreau."
+    SIAM Journal on Imaging Sciences 11.1 (2018): 473-506.
+
+    Parameters
+    ----------
+
+    target : `cuqi.distribution.Distribution`
+        The target distribution to sample from. The target distribution results from
+        a differentiable likelihood and prior of type RestorationPrior.
+    
+    initial_point : ndarray
+        Initial parameters. *Optional*
+
+    scale : float
+        The Langevin diffusion discretization time step (In practice, scale must
+        be smaller than 1/L, where L is the Lipschitz of the gradient of the log
+        target density, logd).
+        
+    smoothing_strength : float
+        This parameter controls the smoothing strength of MYULA.
+
+    callback : callable, *Optional*
+        If set this function will be called after every sample.
+        The signature of the callback function is `callback(sample, sample_index)`,
+        where `sample` is the current sample and `sample_index` is the index of
+        the sample.
+        An example is shown in demos/demo31_callback.py.
+
+    A Deblur example can be found in demos/howtos/myula.py
+    # TODO: update demo once sampler merged
+    """
+    def __init__(self, target=None, scale=1.0, smoothing_strength=0.1, **kwargs):
+        self.smoothing_strength = smoothing_strength
+        super().__init__(target=target, scale=scale, **kwargs)
+
+    @Sampler.target.setter
+    def target(self, value):
+        """ Set the target density. Runs validation of the target. """
+        self._target = value
+
+        if self._target is not None:
+            # Create a smoothed target
+            self._smoothed_target = self._create_smoothed_target(value)
+
+            # Validate the target
+            self.validate_target()
+
+    def _create_smoothed_target(self, value):
+        """ Create a smoothed target using a Moreau-Yoshida envelope. """
+        copied_value = deepcopy(value)
+        if isinstance(copied_value.prior, RestorationPrior):
+            copied_value.prior = MoreauYoshidaPrior(
+                copied_value.prior,
+                self.smoothing_strength)
+        return copied_value
+
+    def validate_target(self):
+        # Call ULA target validation
+        super().validate_target()
+
+        # Additional validation for MYULA target
+        if isinstance(self.target.prior, MoreauYoshidaPrior):
+            raise ValueError(("The prior is already smoothed, apply"
+                              " ULA when using a MoreauYoshidaPrior."))
+        if not hasattr(self.target.prior, "restore"):
+            raise NotImplementedError(
+                ("Using MYULA with a prior that does not have a restore method"
+                " is not supported.")
+            )
+
+    def _eval_target_grad(self, x):
+        return self._smoothed_target.gradient(x)
+
+class PnPULA(MYULA):
+    """Plug-and-Play Unadjusted Langevin algorithm (PnP-ULA)
+    (Laumont et al., 2022)
+
+    Samples a smoothed target distribution given its smoothed logpdf gradient based on
+    Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt, where W_t is
+    a `dim`-dimensional standard Brownian motion.
+    It targets a differentiable density (partially) smoothed by a convolution
+    with Gaussian kernel with zero mean and smoothing_strength variance. The
+    smoothed target density can be made arbitrarily closed to the
+    true unsmoothed target density. 
+
+    For more details see: Laumont, R., Bortoli, V. D., Almansa, A., Delon, J.,
+    Durmus, A., & Pereyra, M. (2022). Bayesian imaging using plug & play priors:
+    when Langevin meets Tweedie. SIAM Journal on Imaging Sciences, 15(2), 701-737.
+
+    Parameters
+    ----------
+
+    target : `cuqi.distribution.Distribution`
+        The target distribution to sample. The target distribution result from
+        a differentiable likelihood and prior of type RestorationPrior.
+    
+    initial_point : ndarray
+        Initial parameters. *Optional*
+
+    scale : float
+        The Langevin diffusion discretization time step (In practice, a scale of
+        1/L, where L is the Lipschitz of the gradient of the log target density
+        is recommended but not guaranteed to be the optimal choice).
+        
+    smoothing_strength : float
+        This parameter controls the smoothing strength of PnP-ULA.
+
+
+    callback : callable, *Optional*
+        If set this function will be called after every sample.
+        The signature of the callback function is `callback(sample, sample_index)`,
+        where `sample` is the current sample and `sample_index` is the index of
+        the sample.
+        An example is shown in demos/demo31_callback.py.
+
+    # TODO: update demo once sampler merged
+    """
+    def __init__ (self, target=None, scale=1.0, smoothing_strength=0.1, **kwargs):
+        super().__init__(target=target, scale=scale, 
+                         smoothing_strength=smoothing_strength, **kwargs)

cuqi/experimental/mcmc/_rto.py

@@ -235,8 +235,8 @@ class RegularizedLinearRTO(LinearRTO):
 
     def step(self):
         y = self.b_tild + np.random.randn(len(self.b_tild))
-        sim = FISTA(self.M, y, self.current_point, self.proximal,
-                    maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+        sim = FISTA(self.M, y, self.proximal,
+                    self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
         self.current_point, _ = sim.solve()
         acc = 1
         return acc

cuqi/implicitprior/__init__.py

@@ -1,3 +1,5 @@
 from ._regularizedGaussian import RegularizedGaussian, ConstrainedGaussian, NonnegativeGaussian
 from ._regularizedGMRF import RegularizedGMRF, ConstrainedGMRF, NonnegativeGMRF
 from ._regularizedUnboundedUniform import RegularizedUnboundedUniform
+from ._restorator import RestorationPrior, MoreauYoshidaPrior
+

cuqi/implicitprior/_restorator.py

@@ -0,0 +1,223 @@
+from abc import ABC, abstractmethod
+from cuqi.distribution import Distribution
+import numpy as np
+    
+class RestorationPrior(Distribution):
+    """    
+    This class defines an implicit distribution associated with a restoration operator
+    (eg denoiser). They are several works relating restorations operators with
+    priors, see 
+        -Laumont et al. https://arxiv.org/abs/2103.04715
+        -Hu et al. https://openreview.net/pdf?id=x7d1qXEn1e
+    We cannot sample from this distribution, neither compute its logpdf except in
+    some cases. It allows us to apply algorithms such as MYULA and PnPULA.
+    
+    Parameters
+    ---------- 
+    restorator : callable f(x, restoration_strength)
+        Function f that accepts input x to be restored and returns the
+        restored version of x and information about the restoration operation.
+            
+    restorator_kwargs : dictionary
+        Dictionary containing information about the restorator.
+        It contains keyword argument parameters that will be passed to the
+        restorator f. An example could be algorithm parameters such as the number
+        of iterations or the stopping criteria.
+    
+    potential : callable function, optional
+        The potential corresponds to the negative logpdf when it is accessible.
+        This function is a mapping from the parameter domain to the real set. 
+        It can be provided if the user knows how to relate it to the restorator.
+        Ex: restorator is the proximal operator of the total variation (TV), then
+        potential is the TV function. 
+    """
+    def __init__(self, restorator, restorator_kwargs
+                =None, potential=None, **kwargs):
+        if restorator_kwargs is None:
+            restorator_kwargs = {}
+        self.restorator = restorator 
+        self.restorator_kwargs = restorator_kwargs
+        self.potential = potential
+        super().__init__(**kwargs)
+
+    def restore(self, x, restoration_strength):
+        """This function allows us to restore the input x and returns the
+        restored version of x.
+        
+        Parameters
+        ---------- 
+        x : ndarray
+            parameter we want to restore.
+        
+        restoration_strength: positive float
+            Strength of the restoration operation. In the case where the
+            restorator is a denoiser, this parameter might correspond to the
+            noise level.
+        """
+        solution, info = self.restorator(x, restoration_strength=restoration_strength,
+                                         **self.restorator_kwargs)
+        self.info = info
+        return solution
+    
+    def logpdf(self, x):
+        """The logpdf function. It returns nan because we don't know the 
+        logpdf of the implicit prior."""
+        if self.potential is None:
+            return np.nan
+        else:
+            return -self.potential(x)
+    
+    def _sample(self, N, rng=None):
+        raise NotImplementedError("The sample method is not implemented for the"
+                                  + "RestorationPrior class.")
+
+    @property
+    def _mutable_vars(self):
+        """ Returns the mutable variables of the distribution. """
+        # Currently mutable variables are not supported for user-defined
+        # distributions.
+        return []
+
+    def get_conditioning_variables(self):
+        """ Returns the conditioning variables of the distribution. """
+        # Currently conditioning variables are not supported for user-defined
+        # distributions.
+        return []
+
+
+class MoreauYoshidaPrior(Distribution):
+    """    
+    This class defines (implicit) smoothed priors for which we can apply 
+    gradient-based algorithms. The smoothing is performed using
+    the Moreau-Yoshida envelope of the target prior potential.
+    
+    In the following we give a detailed explanation of the
+    Moreau-Yoshida smoothing.
+    
+    We consider a density such that - \log\pi(x) = -g(x) with g convex, lsc,
+    proper but not differentiable. Consequently, we cannot apply any
+    algorithm requiring the gradient of g.
+    Idea:
+    We consider the Moreau envelope of g defined as
+    
+    g_{smoothing_strength} (x) = inf_z 0.5*\| x-z \|_2^2/smoothing_strength + g(z).
+    
+    g_{smoothing_strength} has some nice properties
+        - g_{smoothing_strength}(x)-->g(x) as smoothing_strength-->0 for all x
+        - \nabla g_{smoothing_strength} is 1/smoothing_strength-Lipschitz
+        - \nabla g_{smoothing_strength}(x) = (x - prox_g^{smoothing_strength}(x))/smoothing_strength for all x with 
+        
+        prox_g^{smoothing_strength}(x) = argmin_z 0.5*\| x-z \|_2^2/smoothing_strength + g(z) .
+
+    Consequently, we can apply any gradient-based algorithm with
+    g_{smoothing_strength} in lieu of g. These algorithms do not require the
+    full knowledge of g_{smoothing_strength} but only its gradient. The gradient
+    of g_{smoothing_strength} is fully determined by prox_g^{smoothing_strength}
+    and smoothing_strength.
+    It is important as, although there exists an explicit formula for
+    g_{smoothing_strength}, it is rarely used in practice, as it would require
+    us to solve an optimization problem each time we want to 
+    estimate g_{smoothing_strength}. Furthermore, there exist cases where we dont't
+    the regularization g with which the mapping prox_g^{smoothing_strength} is
+    associated.
+
+    Remark (Proximal operators are denoisers):
+        We consider the denoising inverse problem x = u + n, with
+        n \sim \mathcal{N}(0, smoothing_strength I).
+        A mapping solving a denoising inverse problem is called denoiser. It takes
+        the noisy observation x as an input and returns a less noisy version of x
+        which is an estimate of u.
+        We assume a prior density \pi(u) \propto exp(- g(u)).
+        Then the MAP estimate is given by 
+            x_MAP = \argmin_z 0.5 \| x - z \|_2^2/smoothing_strength + g(z) = prox_g^smoothing_strength(x)
+        Then proximal operators are denoisers. 
+    
+    Remark (Denoisers are not necessarily proximal operators): Data-driven
+    denoisers are not necessarily proximal operators
+    (see https://arxiv.org/pdf/2201.13256)
+    
+    Parameters
+    ---------- 
+    prior : RestorationPrior
+        Prior of the RestorationPrior type. In order to stay within the MYULA
+        framework the restorator of RestorationPrior must be a proximal operator.
+        
+    smoothing_strength : float
+        Smoothing strength of the Moreau-Yoshida envelope of the prior potential.
+    """
+
+    def __init__(self, prior:RestorationPrior, smoothing_strength=0.1, 
+                 **kwargs):
+        self.prior = prior
+        self.smoothing_strength = smoothing_strength
+
+        # if kwargs does not contain the geometry,
+        # we set it to the geometry of the prior, if it exists
+        if "geometry" in kwargs:
+            raise ValueError(
+                "The geometry parameter is not supported for the"
+                + "MoreauYoshidaPrior class. The geometry is"
+                + "automatically set to the geometry of the prior.")
+        try:
+            geometry = prior.geometry
+        except:
+            geometry = None
+
+        super().__init__(geometry=geometry, **kwargs)
+
+    @property
+    def geometry(self):
+        return self.prior.geometry
+
+    @geometry.setter
+    def geometry(self, value):
+        self.prior.geometry = value
+
+    @property
+    def smoothing_strength(self):
+        """ smoothing_strength of the distribution"""
+        return self._smoothing_strength
+
+    @smoothing_strength.setter
+    def smoothing_strength(self, value):
+        self._smoothing_strength = value
+        
+    @property
+    def prior(self):
+        """Getter for the MoreauYoshida prior."""
+        return self._prior
+
+    @prior.setter
+    def prior(self, value):
+        self._prior = value
+    
+    def gradient(self, x):
+        """This is the gradient of the regularizer ie gradient of the negative
+        logpdf of the implicit prior."""
+        return -(x - self.prior.restore(x, self.smoothing_strength))/self.smoothing_strength
+        
+    def logpdf(self, x):
+        """The logpdf function. It returns nan because we don't know the 
+        logpdf of the implicit prior."""
+        if self.prior.potential == None:
+            return np.nan
+        else:
+            return -(self.prior.potential(self.prior.restore(x, self.smoothing_strength))*self.smoothing_strength +
+                     0.5*((x-self.prior.restore(x, self.smoothing_strength))**2).sum())
+    
+    def _sample(self, N, rng=None):
+        raise NotImplementedError("The sample method is not implemented for the"
+                                  + f"{self.__class__.__name__} class.")
+
+    @property
+    def _mutable_vars(self):
+        """ Returns the mutable variables of the distribution. """
+        # Currently mutable variables are not supported for user-defined
+        # distributions.
+        return []
+
+    def get_conditioning_variables(self):
+        """ Returns the conditioning variables of the distribution. """
+        # Currently conditioning variables are not supported for user-defined
+        # distributions.
+        return []

cuqi/sampler/_rto.py

@@ -267,8 +267,8 @@ class RegularizedLinearRTO(LinearRTO):
         samples[:, 0] = self.x0
         for s in range(Ns-1):
             y = self.b_tild + np.random.randn(len(self.b_tild))
-            sim = FISTA(self.M, y, samples[:, s], self.proximal,
-                        maxit = self.maxit, stepsize = _stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+            sim = FISTA(self.M, y, self.proximal,
+                        samples[:, s], maxit = self.maxit, stepsize = _stepsize, abstol = self.abstol, adaptive = self.adaptive)         
             samples[:, s+1], _ = sim.solve()
             
             self._print_progress(s+2,Ns) #s+2 is the sample number, s+1 is index assuming x0 is the first sample

cuqi/solver/__init__.py

@@ -7,6 +7,7 @@ from ._solver import (
     LM,
     PDHG,
     FISTA,
+    ADMM,
     ProjectNonnegative,
     ProjectBox,
     ProximalL1

cuqi/solver/_solver.py

@@ -584,8 +584,8 @@ class FISTA(object):
     ----------
     A : ndarray or callable f(x,*args).
     b : ndarray.
-    x0 : ndarray. Initial guess.
     proximal : callable f(x, gamma) for proximal mapping.
+    x0 : ndarray. Initial guess.
     maxit : The maximum number of iterations.
     stepsize : The stepsize of the gradient step.
     abstol : The numerical tolerance for convergence checks.
@@ -606,11 +606,11 @@ class FISTA(object):
         b = rng.standard_normal(m)
         stepsize = 0.99/(sp.linalg.interpolative.estimate_spectral_norm(A)**2)
         x0 = np.zeros(n)
-        fista = FISTA(A, b, x0, proximal = ProximalL1, stepsize = stepsize, maxit = 100, abstol=1e-12, adaptive = True)
+        fista = FISTA(A, b, proximal = ProximalL1, x0, stepsize = stepsize, maxit = 100, abstol=1e-12, adaptive = True)
         sol, _ = fista.solve()
 
     """  
-    def __init__(self, A, b, x0, proximal, maxit=100, stepsize=1e0, abstol=1e-14, adaptive = True):
+    def __init__(self, A, b, proximal, x0, maxit=100, stepsize=1e0, abstol=1e-14, adaptive = True):
         
         self.A = A
         self.b = b
@@ -650,8 +650,157 @@ class FISTA(object):
                 x_new = x_new + ((k-1)/(k+2))*(x_new - x_old)
               
             x = x_new.copy()
+
+class ADMM(object):
+    """Alternating Direction Method of Multipliers for solving regularized linear least squares problems of the form:
+    Minimize ||Ax-b||^2 + sum_i f_i(L_i x),
+    where the sum ranges from 1 to an arbitrary n. See definition of the parameter `penalty_terms` below for more details about f_i and L_i
+
+    Reference:
+    [1] Boyd et al. "Distributed optimization and statistical learning via the alternating direction method of multipliers."Foundations and Trends® in Machine learning, 2011.
+
+    
+    Parameters
+    ----------
+    A : ndarray or callable
+        Represents a matrix or a function that performs matrix-vector multiplications. 
+        When A is a callable, it accepts arguments (x, flag) where:
+        - flag=1 indicates multiplication of A with vector x, that is A @ x.
+        - flag=2 indicates multiplication of the transpose of A with vector x, that is  A.T @ x.
+    b : ndarray.
+    penalty_terms : List of tuples (callable proximal operator of f_i, linear operator L_i)
+        Each callable proximal operator f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
+    x0 : ndarray. Initial guess.
+    penalty_parameter : Trade-off between linear least squares and regularization term in the solver iterates. Denoted as "rho" in [1].
+    maxit : The maximum number of iterations.
+    adaptive : Whether to adaptively update the penalty_parameter each iteration such that the primal and dual residual norms are of the same order of magnitude. Based on [1], Subsection 3.4.1
+    
+    Example
+    -----------
+    .. code-block:: python
     
+        from cuqi.solver import ADMM, ProximalL1, ProjectNonnegative
+        import numpy as np
+
+        rng = np.random.default_rng()
+
+        m, n, k = 10, 5, 4
+        A = rng.standard_normal((m, n))
+        b = rng.standard_normal(m)
+        L = rng.standard_normal((k, n))
+
+        x0 = np.zeros(n)
+        admm = ADMM(A, b, x0, penalty_terms = [(ProximalL1, L), (lambda z, _ : ProjectNonnegative(z), np.eye(n))], tradeoff = 10)
+        sol, _ = admm.solve()
+
+    """  
+
+    def __init__(self, A, b, penalty_terms, x0, penalty_parameter = 10, maxit = 100, inner_max_it = 10, adaptive = True):
+
+        self.A = A
+        self.b = b
+        self.x_cur = x0
+
+        dual_len = [penalty[1].shape[0] for penalty in penalty_terms]
+        self.z_cur = [np.zeros(l) for l in dual_len]
+        self.u_cur = [np.zeros(l) for l in dual_len]
+        self.n = penalty_terms[0][1].shape[1]
+        
+        self.rho = penalty_parameter
+        self.maxit = maxit
+        self.inner_max_it = inner_max_it
+        self.adaptive = adaptive
+
+        self.penalty_terms = penalty_terms
+       
+        self.p = len(self.penalty_terms)
+        self._big_matrix = None
+        self._big_vector = None
+
+    def solve(self):
+        """
+        Solves the regularized linear least squares problem using ADMM in scaled form. Based on [1], Subsection 3.1.1
+        """
+        z_new = self.p*[0]
+        u_new = self.p*[0]
+
+        # Iterating
+        for i in range(self.maxit):
+            self._iteration_pre_processing()
+
+            # Main update (Least Squares)
+            solver = CGLS(self._big_matrix, self._big_vector, self.x_cur, self.inner_max_it)
+            x_new, _ = solver.solve()
+        
+            # Regularization update
+            for j, penalty in enumerate(self.penalty_terms):
+                z_new[j] = penalty[0](penalty[1]@x_new + self.u_cur[j], 1.0/self.rho)
+                
+            res_primal = 0.0
+            # Dual update
+            for j, penalty in enumerate(self.penalty_terms):
+                r_partial = penalty[1]@x_new - z_new[j]
+                res_primal += LA.norm(r_partial)**2
+
+                u_new[j] = self.u_cur[j] + r_partial
+            
+            res_dual = 0.0
+            for j, penalty in enumerate(self.penalty_terms):
+                res_dual += LA.norm(penalty[1].T@(z_new[j] - self.z_cur[j]))**2
+
+            # Adaptive approach based on [1], Subsection 3.4.1
+            if self.adaptive:
+                if res_dual > 1e2*res_primal:
+                    self.rho *= 0.5 # More regularization
+                elif res_primal > 1e2*res_dual:
+                    self.rho *= 2.0 # More data fidelity
+
+            self.x_cur, self.z_cur, self.u_cur = x_new, z_new.copy(), u_new
+            
+        return self.x_cur, i
     
+    def _iteration_pre_processing(self):
+            """ Preprocessing
+            Every iteration of ADMM requires solving a linear least squares system of the form
+                minimize 1/(rho) \|Ax-b\|_2^2 + sum_{i=1}^{p} \|penalty[1]x - (y - u)\|_2^2
+            To solve this, all linear least squares terms are combined into a single big term
+            with matrix big_matrix and data big_vector.
+
+            The matrix only needs to be updated when rho changes, i.e., when the adaptive option is used.
+            The data vector needs to be updated every iteration.
+            """
+
+            self._big_vector = np.hstack([np.sqrt(1/self.rho)*self.b] + [self.z_cur[i] - self.u_cur[i] for i in range(self.p)])
+
+            # Check whether matrix needs to be updated
+            if self._big_matrix is not None and not self.adaptive:
+                return
+
+            # Update big_matrix
+            if callable(self.A):
+                def matrix_eval(x, flag):
+                    if flag == 1:
+                        out1 = np.sqrt(1/self.rho)*self.A(x, 1)
+                        out2 = [penalty[1]@x for penalty in self.penalty_terms]
+                        out  = np.hstack([out1] + out2)
+                    elif flag == 2:
+                        idx_start = len(x)
+                        idx_end = len(x)
+                        out1 = np.zeros(self.n)
+                        for _, t in reversed(self.penalty_terms):
+                            idx_start -= t.shape[0]
+                            out1 += t.T@x[idx_start:idx_end]
+                            idx_end = idx_start
+                        out2 = np.sqrt(1/self.rho)*self.A(x[:idx_end], 2)
+                        out  = out1 + out2     
+                    return out
+                self._big_matrix = matrix_eval
+            else:
+                self._big_matrix = np.vstack([np.sqrt(1/self.rho)*self.A] + [penalty[1] for penalty in self.penalty_terms])
+
+
+
+
 def ProjectNonnegative(x):
     """(Euclidean) projection onto the nonnegative orthant.
     
@@ -678,6 +827,22 @@ def ProjectBox(x, lower = None, upper = None):
     
     return np.minimum(np.maximum(x, lower), upper)
 
+def ProjectHalfspace(x, a, b):
+    """(Euclidean) projection onto the halfspace defined {z|<a,z> <= b}.
+    
+    Parameters
+    ----------
+    x : array_like.
+    a : array_like.
+    b : array_like.
+    """  
+
+    ax_b = np.inner(a,x) - b
+    if ax_b <= 0:
+        return x
+    else:
+        return x - (ax_b/np.inner(a,a))*a
+
 def ProximalL1(x, gamma):
     """(Euclidean) proximal operator of the \|x\|_1 norm.
     Also known as the shrinkage or soft thresholding operator.
@@ -687,4 +852,4 @@ def ProximalL1(x, gamma):
     x : array_like.
     gamma : scale parameter.
     """
-    return np.multiply(np.sign(x), np.maximum(np.abs(x)-gamma, 0))
+    return np.multiply(np.sign(x), np.maximum(np.abs(x)-gamma, 0))