CUQIpy 1.3.0.post0.dev401__py3-none-any.whl → 1.4.0.post0.dev41__py3-none-any.whl

cuqi/sampler/__init__.py

@@ -1,11 +1,123 @@
+"""
+The sampler module of CUQIpy. It has been re-implemented to improve design, flexibility, 
+and extensibility. The old sampler module can be found in :py:mod:`cuqi.legacy.sampler`.
+
+Main changes for users in this implementation
+---------------------------------------------
+
+1. Sampling API
+   ^^^^^^^^^^^^
+
+   Previously one would call the `.sample` or `sample_adapt` methods of a sampler instance at :py:mod:`cuqi.legacy.sampler` to sample from a target distribution and store the samples as the output as follows:
+
+   .. code-block:: python
+
+      from cuqi.legacy.sampler import MH
+      from cuqi.distribution import DistributionGallery
+
+      # Target distribution
+      target = DistributionGallery("donut")
+
+      # Set up sampler
+      sampler = MH(target)
+
+      # Sample from the target distribution (Alternatively calling sample with explicit scale parameter set in sampler)
+      samples = sampler.sample_adapt(Ns=100, Nb=100) # Burn-in (Nb) removed by default
+
+   This has now changed to to a more object-oriented API which provides more flexibility and control over the sampling process.
+
+   For example one can now more explicitly control when the sampler is tuned (warmup) and when it is sampling with fixed parameters.
+
+   .. code-block:: python
+
+      from cuqi.sampler import MH
+      from cuqi.distribution import DistributionGallery
+
+      # Target distribution
+      target = DistributionGallery("donut")
+
+      # Set up sampler
+      sampler = MH(target)
+
+      # Sample from the target distribution
+      sampler.warmup(Nb=100)  # Explicit warmup (tuning) of sampler
+      sampler.sample(Ns=100)  # Sampling with fixed parameters
+      samples = sampler.get_samples().burnthin(Nb=100) # Getting samples and removing burn-in from warmup
+
+   Importantly, the removal of burn-in from e.g. warmup is now a separate step that is done after the sampling process is complete.
+
+2. Sampling API for BayesianProblem
+   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+   :py:class:`cuqi.problem.BayesianProblem` continues to have the same API for `sample_posterior` and the `UQ` method.
+
+   There is a flag `legacy` that can be set to `True` to use the legacy MCMC samplers.
+
+   By default, the flag is set to `False` and the samplers in `cuqi.sampler` are used.
+
+   For this more high-level interface, burn-in is automatically removed from the samples as was the case before.
+
+
+3. More options for Gibbs sampling
+   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+   There are now more options for Gibbs sampling. Previously it was only possible to sample with Gibbs for samplers  :py:class:`cuqi.legacy.sampler.LinearRTO`, :py:class:`cuqi.legacy.sampler.RegularizedLinearRTO`, :py:class:`cuqi.legacy.sampler.Conjugate`, and :py:class:`cuqi.legacy.sampler.ConjugateApprox`.
+
+   Now, it is possible to define a Gibbs sampling scheme using any sampler from the :py:mod:`cuqi.sampler` module.
+
+   **Example using a NUTS-within-Gibbs scheme for a 1D deconvolution problem:**
+
+   .. code-block:: python
+
+      import cuqi
+      import numpy as np
+      from cuqi.distribution import Gamma, Gaussian, GMRF, JointDistribution
+      from cuqi.sampler import NUTS, HybridGibbs, Conjugate
+      from cuqi.testproblem import Deconvolution1D
+
+      # Forward problem
+      A, y_data, info = Deconvolution1D(dim=128, phantom='sinc', noise_std=0.001).get_components()
+
+      # Bayesian Inverse Problem
+      s = Gamma(1, 1e-4)
+      x = GMRF(np.zeros(A.domain_dim), 50)
+      y = Gaussian(A @ x, lambda s: 1 / s)
+
+      # Posterior
+      target = JointDistribution(y, x, s)(y=y_data)
+
+      # Gibbs sampling strategy. Note we can define initial_points and various parameters for each sampler
+      sampling_strategy = {
+          "x": NUTS(max_depth=10, initial_point=np.zeros(A.domain_dim)),
+          "s": Conjugate()
+      }
+
+      # Here we do 10 internal steps with NUTS for each Gibbs step
+      num_sampling_steps = {
+          "x": 10,
+          "s": 1
+      }
+
+      sampler = HybridGibbs(target, sampling_strategy, num_sampling_steps)
+
+      sampler.warmup(50)
+      sampler.sample(200)
+      samples = sampler.get_samples().burnthin(Nb=50)
+
+      samples["x"].plot_ci(exact=info.exactSolution)
+"""
+
+
+
 from ._sampler import Sampler, ProposalBasedSampler
