Merge pull request #660 from jeverink/Fix-HybridGibbs-ordering

Fix hybrid gibbs ordering

Author: amal-ghamdi

cuqi/experimental/mcmc/_gibbs.py

@@ -42,6 +42,10 @@ class HybridGibbs:
     fully stateful at this point. This means samplers like NUTS will lose
     their internal state between Gibbs steps.
 
+    The order in which the conditionals are sampled is the order of the
+    variables in the sampling strategy, unless a different sampling order
+    is specified by the parameter `scan_order`
+
     Parameters
     ----------
     target : cuqi.distribution.JointDistribution
@@ -58,6 +62,11 @@ class HybridGibbs:
         will call its step method in each Gibbs step.
         Default is 1 for all variables.
 
+    scan_order : list or str, *optional*
+        Order in which the conditional distributions are sampled.
+        If set to "random", use a random ordering at each step.
+        If not specified, it will be the order in the sampling_strategy.
+
     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)`.
@@ -107,7 +116,7 @@ class HybridGibbs:
             
     """
 
-    def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampler], num_sampling_steps: Dict[str, int] = None, callback=None):
+    def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampler], num_sampling_steps: Dict[str, int] = None, scan_order = None, callback=None):
 
         # Store target and allow conditioning to reduce to a single density
         self.target = target() # Create a copy of target distribution (to avoid modifying the original)
@@ -121,6 +130,13 @@ def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampl
         # Store parameter names
         self.par_names = self.target.get_parameter_names()
 
+        # Store the scan order
+        self._scan_order = scan_order
+
+        # Check that the parameters of the target align with the sampling_strategy and scan_order
+        if set(self.par_names) != set(self.scan_order):
+            raise ValueError("Parameter names in JointDistribution do not equal the names in the scan order.")
+
         # Initialize sampler (after target is set)
         self._initialize()
 
@@ -148,6 +164,16 @@ def _initialize(self):
         # Validate all targets for samplers.
         self.validate_targets()
 
+    @property
+    def scan_order(self):
+        if self._scan_order is None:
+            return list(self.samplers.keys())
+        if self._scan_order == "random":
+            arr = list(self.samplers.keys())
+            np.random.shuffle(arr) # Shuffle works in-place
+            return arr
+        return self._scan_order
+
     # ------------ Public methods ------------
     def validate_targets(self):
         """ Validate each of the conditional targets used in the Gibbs steps """
@@ -217,7 +243,7 @@ def step(self):
         """ Sequentially go through all parameters and sample them conditionally on each other """
 
         # Sample from each conditional distribution
-        for par_name in self.par_names:
+        for par_name in self.scan_order:
 
             # Set target for current parameter
             self._set_target(par_name)

tests/test_implicit_priors.py

@@ -210,8 +210,8 @@ def test_regression_increasing():
     posterior = joint(y=y_obs)
 
     sampling_strategy = {
+        'd': cuqi.experimental.mcmc.Conjugate(),
         'x': cuqi.experimental.mcmc.RegularizedLinearRTO(maxit=50, penalty_parameter=20, adaptive = False),
-        'd': cuqi.experimental.mcmc.Conjugate()
                         }
     sampler = cuqi.experimental.mcmc.HybridGibbs(posterior, sampling_strategy)
 
@@ -241,8 +241,8 @@ def test_regression_convex():
     posterior = joint(y=y_obs)
 
     sampling_strategy = {
-        'x': cuqi.experimental.mcmc.RegularizedLinearRTO(maxit=50, penalty_parameter=20, adaptive = False),
-        'd': cuqi.experimental.mcmc.Conjugate()
+        'd': cuqi.experimental.mcmc.Conjugate(),
+        'x': cuqi.experimental.mcmc.RegularizedLinearRTO(maxit=50, penalty_parameter=20, adaptive = False)
                         }
     sampler = cuqi.experimental.mcmc.HybridGibbs(posterior, sampling_strategy)
 

tests/zexperimental/test_mcmc.py

@@ -1574,4 +1574,26 @@ def test_all_samplers_that_should_be_tested_for_callback_are_in_the_tested_list(
     for cls in callback_testing_sampler_classes:
         assert cls in tested_classes, f"Sampler {cls} is not tested for callback."
 
-# ============= End testing sampler callback =============
+# ============= End testing sampler callback =============
+def test_gibbs_random_scan_order():
+    target = HybridGibbs_target_1()
+    sampling_strategy={
+                "x": cuqi.experimental.mcmc.LinearRTO(),
+                "s": cuqi.experimental.mcmc.Conjugate(),
+            }
+    
+    sampler = cuqi.experimental.mcmc.HybridGibbs(target, sampling_strategy, scan_order='random')
+    np.random.seed(0)
+    scan_order1 = sampler.scan_order
+    scan_order2 = sampler.scan_order
+    assert scan_order1 != scan_order2
+
+def test_gibbs_scan_order():
+    target = HybridGibbs_target_1()
+    sampling_strategy={
+                "x": cuqi.experimental.mcmc.LinearRTO(),
+                "s": cuqi.experimental.mcmc.Conjugate(),
+            }
+    
+    sampler = cuqi.experimental.mcmc.HybridGibbs(target, sampling_strategy, scan_order=['x', 's'])
+    assert sampler.scan_order == ['x', 's']