def test_as_shape_int(shape_arg, ndim):
if ndim:
shape = pnut.as_shape(shape_arg, ndim=1)
else:
shape = pnut.as_shape(shape_arg)
assert isinstance(shape, tuple)
assert len(shape) == 1
assert all(isinstance(entry, int) for entry in shape)
assert shape[0] == shape_arg
def test_as_shape_iterable(shape_arg, ndim):
if ndim:
shape = pnut.as_shape(shape_arg, ndim=len(shape_arg))
else:
shape = pnut.as_shape(shape_arg)
assert isinstance(shape, tuple)
assert len(shape) == len(shape_arg)
assert all(isinstance(entry, int) for entry in shape)
assert all(entry_shape == entry_shape_arg
for entry_shape, entry_shape_arg in zip(shape_arg, shape))
def __init__(
self,
input_shape: ShapeLike,
output_shape: ShapeLike = (),
):
self._input_shape = _pn_utils.as_shape(input_shape)
self._input_ndim = len(self._input_shape)
if self._input_ndim > 1:
raise ValueError(
"Currently, we only support kernels with at most 1 input dimension."
)
self._output_shape = _pn_utils.as_shape(output_shape)
self._output_ndim = len(self._output_shape)
def _get_shape_kernelmatrix(
self,
x0_shape: ShapeArgType,
x1_shape: ShapeArgType,
) -> ShapeType:
"""Determine the shape of the kernel matrix based on the given arguments.
Determine the correct shape of the covariance function evaluated at the given
input arguments. If inputs are vectors the output is a numpy scalar if the
output dimension of the kernel is 1, otherwise *shape=(output_dim,
output_dim)*. If inputs represent multiple observations, then the resulting
matrix has *shape=(n0, n1) or (n0, n1, output_dim, output_dim)*.
Parameters
----------
x0_shape :
Shape of the first input to the covariance function.
x1_shape :
Shape of the second input to the covariance function.
"""
if len(x0_shape) <= 1 and len(x1_shape) <= 1:
if self.output_dim == 1:
kern_shape = 0
else:
kern_shape = ()
else:
kern_shape = (x0_shape[0], x1_shape[0])
if self.output_dim > 1:
kern_shape += (
self.output_dim,
self.output_dim,
)
return _utils.as_shape(kern_shape)
def sample(self, locations=None, size=()):
# In the present setting, only works for sampling from the smoothing posterior.
size = utils.as_shape(size)
if locations is None:
locations = self.locations
random_vars = self.filtering_posterior.state_rvs
else:
random_vars = self.filtering_posterior(locations)
# Inform the final point in the list about all the data by
# conditioning on the final state rv
if locations[-1] < self.locations[-1]:
final_sample = self.state_rvs[-1].sample()
random_vars[-1], _ = self.transition.backward_realization(
final_sample,
random_vars[-1],
t=locations[-1],
dt=self.locations[-1] - locations[-1],
)
if size == ():
return np.array(
self.transition.jointly_sample_list_backward(
locations=locations, rv_list=random_vars
)
)
return np.array(
[self.sample(locations=locations, size=size[1:]) for _ in range(size[0])]
)
def test_sampling_shapes_1d(locs, size):
"""Make the sampling tests for a 1d posterior."""
locations = np.linspace(0, 2 * np.pi, 100)
data = 0.5 * np.random.randn(100) + np.sin(locations)
prior = statespace.IBM(0, 1)
measmod = statespace.DiscreteLTIGaussian(state_trans_mat=np.eye(1),
shift_vec=np.zeros(1),
proc_noise_cov_mat=np.eye(1))
initrv = randvars.Normal(np.zeros(1), np.eye(1))
kalman = filtsmooth.Kalman(prior, measmod, initrv)
regression_problem = problems.RegressionProblem(observations=data,
locations=locations)
posterior, _ = kalman.filtsmooth(regression_problem)
size = utils.as_shape(size)
if locs is None:
base_measure_reals = np.random.randn(
*(size + posterior.locations.shape + (1, )))
samples = posterior.transform_base_measure_realizations(
base_measure_reals, t=posterior.locations)
else:
locs = np.union1d(locs, posterior.locations)
base_measure_reals = np.random.randn(*(size + (len(locs), )) + (1, ))
samples = posterior.transform_base_measure_realizations(
base_measure_reals, t=locs)
assert samples.shape == base_measure_reals.shape
def reshape(self, newshape: ShapeArgType) -> "RandomVariable":
"""Give a new shape to a random variable.
Parameters
----------
newshape :
New shape for the random variable. It must be compatible with the original
shape.
"""
newshape = _utils.as_shape(newshape)
return RandomVariable(
shape=newshape,
dtype=self.dtype,
random_state=_utils.derive_random_seed(self.random_state),
sample=lambda size: self.sample(size).reshape(size + newshape),
mode=lambda: self.mode.reshape(newshape),
median=lambda: self.median.reshape(newshape),
mean=lambda: self.mean.reshape(newshape),
cov=lambda: self.cov,
var=lambda: self.var.reshape(newshape),
std=lambda: self.std.reshape(newshape),
entropy=lambda: self.entropy,
as_value_type=self.__as_value_type,
)
def _sample(self, size: ShapeArgType = ()) -> _ValueType:
size = _utils.as_shape(size)
if size == ():
return self._support.copy()
else:
return np.tile(self._support, reps=size + (1, ) * self.ndim)
def _sample_at_input(
self,
rng: np.random.Generator,
args: InputType,
size: ShapeLike = (),
) -> OutputType:
size = utils.as_shape(size)
args = np.atleast_1d(args)
if args.ndim > 1:
raise ValueError(f"Invalid args shape {args.shape}")
base_measure_realizations = scipy.stats.norm.rvs(
size=(size + args.shape + self.initrv.shape), random_state=rng)
if size == ():
return np.array(
self.transition.
jointly_transform_base_measure_realization_list_forward(
base_measure_realizations=base_measure_realizations,
t=args,
initrv=self.initrv,
_diffusion_list=np.ones_like(args[:-1]),
))
return np.stack([
self.transition.
jointly_transform_base_measure_realization_list_forward(
base_measure_realizations=base_real,
t=args,
initrv=self.initrv,
_diffusion_list=np.ones_like(args[:-1]),
) for base_real in base_measure_realizations
])
def test_sampling_shapes_1d(locs, size):
"""Make the sampling tests for a 1d posterior."""
locations = np.linspace(0, 2 * np.pi, 100)
data = 0.5 * np.random.randn(100) + np.sin(locations)
prior = randprocs.markov.integrator.IntegratedWienerTransition(0, 1)
measmod = randprocs.markov.discrete.LTIGaussian(
state_trans_mat=np.eye(1),
shift_vec=np.zeros(1),
proc_noise_cov_mat=np.eye(1))
initrv = randvars.Normal(np.zeros(1), np.eye(1))
prior_process = randprocs.markov.MarkovProcess(transition=prior,
initrv=initrv,
initarg=locations[0])
kalman = filtsmooth.gaussian.Kalman(prior_process)
regression_problem = problems.TimeSeriesRegressionProblem(
observations=data, measurement_models=measmod, locations=locations)
posterior, _ = kalman.filtsmooth(regression_problem)
size = utils.as_shape(size)
if locs is None:
base_measure_reals = np.random.randn(
*(size + posterior.locations.shape + (1, )))
samples = posterior.transform_base_measure_realizations(
base_measure_reals, t=posterior.locations)
else:
locs = np.union1d(locs, posterior.locations)
base_measure_reals = np.random.randn(*(size + (len(locs), )) + (1, ))
samples = posterior.transform_base_measure_realizations(
base_measure_reals, t=locs)
assert samples.shape == base_measure_reals.shape
def __init__(
self,
input_dim: IntArgType,
shape: ShapeArgType = (),
):
self._input_dim = int(input_dim)
self._shape = _pn_utils.as_shape(shape)
def _sample(self, rng: np.random.Generator,
size: ShapeLike = ()) -> ValueType:
size = _utils.as_shape(size)
if size == ():
return self._support.copy()
return np.tile(self._support, reps=size + (1, ) * self.ndim)
def sample(self, size: ShapeArgType = ()) -> _ValueType:
"""Draw realizations from a random variable.
Parameters
----------
size :
Size of the drawn sample of realizations.
"""
if self.__sample is None:
raise NotImplementedError("No sampling method provided.")
return self.__sample(_utils.as_shape(size))
def sample(self, t=None, size=()):
# this has its own recursion because of the tedious undoing of preconditioning....
size = utils.as_shape(size)
# implement only single samples, rest via recursion
if size != ():
return np.array(
[self.sample(t=t, size=size[1:]) for _ in range(size[0])])
samples = self._kalman_posterior.sample(locations=t, size=size)
return np.array(self._project_rv_list(samples))
def __init__(
self,
shape: ShapeArgType,
dtype: DTypeArgType,
random_state: RandomStateArgType = None,
parameters: Optional[Dict[str, Any]] = None,
sample: Optional[Callable[[ShapeType], _ValueType]] = None,
in_support: Optional[Callable[[_ValueType], bool]] = None,
cdf: Optional[Callable[[_ValueType], np.float_]] = None,
logcdf: Optional[Callable[[_ValueType], np.float_]] = None,
quantile: Optional[Callable[[FloatArgType], _ValueType]] = None,
mode: Optional[Callable[[], _ValueType]] = None,
median: Optional[Callable[[], _ValueType]] = None,
mean: Optional[Callable[[], _ValueType]] = None,
cov: Optional[Callable[[], _ValueType]] = None,
var: Optional[Callable[[], _ValueType]] = None,
std: Optional[Callable[[], _ValueType]] = None,
entropy: Optional[Callable[[], np.float_]] = None,
as_value_type: Optional[Callable[[Any], _ValueType]] = None,
):
# pylint: disable=too-many-arguments,too-many-locals
"""Create a new random variable."""
self.__shape = _utils.as_shape(shape)
# Data Types
self.__dtype = np.dtype(dtype)
self.__median_dtype = RandomVariable.infer_median_dtype(self.__dtype)
self.__moment_dtype = RandomVariable.infer_moment_dtype(self.__dtype)
self._random_state = _utils.as_random_state(random_state)
# Probability distribution of the random variable
self.__parameters = parameters.copy() if parameters is not None else {}
self.__sample = sample
self.__in_support = in_support
self.__cdf = cdf
self.__logcdf = logcdf
self.__quantile = quantile
# Properties of the random variable
self.__mode = mode
self.__median = median
self.__mean = mean
self.__cov = cov
self.__var = var
self.__std = std
self.__entropy = entropy
# Utilities
self.__as_value_type = as_value_type
def sample(self, locations=None, size=()):
size = utils.as_shape(size)
if locations is None:
locations = self.locations
random_vars = self.state_rvs
else:
random_vars = self.__call__(locations)
if size == ():
return self._single_sample_path(
locations=locations, random_vars=random_vars
)
return np.array(
[self.sample(locations=locations, size=size[1:]) for _ in range(size[0])]
)
def sample(self, size: ShapeArgType = ()) -> _ValueType:
"""
Draw realizations from a random variable.
Parameters
----------
size : tuple
Size of the drawn sample of realizations.
Returns
-------
sample : array-like
Sample of realizations with the given ``size`` and the inherent ``shape``.
"""
if self.__sample is None:
raise NotImplementedError("No sampling method provided.")
return self.__sample(size=_utils.as_shape(size))
def sample(self, rng: np.random.Generator, size: ShapeArgType = ()) -> _ValueType:
"""Draw realizations from a random variable.
Parameters
----------
rng
Random number generator used for sampling.
size
Size of the drawn sample of realizations.
"""
if self.__sample is None:
raise NotImplementedError("No sampling method provided.")
if not isinstance(rng, np.random.Generator):
msg = "Random number generators must be of type np.random.Generator."
raise TypeError(msg)
return self.__sample(rng=rng, size=_utils.as_shape(size))
def __init__(
self,
input_shape: ShapeLike,
lengthscales: Union[np.ndarray, ScalarLike],
nus: Union[np.ndarray, ScalarLike],
):
input_shape = _utils.as_shape(input_shape)
if input_shape == () and not (np.isscalar(lengthscales)
and np.isscalar(nus)):
raise ValueError(
f"'lengthscales' and 'nus' must be scalar if 'input_shape' is "
f"{input_shape}.")
input_dim = 1 if input_shape == () else input_shape[0]
# If only single scalar lengthcsale or nu is given, use this in every dimension
def expand_array(x, ndim):
return np.full((ndim, ), _utils.as_numpy_scalar(x))
if isinstance(lengthscales, np.ndarray):
if lengthscales.shape == ():
lengthscales = expand_array(lengthscales, input_dim)
if isinstance(nus, np.ndarray):
if nus.shape == ():
nus = expand_array(nus, input_dim)
# also expand if scalars are given
if np.isscalar(lengthscales):
lengthscales = expand_array(lengthscales, input_dim)
if np.isscalar(nus):
nus = expand_array(nus, input_dim)
univariate_materns = []
for dim in range(input_dim):
univariate_materns.append(
Matern(input_shape=(),
lengthscale=lengthscales[dim],
nu=nus[dim]))
self.univariate_materns = univariate_materns
self.nus = nus
self.lengthscales = lengthscales
super().__init__(input_shape=input_shape)
def sample(self, locations=None, size=()):
# In the present setting, only works for sampling from the smoothing posterior.
size = utils.as_shape(size)
if locations is None:
locations = self.locations
random_vars = self.filtering_posterior.state_rvs
else:
random_vars = self.filtering_posterior(locations)
if size == ():
return np.array(
self.transition.jointly_sample_list_backward(
locations=locations, rv_list=random_vars))
return np.array([
self.sample(locations=locations, size=size[1:])
for _ in range(size[0])
])
def sample(
self,
rng: np.random.Generator,
t: Optional[ArrayLike] = None,
size: Optional[ShapeLike] = (),
) -> np.ndarray:
size = utils.as_shape(size)
single_rv_shape = self.states[0].shape
single_rv_ndim = self.states[0].ndim
# Early exit if no dense output is required
if t is None:
base_measure_realizations = stats.norm.rvs(
size=(size + self.locations.shape + single_rv_shape),
random_state=rng,
)
return self.transform_base_measure_realizations(
base_measure_realizations=base_measure_realizations,
t=self.locations)
# Compute the union (as sets) of t and self.locations
# This allows that samples "always pass" the grid points.
all_locations = np.union1d(t, self.locations)
slice_these_out = np.where(np.isin(all_locations, t))[0]
base_measure_realizations = stats.norm.rvs(
size=(size + all_locations.shape + single_rv_shape),
random_state=rng,
)
samples = self.transform_base_measure_realizations(
base_measure_realizations=base_measure_realizations,
t=all_locations)
new_samples = np.take(samples,
indices=slice_these_out,
axis=-(single_rv_ndim + 1))
return new_samples
def test_finite_evaluation_is_normal(gaussian_process: randprocs.GaussianProcess):
"""A Gaussian process evaluated at a finite set of inputs is a Gaussian random
variable."""
x = np.random.normal(size=(5,) + utils.as_shape(gaussian_process.input_dim))
assert isinstance(gaussian_process(x), randvars.Normal)
def test_sample(categ, size, rng):
samples = categ.sample(rng=rng, size=size)
expected_shape = utils.as_shape(size) + categ.shape
assert samples.shape == expected_shape
def test_as_shape_wrong_type(shape_arg):
with pytest.raises(TypeError):
pnut.as_shape(shape_arg)
def test_as_shape_wrong_ndim(shape_arg, ndim):
with pytest.raises(TypeError):
pnut.as_shape(shape_arg, ndim=ndim)