-from ._conjugate import Conjugate
-from ._conjugate_approx import ConjugateApprox
-from ._cwmh import CWMH
-from ._gibbs import Gibbs
-from ._hmc import NUTS
-from ._langevin_algorithm import ULA, MALA
-from ._laplace_approximation import UGLA
+from ._langevin_algorithm import ULA, MALA, MYULA, PnPULA
 from ._mh import MH
-from ._pcn import pCN
+from ._pcn import PCN
 from ._rto import LinearRTO, RegularizedLinearRTO
+from ._cwmh import CWMH
+from ._laplace_approximation import UGLA
+from ._hmc import NUTS
+from ._gibbs import HybridGibbs
+from ._conjugate import Conjugate
+from ._conjugate_approx import ConjugateApprox
+from ._direct import Direct

cuqi/sampler/_conjugate.py

@@ -1,55 +1,396 @@
-from cuqi.distribution import Posterior, Gaussian, Gamma, GMRF
-from cuqi.implicitprior import RegularizedGaussian, RegularizedGMRF
 import numpy as np
+from abc import ABC, abstractmethod
+import math
+from cuqi.sampler import Sampler
+from cuqi.distribution import Posterior, Gaussian, Gamma, GMRF, ModifiedHalfNormal
+from cuqi.implicitprior import RegularizedGaussian, RegularizedGMRF, RegularizedUnboundedUniform
+from cuqi.utilities import get_non_default_args, count_nonzero, count_within_bounds, count_constant_components_1D, count_constant_components_2D, piecewise_linear_1D_DoF
+from cuqi.geometry import Continuous1D, Continuous2D, Image2D
 
-class Conjugate: # TODO: Subclass from Sampler once updated
+class Conjugate(Sampler):
     """ Conjugate sampler
 
-    Sampler for sampling a posterior distribution where the likelihood and prior are conjugate.
+    Sampler for sampling a posterior distribution which is a so-called "conjugate" distribution, i.e., where the likelihood and prior are conjugate to each other - denoted as a conjugate pair.   
 
     Currently supported conjugate pairs are:
-    - (Gaussian, Gamma)
-    - (GMRF, Gamma)
-    - (RegularizedGaussian, Gamma) with nonnegativity constraints only
+    - (Gaussian, Gamma) where Gamma is defined on the precision parameter of the Gaussian
+    - (GMRF, Gamma) where Gamma is defined on the precision parameter of the GMRF
+    - (RegularizedGaussian, Gamma) with preset constraints only and Gamma is defined on the precision parameter of the RegularizedGaussian
+    - (RegularizedGMRF, Gamma) with preset constraints only and Gamma is defined on the precision parameter of the RegularizedGMRF
+    - (RegularizedGaussian, ModifiedHalfNormal) with most of the preset constraints and regularization
+    - (RegularizedGMRF, ModifiedHalfNormal) with most of the preset constraints and regularization
 
-    For more information on conjugate pairs, see https://en.wikipedia.org/wiki/Conjugate_prior.
+    Currently the Gamma and ModifiedHalfNormal distribution must be univariate.
 
-    For implicit regularized Gaussians see:
+    A conjugate pair defines implicitly a so-called conjugate distribution which can be sampled from directly.
+
+    The conjugate parameter is the parameter that both the likelihood and prior PDF depend on.
+
+    For more information on conjugacy and conjugate distributions see https://en.wikipedia.org/wiki/Conjugate_prior.
+
+    For implicit regularized Gaussians and the corresponding conjugacy relations, see:
     
-    [1] Everink, Jasper M., Yiqiu Dong, and Martin S. Andersen. "Bayesian inference with projected densities." SIAM/ASA Journal on Uncertainty Quantification 11.3 (2023): 1025-1043.
+    Section 3.3 from [1] Everink, Jasper M., Yiqiu Dong, and Martin S. Andersen. "Bayesian inference with projected densities." SIAM/ASA Journal on Uncertainty Quantification 11.3 (2023): 1025-1043.
+    Section 4 from [2] Everink, Jasper M., Yiqiu Dong, and Martin S. Andersen. "Sparse Bayesian inference with regularized Gaussian distributions." Inverse Problems 39.11 (2023): 115004.
 
     """
 
