Merge pull request #611 from CUQI-DTU/add_rv_to_distribution

Add rv method to Distribution class and update artifact version

Author: chaozg

.github/workflows/docs.yml

@@ -26,7 +26,7 @@ jobs:
       # If PR-based build. Upload artifact of the documentation for download and review
       - name: Upload artifact
         if: ${{ github.event_name == 'pull_request' }}
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: docs_build
           path: ./public

cuqi/distribution/_distribution.py

@@ -406,3 +406,9 @@ def __repr__(self) -> str:
             return "CUQI {}. Conditioning variables {}.".format(self.__class__.__name__,self.get_conditioning_variables())
         else:
             return "CUQI {}.".format(self.__class__.__name__)
+
+    @property
+    def rv(self):
+        """ Return a random variable object representing the distribution. """
+        from cuqi.experimental.algebra import RandomVariable
+        return RandomVariable(self)

demos/demo39_check_RandomVariable.py

@@ -0,0 +1,79 @@
+# %%
+# This demo shows the use of RandomVariable (RV) as prior in posteior sampling.
+# Specifically, we test the use of RV as prior through 
+# 1) BayesianProblem,
+# 2) JointDistribution with MALA,
+# 3) JointDistribution with Gibbs (NUTS + Conjugate) and
+# 4) JointDistribution with Regularized LinearRTO.
+
+from cuqi.testproblem import Deconvolution1D
+from cuqi.distribution import Gaussian, Gamma, GMRF
+from cuqi.problem import BayesianProblem
+from cuqi.distribution import JointDistribution
+from cuqi.implicitprior import NonnegativeGMRF
+from cuqi.experimental.mcmc import HybridGibbs, Conjugate, MALA, NUTS, RegularizedLinearRTO
+import numpy as np
+import matplotlib.pyplot as plt
+
+np.random.seed(1912)
+
+# %% Forward model
+n = 128
+A, y_obs, info = Deconvolution1D(dim=n, phantom='square').get_components()
+
+# %% 1) Bayesian problem
+x = GMRF(np.zeros(A.domain_dim), 100).rv
+y = Gaussian(A @ x, 1/10000).rv
+
+# Combine into a Bayesian Problem and perform UQ
+BP = BayesianProblem(y, x)
+BP.set_data(y=y_obs)
+BP.UQ(exact=info.exactSolution)
+
+# %% 2) Joint distribution with MALA
+target = JointDistribution(y, x)(y=y_obs)
+sampler = MALA(target, scale=0.0002, initial_point=np.zeros(n))
+sampler.sample(10000)
+samples = sampler.get_samples()
+plt.figure()
+samples.plot_ci(exact=info.exactSolution)
+
+# %% 3) Hybrid Gibbs (NUTS + Conjugate)
+d = Gamma(1, 1e-4).rv
+x = GMRF(np.zeros(A.domain_dim), d).rv
+y = Gaussian(A @ x, 1/10000).rv
+
+target = JointDistribution(y, x, d)(y=y_obs)
+
+# Sampling strategy
+sampling_strategy = {
+    "x" : NUTS(),
+    "d" : Conjugate()
+}
+
+# Gibbs sampler
+sampler = HybridGibbs(target, sampling_strategy)
+# Run sampler
+sampler.warmup(200)
+sampler.sample(1000)
+samples = sampler.get_samples()
+
+# Plot
+plt.figure()
+samples["x"].plot_ci(exact=info.exactSolution)
+
+# %% 4) Regularized GMRF
+x = NonnegativeGMRF(np.zeros(A.domain_dim), 50).rv
+y = Gaussian(A @ x, 1/10000).rv
+
+target = JointDistribution(y, x)(y=y_obs)
+
+# Regularized Linear RTO sampler
+sampler = RegularizedLinearRTO(target)
+# Run sampler
+sampler.sample(1000)
+samples = sampler.get_samples()
+
+# Plot
+plt.figure()
+samples.plot_ci(exact=info.exactSolution)

demos/howtos/RandomVariablesAndAlgebra.py

@@ -12,16 +12,16 @@
 #%%
 # Defining Random Variables
 # -------------------------
