CUQIpy 1.3.0__py3-none-any.whl → 1.4.0.post0.dev61__py3-none-any.whl

cuqi/sampler/_langevin_algorithm.py

@@ -1,13 +1,18 @@
 import numpy as np
 import cuqi
 from cuqi.sampler import Sampler
+from cuqi.implicitprior import RestorationPrior, MoreauYoshidaPrior
+from cuqi.array import CUQIarray
+from copy import copy
 
-class ULA(Sampler):
+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.
@@ -19,22 +24,17 @@ class ULA(Sampler):
         The target distribution to sample. Must have logd and gradient method. Custom logpdfs 
         and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
     
-    x0 : ndarray
+    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).
 
-    dim : int
-        Dimension of parameter space. Required if target logpdf and gradient are callable 
-        functions. *Optional*.
-
-    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.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
 
     Example
@@ -58,64 +58,91 @@ class ULA(Sampler):
         sampler = cuqi.sampler.ULA(target, scale=1/dim**2)
 
         # Sample
-        samples = sampler.sample(2000)
+        sampler.sample(2000)
 
     A Deblur example can be found in demos/demo27_ULA.py
+    # TODO: update demo once sampler merged
     """
-    def __init__(self, target, scale, x0=None, dim=None, rng=None, **kwargs):
-        super().__init__(target, x0=x0, dim=dim, **kwargs)
-        self.scale = scale
-        self.rng = rng
-
-    def _sample_adapt(self, N, Nb):
-        return self._sample(N, Nb)
-
-    def _sample(self, N, Nb):    
-        # allocation
-        Ns = Nb+N
-        samples = np.empty((self.dim, Ns))
-        target_eval = np.empty(Ns)
-        g_target_eval = np.empty((self.dim, Ns))
-        acc = np.zeros(Ns)
-
-        # initial state
-        samples[:, 0] = self.x0
-        target_eval[0], g_target_eval[:,0] = self.target.logd(self.x0), self.target.gradient(self.x0)
-        acc[0] = 1
-
-        # ULA
-        for s in range(Ns-1):
-            samples[:, s+1], target_eval[s+1], g_target_eval[:,s+1], acc[s+1] = \
-                self.single_update(samples[:, s], target_eval[s], g_target_eval[:,s])            
-            self._print_progress(s+2,Ns) #s+2 is the sample number, s+1 is index assuming x0 is the first sample
-            self._call_callback(samples[:, s+1], s+1)
-    
-        # apply burn-in 
-        samples = samples[:, Nb:]
-        target_eval = target_eval[Nb:]
-        acc = acc[Nb:]
-        return samples, target_eval, np.mean(acc)
 
-    def single_update(self, x_t, target_eval_t, g_target_eval_t):
-        # approximate Langevin diffusion
-        xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample(rng=self.rng)
-        x_star = x_t + 0.5*self.scale*g_target_eval_t + xi
-        logpi_eval_star, g_logpi_star = self.target.logd(x_star), self.target.gradient(x_star)
+    _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_grad = self._eval_target_grad(self.current_point)
+
+    def validate_target(self):
+        try:
+            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.,
+        current_point, current_target_eval, and current_target_grad_eval.
+
+        Parameters
+        ----------
+        x_star : 
+            The proposed state
+
+        target_eval_star: 
+            The log likelihood evaluated at x_star
+
+        target_grad_star: 
+            The gradient of log likelihood evaluated at x_star
+
+        Returns
+        -------
+        scalar
+            1 (accepted)
+        """
+
+        self.current_point = x_star
+        self.current_target_grad = target_grad_star
+        acc = 1
 
-        # msg
-        if np.isnan(logpi_eval_star):
-            raise NameError('NaN potential func. Consider using smaller scale parameter')
+        return acc
 
-        return x_star, logpi_eval_star, g_logpi_star, 1 # sample always accepted without Metropolis correction
+    def step(self):
+        # propose state
+        xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample()
+        x_star = self.current_point + 0.5*self.scale*self.current_target_grad + xi
 