-    def __init__(self, target: Posterior):
-        if not isinstance(target.likelihood.distribution, (Gaussian, GMRF, RegularizedGaussian, RegularizedGMRF)):
-            raise ValueError("Conjugate sampler only works with a Gaussian-type likelihood function")
-        if not isinstance(target.prior, Gamma):
-            raise ValueError("Conjugate sampler only works with Gamma prior")
-        if not target.prior.dim == 1:
-            raise ValueError("Conjugate sampler only works with univariate Gamma prior")
-            
-        if isinstance(target.likelihood.distribution, (RegularizedGaussian, RegularizedGMRF)) and (target.likelihood.distribution.preset["constraint"] not in ["nonnegativity"] or target.likelihood.distribution.preset["regularization"] is not None) :
-               raise ValueError("Conjugate sampler only works implicit regularized Gaussian likelihood with nonnegativity constraints")
+    def _initialize(self):
+        pass
+
+    @Sampler.target.setter # Overwrite the target setter to set the conjugate pair
+    def target(self, value):
+        """ Set the target density. Runs validation of the target. """
+        self._target = value
+        if self._target is not None:
+            self._set_conjugatepair()
+            self.validate_target()
+
+    def validate_target(self):
+        self._ensure_target_is_posterior()
+        self._conjugatepair.validate_target()
+        
+    def step(self):
+        self.current_point = self._conjugatepair.sample()
+        return 1 # Returns acceptance rate of 1
+
+    def tune(self, skip_len, update_count):
+        pass # No tuning required for conjugate sampler
+
+    def _ensure_target_is_posterior(self):
+        """ Ensure that the target is a Posterior distribution. """
+        if not isinstance(self.target, Posterior):
+            raise TypeError("Conjugate sampler requires a target of type Posterior")
+
+    def _set_conjugatepair(self):
+        """ Set the conjugate pair based on the likelihood and prior. This requires target to be set. """
+        self._ensure_target_is_posterior()
+        if isinstance(self.target.likelihood.distribution, (Gaussian, GMRF)) and isinstance(self.target.prior, Gamma):
+            self._conjugatepair = _GaussianGammaPair(self.target)
+        elif isinstance(self.target.likelihood.distribution, RegularizedUnboundedUniform) and isinstance(self.target.prior, Gamma):
+            # Check RegularizedUnboundedUniform before RegularizedGaussian and RegularizedGMRF due to the first inheriting from the second.
+            self._conjugatepair = _RegularizedUnboundedUniformGammaPair(self.target)
+        elif isinstance(self.target.likelihood.distribution, (RegularizedGaussian, RegularizedGMRF)) and isinstance(self.target.prior, Gamma):
+            self._conjugatepair = _RegularizedGaussianGammaPair(self.target)
+        elif isinstance(self.target.likelihood.distribution, (RegularizedGaussian, RegularizedGMRF)) and isinstance(self.target.prior, ModifiedHalfNormal):
+            self._conjugatepair = _RegularizedGaussianModifiedHalfNormalPair(self.target)
+        else:
+            raise ValueError(f"Conjugacy is not defined for likelihood {type(self.target.likelihood.distribution)} and prior {type(self.target.prior)}, in CUQIpy")
+
+    def conjugate_distribution(self):
+        return self._conjugatepair.conjugate_distribution()
+
+    def __repr__(self):
+        msg = super().__repr__()
+        if hasattr(self, "_conjugatepair"):
+            msg += f"\n Conjugate pair:\n\t {type(self._conjugatepair).__name__.removeprefix('_')}"
+        return msg
         