-# Random variables can be defined using the RandomVariable class. The RandomVariable
-# class requires a distribution object to be passed as an argument. The distribution
-# object can be any distribution object from the `cuqi.distribution` module.
+
+# Random variables can be defined by either initialising the RandomVariable class
+# with a distribution object or by retrieving the `rv` attribute of a distribution.
+# The distribution object can be any distribution from the `cuqi.distribution` module.
 
 from cuqi.distribution import Normal
 from cuqi.experimental.algebra import RandomVariable
 
-
 x = RandomVariable(Normal(0, 1))
-y = RandomVariable(Normal(0, 1))
+y = Normal(0, 1).rv
 
 # %%
 # Recording Algebraic Operations
@@ -66,33 +66,59 @@
 from cuqi.distribution import Gaussian, Gamma, GMRF
 from cuqi.experimental.algebra import RandomVariable
 from cuqi.problem import BayesianProblem
+from cuqi.distribution import JointDistribution
+from cuqi.experimental.mcmc import HybridGibbs, LinearRTO, Conjugate, ConjugateApprox
 import numpy as np
+import matplotlib.pyplot as plt
 
 # Forward model
 A, y_obs, info = Deconvolution1D().get_components()
 
 # Bayesian Problem (defined using Random Variables)
-d = RandomVariable(Gamma(1, 1e-4))
-s = RandomVariable(Gamma(1, 1e-4))
-x = RandomVariable(GMRF(np.zeros(A.domain_dim), d))
-y = RandomVariable(Gaussian(A @ x, 1/s))
+d = Gamma(1, 1e-4).rv
+s = Gamma(1, 1e-4).rv
+x = GMRF(np.zeros(A.domain_dim), d).rv
+y = Gaussian(A @ x, 1/s).rv
 
 # Combine into a Bayesian Problem and perform UQ
 BP = BayesianProblem(y, x, s, d)
 BP.set_data(y=y_obs)
 BP.UQ(exact={"x": info.exactSolution})
 
+# Random variables can also be used to define JointDistribution. Here we solve the same
+# problem above by explictly forming a target distribution and then drawing samples with
+# the HybridGibbs sampler.
+target = JointDistribution(y, x, s, d)(y=y_obs)
+
+# Sampling strategy
+sampling_strategy = {
+    "x" : LinearRTO(),
+    "s" : Conjugate(),
+    "d" : Conjugate()
+}
+
+# Gibbs sampler
+sampler = HybridGibbs(target, sampling_strategy)
+
+# Run sampler
+sampler.warmup(200)
+sampler.sample(1000)
+samples = sampler.get_samples()
+
+# Plot
+plt.figure()
+samples["x"].plot_ci(exact=info.exactSolution)
+
 # %%
 # Conditioning on random variables (example 1)
-from cuqi.distribution import Gaussian
-from cuqi.experimental.algebra import RandomVariable
-
-x = RandomVariable(Gaussian(0, lambda s: s))
-y = RandomVariable(Gaussian(0, lambda d: d))
+s = Gaussian(0, 1).rv
+x = Gaussian(0, s).rv
+y = Gaussian(0, lambda d: d).rv
 
 z = x+y
 
-z.condition(x=1)
+z.condition(s=1)
+z.condition(d=2)
 
 # %%
 # Or conditioning on the variables s, or d
@@ -110,10 +136,10 @@
 A, y_obs, info = Deconvolution1D(dim=4).get_components()
 
 # Bayesian Problem (defined using Random Variables)
-d = RandomVariable(Gamma(1, 1e-4))
-s = RandomVariable(Gamma(1, 1e-4))
-x = RandomVariable(GMRF(np.zeros(A.domain_dim), d))
-y = RandomVariable(Gaussian(A @ x, 1/s))
+d = Gamma(1, 1e-4).rv
+s = Gamma(1, 1e-4).rv
+x = GMRF(np.zeros(A.domain_dim), d).rv
+y = Gaussian(A @ x, 1/s).rv
 
 
 z = x+y

tests/zexperimental/test_mcmc.py

@@ -1255,6 +1255,142 @@ def test_UGLA_with_AffineModel_is_equivalent_to_LinearModel_and_shifted_data():
     # Check that the samples are the same
     assert np.allclose(samples_linear.samples, samples_affine.samples)
 
