Merge pull request #707 from CUQI-DTU/gibbs_online_thinning

Add online thinning to HybridGibbs

Author: chaozg

cuqi/sampler/_gibbs.py

@@ -180,33 +180,39 @@ def validate_targets(self):
         for sampler in self.samplers.values():
             sampler.validate_target()
 
-    def sample(self, Ns) -> 'HybridGibbs':
+    def sample(self, Ns, Nt=1) -> 'HybridGibbs':
         """ Sample from the joint distribution using Gibbs sampling
 
         Parameters
         ----------
         Ns : int
             The number of samples to draw.
-
+        Nt : int, optional, default=1
+            The thinning interval. If Nt >= 1, every Nt'th sample is stored. The larger Nt, the fewer samples are stored.
+            
         """
         for idx in tqdm(range(Ns), "Sample: "):
 
             self.step()
 
-            self._store_samples()
+            if (Nt > 0) and ((idx + 1) % Nt == 0):
+                self._store_samples()
 
             # Call callback function if specified
             self._call_callback(idx, Ns)
 
         return self
 
-    def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
+    def warmup(self, Nb, Nt=1, tune_freq=0.1) -> 'HybridGibbs':
         """ Warmup (tune) the samplers in the Gibbs sampling scheme
 
         Parameters
         ----------
         Nb : int
             The number of samples to draw during warmup.
+        
+        Nt : int, optional, default=1
+            The thinning interval. If Nt >= 1, every Nt'th sample is stored. The larger Nt, the fewer samples are stored.
 
         tune_freq : float, optional
             Frequency of tuning the samplers. Tuning is performed every tune_freq*Nb steps.
@@ -221,9 +227,10 @@ def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
 
             # Tune the sampler at tuning intervals (matching behavior of Sampler class)
             if (idx + 1) % tune_interval == 0:
-                self.tune(tune_interval, idx // tune_interval) 
-                
-            self._store_samples()
+                self.tune(tune_interval, idx // tune_interval)
+
+            if (Nt > 0) and ((idx + 1) % Nt == 0):
+                self._store_samples()
 
             # Call callback function if specified
             self._call_callback(idx, Nb)

tests/zexperimental/test_mcmc.py

@@ -1656,6 +1656,61 @@ def test_online_thinning_with_mala_and_rto():
     assert np.allclose(samples_rto_1.samples[:,1], samples_rto_2.samples[:,9], rtol=1e-8)
     assert np.allclose(samples_rto_1.samples[:,2], samples_rto_2.samples[:,14], rtol=1e-8)
 
+def test_online_thinning_with_hybrid_gibbs():
+
+    # example adapted from https://cuqi-dtu.github.io/CUQI-Book/chapter04/gibbs.html
+
+    # Model and data
+    A, y_data, _ = cuqi.testproblem.Deconvolution1D(phantom='sinc', noise_std=0.005, PSF_param=6).get_components()
+
+    # Get dimension of signal
+    n = A.domain_dim
+
+    d = cuqi.distribution.Gamma(1, 1e-4)
+    s = cuqi.distribution.Gamma(1, 1e-4)
+    x = cuqi.distribution.GMRF(np.zeros(n), lambda d: d)
+    y = cuqi.distribution.Gaussian(A@x, cov=lambda s: 1/s)
+
+    # Create joint distribution
+    joint = cuqi.distribution.JointDistribution(y, x, d, s)
+
+    # Define posterior by conditioning on the data
+    posterior = joint(y=y_data)
+
+    # Define sampling strategies
+    sampling_strategy_1 = {
+        'x': cuqi.sampler.LinearRTO(),
+        'd': cuqi.sampler.Conjugate(),
+        's': cuqi.sampler.Conjugate()
+    }
+    sampling_strategy_2 = {
+        'x': cuqi.sampler.LinearRTO(),
+        'd': cuqi.sampler.Conjugate(),
+        's': cuqi.sampler.Conjugate()
+    }
+
+    # Define Gibbs samplers
+    sampler_1 = cuqi.sampler.HybridGibbs(posterior, sampling_strategy_1)
+    sampler_2 = cuqi.sampler.HybridGibbs(posterior, sampling_strategy_2)
+
+    # Run sampler with different online thinnning Nt
+    np.random.seed(0)
+    samples_1 = sampler_1.sample(20, Nt=5).get_samples()
+    np.random.seed(0)
+    samples_2 = sampler_2.sample(20).get_samples() # by default Nt=1
+
+    # Compare samples
+    assert np.allclose(samples_1['d'].samples[:, 0], samples_2['d'].samples[:, 4], rtol=1e-5)
+    assert np.allclose(samples_1['d'].samples[:, 1], samples_2['d'].samples[:, 9], rtol=1e-5)
+    assert np.allclose(samples_1['d'].samples[:, 2], samples_2['d'].samples[:, 14], rtol=1e-5)
+    assert np.allclose(samples_1['s'].samples[:, 0], samples_2['s'].samples[:, 4], rtol=1e-5)
+    assert np.allclose(samples_1['s'].samples[:, 1], samples_2['s'].samples[:, 9], rtol=1e-5)
+    assert np.allclose(samples_1['s'].samples[:, 2], samples_2['s'].samples[:, 14], rtol=1e-5)
+    assert np.allclose(samples_1['x'].samples[:, 0], samples_2['x'].samples[:, 4], rtol=1e-5)
+    assert np.allclose(samples_1['x'].samples[:, 1], samples_2['x'].samples[:, 9], rtol=1e-5)
+    assert np.allclose(samples_1['x'].samples[:, 2], samples_2['x'].samples[:, 14], rtol=1e-5)
+
+
 @pytest.mark.parametrize("step_size", [None, 0.1])
 @pytest.mark.parametrize("num_sampling_steps_x", [1, 5])
 @pytest.mark.parametrize("nb", [5, 20])