+class _ConjugatePair(ABC):
+    """ Abstract base class for conjugate pairs (likelihood, prior) used in the Conjugate sampler. """
+
+    def __init__(self, target):
         self.target = target
 
-    def step(self, x=None):
+    @abstractmethod
+    def validate_target(self):
+        """ Validate the target distribution for the conjugate pair. """
+        pass
+
+    @abstractmethod
+    def conjugate_distribution(self):
+        """ Returns the posterior distribution in the form of a CUQIpy distribution """
+        pass
+
+    def sample(self):
+        """ Sample from the conjugate distribution. """
+        return self.conjugate_distribution().sample()
+
+
+class _GaussianGammaPair(_ConjugatePair):
+    """ Implementation for the Gaussian-Gamma conjugate pair."""
+
+    def validate_target(self):
+        if self.target.prior.dim != 1:
+            raise ValueError("Gaussian-Gamma conjugacy only works with univariate Gamma prior")
+
+        key_value_pairs = _get_conjugate_parameter(self.target)
+        if len(key_value_pairs) != 1:
+            raise ValueError(f"Multiple references to conjugate parameter {self.target.prior.name} found in likelihood. Only one occurance is supported.")
+        for key, value in key_value_pairs:
+            if key == "cov":
+                if not _check_conjugate_parameter_is_scalar_linear_reciprocal(value):
+                    raise ValueError("Gaussian-Gamma conjugate pair defined via covariance requires cov: lambda x : s/x for the conjugate parameter")
+            elif key == "prec":
+                if not _check_conjugate_parameter_is_scalar_linear(value):
+                    raise ValueError("Gaussian-Gamma conjugate pair defined via precision requires prec: lambda x : s*x for the conjugate parameter")
+            else:
+                raise ValueError(f"RegularizedGaussian-ModifiedHalfNormal conjugacy does not support the conjugate parameter {self.target.prior.name} in the {key} attribute. Only cov and prec")
+
+    def conjugate_distribution(self):
         # Extract variables
-        b = self.target.likelihood.data                                 #mu
-        m = self._calc_m_for_Gaussians(b)                               #n
-        Ax = self.target.likelihood.distribution.mean                   #x_i
-        L = self.target.likelihood.distribution(np.array([1])).sqrtprec #L
-        alpha = self.target.prior.shape                                 #alpha
-        beta = self.target.prior.rate                                   #beta
+        b = self.target.likelihood.data                                 # mu
+        m = len(b)                                                      # n
+        Ax = self.target.likelihood.distribution.mean                   # x_i
+        L = self.target.likelihood.distribution(np.array([1])).sqrtprec # L
+        alpha = self.target.prior.shape                                 # alpha
+        beta = self.target.prior.rate                                   # beta
 
         # Create Gamma distribution and sample