+# ============ Test for sampling with RandomVariable prior against Distribution prior ============
+samplers_for_rv_against_dist = [cuqi.experimental.mcmc.MALA, 
+                                cuqi.experimental.mcmc.ULA,
+                                cuqi.experimental.mcmc.MH,
+                                cuqi.experimental.mcmc.PCN,
+                                cuqi.experimental.mcmc.CWMH,
+                                cuqi.experimental.mcmc.NUTS,
+                                cuqi.experimental.mcmc.LinearRTO]
+
+@pytest.mark.parametrize("sampler", samplers_for_rv_against_dist)
+def test_RandomVariable_prior_against_Distribution_prior(sampler: cuqi.experimental.mcmc.Sampler):
+    """ Test RandomVariable prior is equivalent to Distribution prior for 
+        MALA, ULA, MH, PCN, CWMH, NUTS and LinearRTO. 
+    """
+
+    # Set dim
+    dim = 32
+
+    # Extract model and data
+    A, y_data, info = cuqi.testproblem.Deconvolution1D(dim=32, phantom='square').get_components()
+
+    # Set up RandomVariable prior and do posterior sampling
+    np.random.seed(0)
+    x_rv = cuqi.distribution.Gaussian(0.5*np.ones(dim), 0.1).rv
+    y_rv = cuqi.distribution.Gaussian(A@x_rv, 0.001).rv
+    joint_rv = cuqi.distribution.JointDistribution(x_rv, y_rv)(y_rv=y_data)
+    sampler_rv = sampler(joint_rv)
+    sampler_rv.sample(10)
+    samples_rv = sampler_rv.get_samples()
+    
+    # Set up Distribution prior and do posterior sampling
+    np.random.seed(0)
+    x_dist = cuqi.distribution.Gaussian(0.5*np.ones(dim), 0.1)
+    y_dist = cuqi.distribution.Gaussian(A@x_dist, 0.001)
+    joint_dist = cuqi.distribution.JointDistribution(x_dist, y_dist)(y_dist=y_data)
+    sampler_dist = sampler(joint_dist)
+    sampler_dist.sample(10)
+    samples_dist = sampler_dist.get_samples()
+
+    assert np.allclose(samples_rv.samples, samples_dist.samples)
+
+def test_RandomVariable_prior_against_Distribution_prior_regularized_RTO():
+    """ Test RandomVariable prior is equivalent to Distribution prior for 
+        RegularizedLinearRTO. 
+    """
+
+    # Set dim
+    dim = 32
+
+    # Extract model and data
+    A, y_data, info = cuqi.testproblem.Deconvolution1D(dim=32, phantom='square').get_components()
+
+    # Set up RandomVariable prior and do posterior sampling
+    np.random.seed(0)
+    x_rv = cuqi.implicitprior.RegularizedGaussian(0.5*np.ones(dim), 0.1, constraint="nonnegativity").rv
+    y_rv = cuqi.distribution.Gaussian(A@x_rv, 0.001).rv
+    joint_rv = cuqi.distribution.JointDistribution(x_rv, y_rv)(y_rv=y_data)
+    sampler_rv = cuqi.experimental.mcmc.RegularizedLinearRTO(joint_rv)
+    sampler_rv.sample(10)
+    samples_rv = sampler_rv.get_samples()
+    
+    # Set up Distribution prior and do posterior sampling
+    np.random.seed(0)
+    x_dist = cuqi.implicitprior.RegularizedGaussian(0.5*np.ones(dim), 0.1, constraint="nonnegativity")
+    y_dist = cuqi.distribution.Gaussian(A@x_dist, 0.001)
+    joint_dist = cuqi.distribution.JointDistribution(x_dist, y_dist)(y_dist=y_data)
+    sampler_dist = cuqi.experimental.mcmc.RegularizedLinearRTO(joint_dist)
+    sampler_dist.sample(10)
+    samples_dist = sampler_dist.get_samples()
+
+    assert np.allclose(samples_rv.samples, samples_dist.samples)
+
+def test_RandomVariable_prior_against_Distribution_prior_UGLA_Conjugate_ConjugateApprox_HybridGibbs():
+    """ Test RandomVariable prior is equivalent to Distribution prior for 
+        UGLA, Conjugate, ConjugateApprox and HybridGibbs samplers.
+    """
+
+    # Forward problem
+    A, y_data, info = cuqi.testproblem.Deconvolution1D(dim=28, phantom='square', noise_std=0.001).get_components()
+
+    # Random seed
+    np.random.seed(0)
+
+    # Bayesian Inverse Problem
+    d = cuqi.distribution.Gamma(1, 1e-4)
+    s = cuqi.distribution.Gamma(1, 1e-4)
+    x = cuqi.distribution.LMRF(0, lambda d: 1/d, geometry=A.domain_geometry)
+    y = cuqi.distribution.Gaussian(A@x, lambda s: 1/s)
+
+    # Posterior
+    target = cuqi.distribution.JointDistribution(y, x, s, d)(y=y_data)
+
+    # Sampling strategy
+    sampling_strategy = {
+        "x" : cuqi.experimental.mcmc.UGLA(),
+        "s" : cuqi.experimental.mcmc.Conjugate(),
+        "d" : cuqi.experimental.mcmc.ConjugateApprox()
+    }
+
+    # Gibbs sampler
+    sampler = cuqi.experimental.mcmc.HybridGibbs(target, sampling_strategy)
+
+    # Run sampler
+    sampler.warmup(50)
+    sampler.sample(200)
+    samples = sampler.get_samples()
+
+    # Random seed
+    np.random.seed(0)
+
+    # Bayesian Inverse Problem
+    d_rv = cuqi.distribution.Gamma(1, 1e-4).rv
+    s_rv = cuqi.distribution.Gamma(1, 1e-4).rv
+    x_rv = cuqi.distribution.LMRF(0, lambda d_rv: 1/d_rv, geometry=A.domain_geometry).rv
+    y_rv = cuqi.distribution.Gaussian(A@x_rv, lambda s_rv: 1/s_rv).rv
+
+    # Posterior
+    target_rv = cuqi.distribution.JointDistribution(y_rv, x_rv, s_rv, d_rv)(y_rv=y_data)
+
+    # Sampling strategy
+    sampling_strategy_rv = {
+        "x_rv" : cuqi.experimental.mcmc.UGLA(),
+        "s_rv" : cuqi.experimental.mcmc.Conjugate(),
+        "d_rv" : cuqi.experimental.mcmc.ConjugateApprox()
+    }
+
+    # Gibbs sampler
+    sampler_rv = cuqi.experimental.mcmc.HybridGibbs(target_rv, sampling_strategy_rv)
+
+    # Run sampler
+    sampler_rv.warmup(50)
+    sampler_rv.sample(200)
+    samples_rv = sampler_rv.get_samples()
+
+    assert np.allclose(samples['x'].samples, samples_rv['x_rv'].samples)
+
 def Conjugate_GaussianGammaPair():
     """ Unit test whether Conjugacy Pair (Gaussian, Gamma) constructs the right distribution """
     x = cuqi.distribution.Gamma(1.0, 2.0)
