Fix ProposalBasedSampler for JointDistribution

cuqi/distribution/_joint_distribution.py

@@ -5,6 +5,7 @@
 from cuqi.distribution import Distribution, Posterior
 from cuqi.likelihood import Likelihood
 from cuqi.geometry import Geometry, _DefaultGeometry1D
+from cuqi.samples import Samples, JointSamples
 import cuqi
 import numpy as np # for splitting array. Can avoid.
 
@@ -311,16 +312,7 @@ def geometry(self):
 
     def logd(self, stacked_input):
         """ Return the un-normalized log density function stacked joint density. """
-
-        # Split the stacked input into individual inputs and call superclass
-        split_indices = np.cumsum(super().dim)  # list(accumulate(super().dim))
-        inputs = np.split(stacked_input, split_indices[:-1])
-        names = self.get_parameter_names()
-
-        # Create keyword arguments
-        kwargs = dict(zip(names, inputs))
-
-        return super().logd(**kwargs)
+        return super().logd(**self._unstack_input(stacked_input))
 
     def logpdf(self, stacked_input):
         return self.logd(stacked_input)
@@ -330,6 +322,24 @@ def _sample(self, Ns=1):
 
     def __repr__(self):
         return "_Stacked"+super().__repr__()
+
+    def _unstack_input(self, stacked_input):
+        # Split the stacked input into individual inputs and call superclass
+        split_indices = np.cumsum(super().dim)  # list(accumulate(super().dim))
+        inputs = np.split(stacked_input, split_indices[:-1])
+        names = self.get_parameter_names()
+
+        # Create keyword arguments
+        return dict(zip(names, inputs))
+
+    def _unstack_samples(self, stacked_samples):
+        split_indices = np.cumsum(super().dim)
+        split_samples = np.split(stacked_samples.samples, split_indices[:-1], axis = 0)
+        names = self.get_parameter_names()
+        geometries = [dist.geometry for dist in self._distributions]
+
+        unstacked_samples = [Samples(samples, geometry) for samples, geometry in zip(split_samples, geometries)]
+        return JointSamples(zip(names, unstacked_samples))
 
 
 class MultipleLikelihoodPosterior(JointDistribution, Distribution):

cuqi/experimental/mcmc/_cwmh.py

@@ -93,10 +93,6 @@ def scale(self, value):
         self._scale = value
 
     def validate_target(self):
-        if not isinstance(self.target, cuqi.density.Density):
-            raise ValueError(
-                "Target should be an instance of "+\
-                f"{cuqi.density.Density.__class__.__name__}")
         # Fail when there is no log density, which is currently assumed to be the case in case NaN is returned.
         if np.isnan(self.target.logd(self._get_default_initial_point(self.dim))):
             raise ValueError("Target does not have valid logd")

cuqi/experimental/mcmc/_sampler.py

@@ -388,6 +388,8 @@ def _ensure_initialized(self):
 
     def _get_default_initial_point(self, dim):
         """ Return the default initial point for the sampler. Defaults to an array of ones. """
+        if isinstance(dim, list):
+            return [np.ones(d) for d in dim]
         return np.ones(dim)
 
     def __repr__(self):
@@ -475,6 +477,26 @@ def initialize(self):
 
         self._is_initialized = True
 
+    @Sampler.target.setter
+    def target(self, value):
+        """ Set the target density. Runs validation of the target. """
+        self._target = value
+
+        if isinstance(self._target, cuqi.distribution.JointDistribution):
+            self._target = self._target._as_stacked()
+
+        if self._target is not None:
+            self.validate_target()
+
+    def get_samples(self) -> Samples:
+        """ Return the samples. The internal data-structure for the samples is a dynamic list so this creates a copy. """
+
+        samples = Samples(np.array(self._samples).T, self.target.geometry)
+        if isinstance(self.target, cuqi.distribution.JointDistribution):
+            return self.target._unstack_samples(samples)
+        return samples
+
+
     @abstractmethod
     def validate_proposal(self):
         """ Validate the proposal distribution. """