-        dist = Gamma(shape=m/2+alpha,rate=.5*np.linalg.norm(L@(Ax-b))**2+beta)
+        return Gamma(shape=m/2 + alpha, rate=.5 * np.linalg.norm(L @ (Ax - b))**2 + beta)
+
+
+class _RegularizedGaussianGammaPair(_ConjugatePair):
+    """Implementation for the Regularized Gaussian-Gamma conjugate pair using the conjugacy rules from [1], Section 3.3."""
+
+    def validate_target(self):
+        if self.target.prior.dim != 1:
+            raise ValueError("RegularizedGaussian-Gamma conjugacy only works with univariate ModifiedHalfNormal prior")
+
+        # Raises error if preset is not supported
+        _compute_sparsity_level(self.target)
+
+        key_value_pairs = _get_conjugate_parameter(self.target)
+        if len(key_value_pairs) != 1:
+            raise ValueError(f"Multiple references to conjugate parameter {self.target.prior.name} found in likelihood. Only one occurance is supported.")
+        for key, value in key_value_pairs:
+            if key == "cov":
+                if not _check_conjugate_parameter_is_scalar_linear_reciprocal(value):
+                    raise ValueError("Regularized Gaussian-Gamma conjugacy defined via covariance requires cov: lambda x : s/x for the conjugate parameter")
+            elif key == "prec":
+                if not _check_conjugate_parameter_is_scalar_linear(value):
+                    raise ValueError("Regularized Gaussian-Gamma conjugacy defined via precision requires prec: lambda x : s*x for the conjugate parameter")
+            else:
+                raise ValueError(f"RegularizedGaussian-ModifiedHalfNormal conjugacy does not support the conjugate parameter {self.target.prior.name} in the {key} attribute. Only cov and prec")
+
+    def conjugate_distribution(self):
+        # Extract variables
+        b = self.target.likelihood.data                                 # mu
+        m = _compute_sparsity_level(self.target)
+        Ax = self.target.likelihood.distribution.mean                   # x_i
+        L = self.target.likelihood.distribution(np.array([1])).sqrtprec # L
+        alpha = self.target.prior.shape                                 # alpha
+        beta = self.target.prior.rate                                   # beta
+
+        # Create Gamma distribution and sample
+        return Gamma(shape=m/2 + alpha, rate=.5 * np.linalg.norm(L @ (Ax - b))**2 + beta)
+    
+    
+class _RegularizedUnboundedUniformGammaPair(_ConjugatePair):
+    """Implementation for the RegularizedUnboundedUniform-ModifiedHalfNormal conjugate pair using the conjugacy rules from [2], Section 4."""
+
+    def validate_target(self):
+        if self.target.prior.dim != 1:
+            raise ValueError("RegularizedUnboundedUniform-Gamma conjugacy only works with univariate Gamma prior")
+        
+        # Raises error if preset is not supported
+        _compute_sparsity_level(self.target)
+
+        key_value_pairs = _get_conjugate_parameter(self.target)
+        if len(key_value_pairs) != 1:
+            raise ValueError(f"Multiple references to conjugate parameter {self.target.prior.name} found in likelihood. Only one occurance is supported.")
+        for key, value in key_value_pairs:
+            if key == "strength":
+                if not _check_conjugate_parameter_is_scalar_linear(value):
+                    raise ValueError("RegularizedUnboundedUniform-Gamma conjugacy defined via strength requires strength: lambda x : s*x for the conjugate parameter")
+            else:
+                raise ValueError(f"RegularizedUnboundedUniform-Gamma conjugacy does not support the conjugate parameter {self.target.prior.name} in the {key} attribute. Only strength is supported")
+
+    def conjugate_distribution(self):
+        # Extract prior variables
+        alpha = self.target.prior.shape
+        beta = self.target.prior.rate
+
+        # Compute likelihood quantities
+        x = self.target.likelihood.data
+        m = _compute_sparsity_level(self.target)
 
-        return dist.sample()
+        reg_op = self.target.likelihood.distribution._regularization_oper
+        reg_strength = self.target.likelihood.distribution(np.array([1])).strength
+        fx = reg_strength*np.linalg.norm(reg_op@x, ord = 1)
 