+        # evaluate target
+        target_eval_star = self._eval_target_logd(x_star)
+        target_grad_star = self._eval_target_grad(x_star)
 
-class MALA(ULA):
+        # accept or reject proposal
+        acc = self._accept_or_reject(x_star, target_eval_star, target_grad_star)
+
+        return acc
+
+    def tune(self, skip_len, update_count):
+        pass
+
+
+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.
@@ -127,21 +154,17 @@ class MALA(ULA):
         The target distribution to sample. Must have logpdf and gradient method. Custom logpdfs 
         and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
     
-    x0 : ndarray
+    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).
 
-    dim : int
-        Dimension of parameter space. Required if target logpdf and gradient are callable 
-        functions. *Optional*.
-
-    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.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
 
     Example
@@ -165,34 +188,195 @@ class MALA(ULA):
         sampler = cuqi.sampler.MALA(target, scale=1/5**2)
 
         # Sample
-        samples = sampler.sample(2000)
+        sampler.sample(2000)
 
     A Deblur example can be found in demos/demo28_MALA.py
+    # TODO: update demo once sampler merged
     """
-    def __init__(self, target, scale, x0=None, dim=None, rng=None, **kwargs):
-        super().__init__(target, scale, x0=x0, dim=dim, rng=rng, **kwargs)
-
-    def single_update(self, x_t, target_eval_t, g_target_eval_t):
-        # approximate Langevin diffusion
-        xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample(rng=self.rng)
-        x_star = x_t + (self.scale/2)*g_target_eval_t + xi
-        logpi_eval_star, g_logpi_star = self.target.logd(x_star), self.target.gradient(x_star)
-
-        # Metropolis step
-        log_target_ratio = logpi_eval_star - target_eval_t
-        log_prop_ratio = self.log_proposal(x_t, x_star, g_logpi_star) \
-            - self.log_proposal(x_star, x_t,  g_target_eval_t)
-        log_alpha = min(0, log_target_ratio + log_prop_ratio)
 
-        # accept/reject
-        log_u = np.log(cuqi.distribution.Uniform(low=0, high=1).sample(rng=self.rng))
-        if (log_u <= log_alpha) and (np.isnan(logpi_eval_star) == False):
-            return x_star, logpi_eval_star, g_logpi_star, 1
-        else:
-            return x_t.copy(), target_eval_t, g_target_eval_t.copy(), 0
+    _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.
+
+        Parameters
+        ----------
+        x_star : 
+            The proposed state
+
+        target_eval_star: 
+            The log likelihood evaluated at x_star
+
+        target_grad_star: 
+            The gradient of log likelihood evaluated at x_star
+
+        Returns
+        -------
+        scaler
+            1 if accepted, 0 otherwise
+        """
+        log_target_ratio = target_eval_star - self.current_target_logd
+        log_prop_ratio = self._log_proposal(self.current_point, x_star, target_grad_star) \
+            - self._log_proposal(x_star, self.current_point,  self.current_target_grad)
+        log_alpha = min(0, log_target_ratio + log_prop_ratio)
 
-    def log_proposal(self, theta_star, theta_k, g_logpi_k):
+        # accept/reject with Metropolis
+        acc = 0
+        log_u = np.log(np.random.rand())
+        if (log_u <= log_alpha) and \
+           (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
+        return acc
+
+    def tune(self, skip_len, update_count):
+        pass
+
+    def _log_proposal(self, theta_star, theta_k, g_logpi_k):
         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
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
+
+    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 = copy(value)
+        if isinstance(copied_value.prior, RestorationPrior):
+            # Acceess the prior name
+            name = value.prior.name
+            copied_value.prior = MoreauYoshidaPrior(
+                copied_value.prior,
+                self.smoothing_strength,
+                name=name)
+        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
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
+
+    # 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)