diff --git a/cuqi/distribution/_joint_distribution.py b/cuqi/distribution/_joint_distribution.py index 682e962735..e952138108 100644 --- a/cuqi/distribution/_joint_distribution.py +++ b/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): diff --git a/cuqi/experimental/mcmc/_cwmh.py b/cuqi/experimental/mcmc/_cwmh.py index 84ee1e1068..ae48f3748b 100644 --- a/cuqi/experimental/mcmc/_cwmh.py +++ b/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") diff --git a/cuqi/experimental/mcmc/_sampler.py b/cuqi/experimental/mcmc/_sampler.py index 90e11f2eeb..be110235bb 100644 --- a/cuqi/experimental/mcmc/_sampler.py +++ b/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. """