-    def _calc_m_for_Gaussians(self, b):
-        """ Helper method to calculate m parameter for Gaussian-Gamma conjugate pair. """
-        if isinstance(self.target.likelihood.distribution, (Gaussian, GMRF)):
-            return len(b)
-        elif isinstance(self.target.likelihood.distribution, (RegularizedGaussian, RegularizedGMRF)):
-            return np.count_nonzero(b) # See 
+        # Create Gamma distribution
+        return Gamma(shape=m/2 + alpha, rate=fx + beta)
+    
+class _RegularizedGaussianModifiedHalfNormalPair(_ConjugatePair):
+    """Implementation for the Regularized Gaussian-ModifiedHalfNormal conjugate pair using the conjugacy rules from [2], Section 4."""
+
+    def validate_target(self):
+        if self.target.prior.dim != 1:
+            raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy only works with univariate ModifiedHalfNormal prior")
+        
+        # Raises error if preset is not supported
+        _compute_sparsity_level(self.target)
+
+        key_value_pairs = _get_conjugate_parameter(self.target)
+        if len(key_value_pairs) != 2:
+            raise ValueError(f"Incorrect number of references to conjugate parameter {self.target.prior.name} found in likelihood. Found {len(key_value_pairs)} times, but needs to occur in prec or cov, and in strength")
+        for key, value in key_value_pairs:
+            if key == "strength":
+                if not _check_conjugate_parameter_is_scalar_linear(value):
+                    raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy defined via strength requires strength: lambda x : s*x for the conjugate parameter")
+            elif key == "prec":
+                if not _check_conjugate_parameter_is_scalar_quadratic(value):
+                    raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy defined via precision requires prec: lambda x : s*x for the conjugate parameter")
+            elif key == "cov":
+                if not _check_conjugate_parameter_is_scalar_quadratic_reciprocal(value):
+                    raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy  defined via covariance requires cov: lambda x : s/x for the conjugate parameter")
+            else:
+                raise ValueError(f"RegularizedGaussian-ModifiedHalfNormal conjugacy does not support the conjugate parameter {self.target.prior.name} in the {key} attribute. Only cov, prec and strength are supported")
+
+
+    def conjugate_distribution(self):
+        # Extract prior variables
+        alpha = self.target.prior.alpha
+        beta = self.target.prior.beta
+        gamma = self.target.prior.gamma
+
+        # Compute likelihood variables
+        x = self.target.likelihood.data
+        mu = self.target.likelihood.distribution.mean
+        L = self.target.likelihood.distribution(np.array([1])).sqrtprec
+        
+        m = _compute_sparsity_level(self.target)
+
+        reg_op = self.target.likelihood.distribution._regularization_oper
+        reg_strength = self.target.likelihood.distribution(np.array([1])).strength
+        fx = reg_strength*np.linalg.norm(reg_op@x, ord = 1)
+
+        # Compute parameters of conjugate distribution
+        conj_alpha = m + alpha
+        conj_beta = 0.5*np.linalg.norm(L @ (mu - x))**2 + beta
+        conj_gamma = -fx + gamma
+
+        # Create conjugate distribution
+        return ModifiedHalfNormal(conj_alpha, conj_beta, conj_gamma)
+    
+
+def _compute_sparsity_level(target):
+    """Computes the sparsity level in accordance with Section 4 from [2],
+    this can be interpreted as the number of degrees of freedom, that is,
+    the number of components n minus the dimension the of the subdifferential of the regularized.
+    """
+    x = target.likelihood.data
+
+    constraint = target.likelihood.distribution.preset["constraint"]
+    regularization = target.likelihood.distribution.preset["regularization"]
+
+    # There is no reference for some of these conjugacy rules
+    if constraint == "nonnegativity":
+        if regularization in [None, "l1"]:
+            # Number of non-zero components in x
+            return count_nonzero(x)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of non-zero constant components in x
+            return count_constant_components_1D(x, lower = 0.0)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            # Number of non-zero constant components in x
+            return count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x), lower = 0.0)
+    elif constraint == "box":
+        bounds = target.likelihood.distribution._box_bounds
+        if regularization is None:
+            # Number of components in x that are strictly between the lower and upper bound
+            return count_within_bounds(x, bounds[0], bounds[1])
+        elif regularization == "l1":
+            # Number of components in x that are strictly between the lower and upper bound and are not zero
+            return count_within_bounds(x, bounds[0], bounds[1], exception = 0.0)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of constant components in x between are strictly between the lower and upper bound
+            return count_constant_components_1D(x, lower = bounds[0], upper = bounds[1])
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            # Number of constant components in x between are strictly between the lower and upper bound
+            return count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x), lower = bounds[0], upper = bounds[1])
+    elif constraint in ["increasing", "decreasing"]:
+        if regularization is None:
+            # Number of constant components in x
+            return count_constant_components_1D(x)
+        elif regularization == "l1":
+            # Number of constant components in x that are not zero
+            return count_constant_components_1D(x, exception = 0.0)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of constant components in x
+            return count_constant_components_1D(x)
+        # Increasing and decreasing cannot be done in 2D
+    elif constraint in ["convex", "concave"]:
+        if regularization is None:
+            # Number of piecewise linear components in x
+            return piecewise_linear_1D_DoF(x)
+        elif regularization == "l1":
+            # Number of piecewise linear components in x that are not zero
+            return piecewise_linear_1D_DoF(x, exception_zero = True)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of piecewise linear components in x that are not flat
+            return piecewise_linear_1D_DoF(x, exception_flat = True)
+        # convex and concave has only been implemented in 1D
+    elif constraint == None: 
+        if regularization == "l1":
+            # Number of non-zero components in x
+            return count_nonzero(x)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of non-zero constant components in x
+            return count_constant_components_1D(x)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            # Number of non-zero constant components in x
+            return count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x))
+
+    raise ValueError("RegularizedGaussian preset constraint and regularization choice is currently not supported with conjugacy.")
+
+
+def _get_conjugate_parameter(target):
+    """Extract the conjugate parameter name (e.g. d), and returns the mutable variable that is defined by the conjugate parameter, e.g. cov and its value e.g. lambda d:1/d"""
+    par_name = target.prior.name
+    mutable_likelihood_vars  = target.likelihood.distribution.get_mutable_variables()
+
+    found_parameter_pairs = []
+
+    for var_key in mutable_likelihood_vars:
+        attr = getattr(target.likelihood.distribution, var_key)
+        if callable(attr) and par_name in get_non_default_args(attr):
+            found_parameter_pairs.append((var_key, attr))
+    if len(found_parameter_pairs) == 0:
+        raise ValueError(f"Unable to find conjugate parameter {par_name} in likelihood function for conjugate sampler with target {target}")
+    return found_parameter_pairs
+
+def _check_conjugate_parameter_is_scalar_identity(f):
+    """Tests whether a function (scalar to scalar) is the identity (lambda x: x)."""
+    test_values = [1.0, 10.0, 100.0]
+    return all(np.allclose(f(x), x) for x in test_values)
+
+def _check_conjugate_parameter_is_scalar_reciprocal(f):
+    """Tests whether a function (scalar to scalar) is the reciprocal (lambda x : 1.0/x)."""
+    return all(math.isclose(f(x), 1.0 / x) for x in [1.0, 10.0, 100.0])
+
+def _check_conjugate_parameter_is_scalar_linear(f):
+    """
+        Tests whether a function (scalar to scalar) is linear (lambda x: s*x for some s).
+        The tests checks whether the function is zero and some finite differences are constant.
+    """
+    test_values = [1.0, 10.0, 100.0]
+    h = 1e-2
+    finite_diffs = [(f(x + h*x)-f(x))/(h*x) for x in test_values]
+    return np.isclose(f(0.0), 0.0) and all(np.allclose(c, finite_diffs[0]) for c in finite_diffs[1:])
+
+def _check_conjugate_parameter_is_scalar_linear_reciprocal(f):
+    """
+        Tests whether a function (scalar to scalar) is a constant times the inverse of the input (lambda x: s/x for some s).
+        The tests checks whether the the reciprocal of the function has constant finite differences.
+    """
+    g = lambda x : 1.0/f(x)
+    test_values = [1.0, 10.0, 100.0]
+    h = 1e-2
+    finite_diffs = [(g(x + h*x)-g(x))/(h*x) for x in test_values]
+    return all(np.allclose(c, finite_diffs[0]) for c in finite_diffs[1:])
+
+def _check_conjugate_parameter_is_scalar_quadratic(f):
+    """
+        Tests whether a function (scalar to scalar) is linear (lambda x: s*x**2 for some s).
+        The tests checks whether the function divided by the parameter is linear
+    """
+    return _check_conjugate_parameter_is_scalar_linear(lambda x: f(x)/x if x != 0.0 else f(0.0))
+
+def _check_conjugate_parameter_is_scalar_quadratic_reciprocal(f):
+    """
+        Tests whether a function (scalar to scalar) is linear (lambda x: s*x**-2 for some s).
+        The tests checks whether the function divided by the parameter is the reciprical of a linear function.
+    """
+    return _check_conjugate_parameter_is_scalar_linear_reciprocal(lambda x: f(x)/x)