@@ -1373,4 +1509,4 @@ def test_RegularizedLinearRTO_ScipyLinearLSQ_option_invalid():
     posterior = joint(y=y_data)
 
     with pytest.raises(ValueError, match="ScipyLinearLSQ"):
-        sampler = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver = "ScipyLinearLSQ")
+        sampler = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver = "ScipyLinearLSQ")

tests/zexperimental/test_randomvariable.py

@@ -328,6 +328,21 @@ def test_RV_condition_maintains_parameter_name_order():
     assert z.condition(d=1).parameter_names == ['x', 'y']
     assert z.condition(d=1, s=1).parameter_names == ['x', 'y']
 
+def test_equivalent_ways_to_create_RV_from_distribution():
+    x = RandomVariable(cuqi.distribution.Gaussian(0, 1))
+    y = cuqi.distribution.Gaussian(0, 1).rv
+
+    assert x.dim == y.dim
+    assert x.distribution.mean == y.distribution.mean
+    assert x.distribution.cov == y.distribution.cov
+
+    x = RandomVariable(cuqi.distribution.Gaussian(0, lambda s: s))
+    y = cuqi.distribution.Gaussian(0, lambda s: s).rv
+
+    assert x.dim == y.dim
+    assert x.distribution.mean == y.distribution.mean
+    assert x.condition(s=1).distribution.cov == y.condition(s=1).distribution.cov
+
 def test_RV_sampling_unable_if_conditioning_variables_from_lambda():
     x = RandomVariable(cuqi.distribution.Gaussian(0, lambda s: s))
 
@@ -354,3 +369,4 @@ def test_RV_sample_against_distribution_sample():
     dist_samples = x.distribution.sample(10)
 
     assert np.allclose(rv_samples.samples, dist_samples.samples)
+