def __init__(self, nu=1.5, dt=1.0, obs_dim=2, num_gps=1,
length_scale_init=None, kernel_scale_init=None,
obs_noise_scale_init=None):
self.nu = nu
self.dt = dt
assert obs_dim > 1, "If obs_dim==1 you should use IndependentMaternGP"
self.obs_dim = obs_dim
self.num_gps = num_gps
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim,)
self.dt = dt
self.obs_dim = obs_dim
self.num_gps = num_gps
super().__init__()
self.kernel = MaternKernel(nu=nu, num_gps=num_gps,
length_scale_init=length_scale_init,
kernel_scale_init=kernel_scale_init)
self.full_state_dim = num_gps * self.kernel.state_dim
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.A = nn.Parameter(0.3 * torch.randn(self.num_gps, self.obs_dim))
def __init__(
self,
nu=1.5,
dt=1.0,
obs_dim=1,
length_scale_init=None,
kernel_scale_init=None,
obs_noise_scale_init=None,
):
self.nu = nu
self.dt = dt
self.obs_dim = obs_dim
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim, )
super().__init__()
self.kernel = MaternKernel(
nu=nu,
num_gps=obs_dim,
length_scale_init=length_scale_init,
kernel_scale_init=kernel_scale_init,
)
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
obs_matrix = [1.0] + [0.0] * (self.kernel.state_dim - 1)
self.register_buffer("obs_matrix",
torch.tensor(obs_matrix).unsqueeze(-1))
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
self._event_dims = {}
self._cond_indep_stacks = {}
self.locs = PyroModule()
self.scales = PyroModule()
# Initialize guide params
for name, site in self.prototype_trace.iter_stochastic_nodes():
# Collect unconstrained event_dims, which may differ from constrained event_dims.
with helpful_support_errors(site):
init_loc = biject_to(site["fn"].support).inv(site["value"].detach()).detach()
event_dim = site["fn"].event_dim + init_loc.dim() - site["value"].dim()
self._event_dims[name] = event_dim
# Collect independence contexts.
self._cond_indep_stacks[name] = site["cond_indep_stack"]
# If subsampling, repeat init_value to full size.
for frame in site["cond_indep_stack"]:
full_size = getattr(frame, "full_size", frame.size)
if full_size != frame.size:
dim = frame.dim - event_dim
init_loc = periodic_repeat(init_loc, full_size, dim).contiguous()
init_scale = torch.full_like(init_loc, self._init_scale)
_deep_setattr(self.locs, name, PyroParam(init_loc, constraints.real, event_dim))
_deep_setattr(self.scales, name,
PyroParam(init_scale, self.scale_constraint, event_dim))
def __init__(
self,
n_input: int,
n_conditions: int,
lam_scale: float,
bias_scale: float,
alpha: float = 1.0,
):
super().__init__()
self.n_input = n_input
self.n_conditions = n_conditions
# weight on monotonic constraint
self.register_buffer("alpha", torch.as_tensor(alpha))
# scale of priors on weights and bias
self.register_buffer("lam_p_scale", torch.as_tensor(lam_scale))
self.register_buffer("bias_p_scale", torch.as_tensor(bias_scale))
# parameters for guide
self.weight_loc = nn.Parameter(
torch.nn.init.normal_(torch.Tensor(n_input)))
self.weight_scale = PyroParam(torch.full((n_input, ), lam_scale),
constraint=constraints.positive)
self.bias_loc = nn.Parameter(
torch.nn.init.normal_(torch.Tensor(n_conditions)))
self.bias_scale = PyroParam(torch.full((n_conditions, ), bias_scale),
constraint=constraints.positive)
def __init__(self,
nu=1.5,
num_gps=1,
length_scale_init=None,
kernel_scale_init=None):
if nu not in [0.5, 1.5, 2.5]:
raise NotImplementedError(
"The only supported values of nu are 0.5, 1.5 and 2.5")
self.nu = nu
self.state_dim = {0.5: 1, 1.5: 2, 2.5: 3}[nu]
self.num_gps = num_gps
if length_scale_init is None:
length_scale_init = torch.ones(num_gps)
assert length_scale_init.shape == (num_gps, )
if kernel_scale_init is None:
kernel_scale_init = torch.ones(num_gps)
assert kernel_scale_init.shape == (num_gps, )
super().__init__()
self.length_scale = PyroParam(length_scale_init,
constraint=constraints.positive)
self.kernel_scale = PyroParam(kernel_scale_init,
constraint=constraints.positive)
if self.state_dim > 1:
for x in range(self.state_dim):
for y in range(self.state_dim):
mask = torch.zeros(self.state_dim, self.state_dim)
mask[x, y] = 1.0
self.register_buffer("mask{}{}".format(x, y), mask)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
# Initialize guide params
self.loc = nn.Parameter(self._init_loc())
self.scale_tril = PyroParam(
eye_like(self.loc, self.latent_dim) * self._init_scale,
constraints.lower_cholesky)
def init_mvn_guide(self):
""" Initialize multivariate normal guide
"""
init_loc = torch.full((self.n_params, ), 0.0)
init_scale = eye_like(init_loc, self.n_params) * 0.1
_deep_setattr(self, "mvn.loc", PyroParam(init_loc, constraints.real))
_deep_setattr(self, "mvn.scale_tril",
PyroParam(init_scale, constraints.lower_cholesky))
class PartialMultivariateNormalSamplingGroup(LocatedSamplingGroupWithPrior):
def __init__(self,
sites,
name='',
diag=_nomatch,
init_scale_full: tp.Union[torch.Tensor, float] = 1.,
init_scale_diag: tp.Union[torch.Tensor, float] = 1.,
*args,
**kwargs):
self.diag_pattern = re.compile(diag)
self.sites_full, self.sites_diag = ({
site['name']: site
for site in _
} for _ in partition(lambda _: self.diag_pattern.match(_['name']),
sites))
super().__init__(
dict_union(self.sites_full, self.sites_diag).values(), name, *args,
**kwargs)
self.size_full, self.size_diag = (sum(self.sizes[site] for site in _)
for _ in (self.sites_full,
self.sites_diag))
jac = self.jacobian(self.loc)
self.scale_full = PyroParam(self._scale_matrix(init_scale_full,
jac[:self.size_full]),
event_dim=2,
constraint=constraints.lower_cholesky)
self.scale_cross = PyroParam(self.loc.new_zeros(
torch.Size((self.size_diag, self.size_full))),
event_dim=2)
self.scale_diag = PyroParam(self._scale_diagonal(
init_scale_diag, jac[self.size_full:]),
event_dim=1,
constraint=constraints.positive)
self.guide_z_aux = PyroSample(
dist.Normal(self.loc.new_zeros(()),
1.).expand(self.event_shape).to_event(1))
@property
def half_log_det(self):
return self.scale_full.diagonal(
dim1=-2, dim2=-1).log().sum(-1) + self.scale_diag.log().sum(-1)
def prior(self):
z_aux = self.guide_z_aux
zfull, zdiag = z_aux[..., :self.size_full,
None], z_aux[..., self.size_full:]
return dist.Delta(self.loc + torch.cat(
((self.scale_full @ zfull).squeeze(-1),
(self.scale_cross @ zfull).squeeze(-1) + self.scale_diag * zdiag),
dim=-1),
log_density=-self.half_log_det,
event_dim=1)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
# Initialize guide params
for name, site in self.prototype_trace.iter_stochastic_nodes():
value = PyroParam(site["value"].detach(), constraint=site["fn"].support)
_deep_setattr(self, name, value)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
# Initialize guide params
self.loc = nn.Parameter(self._init_loc())
self.scale = PyroParam(
self.loc.new_full((self.latent_dim, ), self._init_scale),
constraints.positive)
def __init__(
self,
nu=1.5,
dt=1.0,
obs_dim=1,
linearly_coupled=False,
length_scale_init=None,
obs_noise_scale_init=None,
):
if nu != 1.5:
raise NotImplementedError("The only supported value of nu is 1.5")
self.dt = dt
self.obs_dim = obs_dim
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim, )
super().__init__()
self.kernel = MaternKernel(nu=nu,
num_gps=obs_dim,
length_scale_init=length_scale_init)
self.full_state_dim = self.kernel.state_dim * obs_dim
# we demote self.kernel.kernel_scale from being a nn.Parameter
# since the relevant scales are now encoded in the wiener noise matrix
del self.kernel.kernel_scale
self.kernel.register_buffer("kernel_scale", torch.ones(obs_dim))
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.wiener_noise_tril = PyroParam(
torch.eye(obs_dim) + 0.03 * torch.randn(obs_dim, obs_dim).tril(-1),
constraint=constraints.lower_cholesky,
)
if linearly_coupled:
self.obs_matrix = nn.Parameter(
0.3 * torch.randn(self.obs_dim, self.obs_dim))
else:
obs_matrix = torch.zeros(self.full_state_dim, obs_dim)
for i in range(obs_dim):
obs_matrix[self.kernel.state_dim * i, i] = 1.0
self.register_buffer("obs_matrix", obs_matrix)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
# Initialize guide params
self.loc = nn.Parameter(self._init_loc())
if self.rank is None:
self.rank = int(round(self.latent_dim ** 0.5))
self.scale = PyroParam(
self.loc.new_full((self.latent_dim,), 0.5 ** 0.5 * self._init_scale),
constraint=constraints.positive)
self.cov_factor = nn.Parameter(
self.loc.new_empty(self.latent_dim, self.rank).normal_(0, 1 / self.rank ** 0.5))
class AutoMultivariateNormal(AutoContinuous):
"""
This implementation of :class:`AutoContinuous` uses a Cholesky
factorization of a Multivariate Normal distribution to construct a guide
over the entire latent space. The guide does not depend on the model's
``*args, **kwargs``.
Usage::
guide = AutoMultivariateNormal(model)
svi = SVI(model, guide, ...)
By default the mean vector is initialized by ``init_loc_fn()`` and the
Cholesky factor is initialized to the identity times a small factor.
:param callable model: A generative model.
:param callable init_loc_fn: A per-site initialization function.
See :ref:`autoguide-initialization` section for available functions.
:param float init_scale: Initial scale for the standard deviation of each
(unconstrained transformed) latent variable.
"""
# TODO consider switching to constraints.softplus_lower_cholesky
# See https://github.com/pyro-ppl/numpyro/issues/855
scale_tril_constraint = constraints.lower_cholesky
def __init__(self, model, init_loc_fn=init_to_median, init_scale=0.1):
if not isinstance(init_scale, float) or not (init_scale > 0):
raise ValueError("Expected init_scale > 0. but got {}".format(init_scale))
self._init_scale = init_scale
super().__init__(model, init_loc_fn=init_loc_fn)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
# Initialize guide params
self.loc = nn.Parameter(self._init_loc())
self.scale_tril = PyroParam(eye_like(self.loc, self.latent_dim) * self._init_scale,
self.scale_tril_constraint)
def get_base_dist(self):
return dist.Normal(torch.zeros_like(self.loc), torch.zeros_like(self.loc)).to_event(1)
def get_transform(self, *args, **kwargs):
return dist.transforms.LowerCholeskyAffine(self.loc, scale_tril=self.scale_tril)
def get_posterior(self, *args, **kwargs):
"""
Returns a MultivariateNormal posterior distribution.
"""
return dist.MultivariateNormal(self.loc, scale_tril=self.scale_tril)
def _loc_scale(self, *args, **kwargs):
return self.loc, self.scale_tril.diag()
def __init__(self,
sites,
name='',
init_scale: tp.Union[torch.Tensor, float] = 1.,
*args,
**kwargs):
super().__init__(sites, name, *args, **kwargs)
self.scale = PyroParam(self._scale_diagonal(init_scale,
self.jacobian(self.loc)),
event_dim=1,
constraint=constraints.positive)
def __init__(self,
sites,
name='',
diag=_nomatch,
init_scale_full: tp.Union[torch.Tensor, float] = 1.,
init_scale_diag: tp.Union[torch.Tensor, float] = 1.,
*args,
**kwargs):
self.diag_pattern = re.compile(diag)
self.sites_full, self.sites_diag = ({
site['name']: site
for site in _
} for _ in partition(lambda _: self.diag_pattern.match(_['name']),
sites))
super().__init__(
dict_union(self.sites_full, self.sites_diag).values(), name, *args,
**kwargs)
self.size_full, self.size_diag = (sum(self.sizes[site] for site in _)
for _ in (self.sites_full,
self.sites_diag))
jac = self.jacobian(self.loc)
self.scale_full = PyroParam(self._scale_matrix(init_scale_full,
jac[:self.size_full]),
event_dim=2,
constraint=constraints.lower_cholesky)
self.scale_cross = PyroParam(self.loc.new_zeros(
torch.Size((self.size_diag, self.size_full))),
event_dim=2)
self.scale_diag = PyroParam(self._scale_diagonal(
init_scale_diag, jac[self.size_full:]),
event_dim=1,
constraint=constraints.positive)
self.guide_z_aux = PyroSample(
dist.Normal(self.loc.new_zeros(()),
1.).expand(self.event_shape).to_event(1))
def __init__(self,
obs_dim=1,
state_dim=2,
obs_noise_scale_init=None,
learnable_observation_loc=False):
self.obs_dim = obs_dim
self.state_dim = state_dim
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim, )
super().__init__()
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.trans_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
self.trans_matrix = nn.Parameter(
torch.eye(state_dim) + 0.03 * torch.randn(state_dim, state_dim))
self.obs_matrix = nn.Parameter(0.3 * torch.randn(state_dim, obs_dim))
self.init_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
if learnable_observation_loc:
self.obs_loc = nn.Parameter(torch.zeros(obs_dim))
else:
self.register_buffer('obs_loc', torch.zeros(obs_dim))
def __init__(self,
sites,
name='',
init_scale: tp.Union[torch.Tensor, float] = 1.,
*args,
**kwargs):
super().__init__(sites, name, *args, **kwargs)
self.scale_tril = PyroParam(self._scale_matrix(init_scale,
self.jacobian(
self.loc)),
event_dim=2,
constraint=constraints.lower_cholesky)
def _setup_prototype(self, *args, **kwargs):
# run the model so we can inspect its structure
model = config_enumerate(self.model)
self.prototype_trace = poutine.block(poutine.trace(model).get_trace)(
*args, **kwargs)
self.prototype_trace = prune_subsample_sites(self.prototype_trace)
if self.master is not None:
self.master()._check_prototype(self.prototype_trace)
self._discrete_sites = []
self._cond_indep_stacks = {}
self._plates = {}
for name, site in self.prototype_trace.iter_stochastic_nodes():
if site["infer"].get("enumerate") != "parallel":
raise NotImplementedError(
'Expected sample site "{}" to be discrete and '
'configured for parallel enumeration'.format(name))
# collect discrete sample sites
fn = site["fn"]
Dist = type(fn)
if Dist in (dist.Bernoulli, dist.Categorical,
dist.OneHotCategorical):
params = [("probs", fn.probs.detach().clone(),
fn.arg_constraints["probs"])]
else:
raise NotImplementedError("{} is not supported".format(
Dist.__name__))
self._discrete_sites.append((site, Dist, params))
# collect independence contexts
self._cond_indep_stacks[name] = site["cond_indep_stack"]
for frame in site["cond_indep_stack"]:
if frame.vectorized:
self._plates[frame.name] = frame
else:
raise NotImplementedError(
"AutoDiscreteParallel does not support sequential pyro.plate"
)
# Initialize guide params
for site, Dist, param_spec in self._discrete_sites:
name = site["name"]
for param_name, param_init, param_constraint in param_spec:
_deep_setattr(
self, "{}_{}".format(name, param_name),
PyroParam(param_init, constraint=param_constraint))
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
# Initialize guide params
for name, site in self.prototype_trace.iter_stochastic_nodes():
value = site["value"].detach()
event_dim = site["fn"].event_dim
# If subsampling, repeat init_value to full size.
for frame in site["cond_indep_stack"]:
full_size = getattr(frame, "full_size", frame.size)
if full_size != frame.size:
dim = frame.dim - event_dim
value = periodic_repeat(value, full_size, dim).contiguous()
value = PyroParam(value, site["fn"].support, event_dim)
_deep_setattr(self, name, value)
def __init__(
self,
obs_dim=1,
state_dim=2,
nu=1.5,
obs_noise_scale_init=None,
length_scale_init=None,
kernel_scale_init=None,
learnable_observation_loc=False,
):
self.obs_dim = obs_dim
self.state_dim = state_dim
self.nu = nu
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim, )
super().__init__()
self.kernel = MaternKernel(
nu=nu,
num_gps=obs_dim,
length_scale_init=length_scale_init,
kernel_scale_init=kernel_scale_init,
)
self.dt = 1.0
self.full_state_dim = self.kernel.state_dim * obs_dim + state_dim
self.full_gp_state_dim = self.kernel.state_dim * obs_dim
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.trans_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
self.z_trans_matrix = nn.Parameter(
torch.eye(state_dim) + 0.03 * torch.randn(state_dim, state_dim))
self.z_obs_matrix = nn.Parameter(0.3 * torch.randn(state_dim, obs_dim))
self.init_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
gp_obs_matrix = torch.zeros(self.kernel.state_dim * obs_dim, obs_dim)
for i in range(obs_dim):
gp_obs_matrix[self.kernel.state_dim * i, i] = 1.0
self.register_buffer("gp_obs_matrix", gp_obs_matrix)
self.obs_selector = torch.tensor(
[self.kernel.state_dim * d for d in range(obs_dim)],
dtype=torch.long)
if learnable_observation_loc:
self.obs_loc = nn.Parameter(torch.zeros(obs_dim))
else:
self.register_buffer("obs_loc", torch.zeros(obs_dim))
def __init__(self, create_plates=None):
super().__init__()
# we define parameters here; make sure that the shape is aligned
# with the shapes of sample sites in model.
self.ma_weight_loc = PyroParam(torch.zeros(10, 1, 1, 2, 3, 7),
event_dim=3)
self.ma_weight_scale = PyroParam(torch.ones(10, 1, 1, 2, 3, 7) * 0.1,
dist.constraints.positive,
event_dim=3)
self.snap_weight_loc = PyroParam(torch.zeros(10, 1, 1, 2, 7),
event_dim=2)
self.snap_weight_scale = PyroParam(torch.ones(10, 1, 1, 2, 7) * 0.1,
dist.constraints.positive,
event_dim=2)
self.seasonal_loc = PyroParam(torch.zeros(10, 1, 7, 2, 7), event_dim=2)
self.seasonal_scale = PyroParam(torch.ones(10, 1, 7, 2, 7) * 0.1,
dist.constraints.positive,
event_dim=2)
self.create_plates = create_plates
class GenericLGSSMWithGPNoiseModel(TimeSeriesModel):
"""
A generic Linear Gaussian State Space Model parameterized with arbitrary time invariant
transition and observation dynamics together with separate Gaussian Process noise models
for each output dimension. In more detail, the generative process is:
:math:`y_i(t) = \\sum_j A_{ij} z_j(t) + f_i(t) + \\epsilon_i(t)`
where the latent variables :math:`{\\bf z}(t)` follow generic time invariant Linear Gaussian dynamics
and the :math:`f_i(t)` are Gaussian Processes with Matern kernels.
The targets are (implicitly) assumed to be evenly spaced in time. In particular a timestep of
:math:`dt=1.0` for the continuous-time GP dynamics corresponds to a single discrete step of
the :math:`{\\bf z}`-space dynamics. Training and inference are logarithmic in the length of
the time series T.
:param int obs_dim: The dimension of the targets at each time step.
:param int state_dim: The dimension of the :math:`{\\bf z}` latent state at each time step.
:param float nu: The order of the Matern kernel; one of 0.5, 1.5 or 2.5.
:param torch.Tensor length_scale_init: optional initial values for the kernel length scale
given as a ``obs_dim``-dimensional tensor
:param torch.Tensor kernel_scale_init: optional initial values for the kernel scale
given as a ``obs_dim``-dimensional tensor
:param torch.Tensor obs_noise_scale_init: optional initial values for the observation noise scale
given as a ``obs_dim``-dimensional tensor
:param bool learnable_observation_loc: whether the mean of the observation model should be learned or not;
defaults to False.
"""
def __init__(self,
obs_dim=1,
state_dim=2,
nu=1.5,
obs_noise_scale_init=None,
length_scale_init=None,
kernel_scale_init=None,
learnable_observation_loc=False):
self.obs_dim = obs_dim
self.state_dim = state_dim
self.nu = nu
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim, )
super().__init__()
self.kernel = MaternKernel(nu=nu,
num_gps=obs_dim,
length_scale_init=length_scale_init,
kernel_scale_init=kernel_scale_init)
self.dt = 1.0
self.full_state_dim = self.kernel.state_dim * obs_dim + state_dim
self.full_gp_state_dim = self.kernel.state_dim * obs_dim
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.trans_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
self.z_trans_matrix = nn.Parameter(
torch.eye(state_dim) + 0.03 * torch.randn(state_dim, state_dim))
self.z_obs_matrix = nn.Parameter(0.3 * torch.randn(state_dim, obs_dim))
self.init_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
gp_obs_matrix = torch.zeros(self.kernel.state_dim * obs_dim, obs_dim)
for i in range(obs_dim):
gp_obs_matrix[self.kernel.state_dim * i, i] = 1.0
self.register_buffer("gp_obs_matrix", gp_obs_matrix)
self.obs_selector = torch.tensor(
[self.kernel.state_dim * d for d in range(obs_dim)],
dtype=torch.long)
if learnable_observation_loc:
self.obs_loc = nn.Parameter(torch.zeros(obs_dim))
else:
self.register_buffer('obs_loc', torch.zeros(obs_dim))
def _get_obs_matrix(self):
# (obs_dim + state_dim, obs_dim) => (gp_state_dim * obs_dim + state_dim, obs_dim)
return torch.cat([self.gp_obs_matrix, self.z_obs_matrix], dim=0)
def _get_init_dist(self):
loc = self.z_trans_matrix.new_zeros(self.full_state_dim)
covar = self.z_trans_matrix.new_zeros(self.full_state_dim,
self.full_state_dim)
covar[:self.full_gp_state_dim, :self.
full_gp_state_dim] = block_diag_embed(
self.kernel.stationary_covariance())
covar[self.full_gp_state_dim:,
self.full_gp_state_dim:] = self.init_noise_scale_sq.diag_embed()
return MultivariateNormal(loc, covar)
def _get_obs_dist(self):
return dist.Normal(self.obs_loc, self.obs_noise_scale).to_event(1)
def _get_dist(self):
"""
Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds
to :class:`GenericLGSSMWithGPNoiseModel`.
"""
gp_trans_matrix, gp_process_covar = self.kernel.transition_matrix_and_covariance(
dt=self.dt)
trans_covar = self.z_trans_matrix.new_zeros(self.full_state_dim,
self.full_state_dim)
trans_covar[:self.full_gp_state_dim, :self.
full_gp_state_dim] = block_diag_embed(gp_process_covar)
trans_covar[
self.full_gp_state_dim:,
self.full_gp_state_dim:] = self.trans_noise_scale_sq.diag_embed()
trans_dist = MultivariateNormal(
trans_covar.new_zeros(self.full_state_dim), trans_covar)
full_trans_mat = trans_covar.new_zeros(self.full_state_dim,
self.full_state_dim)
full_trans_mat[:self.full_gp_state_dim, :self.
full_gp_state_dim] = block_diag_embed(gp_trans_matrix)
full_trans_mat[self.full_gp_state_dim:,
self.full_gp_state_dim:] = self.z_trans_matrix
return dist.GaussianHMM(self._get_init_dist(), full_trans_mat,
trans_dist, self._get_obs_matrix(),
self._get_obs_dist())
@pyro_method
def log_prob(self, targets):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued ``targets`` at each time step
:returns torch.Tensor: A (scalar) log probability.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self._get_dist().log_prob(targets)
@torch.no_grad()
def _filter(self, targets):
"""
Return the filtering state for the associated state space model.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self._get_dist().filter(targets)
@torch.no_grad()
def _forecast(self,
N_timesteps,
filtering_state,
include_observation_noise=True):
"""
Internal helper for forecasting.
"""
dts = torch.arange(N_timesteps,
dtype=self.z_trans_matrix.dtype,
device=self.z_trans_matrix.device) + 1.0
dts = dts.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
gp_trans_matrix, gp_process_covar = self.kernel.transition_matrix_and_covariance(
dt=dts)
gp_trans_matrix = block_diag_embed(gp_trans_matrix)
gp_process_covar = block_diag_embed(gp_process_covar[..., 0:1, 0:1])
N_trans_matrix = repeated_matmul(self.z_trans_matrix, N_timesteps)
N_trans_obs = torch.matmul(N_trans_matrix, self.z_obs_matrix)
# z-state contribution + gp contribution
predicted_mean1 = torch.matmul(
filtering_state.loc[-self.state_dim:].unsqueeze(-2),
N_trans_obs).squeeze(-2)
predicted_mean2 = torch.matmul(
filtering_state.loc[:self.full_gp_state_dim].unsqueeze(-2),
gp_trans_matrix[..., self.obs_selector]).squeeze(-2)
predicted_mean = predicted_mean1 + predicted_mean2
# first compute the contributions from filtering_state.covariance_matrix: z-space and gp
fs_cov = filtering_state.covariance_matrix
predicted_covar1z = torch.matmul(N_trans_obs.transpose(-1, -2),
torch.matmul(
fs_cov[self.full_gp_state_dim:,
self.full_gp_state_dim:],
N_trans_obs)) # N O O
gp_trans = gp_trans_matrix[..., self.obs_selector]
predicted_covar1gp = torch.matmul(
gp_trans.transpose(-1, -2),
torch.matmul(
fs_cov[:self.full_gp_state_dim:, :self.full_gp_state_dim],
gp_trans))
# next compute the contribution from process noise that is injected at each timestep.
# (we need to do a cumulative sum to integrate across time for the z-state contribution)
z_process_covar = self.trans_noise_scale_sq.diag_embed()
N_trans_obs_shift = torch.cat(
[self.z_obs_matrix.unsqueeze(0), N_trans_obs[0:-1]])
predicted_covar2z = torch.matmul(N_trans_obs_shift.transpose(
-1, -2), torch.matmul(z_process_covar, N_trans_obs_shift)) # N O O
predicted_covar = predicted_covar1z + predicted_covar1gp + gp_process_covar + \
torch.cumsum(predicted_covar2z, dim=0)
if include_observation_noise:
predicted_covar = predicted_covar + self.obs_noise_scale.pow(
2.0).diag_embed()
return predicted_mean, predicted_covar
@pyro_method
def forecast(self, targets, N_timesteps):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued targets at each time step. These
represent the training data that are conditioned on for the purpose of making
forecasts.
:param int N_timesteps: The number of timesteps to forecast into the future from
the final target ``targets[-1]``.
:returns torch.distributions.MultivariateNormal: Returns a predictive MultivariateNormal distribution
with batch shape ``(N_timesteps,)`` and event shape ``(obs_dim,)``
"""
filtering_state = self._filter(targets)
predicted_mean, predicted_covar = self._forecast(
N_timesteps, filtering_state)
return MultivariateNormal(predicted_mean, predicted_covar)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
self._event_dims = {}
self._cond_indep_stacks = {}
self.hidden2locs = PyroModule()
self.hidden2scales = PyroModule()
if "multiple" in self.encoder_mode:
# create module for collecting multiple encoder NN
self.multiple_encoders = PyroModule()
# Initialize guide params
for name, site in self.prototype_trace.iter_stochastic_nodes():
# Collect unconstrained event_dims, which may differ from constrained event_dims.
with helpful_support_errors(site):
init_loc = biject_to(site["fn"].support).inv(site["value"].detach()).detach()
event_dim = site["fn"].event_dim + init_loc.dim() - site["value"].dim()
self._event_dims[name] = event_dim
# Collect independence contexts.
self._cond_indep_stacks[name] = site["cond_indep_stack"]
# determine the number of hidden layers
if "multiple" in self.encoder_mode:
if "multiple" in self.n_hidden.keys():
n_hidden = self.n_hidden["multiple"]
else:
n_hidden = self.n_hidden[name]
elif "single" in self.encoder_mode:
n_hidden = self.n_hidden["single"]
# add linear layer for locs and scales
param_dim = (n_hidden, self.amortised_plate_sites["sites"][name])
init_param = np.random.normal(
np.zeros(param_dim),
(np.ones(param_dim) * self.init_param_scale) / np.sqrt(n_hidden),
).astype("float32")
_deep_setattr(
self.hidden2locs,
name,
PyroParam(torch.tensor(init_param, device=site["value"].device, requires_grad=True)),
)
init_param = np.random.normal(
np.zeros(param_dim),
(np.ones(param_dim) * self.init_param_scale) / np.sqrt(n_hidden),
).astype("float32")
_deep_setattr(
self.hidden2scales,
name,
PyroParam(torch.tensor(init_param, device=site["value"].device, requires_grad=True)),
)
if "multiple" in self.encoder_mode:
# create multiple encoders
if self.encoder_instance is not None:
# copy instances
encoder_ = deepcopy(self.encoder_instance).to(site["value"].device)
# convert to pyro module
to_pyro_module_(encoder_)
_deep_setattr(
self.multiple_encoders,
name,
encoder_,
)
else:
# create instances
_deep_setattr(
self.multiple_encoders,
name,
self.encoder_class(n_in=self.multiple_n_in, n_out=n_hidden, **self.multi_encoder_kwargs).to(
site["value"].device
),
)
class GenericLGSSM(TimeSeriesModel):
"""
A generic Linear Gaussian State Space Model parameterized with arbitrary time invariant
transition and observation dynamics. The targets are (implicitly) assumed to be evenly
spaced in time. Training and inference are logarithmic in the length of the time series T.
:param int obs_dim: The dimension of the targets at each time step.
:param int state_dim: The dimension of latent state at each time step.
:param bool learnable_observation_loc: whether the mean of the observation model should be learned or not;
defaults to False.
"""
def __init__(self,
obs_dim=1,
state_dim=2,
obs_noise_scale_init=None,
learnable_observation_loc=False):
self.obs_dim = obs_dim
self.state_dim = state_dim
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim, )
super().__init__()
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.trans_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
self.trans_matrix = nn.Parameter(
torch.eye(state_dim) + 0.03 * torch.randn(state_dim, state_dim))
self.obs_matrix = nn.Parameter(0.3 * torch.randn(state_dim, obs_dim))
self.init_noise_scale_sq = PyroParam(torch.ones(state_dim),
constraint=constraints.positive)
if learnable_observation_loc:
self.obs_loc = nn.Parameter(torch.zeros(obs_dim))
else:
self.register_buffer('obs_loc', torch.zeros(obs_dim))
def _get_init_dist(self):
loc = self.obs_matrix.new_zeros(self.state_dim)
return MultivariateNormal(loc, self.init_noise_scale_sq.diag_embed())
def _get_obs_dist(self):
return dist.Normal(self.obs_loc, self.obs_noise_scale).to_event(1)
def _get_trans_dist(self):
loc = self.obs_matrix.new_zeros(self.state_dim)
return MultivariateNormal(loc, self.trans_noise_scale_sq.diag_embed())
def _get_dist(self):
"""
Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds to :class:`GenericLGSSM`.
"""
return dist.GaussianHMM(self._get_init_dist(), self.trans_matrix,
self._get_trans_dist(), self.obs_matrix,
self._get_obs_dist())
@pyro_method
def log_prob(self, targets):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued ``targets`` at each time step
:returns torch.Tensor: A (scalar) log probability.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self._get_dist().log_prob(targets)
@torch.no_grad()
def _filter(self, targets):
"""
Return the filtering state for the associated state space model.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self._get_dist().filter(targets)
@torch.no_grad()
def _forecast(self,
N_timesteps,
filtering_state,
include_observation_noise=True):
"""
Internal helper for forecasting.
"""
N_trans_matrix = repeated_matmul(self.trans_matrix, N_timesteps)
N_trans_obs = torch.matmul(N_trans_matrix, self.obs_matrix)
predicted_mean = torch.matmul(filtering_state.loc, N_trans_obs)
# first compute the contribution from filtering_state.covariance_matrix
predicted_covar1 = torch.matmul(N_trans_obs.transpose(-1, -2),
torch.matmul(
filtering_state.covariance_matrix,
N_trans_obs)) # N O O
# next compute the contribution from process noise that is injected at each timestep.
# (we need to do a cumulative sum to integrate across time)
process_covar = self._get_trans_dist().covariance_matrix
N_trans_obs_shift = torch.cat(
[self.obs_matrix.unsqueeze(0), N_trans_obs[:-1]])
predicted_covar2 = torch.matmul(N_trans_obs_shift.transpose(
-1, -2), torch.matmul(process_covar, N_trans_obs_shift)) # N O O
predicted_covar = predicted_covar1 + torch.cumsum(predicted_covar2,
dim=0)
if include_observation_noise:
predicted_covar = predicted_covar + self.obs_noise_scale.pow(
2.0).diag_embed()
return predicted_mean, predicted_covar
@pyro_method
def forecast(self, targets, N_timesteps):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued targets at each time step. These
represent the training data that are conditioned on for the purpose of making
forecasts.
:param int N_timesteps: The number of timesteps to forecast into the future from
the final target ``targets[-1]``.
:returns torch.distributions.MultivariateNormal: Returns a predictive MultivariateNormal distribution
with batch shape ``(N_timesteps,)`` and event shape ``(obs_dim,)``
"""
filtering_state = self._filter(targets)
predicted_mean, predicted_covar = self._forecast(
N_timesteps, filtering_state)
return torch.distributions.MultivariateNormal(predicted_mean,
predicted_covar)
class LinearlyCoupledMaternGP(TimeSeriesModel):
"""
A time series model in which each output dimension is modeled as a linear combination
of shared univariate Gaussian Processes with Matern kernels.
In more detail, the generative process is:
:math:`y_i(t) = \\sum_j A_{ij} f_j(t) + \\epsilon_i(t)`
The targets :math:`y_i` are assumed to be evenly spaced in time. Training and inference
are logarithmic in the length of the time series T.
:param float nu: The order of the Matern kernel; one of 0.5, 1.5 or 2.5.
:param float dt: The time spacing between neighboring observations of the time series.
:param int obs_dim: The dimension of the targets at each time step.
:param int num_gps: The number of independent GPs that are mixed to model the time series.
Typical values might be :math:`\\N_{\\rm gp} \\in [\\D_{\\rm obs} / 2, \\D_{\\rm obs}]`
:param torch.Tensor length_scale_init: optional initial values for the kernel length scale
given as a ``num_gps``-dimensional tensor
:param torch.Tensor kernel_scale_init: optional initial values for the kernel scale
given as a ``num_gps``-dimensional tensor
:param torch.Tensor obs_noise_scale_init: optional initial values for the observation noise scale
given as a ``obs_dim``-dimensional tensor
"""
def __init__(self, nu=1.5, dt=1.0, obs_dim=2, num_gps=1,
length_scale_init=None, kernel_scale_init=None,
obs_noise_scale_init=None):
self.nu = nu
self.dt = dt
assert obs_dim > 1, "If obs_dim==1 you should use IndependentMaternGP"
self.obs_dim = obs_dim
self.num_gps = num_gps
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim,)
self.dt = dt
self.obs_dim = obs_dim
self.num_gps = num_gps
super().__init__()
self.kernel = MaternKernel(nu=nu, num_gps=num_gps,
length_scale_init=length_scale_init,
kernel_scale_init=kernel_scale_init)
self.full_state_dim = num_gps * self.kernel.state_dim
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.A = nn.Parameter(0.3 * torch.randn(self.num_gps, self.obs_dim))
def _get_obs_matrix(self):
# (num_gps, obs_dim) => (state_dim * num_gps, obs_dim)
return self.A.repeat_interleave(self.kernel.state_dim, dim=0) * \
self.A.new_tensor([1.0] + [0.0] * (self.kernel.state_dim - 1)).repeat(self.num_gps).unsqueeze(-1)
def _stationary_covariance(self):
return block_diag_embed(self.kernel.stationary_covariance())
def _get_init_dist(self):
loc = self.A.new_zeros(self.full_state_dim)
return MultivariateNormal(loc, self._stationary_covariance())
def _get_obs_dist(self):
loc = self.A.new_zeros(self.obs_dim)
return dist.Normal(loc, self.obs_noise_scale).to_event(1)
def get_dist(self, duration=None):
"""
Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds
to a :class:`LinearlyCoupledMaternGP`.
:param int duration: Optional size of the time axis ``event_shape[0]``.
This is required when sampling from homogeneous HMMs whose parameters
are not expanded along the time axis.
"""
trans_matrix, process_covar = self.kernel.transition_matrix_and_covariance(dt=self.dt)
trans_matrix = block_diag_embed(trans_matrix)
process_covar = block_diag_embed(process_covar)
loc = self.A.new_zeros(self.full_state_dim)
trans_dist = MultivariateNormal(loc, process_covar)
return dist.GaussianHMM(self._get_init_dist(), trans_matrix, trans_dist,
self._get_obs_matrix(), self._get_obs_dist(), duration=duration)
@pyro_method
def log_prob(self, targets):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued ``targets`` at each time step
:returns torch.Tensor: a (scalar) log probability
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self.get_dist().log_prob(targets)
@torch.no_grad()
def _filter(self, targets):
"""
Return the filtering state for the associated state space model.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self.get_dist().filter(targets)
@torch.no_grad()
def _forecast(self, dts, filtering_state, include_observation_noise=True, full_covar=True):
"""
Internal helper for forecasting.
"""
assert dts.dim() == 1
dts = dts.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
trans_mat, process_covar = self.kernel.transition_matrix_and_covariance(dt=dts)
trans_mat = block_diag_embed(trans_mat) # S x full_state_dim x full_state_dim
process_covar = block_diag_embed(process_covar) # S x full_state_dim x full_state_dim
obs_matrix = self._get_obs_matrix() # full_state_dim x obs_dim
trans_obs = torch.matmul(trans_mat, obs_matrix) # S x full_state_dim x obs_dim
predicted_mean = torch.matmul(filtering_state.loc.unsqueeze(-2), trans_obs).squeeze(-2)
predicted_function_covar = torch.matmul(trans_obs.transpose(-1, -2),
torch.matmul(filtering_state.covariance_matrix,
trans_obs))
predicted_function_covar = predicted_function_covar + \
torch.matmul(obs_matrix.transpose(-1, -2), torch.matmul(process_covar, obs_matrix))
if include_observation_noise:
obs_noise = self.obs_noise_scale.pow(2.0).diag_embed()
predicted_function_covar = predicted_function_covar + obs_noise
if not full_covar:
predicted_function_covar = predicted_function_covar.diagonal(dim1=-1, dim2=-2)
return predicted_mean, predicted_function_covar
@pyro_method
def forecast(self, targets, dts):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued targets at each time step. These
represent the training data that are conditioned on for the purpose of making
forecasts.
:param torch.Tensor dts: A 1-dimensional tensor of times to forecast into the future,
with zero corresponding to the time of the final target ``targets[-1]``.
:returns torch.distributions.MultivariateNormal: Returns a predictive MultivariateNormal
distribution with batch shape ``(S,)`` and event shape ``(obs_dim,)``,
where ``S`` is the size of ``dts``.
"""
filtering_state = self._filter(targets)
predicted_mean, predicted_covar = self._forecast(dts, filtering_state)
return MultivariateNormal(predicted_mean, predicted_covar)
class MaternKernel(PyroModule):
"""
Provides the building blocks for representing univariate Gaussian Processes (GPs)
with Matern kernels as state space models.
:param float nu: The order of the Matern kernel (one of 0.5, 1.5 or 2.5)
:param int num_gps: the number of GPs
:param torch.Tensor length_scale_init: optional `num_gps`-dimensional vector of initializers
for the length scale
:param torch.Tensor kernel_scale_init: optional `num_gps`-dimensional vector of initializers
for the kernel scale
**References**
[1] `Kalman Filtering and Smoothing Solutions to Temporal Gaussian Process Regression Models`,
Jouni Hartikainen and Simo Sarkka.
[2] `Stochastic Differential Equation Methods for Spatio-Temporal Gaussian Process Regression`,
Arno Solin.
"""
def __init__(self,
nu=1.5,
num_gps=1,
length_scale_init=None,
kernel_scale_init=None):
if nu not in [0.5, 1.5, 2.5]:
raise NotImplementedError(
"The only supported values of nu are 0.5, 1.5 and 2.5")
self.nu = nu
self.state_dim = {0.5: 1, 1.5: 2, 2.5: 3}[nu]
self.num_gps = num_gps
if length_scale_init is None:
length_scale_init = torch.ones(num_gps)
assert length_scale_init.shape == (num_gps, )
if kernel_scale_init is None:
kernel_scale_init = torch.ones(num_gps)
assert kernel_scale_init.shape == (num_gps, )
super().__init__()
self.length_scale = PyroParam(length_scale_init,
constraint=constraints.positive)
self.kernel_scale = PyroParam(kernel_scale_init,
constraint=constraints.positive)
if self.state_dim > 1:
for x in range(self.state_dim):
for y in range(self.state_dim):
mask = torch.zeros(self.state_dim, self.state_dim)
mask[x, y] = 1.0
self.register_buffer("mask{}{}".format(x, y), mask)
@pyro_method
def transition_matrix(self, dt):
"""
Compute the (exponentiated) transition matrix of the GP latent space.
The resulting matrix has layout (num_gps, old_state, new_state), i.e. this
matrix multiplies states from the right.
See section 5 in reference [1] for details.
:param float dt: the time interval over which the GP latent space evolves.
:returns torch.Tensor: a 3-dimensional tensor of transition matrices of shape
(num_gps, state_dim, state_dim).
"""
if self.nu == 0.5:
rho = self.length_scale.unsqueeze(-1).unsqueeze(-1)
return torch.exp(-dt / rho)
elif self.nu == 1.5:
rho = self.length_scale.unsqueeze(-1).unsqueeze(-1)
dt_rho = dt / rho
trans = (1.0 + root_three * dt_rho) * self.mask00 + \
(-3.0 * dt_rho / rho) * self.mask01 + \
dt * self.mask10 + \
(1.0 - root_three * dt_rho) * self.mask11
return torch.exp(-root_three * dt_rho) * trans
elif self.nu == 2.5:
rho = self.length_scale.unsqueeze(-1).unsqueeze(-1)
dt_rho = root_five * dt / rho
dt_rho_sq = dt_rho.pow(2.0)
dt_rho_cu = dt_rho.pow(3.0)
dt_rho_qu = dt_rho.pow(4.0)
dt_sq = dt**2.0
trans = (1.0 + dt_rho + 0.5 * dt_rho_sq) * self.mask00 + \
(-0.5 * dt_rho_cu / dt) * self.mask01 + \
((0.5 * dt_rho_qu - dt_rho_cu) / dt_sq) * self.mask02 + \
((dt_rho + 1.0) * dt) * self.mask10 + \
(1.0 + dt_rho - dt_rho_sq) * self.mask11 + \
((dt_rho_cu - 3.0 * dt_rho_sq) / dt) * self.mask12 + \
(0.5 * dt_sq) * self.mask20 + \
((1.0 - 0.5 * dt_rho) * dt) * self.mask21 + \
(1.0 - 2.0 * dt_rho + 0.5 * dt_rho_sq) * self.mask22
return torch.exp(-dt_rho) * trans
@pyro_method
def stationary_covariance(self):
"""
Compute the stationary state covariance. See Eqn. 3.26 in reference [2].
:returns torch.Tensor: a 3-dimensional tensor of covariance matrices of shape
(num_gps, state_dim, state_dim).
"""
if self.nu == 0.5:
sigmasq = self.kernel_scale.pow(2).unsqueeze(-1).unsqueeze(-1)
return sigmasq
elif self.nu == 1.5:
sigmasq = self.kernel_scale.pow(2).unsqueeze(-1).unsqueeze(-1)
rhosq = self.length_scale.pow(2).unsqueeze(-1).unsqueeze(-1)
p_infinity = self.mask00 + (3.0 / rhosq) * self.mask11
return sigmasq * p_infinity
elif self.nu == 2.5:
sigmasq = self.kernel_scale.pow(2).unsqueeze(-1).unsqueeze(-1)
rhosq = self.length_scale.pow(2).unsqueeze(-1).unsqueeze(-1)
p_infinity = 0.0
p_infinity = self.mask00 + \
(five_thirds / rhosq) * (self.mask11 - self.mask02 - self.mask20) + \
(25.0 / rhosq.pow(2.0)) * self.mask22
return sigmasq * p_infinity
@pyro_method
def process_covariance(self, A):
"""
Given a transition matrix `A` computed with `transition_matrix` compute the
the process covariance as described in Eqn. 3.11 in reference [2].
:returns torch.Tensor: a batched covariance matrix of shape (num_gps, state_dim, state_dim)
"""
assert A.shape[-3:] == (self.num_gps, self.state_dim, self.state_dim)
p = self.stationary_covariance()
q = p - torch.matmul(A.transpose(-1, -2), torch.matmul(p, A))
return q
@pyro_method
def transition_matrix_and_covariance(self, dt):
"""
Get the transition matrix and process covariance corresponding to a time interval `dt`.
:param float dt: the time interval over which the GP latent space evolves.
:returns tuple: (`transition_matrix`, `process_covariance`) both 3-dimensional tensors of
shape (num_gps, state_dim, state_dim)
"""
trans_matrix = self.transition_matrix(dt)
process_covar = self.process_covariance(trans_matrix)
return trans_matrix, process_covar
def __init__(self, sites, name='', *args, **kwargs):
super().__init__(sites, name, *args, **kwargs)
self.loc = PyroParam(self.init[self.mask], event_dim=1)
class IndependentMaternGP(TimeSeriesModel):
"""
A time series model in which each output dimension is modeled independently
with a univariate Gaussian Process with a Matern kernel. The targets are assumed
to be evenly spaced in time. Training and inference are logarithmic in the length
of the time series T.
:param float nu: The order of the Matern kernel; one of 0.5, 1.5 or 2.5.
:param float dt: The time spacing between neighboring observations of the time series.
:param int obs_dim: The dimension of the targets at each time step.
:param torch.Tensor length_scale_init: optional initial values for the kernel length scale
given as a ``obs_dim``-dimensional tensor
:param torch.Tensor kernel_scale_init: optional initial values for the kernel scale
given as a ``obs_dim``-dimensional tensor
:param torch.Tensor obs_noise_scale_init: optional initial values for the observation noise scale
given as a ``obs_dim``-dimensional tensor
"""
def __init__(self, nu=1.5, dt=1.0, obs_dim=1,
length_scale_init=None, kernel_scale_init=None,
obs_noise_scale_init=None):
self.nu = nu
self.dt = dt
self.obs_dim = obs_dim
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim,)
super().__init__()
self.kernel = MaternKernel(nu=nu, num_gps=obs_dim,
length_scale_init=length_scale_init,
kernel_scale_init=kernel_scale_init)
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
obs_matrix = [1.0] + [0.0] * (self.kernel.state_dim - 1)
self.register_buffer("obs_matrix", torch.tensor(obs_matrix).unsqueeze(-1))
def _get_init_dist(self):
return torch.distributions.MultivariateNormal(self.obs_matrix.new_zeros(self.obs_dim, self.kernel.state_dim),
self.kernel.stationary_covariance().squeeze(-3))
def _get_obs_dist(self):
return dist.Normal(self.obs_matrix.new_zeros(self.obs_dim, 1, 1),
self.obs_noise_scale.unsqueeze(-1).unsqueeze(-1)).to_event(1)
def get_dist(self, duration=None):
"""
Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds
to ``obs_dim``-many independent Matern GPs.
:param int duration: Optional size of the time axis ``event_shape[0]``.
This is required when sampling from homogeneous HMMs whose parameters
are not expanded along the time axis.
"""
trans_matrix, process_covar = self.kernel.transition_matrix_and_covariance(dt=self.dt)
trans_dist = MultivariateNormal(self.obs_matrix.new_zeros(self.obs_dim, 1, self.kernel.state_dim),
process_covar.unsqueeze(-3))
trans_matrix = trans_matrix.unsqueeze(-3)
return dist.GaussianHMM(self._get_init_dist(), trans_matrix, trans_dist,
self.obs_matrix, self._get_obs_dist(), duration=duration)
@pyro_method
def log_prob(self, targets):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued ``targets`` at each time step
:returns torch.Tensor: A 1-dimensional tensor of log probabilities of shape ``(obs_dim,)``
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self.get_dist().log_prob(targets.t().unsqueeze(-1))
@torch.no_grad()
def _filter(self, targets):
"""
Return the filtering state for the associated state space model.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self.get_dist().filter(targets.t().unsqueeze(-1))
@torch.no_grad()
def _forecast(self, dts, filtering_state, include_observation_noise=True):
"""
Internal helper for forecasting.
"""
assert dts.dim() == 1
dts = dts.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
trans_matrix, process_covar = self.kernel.transition_matrix_and_covariance(dt=dts)
trans_matrix = trans_matrix[..., 0:1]
predicted_mean = torch.matmul(filtering_state.loc.unsqueeze(-2), trans_matrix).squeeze(-2)[..., 0]
predicted_function_covar = torch.matmul(trans_matrix.transpose(-1, -2), torch.matmul(
filtering_state.covariance_matrix, trans_matrix))[..., 0, 0] + \
process_covar[..., 0, 0]
if include_observation_noise:
predicted_function_covar = predicted_function_covar + self.obs_noise_scale.pow(2.0)
return predicted_mean, predicted_function_covar
@pyro_method
def forecast(self, targets, dts):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued targets at each time step. These
represent the training data that are conditioned on for the purpose of making
forecasts.
:param torch.Tensor dts: A 1-dimensional tensor of times to forecast into the future,
with zero corresponding to the time of the final target ``targets[-1]``.
:returns torch.distributions.Normal: Returns a predictive Normal distribution with batch shape ``(S,)`` and
event shape ``(obs_dim,)``, where ``S`` is the size of ``dts``.
"""
filtering_state = self._filter(targets)
predicted_mean, predicted_covar = self._forecast(dts, filtering_state)
return torch.distributions.Normal(predicted_mean, predicted_covar.sqrt())
class DependentMaternGP(TimeSeriesModel):
"""
A time series model in which each output dimension is modeled as a univariate Gaussian Process
with a Matern kernel. The different output dimensions become correlated because the Gaussian
Processes are driven by a correlated Wiener process; see reference [1] for details.
If, in addition, `linearly_coupled` is True, additional correlation is achieved through
linear mixing as in :class:`LinearlyCoupledMaternGP`. The targets are assumed to be evenly
spaced in time. Training and inference are logarithmic in the length of the time series T.
:param float nu: The order of the Matern kernel; must be 1.5.
:param float dt: The time spacing between neighboring observations of the time series.
:param int obs_dim: The dimension of the targets at each time step.
:param bool linearly_coupled: Whether to linearly mix the various gaussian processes in the likelihood.
Defaults to False.
:param torch.Tensor length_scale_init: optional initial values for the kernel length scale
given as a ``obs_dim``-dimensional tensor
:param torch.Tensor obs_noise_scale_init: optional initial values for the observation noise scale
given as a ``obs_dim``-dimensional tensor
References
[1] "Dependent Matern Processes for Multivariate Time Series," Alexander Vandenberg-Rodes, Babak Shahbaba.
"""
def __init__(self, nu=1.5, dt=1.0, obs_dim=1, linearly_coupled=False,
length_scale_init=None, obs_noise_scale_init=None):
if nu != 1.5:
raise NotImplementedError("The only supported value of nu is 1.5")
self.dt = dt
self.obs_dim = obs_dim
if obs_noise_scale_init is None:
obs_noise_scale_init = 0.2 * torch.ones(obs_dim)
assert obs_noise_scale_init.shape == (obs_dim,)
super().__init__()
self.kernel = MaternKernel(nu=nu, num_gps=obs_dim,
length_scale_init=length_scale_init)
self.full_state_dim = self.kernel.state_dim * obs_dim
# we demote self.kernel.kernel_scale from being a nn.Parameter
# since the relevant scales are now encoded in the wiener noise matrix
del self.kernel.kernel_scale
self.kernel.register_buffer("kernel_scale", torch.ones(obs_dim))
self.obs_noise_scale = PyroParam(obs_noise_scale_init,
constraint=constraints.positive)
self.wiener_noise_tril = PyroParam(torch.eye(obs_dim) +
0.03 * torch.randn(obs_dim, obs_dim).tril(-1),
constraint=constraints.lower_cholesky)
if linearly_coupled:
self.obs_matrix = nn.Parameter(0.3 * torch.randn(self.obs_dim, self.obs_dim))
else:
obs_matrix = torch.zeros(self.full_state_dim, obs_dim)
for i in range(obs_dim):
obs_matrix[self.kernel.state_dim * i, i] = 1.0
self.register_buffer("obs_matrix", obs_matrix)
def _get_obs_matrix(self):
if self.obs_matrix.size(0) == self.obs_dim:
# (num_gps, obs_dim) => (state_dim * num_gps, obs_dim)
selector = [1.0] + [0.0] * (self.kernel.state_dim - 1)
return self.obs_matrix.repeat_interleave(self.kernel.state_dim, dim=0) * \
self.obs_matrix.new_tensor(selector).repeat(self.obs_dim).unsqueeze(-1)
else:
return self.obs_matrix
def _get_init_dist(self, stationary_covariance):
return torch.distributions.MultivariateNormal(self.obs_matrix.new_zeros(self.full_state_dim),
stationary_covariance)
def _get_obs_dist(self):
return dist.Normal(self.obs_matrix.new_zeros(self.obs_dim),
self.obs_noise_scale).to_event(1)
def _get_wiener_cov(self):
chol = self.wiener_noise_tril
wiener_cov = torch.mm(chol, chol.t()).reshape(self.obs_dim, 1, self.obs_dim, 1)
wiener_cov = wiener_cov * wiener_cov.new_ones(self.kernel.state_dim, 1, self.kernel.state_dim)
return wiener_cov.reshape(self.full_state_dim, self.full_state_dim)
def _stationary_covariance(self):
rho_j = math.sqrt(3.0) / self.kernel.length_scale.unsqueeze(-1).unsqueeze(-1)
rho_i = rho_j.unsqueeze(-1)
block = 2.0 * self.kernel.mask00 + \
(rho_i - rho_j) * (self.kernel.mask01 - self.kernel.mask10) + \
(2.0 * rho_i * rho_j) * self.kernel.mask11
block = block / (rho_i + rho_j).pow(3.0)
block = block.transpose(-2, -3).reshape(self.full_state_dim, self.full_state_dim)
return self._get_wiener_cov() * block
def _get_trans_dist(self, trans_matrix, stationary_covariance):
covar = stationary_covariance - torch.matmul(trans_matrix.transpose(-1, -2),
torch.matmul(stationary_covariance, trans_matrix))
return MultivariateNormal(covar.new_zeros(self.full_state_dim), covar)
def _trans_matrix_distribution_stat_covar(self, dts):
stationary_covariance = self._stationary_covariance()
trans_matrix = self.kernel.transition_matrix(dt=dts)
trans_matrix = block_diag_embed(trans_matrix)
trans_dist = self._get_trans_dist(trans_matrix, stationary_covariance)
return trans_matrix, trans_dist, stationary_covariance
def get_dist(self, duration=None):
"""
Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds to a :class:`DependentMaternGP`
:param int duration: Optional size of the time axis ``event_shape[0]``.
This is required when sampling from homogeneous HMMs whose parameters
are not expanded along the time axis.
"""
trans_matrix, trans_dist, stat_covar = self._trans_matrix_distribution_stat_covar(self.dt)
return dist.GaussianHMM(self._get_init_dist(stat_covar), trans_matrix,
trans_dist, self._get_obs_matrix(), self._get_obs_dist(), duration=duration)
@pyro_method
def log_prob(self, targets):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued ``targets`` at each time step
:returns torch.Tensor: A (scalar) log probability
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self.get_dist().log_prob(targets)
@torch.no_grad()
def _filter(self, targets):
"""
Return the filtering state for the associated state space model.
"""
assert targets.dim() == 2 and targets.size(-1) == self.obs_dim
return self.get_dist().filter(targets)
@torch.no_grad()
def _forecast(self, dts, filtering_state, include_observation_noise=True):
"""
Internal helper for forecasting.
"""
assert dts.dim() == 1
dts = dts.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
trans_matrix, trans_dist, _ = self._trans_matrix_distribution_stat_covar(dts)
obs_matrix = self._get_obs_matrix()
trans_obs = torch.matmul(trans_matrix, obs_matrix)
predicted_mean = torch.matmul(filtering_state.loc.unsqueeze(-2), trans_obs).squeeze(-2)
predicted_function_covar = torch.matmul(trans_obs.transpose(-1, -2),
torch.matmul(filtering_state.covariance_matrix, trans_obs)) + \
torch.matmul(obs_matrix.t(), torch.matmul(trans_dist.covariance_matrix, obs_matrix))
if include_observation_noise:
predicted_function_covar = predicted_function_covar + self.obs_noise_scale.pow(2.0)
return predicted_mean, predicted_function_covar
@pyro_method
def forecast(self, targets, dts):
"""
:param torch.Tensor targets: A 2-dimensional tensor of real-valued targets
of shape ``(T, obs_dim)``, where ``T`` is the length of the time series and ``obs_dim``
is the dimension of the real-valued targets at each time step. These
represent the training data that are conditioned on for the purpose of making
forecasts.
:param torch.Tensor dts: A 1-dimensional tensor of times to forecast into the future,
with zero corresponding to the time of the final target ``targets[-1]``.
:returns torch.distributions.MultivariateNormal: Returns a predictive MultivariateNormal
distribution with batch shape ``(S,)`` and event shape ``(obs_dim,)``, where ``S`` is the size of ``dts``.
"""
filtering_state = self._filter(targets)
predicted_mean, predicted_covar = self._forecast(dts, filtering_state)
return MultivariateNormal(predicted_mean, predicted_covar)