cuqi/sampler/_conjugate_approx.py

@@ -1,31 +1,57 @@
-from cuqi.distribution import Posterior, LMRF, Gamma
 import numpy as np
+from cuqi.sampler import Conjugate
+from cuqi.sampler._conjugate import _ConjugatePair, _get_conjugate_parameter, _check_conjugate_parameter_is_scalar_reciprocal
+from cuqi.distribution import LMRF, Gamma
 import scipy as sp
 
-class ConjugateApprox: # TODO: Subclass from Sampler once updated
+class ConjugateApprox(Conjugate):
     """ Approximate Conjugate sampler
 
     Sampler for sampling a posterior distribution where the likelihood and prior can be approximated
     by a conjugate pair.
 
     Currently supported pairs are:
-    - (LMRF, Gamma): Approximated by (Gaussian, Gamma)
+    - (LMRF, Gamma): Approximated by (Gaussian, Gamma) where Gamma is defined on the inverse of the scale parameter of the LMRF distribution.
 
-    For more information on conjugate pairs, see https://en.wikipedia.org/wiki/Conjugate_prior.
+    Gamma distribution must be univariate.
+
+    LMRF likelihood must have zero mean.
+
+    For more details on conjugacy see :class:`Conjugate`.
 
     """
 
-    
-    def __init__(self, target: Posterior):
-        if not isinstance(target.likelihood.distribution, LMRF):
-            raise ValueError("Conjugate sampler only works with Laplace diff likelihood function")
-        if not isinstance(target.prior, Gamma):
-            raise ValueError("Conjugate sampler only works with Gamma prior")
-        self.target = target
+    def _set_conjugatepair(self):
+        """ Set the conjugate pair based on the likelihood and prior. This requires target to be set. """
+        if isinstance(self.target.likelihood.distribution, LMRF) and isinstance(self.target.prior, Gamma):
+            self._conjugatepair = _LMRFGammaPair(self.target)
+        else:
+            raise ValueError(f"Conjugacy is not defined for likelihood {type(self.target.likelihood.distribution)} and prior {type(self.target.prior)}, in CUQIpy")
+
 
-    def step(self, x=None):
+class _LMRFGammaPair(_ConjugatePair):
+    """ Implementation of the conjugate pair (LMRF, Gamma) """
+
+    def validate_target(self):
+        if not self.target.prior.dim == 1:
+            raise ValueError("Approximate conjugate sampler only works with univariate Gamma prior")
+        
+        if np.sum(self.target.likelihood.distribution.location) != 0:
+            raise ValueError("Approximate conjugate sampler only works with zero mean LMRF likelihood")
+        
+        key_value_pairs = _get_conjugate_parameter(self.target)
+        if len(key_value_pairs) != 1:
+            raise ValueError(f"Multiple references to conjugate parameter {self.target.prior.name} found in likelihood. Only one occurance is supported.")
+        for key, value in key_value_pairs:
+            if key == "scale":
+                if not _check_conjugate_parameter_is_scalar_reciprocal(value):
+                    raise ValueError("Approximate conjugate sampler only works with Gamma prior on the inverse of the scale parameter of the LMRF likelihood")
+            else:
+                raise ValueError(f"No approximate conjugacy defined for likelihood {type(self.target.likelihood.distribution)} and prior {type(self.target.prior)}, in CUQIpy")
+        
+    def conjugate_distribution(self):
         # Extract variables
-        # Here we approximate the Laplace diff with a Gaussian
+        # Here we approximate the LMRF with a Gaussian
 
         # Extract diff_op from target likelihood
         D = self.target.likelihood.distribution._diff_op
@@ -47,6 +73,4 @@ class ConjugateApprox: # TODO: Subclass from Sampler once updated
         beta = self.target.prior.rate                   #beta
 
         # Create Gamma distribution and sample
-        dist = Gamma(shape=d+alpha, rate=np.linalg.norm(Lx)**2+beta)
-
-        return dist.sample()
+        return Gamma(shape=d+alpha, rate=np.linalg.norm(Lx)**2+beta)