def testRWM2DMixNormal(self):
"""Sampling from a 2-D Mixture Normal Distribution."""
dtype = np.float32
# By symmetry, target has mean [0, 0]
# Therefore, Var = E[X^2] = E[E[X^2 | c]], where c is the component.
# Now..., for the first component,
# E[X1^2] = Var[X1] + Mean[X1]^2
# = 0.3^2 + 1^2,
# and similarly for the second. As a result,
# Var[mixture] = 1.09.
target = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=[0.5, 0.5]),
components_distribution=tfd.MultivariateNormalDiag(
loc=[[-1., -1], [1., 1.]],
scale_identity_multiplier=[0.3, 0.3]))
inverse_temperatures = 10.**tf.linspace(0., -2., 4)
step_sizes = tf.constant([0.3, 0.6, 1.2, 2.4])
def make_kernel_fn(target_log_prob_fn, seed):
kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_log_prob_fn,
seed=seed,
step_size=step_sizes[make_kernel_fn.idx],
num_leapfrog_steps=2)
make_kernel_fn.idx += 1
return kernel
# TODO(b/124770732): Remove this hack.
make_kernel_fn.idx = 0
remc = tfp.mcmc.ReplicaExchangeMC(
target_log_prob_fn=tf.function(target.log_prob, autograph=False),
# Verified that test fails if inverse_temperatures = [1.]
inverse_temperatures=inverse_temperatures,
make_kernel_fn=make_kernel_fn,
seed=_set_seed())
def _trace_log_accept_ratio(state, results):
del state
return [
r.log_accept_ratio for r in results.sampled_replica_results
]
num_results = 1000
samples, log_accept_ratios = tfp.mcmc.sample_chain(
num_results=num_results,
# Start at one of the modes, in order to make mode jumping necessary
# if we want to pass test.
current_state=np.ones(2, dtype=dtype),
kernel=remc,
num_burnin_steps=500,
trace_fn=_trace_log_accept_ratio,
parallel_iterations=1) # For determinism.
self.assertAllEqual((num_results, 2), samples.shape)
log_accept_ratios = [
tf.reduce_mean(input_tensor=tf.exp(tf.minimum(0., lar)))
for lar in log_accept_ratios
]
sample_mean = tf.reduce_mean(input_tensor=samples, axis=0)
sample_std = tf.sqrt(
tf.reduce_mean(input_tensor=tf.math.squared_difference(
samples, sample_mean),
axis=0))
[sample_mean_, sample_std_, log_accept_ratios_
] = self.evaluate([sample_mean, sample_std, log_accept_ratios])
tf1.logging.vlog(1, 'log_accept_ratios: %s eager: %s',
log_accept_ratios_, tf.executing_eagerly())
self.assertAllClose(sample_mean_, [0., 0.], atol=0.3, rtol=0.3)
self.assertAllClose(sample_std_,
[np.sqrt(1.09), np.sqrt(1.09)],
atol=0.1,
rtol=0.1)
class MarkovChainBijectorTest(test_util.TestCase):
# pylint: disable=g-long-lambda
@parameterized.named_parameters(
dict(testcase_name='deterministic_prior',
prior_fn=lambda: tfd.Deterministic([-100., 0., 100.]),
transition_fn=lambda _, x: tfd.Normal(loc=x, scale=1.)),
dict(testcase_name='deterministic_transition',
prior_fn=lambda: tfd.Normal(loc=[-100., 0., 100.], scale=1.),
transition_fn=lambda _, x: tfd.Deterministic(x)),
dict(testcase_name='fully_deterministic',
prior_fn=lambda: tfd.Deterministic([-100., 0., 100.]),
transition_fn=lambda _, x: tfd.Deterministic(x)),
dict(testcase_name='mvn_diag',
prior_fn=(lambda: tfd.MultivariateNormalDiag(loc=[[2.], [2.]],
scale_diag=[1.])),
transition_fn=lambda _, x: tfd.VectorDeterministic(x)),
dict(testcase_name='docstring_dirichlet',
prior_fn=lambda: tfd.JointDistributionNamedAutoBatched(
{'probs': tfd.Dirichlet([1., 1.])}),
transition_fn=lambda _, x: tfd.JointDistributionNamedAutoBatched(
{
'probs':
tfd.MultivariateNormalDiag(loc=x['probs'],
scale_diag=[0.1, 0.1])
},
batch_ndims=ps.rank(x['probs']))),
dict(testcase_name='uniform_step',
prior_fn=lambda: tfd.Exponential(tf.ones([4, 1])),
transition_fn=lambda _, x: tfd.Uniform(low=x, high=x + 1.)),
dict(testcase_name='joint_distribution',
prior_fn=lambda: tfd.JointDistributionNamedAutoBatched(
batch_ndims=2,
model={
'a':
tfd.Gamma(tf.zeros([5]), 1.),
'b':
lambda a: (tfb.Reshape(event_shape_in=[4, 3],
event_shape_out=[2, 3, 2])
(tfd.Independent(tfd.Normal(
loc=tf.zeros([5, 4, 3]),
scale=a[..., tf.newaxis, tf.newaxis]),
reinterpreted_batch_ndims=2)))
}),
transition_fn=lambda _, x: tfd.JointDistributionNamedAutoBatched(
batch_ndims=ps.rank_from_shape(x['a'].shape),
model={
'a':
tfd.Normal(loc=x['a'], scale=1.),
'b':
lambda a: tfd.Deterministic(x['b'] + a[
..., tf.newaxis, tf.newaxis, tf.newaxis])
})),
dict(testcase_name='nested_chain',
prior_fn=lambda: tfd.
MarkovChain(initial_state_prior=tfb.Split(2)
(tfd.MultivariateNormalDiag(0., [1., 2.])),
transition_fn=lambda _, x: tfb.Split(2)
(tfd.MultivariateNormalDiag(x[0], [1., 2.])),
num_steps=6),
transition_fn=(
lambda _, x: tfd.JointDistributionSequentialAutoBatched(
[
tfd.MultivariateNormalDiag(x[0], [1.]),
tfd.MultivariateNormalDiag(x[1], [1.])
],
batch_ndims=ps.rank(x[0])))))
# pylint: enable=g-long-lambda
def test_default_bijector(self, prior_fn, transition_fn):
chain = tfd.MarkovChain(initial_state_prior=prior_fn(),
transition_fn=transition_fn,
num_steps=7)
y = self.evaluate(chain.sample(seed=test_util.test_seed()))
bijector = chain.experimental_default_event_space_bijector()
self.assertAllEqual(chain.batch_shape_tensor(),
bijector.experimental_batch_shape_tensor())
x = bijector.inverse(y)
yy = bijector.forward(tf.nest.map_structure(
tf.identity, x)) # Bypass bijector cache.
self.assertAllCloseNested(y, yy)
chain_event_ndims = tf.nest.map_structure(ps.rank_from_shape,
chain.event_shape_tensor())
self.assertAllEqualNested(bijector.inverse_min_event_ndims,
chain_event_ndims)
ildj = bijector.inverse_log_det_jacobian(
tf.nest.map_structure(tf.identity, y), # Bypass bijector cache.
event_ndims=chain_event_ndims)
if not bijector.is_constant_jacobian:
self.assertAllEqual(ildj.shape, chain.batch_shape)
fldj = bijector.forward_log_det_jacobian(
tf.nest.map_structure(tf.identity, x), # Bypass bijector cache.
event_ndims=bijector.inverse_event_ndims(chain_event_ndims))
self.assertAllClose(ildj, -fldj)
# Verify that event shapes are passed through and flattened/unflattened
# correctly.
inverse_event_shapes = bijector.inverse_event_shape(chain.event_shape)
x_event_shapes = tf.nest.map_structure(
lambda t, nd: t.shape[ps.rank(t) - nd:], x,
bijector.forward_min_event_ndims)
self.assertAllEqualNested(inverse_event_shapes, x_event_shapes)
forward_event_shapes = bijector.forward_event_shape(
inverse_event_shapes)
self.assertAllEqualNested(forward_event_shapes, chain.event_shape)
# Verify that the outputs of other methods have the correct structure.
inverse_event_shape_tensors = bijector.inverse_event_shape_tensor(
chain.event_shape_tensor())
self.assertAllEqualNested(inverse_event_shape_tensors, x_event_shapes)
forward_event_shape_tensors = bijector.forward_event_shape_tensor(
inverse_event_shape_tensors)
self.assertAllEqualNested(forward_event_shape_tensors,
chain.event_shape_tensor())
def test_posterior_marginals_high_rank(self, rank_o, rank_t, rank_i, rank_s):
def increase_rank(n, x):
# By choosing prime number dimensions we make it less
# likely that a test will pass for accidental reasons.
primes = [3, 5, 7]
for i in range(n):
x = primes[i] * [x]
return x
observation_locs_data = tf.identity(
increase_rank(rank_o, tf.eye(4, dtype=self.dtype)))
observation_scales_data = tf.constant(
[0.25, 0.25, 0.25, 0.25],
dtype=self.dtype)
transition_matrix_data = tf.constant(
increase_rank(rank_t, [[0.8, 0.1, 0.1, 0.0],
[0.1, 0.8, 0.0, 0.1],
[0.1, 0.0, 0.8, 0.1],
[0.0, 0.1, 0.1, 0.8]]),
dtype=self.dtype)
initial_prob_data = tf.constant(
increase_rank(rank_i, [0.25, 0.25, 0.25, 0.25]),
dtype=self.dtype)
(initial_prob, transition_matrix,
observation_locs, observation_scales) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_locs_data, observation_scales_data])
observations = tf.constant(
increase_rank(rank_s,
[[[0.91, 0.11], [0.21, 0.09]],
[[0.11, 0.97], [0.12, 0.08]],
[[0.01, 0.12], [0.92, 0.11]],
[[0.02, 0.11], [0.77, 0.11]],
[[0.81, 0.15], [0.21, 0.03]],
[[0.01, 0.13], [0.23, 0.91]],
[[0.11, 0.12], [0.23, 0.79]],
[[0.13, 0.11], [0.91, 0.29]]]),
dtype=self.dtype)
observation_distribution = tfp.distributions.TransformedDistribution(
tfd.MultivariateNormalDiag(observation_locs,
scale_diag=observation_scales),
tfp.bijectors.Reshape((2, 2)))
[num_steps] = self.make_placeholders([8])
model = tfd.HiddenMarkovModel(
tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
observation_distribution,
num_steps=num_steps,
validate_args=True)
inferred_probs = self.evaluate(
model.posterior_marginals(observations).probs_parameter())
rank_e = max(rank_o, rank_t, rank_i, rank_s)
expected_probs = increase_rank(rank_e,
[[0.99994, 0.00000, 0.00006, 0.00000],
[0.45137, 0.01888, 0.52975, 0.00000],
[0.00317, 0.00002, 0.98112, 0.01570],
[0.00000, 0.00001, 0.99998, 0.00001],
[0.00495, 0.00001, 0.94214, 0.05289],
[0.00000, 0.00083, 0.00414, 0.99503],
[0.00000, 0.00000, 0.00016, 0.99984],
[0.00000, 0.00000, 0.99960, 0.00039]])
self.assertAllClose(inferred_probs, expected_probs, rtol=0., atol=1e-4)
def __init__(self,
num_timesteps,
level_scale,
slope_scale,
initial_state_prior,
observation_noise_scale=0.,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a state space model implementing a local linear trend.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
level_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
level transitions.
slope_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
slope transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states; must
have event shape `[2]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
initial_step: Optional scalar `int` `Tensor` specifying the starting
timestep.
Default value: 0.
validate_args: Python `bool`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: "LocalLinearTrendStateSpaceModel".
"""
with tf.name_scope(name, 'LocalLinearTrendStateSpaceModel',
[level_scale, slope_scale]) as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
level_scale = tf.convert_to_tensor(
level_scale, name='level_scale', dtype=dtype)
slope_scale = tf.convert_to_tensor(
slope_scale, name='slope_scale', dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
observation_noise_scale, name='observation_noise_scale', dtype=dtype)
# Explicitly broadcast all parameters to the same batch shape. This
# allows us to use `tf.stack` for a compact model specification.
broadcast_batch_shape = dist_util.get_broadcast_shape(
level_scale, slope_scale)
broadcast_ones = tf.ones(broadcast_batch_shape, dtype=dtype)
self._level_scale = level_scale
self._slope_scale = slope_scale
self._observation_noise_scale = observation_noise_scale
# Construct a linear Gaussian state space model implementing the
# local linear trend model. See "Mathematical Details" in the
# class docstring for further explanation.
super(LocalLinearTrendStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.constant(
[[1., 1.], [0., 1.]], dtype=dtype, name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=tf.stack(
[level_scale * broadcast_ones, slope_scale * broadcast_ones],
axis=-1),
name='transition_noise'),
observation_matrix=tf.constant(
[[1., 0.]], dtype=dtype, name='observation_matrix'),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def __init__(self,
num_timesteps,
period,
frequency_multipliers,
drift_scale,
initial_state_prior,
observation_noise_scale=0.,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a smooth seasonal state space model.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
period: positive scalar `float` `Tensor` giving the number of timesteps
required for the longest cyclic effect to repeat.
frequency_multipliers: One-dimensional `float` `Tensor` listing the
frequencies (cyclic components) included in the model, as multipliers of
the base/fundamental frequency `2. * pi / period`. Each component is
specified by the number of times it repeats per period, and adds two
latent dimensions to the model. A smooth seasonal model that can
represent any periodic function is given by `frequency_multipliers = [1,
2, ..., floor(period / 2)]`. However, it is often desirable to enforce a
smoothness assumption (and reduce the computational burden) by dropping
some of the higher frequencies.
drift_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
latent state transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states. Must have
event shape `[num_features]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
Default value: `0.`.
initial_step: scalar `int` `Tensor` specifying the starting timestep.
Default value: `0`.
validate_args: Python `bool`. Whether to validate input with asserts. If
`validate_args` is `False`, and the inputs are invalid, correct behavior
is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: 'SmoothSeasonalStateSpaceModel'.
"""
with tf.name_scope(name or 'SmoothSeasonalStateSpaceModel') as name:
dtype = dtype_util.common_dtype([
period, frequency_multipliers, drift_scale, initial_state_prior
])
period = tf.convert_to_tensor(value=period,
name='period',
dtype=dtype)
frequency_multipliers = tf.convert_to_tensor(
value=frequency_multipliers,
name='frequency_multipliers',
dtype=dtype)
drift_scale = tf.convert_to_tensor(value=drift_scale,
name='drift_scale',
dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
num_frequencies = static_num_frequencies(frequency_multipliers)
observation_matrix = tf.tile(tf.constant([[1., 0.]], dtype=dtype),
multiples=[1, num_frequencies])
transition_matrix = build_smooth_seasonal_transition_matrix(
period=period,
frequency_multipliers=frequency_multipliers,
dtype=dtype)
self._drift_scale = drift_scale
self._observation_noise_scale = observation_noise_scale
self._period = period
self._frequency_multipliers = frequency_multipliers
super(SmoothSeasonalStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=transition_matrix,
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=(drift_scale[..., tf.newaxis] *
tf.ones([2 * num_frequencies], dtype=dtype)),
name='transition_noise'),
observation_matrix=observation_matrix,
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def __init__(self,
num_seasons,
num_steps_per_season=1,
drift_scale_prior=None,
initial_effect_prior=None,
observed_time_series=None,
name=None):
"""Specify a seasonal effects model.
Args:
num_seasons: Scalar Python `int` number of seasons.
num_steps_per_season: Python `int` number of steps in each
season. This may be either a scalar (shape `[]`), in which case all
seasons have the same length, or a NumPy array of shape `[num_seasons]`.
Default value: 1.
drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `drift_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_effect_prior: optional `tfd.Distribution` instance specifying a
normal prior on the initial effect of each season. This may be either
a scalar `tfd.Normal` prior, in which case it applies independently to
every season, or it may be multivariate normal (e.g.,
`tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
which case it specifies a joint prior across all seasons. If `None`, a
heuristic default prior is constructed based on the provided
`observed_time_series`.
Default value: `None`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series).
Default value: `None`.
name: the name of this model component.
Default value: 'Seasonal'.
"""
with tf.name_scope(name, 'Seasonal', values=[observed_time_series]) as name:
observed_stddev, observed_initial = (
sts_util.empirical_statistics(observed_time_series)
if observed_time_series is not None else (1., 0.))
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if drift_scale_prior is None:
drift_scale_prior = tfd.LogNormal(
loc=tf.math.log(.01 * observed_stddev), scale=3.)
if initial_effect_prior is None:
initial_effect_prior = tfd.Normal(
loc=observed_initial,
scale=tf.abs(observed_initial) + observed_stddev)
self._num_seasons = num_seasons
self._num_steps_per_season = num_steps_per_season
tf.debugging.assert_same_float_dtype(
[drift_scale_prior, initial_effect_prior])
if isinstance(initial_effect_prior, tfd.Normal):
self._initial_state_prior = tfd.MultivariateNormalDiag(
loc=tf.stack([initial_effect_prior.mean()] * num_seasons, axis=-1),
scale_diag=tf.stack([initial_effect_prior.stddev()] * num_seasons,
axis=-1))
else:
self._initial_state_prior = initial_effect_prior
super(Seasonal, self).__init__(
parameters=[
Parameter('drift_scale', drift_scale_prior, tfb.Softplus()),
],
latent_size=num_seasons,
name=name)
def test_transform_joint_to_joint(self, split_sizes):
dist_batch_shape = tf.nest.pack_sequence_as(split_sizes, [
tensorshape_util.constant_value_as_shape(s)
for s in [[2, 3], [2, 1], [1, 3]]
])
bijector_batch_shape = [1, 3]
# Build a joint distribution with parts of the specified sizes.
seed = test_util.test_seed_stream()
component_dists = tf.nest.map_structure(
lambda size, batch_shape: tfd.MultivariateNormalDiag( # pylint: disable=g-long-lambda
loc=tf.random.normal(batch_shape + [size], seed=seed()),
scale_diag=tf.random.uniform(minval=1.,
maxval=2.,
shape=batch_shape + [size],
seed=seed())),
split_sizes,
dist_batch_shape)
if isinstance(split_sizes, dict):
base_dist = tfd.JointDistributionNamed(component_dists)
else:
base_dist = tfd.JointDistributionSequential(component_dists)
# Transform the distribution by applying a separate bijector to each part.
bijectors = [
tfb.Exp(),
tfb.Scale(
tf.random.uniform(minval=1.,
maxval=2.,
shape=bijector_batch_shape,
seed=seed())),
tfb.Reshape([2, 1])
]
bijector = ToyZipMap(tf.nest.pack_sequence_as(split_sizes, bijectors))
# Transform a joint distribution that has different batch shape components
transformed_dist = tfd.TransformedDistribution(base_dist, bijector)
self.assertAllEqualNested(
transformed_dist.event_shape,
bijector.forward_event_shape(base_dist.event_shape))
self.assertAllEqualNested(
*self.evaluate((transformed_dist.event_shape_tensor(),
bijector.forward_event_shape_tensor(
base_dist.event_shape_tensor()))))
# Test that the batch shape components of the input are the same as those of
# the output.
self.assertAllEqualNested(transformed_dist.batch_shape,
dist_batch_shape)
self.assertAllEqualNested(
self.evaluate(transformed_dist.batch_shape_tensor()),
dist_batch_shape)
self.assertAllEqualNested(dist_batch_shape, base_dist.batch_shape)
# Check transformed `log_prob` against the base distribution.
sample_shape = [3]
sample = base_dist.sample(sample_shape, seed=seed())
x = tf.nest.map_structure(tf.zeros_like, sample)
y = bijector.forward(x)
base_logprob = base_dist.log_prob(x)
event_ndims = tf.nest.map_structure(lambda s: s.ndims,
transformed_dist.event_shape)
ildj = bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims)
(transformed_logprob, base_logprob_plus_ildj,
log_transformed_prob) = self.evaluate([
transformed_dist.log_prob(y), base_logprob + ildj,
tf.math.log(transformed_dist.prob(y))
])
self.assertAllClose(base_logprob_plus_ildj, transformed_logprob)
self.assertAllClose(transformed_logprob, log_transformed_prob)
# Test that `.sample()` works and returns a result of the expected structure
# and shape.
y_sampled = transformed_dist.sample(sample_shape, seed=seed())
self.assertAllEqual(
tf.nest.map_structure(lambda y: y.shape, y),
tf.nest.map_structure(lambda y: y.shape, y_sampled))
def __init__(self,
num_seasons,
num_steps_per_season=1,
allow_drift=True,
drift_scale_prior=None,
initial_effect_prior=None,
constrain_mean_effect_to_zero=True,
observed_time_series=None,
name=None):
"""Specify a seasonal effects model.
Args:
num_seasons: Scalar Python `int` number of seasons.
num_steps_per_season: Python `int` number of steps in each
season. This may be either a scalar (shape `[]`), in which case all
seasons have the same length, or a NumPy array of shape `[num_seasons]`,
in which seasons have different length, but remain constant around
different cycles, or a NumPy array of shape `[num_cycles, num_seasons]`,
in which num_steps_per_season for each season also varies in different
cycle (e.g., a 4 years cycle with leap day).
Default value: 1.
allow_drift: optional Python `bool` specifying whether the seasonal
effects can drift over time. Setting this to `False`
removes the `drift_scale` parameter from the model. This is
mathematically equivalent to
`drift_scale_prior = tfd.Deterministic(0.)`, but removing drift
directly is preferred because it avoids the use of a degenerate prior.
Default value: `True`.
drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `drift_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_effect_prior: optional `tfd.Distribution` instance specifying a
normal prior on the initial effect of each season. This may be either
a scalar `tfd.Normal` prior, in which case it applies independently to
every season, or it may be multivariate normal (e.g.,
`tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
which case it specifies a joint prior across all seasons. If `None`, a
heuristic default prior is constructed based on the provided
`observed_time_series`.
Default value: `None`.
constrain_mean_effect_to_zero: if `True`, use a model parameterization
that constrains the mean effect across all seasons to be zero. This
constraint is generally helpful in identifying the contributions of
different model components and can lead to more interpretable
posterior decompositions. It may be undesirable if you plan to directly
examine the latent space of the underlying state space model.
Default value: `True`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series). May
optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
a mask `Tensor` to specify timesteps with missing observations.
Default value: `None`.
name: the name of this model component.
Default value: 'Seasonal'.
"""
with tf.name_scope(name or 'Seasonal') as name:
_, observed_stddev, observed_initial = (
sts_util.empirical_statistics(observed_time_series)
if observed_time_series is not None else (0., 1., 0.))
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if drift_scale_prior is None:
drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
.01 * observed_stddev),
scale=3.)
if initial_effect_prior is None:
initial_effect_prior = tfd.Normal(
loc=observed_initial,
scale=tf.abs(observed_initial) + observed_stddev)
dtype = tf.debugging.assert_same_float_dtype(
[drift_scale_prior, initial_effect_prior])
if isinstance(initial_effect_prior, tfd.Normal):
initial_state_prior = tfd.MultivariateNormalDiag(
loc=tf.stack([initial_effect_prior.mean()] * num_seasons,
axis=-1),
scale_diag=tf.stack([initial_effect_prior.stddev()] *
num_seasons,
axis=-1))
else:
initial_state_prior = initial_effect_prior
if constrain_mean_effect_to_zero:
# Transform the prior to the residual parameterization used by
# `ConstrainedSeasonalStateSpaceModel`, imposing a zero-sum constraint.
# This doesn't change the marginal prior on individual effects, but
# does introduce dependence between the effects.
(effects_to_residuals,
_) = build_effects_to_residuals_matrix(num_seasons,
dtype=dtype)
effects_to_residuals_linop = tf.linalg.LinearOperatorFullMatrix(
effects_to_residuals
) # Use linop so that matmul broadcasts.
initial_state_prior_loc = effects_to_residuals_linop.matvec(
initial_state_prior.mean())
initial_state_prior_scale_linop = effects_to_residuals_linop.matmul(
initial_state_prior.scale) # returns LinearOperator
initial_state_prior = tfd.MultivariateNormalFullCovariance(
loc=initial_state_prior_loc,
covariance_matrix=initial_state_prior_scale_linop.matmul(
initial_state_prior_scale_linop.to_dense(),
adjoint_arg=True))
self._constrain_mean_effect_to_zero = constrain_mean_effect_to_zero
self._initial_state_prior = initial_state_prior
self._num_seasons = num_seasons
self._num_steps_per_season = num_steps_per_season
parameters = []
if allow_drift:
parameters.append(
Parameter(
'drift_scale', drift_scale_prior,
tfb.Chain([
tfb.AffineScalar(scale=observed_stddev),
tfb.Softplus()
])))
self._allow_drift = allow_drift
super(Seasonal, self).__init__(
parameters,
latent_size=(num_seasons -
1 if self.constrain_mean_effect_to_zero else
num_seasons),
name=name)
def __init__(self,
num_timesteps,
level_scale,
slope_mean,
slope_scale,
autoregressive_coef,
initial_state_prior,
observation_noise_scale=0.,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a state space model implementing a semi-local linear trend.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
level_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
level transitions.
slope_mean: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the expected long-term mean of
the latent slope.
slope_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
slope transitions.
autoregressive_coef: Scalar (any additional dimensions are treated as
batch dimensions) `float` `Tensor` defining the AR1 process on the
latent slope.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states; must
have event shape `[2]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
initial_step: Optional scalar `int` `Tensor` specifying the starting
timestep.
Default value: 0.
validate_args: Python `bool`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: "SemiLocalLinearTrendStateSpaceModel".
"""
with tf.name_scope(name or 'SemiLocalLinearTrendStateSpaceModel') as name:
dtype = initial_state_prior.dtype
level_scale = tf.convert_to_tensor(
value=level_scale, dtype=dtype, name='level_scale')
slope_mean = tf.convert_to_tensor(
value=slope_mean, dtype=dtype, name='slope_mean')
slope_scale = tf.convert_to_tensor(
value=slope_scale, dtype=dtype, name='slope_scale')
autoregressive_coef = tf.convert_to_tensor(
value=autoregressive_coef, dtype=dtype, name='autoregressive_coef')
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
dtype=dtype,
name='observation_noise_scale')
self._level_scale = level_scale
self._slope_mean = slope_mean
self._slope_scale = slope_scale
self._autoregressive_coef = autoregressive_coef
self._observation_noise_scale = observation_noise_scale
super(SemiLocalLinearTrendStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=semilocal_linear_trend_transition_matrix(
autoregressive_coef),
transition_noise=semilocal_linear_trend_transition_noise(
level_scale, slope_mean, slope_scale, autoregressive_coef),
observation_matrix=tf.constant(
[[1., 0.]], dtype=dtype),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis]),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def __init__(
self,
ndims=2,
curvature=0.03,
name='banana',
pretty_name='Banana',
):
"""Construct the banana model.
Args:
ndims: Python integer. Dimensionality of the distribution. Must be at
least 2.
curvature: Python float. Controls the strength of the curvature of
the distribution.
name: Python `str` name prefixed to Ops created by this class.
pretty_name: A Python `str`. The pretty name of this model.
Raises:
ValueError: If ndims < 2.
"""
if ndims < 2:
raise ValueError('ndims must be at least 2, saw: {}'.format(ndims))
with tf.name_scope(name):
def bijector_fn(x):
"""Banana transform."""
batch_shape = ps.shape(x)[:-1]
shift = tf.concat(
[
tf.zeros(ps.concat([batch_shape, [1]], axis=0)),
curvature * (tf.square(x[..., :1]) - 100),
tf.zeros(ps.concat([batch_shape, [ndims - 2]], axis=0)),
],
axis=-1,
)
return tfb.Shift(shift)
mg = tfd.MultivariateNormalDiag(
loc=tf.zeros(ndims), scale_diag=[10.] + [1.] * (ndims - 1))
banana = tfd.TransformedDistribution(
mg, bijector=tfb.MaskedAutoregressiveFlow(bijector_fn=bijector_fn))
sample_transformations = {
'identity':
model.Model.SampleTransformation(
fn=lambda params: params,
pretty_name='Identity',
# The second dimension is a sum of scaled Chi2 and normal
# distribution.
# Mean of Chi2 with one degree of freedom is 1, but since the
# first element has variance of 100, it cancels with the shift
# (which is why the shift is there).
ground_truth_mean=np.zeros(ndims),
# Variance of Chi2 with one degree of freedom is 2.
ground_truth_standard_deviation=np.array(
[10.] + [np.sqrt(1. + 2 * curvature**2 * 10.**4)] +
[1.] * (ndims - 2)),
)
}
self._banana = banana
super(Banana, self).__init__(
default_event_space_bijector=tfb.Identity(),
event_shape=banana.event_shape,
dtype=banana.dtype,
name=name,
pretty_name=pretty_name,
sample_transformations=sample_transformations,
)
def __init__(
self,
num_timesteps,
num_seasons,
drift_scale,
initial_state_prior,
observation_noise_scale=1e-4, # Avoid degeneracy.
num_steps_per_season=1,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None): # pylint: disable=g-doc-args
"""Build a seasonal effect state space model with a zero-sum constraint.
{seasonal_init_args}
"""
parameters = dict(locals())
with tf.name_scope(name
or 'ConstrainedSeasonalStateSpaceModel') as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
drift_scale = tf.convert_to_tensor(value=drift_scale,
name='drift_scale',
dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
# Coerce `num_steps_per_season` to a canonical form, an array of
# `num_seasons` integers.
num_steps_per_season = np.squeeze(np.asarray(num_steps_per_season))
if num_steps_per_season.ndim == 0: # scalar case
num_steps_per_season = np.tile(num_steps_per_season,
num_seasons)
elif ((num_steps_per_season.ndim <= 2) # 1D and 2D case
and (num_steps_per_season.shape[-1] != num_seasons)):
raise ValueError(
'num_steps_per_season must either be scalar (shape [])'
' or have the last dimension equal to [num_seasons] = '
'[{}] (saw: shape {})'.format(num_seasons,
num_steps_per_season.shape))
is_last_day_of_season = build_is_last_day_of_season(
num_steps_per_season)
[effects_to_residuals, residuals_to_effects
] = build_effects_to_residuals_matrix(num_seasons, dtype=dtype)
seasonal_transition_matrix = build_seasonal_transition_matrix(
num_seasons=num_seasons,
is_last_day_of_season=is_last_day_of_season,
dtype=dtype,
basis_change_matrix=effects_to_residuals,
basis_change_matrix_inv=residuals_to_effects)
seasonal_transition_noise = build_constrained_seasonal_transition_noise(
drift_scale, num_seasons, is_last_day_of_season)
observation_matrix = tf.concat([
tf.ones([1, 1], dtype=dtype),
tf.zeros([1, num_seasons - 1], dtype=dtype)
],
axis=-1)
observation_matrix = tf.matmul(observation_matrix,
residuals_to_effects)
self._drift_scale = drift_scale
self._observation_noise_scale = observation_noise_scale
self._num_seasons = num_seasons
self._num_steps_per_season = num_steps_per_season
super(ConstrainedSeasonalStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=seasonal_transition_matrix,
transition_noise=seasonal_transition_noise,
observation_matrix=observation_matrix,
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis]),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
self._parameters = parameters
def __init__(self,
num_timesteps,
level_scale,
initial_state_prior,
observation_noise_scale=0.,
name=None,
**linear_gaussian_ssm_kwargs):
"""Build a state space model implementing a local level.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
level_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
level transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states. Must have
event shape `[1]` (as `tfd.LinearGaussianStateSpaceModel` requires a
rank-1 event shape).
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
name: Python `str` name prefixed to ops created by this class.
Default value: "LocalLevelStateSpaceModel".
**linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
to the base `tfd.LinearGaussianStateSpaceModel` constructor.
"""
parameters = dict(locals())
parameters.update(linear_gaussian_ssm_kwargs)
del parameters['linear_gaussian_ssm_kwargs']
with tf.name_scope(name or 'LocalLevelStateSpaceModel') as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
level_scale = tf.convert_to_tensor(value=level_scale,
name='level_scale',
dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
self._level_scale = level_scale
self._observation_noise_scale = observation_noise_scale
# Construct a linear Gaussian state space model implementing the
# local level model. See "Mathematical Details" in the
# class docstring for further explanation.
super(LocalLevelStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.constant([[1.]],
dtype=dtype,
name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=level_scale[..., tf.newaxis],
name='transition_noise'),
observation_matrix=tf.constant([[1.]],
dtype=dtype,
name='observation_matrix'),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
name=name,
**linear_gaussian_ssm_kwargs)
self._parameters = parameters
def observation_noise_fn(t):
current_slice = tf.gather(timeseries, t)
return tfd.MultivariateNormalDiag(
loc=current_slice, scale_diag=tf.zeros_like(current_slice))
def __init__(self,
num_timesteps,
design_matrix,
drift_scale,
initial_state_prior,
observation_noise_scale=0.,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""State space model for a dynamic linear regression.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
design_matrix: float `Tensor` of shape `concat([batch_shape,
[num_timesteps, num_features]])`.
drift_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
latent state transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states. Must have
event shape `[num_features]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
Default value: `0.`.
initial_step: scalar `int` `Tensor` specifying the starting timestep.
Default value: `0`.
validate_args: Python `bool`. Whether to validate input with asserts. If
`validate_args` is `False`, and the inputs are invalid, correct behavior
is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: 'DynamicLinearRegressionStateSpaceModel'.
"""
with tf1.name_scope(name,
'DynamicLinearRegressionStateSpaceModel',
values=[drift_scale]) as name:
dtype = dtype_util.common_dtype(
[design_matrix, drift_scale, initial_state_prior])
design_matrix = tf.convert_to_tensor(value=design_matrix,
name='design_matrix',
dtype=dtype)
design_matrix_with_time_in_first_dim = distribution_util.move_dimension(
design_matrix, -2, 0)
drift_scale = tf.convert_to_tensor(value=drift_scale,
name='drift_scale',
dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
num_features = prefer_static.shape(design_matrix)[-1]
def observation_matrix_fn(t):
observation_matrix = tf.linalg.LinearOperatorFullMatrix(
tf.gather(design_matrix_with_time_in_first_dim,
t)[..., tf.newaxis, :],
name='observation_matrix')
return observation_matrix
self._drift_scale = drift_scale
self._observation_noise_scale = observation_noise_scale
super(DynamicLinearRegressionStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.linalg.LinearOperatorIdentity(
num_rows=num_features,
dtype=dtype,
name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=(drift_scale[..., tf.newaxis] *
tf.ones([num_features], dtype=dtype)),
name='transition_noise'),
observation_matrix=observation_matrix_fn,
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def test_transform_joint_to_joint(self, split_sizes):
dist_batch_shape = tf.nest.pack_sequence_as(split_sizes, [
tensorshape_util.constant_value_as_shape(s)
for s in [[2, 3], [2, 1], [1, 3]]
])
bijector_batch_shape = [1, 3]
# Build a joint distribution with parts of the specified sizes.
seed = test_util.test_seed_stream()
component_dists = tf.nest.map_structure(
lambda size, batch_shape: tfd.MultivariateNormalDiag( # pylint: disable=g-long-lambda
loc=tf.random.normal(batch_shape + [size], seed=seed()),
scale_diag=tf.random.uniform(minval=1.,
maxval=2.,
shape=batch_shape + [size],
seed=seed())),
split_sizes,
dist_batch_shape)
if isinstance(split_sizes, dict):
base_dist = tfd.JointDistributionNamed(component_dists)
else:
base_dist = tfd.JointDistributionSequential(component_dists)
# Transform the distribution by applying a separate bijector to each part.
bijectors = [
tfb.Exp(),
tfb.Scale(
tf.random.uniform(minval=1.,
maxval=2.,
shape=bijector_batch_shape,
seed=seed())),
tfb.Reshape([2, 1])
]
bijector = tfb.JointMap(tf.nest.pack_sequence_as(
split_sizes, bijectors),
validate_args=True)
# Transform a joint distribution that has different batch shape components
transformed_dist = tfd.TransformedDistribution(base_dist, bijector)
self.assertRegex(
str(transformed_dist),
'{}.*batch_shape.*event_shape.*dtype'.format(
transformed_dist.name))
self.assertAllEqualNested(
transformed_dist.event_shape,
bijector.forward_event_shape(base_dist.event_shape))
self.assertAllEqualNested(
*self.evaluate((transformed_dist.event_shape_tensor(),
bijector.forward_event_shape_tensor(
base_dist.event_shape_tensor()))))
# Test that the batch shape components of the input are the same as those of
# the output.
self.assertAllEqualNested(transformed_dist.batch_shape,
dist_batch_shape)
self.assertAllEqualNested(
self.evaluate(transformed_dist.batch_shape_tensor()),
dist_batch_shape)
self.assertAllEqualNested(dist_batch_shape, base_dist.batch_shape)
# Check transformed `log_prob` against the base distribution.
sample_shape = [3]
sample = base_dist.sample(sample_shape, seed=seed())
x = tf.nest.map_structure(tf.zeros_like, sample)
y = bijector.forward(x)
base_logprob = base_dist.log_prob(x)
event_ndims = tf.nest.map_structure(lambda s: s.ndims,
transformed_dist.event_shape)
ildj = bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims)
(transformed_logprob, base_logprob_plus_ildj,
log_transformed_prob) = self.evaluate([
transformed_dist.log_prob(y), base_logprob + ildj,
tf.math.log(transformed_dist.prob(y))
])
self.assertAllClose(base_logprob_plus_ildj, transformed_logprob)
self.assertAllClose(transformed_logprob, log_transformed_prob)
# Test that `.sample()` works and returns a result of the expected structure
# and shape.
y_sampled = transformed_dist.sample(sample_shape, seed=seed())
self.assertAllEqual(
tf.nest.map_structure(lambda y: y.shape, y),
tf.nest.map_structure(lambda y: y.shape, y_sampled))
# Test that a `Restructure` bijector applied to a `JointDistribution` works
# as expected.
num_components = len(split_sizes)
input_keys = (split_sizes.keys() if isinstance(split_sizes, dict) else
range(num_components))
output_keys = [str(i) for i in range(num_components)]
output_structure = {k: v for k, v in zip(output_keys, input_keys)}
restructure = tfb.Restructure(output_structure)
restructured_dist = tfd.TransformedDistribution(base_dist,
bijector=restructure,
validate_args=True)
# Check that attributes of the restructured distribution have the same
# nested structure as the `output_structure` of the bijector. Pass a no-op
# as the `assert_fn` since the contents of the structures are not
# required to be the same.
noop_assert_fn = lambda *_: None
self.assertAllAssertsNested(noop_assert_fn,
restructured_dist.event_shape,
output_structure)
self.assertAllAssertsNested(noop_assert_fn,
restructured_dist.batch_shape,
output_structure)
self.assertAllAssertsNested(
noop_assert_fn,
self.evaluate(restructured_dist.event_shape_tensor()),
output_structure)
self.assertAllAssertsNested(
noop_assert_fn,
self.evaluate(restructured_dist.batch_shape_tensor()),
output_structure)
self.assertAllAssertsNested(
noop_assert_fn,
self.evaluate(
restructured_dist.sample(seed=test_util.test_seed())))
def __init__(self,
level_scale_prior=None,
slope_mean_prior=None,
slope_scale_prior=None,
autoregressive_coef_prior=None,
initial_level_prior=None,
initial_slope_prior=None,
observed_time_series=None,
constrain_ar_coef_stationary=True,
constrain_ar_coef_positive=False,
name=None):
"""Specify a semi-local linear trend model.
Args:
level_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `level_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
slope_mean_prior: optional `tfd.Distribution` instance specifying a prior
on the `slope_mean` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `slope_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
autoregressive_coef_prior: optional `tfd.Distribution` instance specifying
a prior on the `autoregressive_coef` parameter. If `None`, the default
prior is a standard `Normal(0., 1.)`. Note that the prior may be
implicitly truncated by `constrain_ar_coef_stationary` and/or
`constrain_ar_coef_positive`.
Default value: `None`.
initial_level_prior: optional `tfd.Distribution` instance specifying a
prior on the initial level. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_slope_prior: optional `tfd.Distribution` instance specifying a
prior on the initial slope. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series). May
optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
a mask `Tensor` to specify timesteps with missing observations.
Default value: `None`.
constrain_ar_coef_stationary: if `True`, perform inference using a
parameterization that restricts `autoregressive_coef` to the interval
`(-1, 1)`, or `(0, 1)` if `force_positive_ar_coef` is also `True`,
corresponding to stationary processes. This will implicitly truncates
the support of `autoregressive_coef_prior`.
Default value: `True`.
constrain_ar_coef_positive: if `True`, perform inference using a
parameterization that restricts `autoregressive_coef` to be positive,
or in `(0, 1)` if `constrain_ar_coef_stationary` is also `True`. This
will implicitly truncate the support of `autoregressive_coef_prior`.
Default value: `False`.
name: the name of this model component.
Default value: 'SemiLocalLinearTrend'.
"""
with tf.name_scope(name or 'SemiLocalLinearTrend') as name:
if observed_time_series is not None:
_, observed_stddev, observed_initial = sts_util.empirical_statistics(
observed_time_series)
else:
observed_stddev, observed_initial = 1., 0.
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if level_scale_prior is None:
level_scale_prior = tfd.LogNormal(
loc=tf.math.log(.01 * observed_stddev), scale=2.)
if slope_mean_prior is None:
slope_mean_prior = tfd.Normal(loc=0.,
scale=observed_stddev)
if slope_scale_prior is None:
slope_scale_prior = tfd.LogNormal(
loc=tf.math.log(.01 * observed_stddev), scale=2.)
if autoregressive_coef_prior is None:
autoregressive_coef_prior = tfd.Normal(
loc=0., scale=tf.ones_like(observed_initial))
if initial_level_prior is None:
initial_level_prior = tfd.Normal(
loc=observed_initial,
scale=tf.abs(observed_initial) + observed_stddev)
if initial_slope_prior is None:
initial_slope_prior = tfd.Normal(loc=0., scale=observed_stddev)
self._initial_state_prior = tfd.MultivariateNormalDiag(
loc=tf.stack(
[initial_level_prior.mean(),
initial_slope_prior.mean()
], axis=-1),
scale_diag=tf.stack([
initial_level_prior.stddev(),
initial_slope_prior.stddev()
], axis=-1))
# Constrain the support of the autoregressive coefficient.
if constrain_ar_coef_stationary and constrain_ar_coef_positive:
autoregressive_coef_bijector = tfb.Sigmoid() # support in (0, 1)
elif constrain_ar_coef_positive:
autoregressive_coef_bijector = tfb.Softplus() # support in (0, infty)
elif constrain_ar_coef_stationary:
autoregressive_coef_bijector = tfb.Tanh() # support in (-1, 1)
else:
autoregressive_coef_bijector = tfb.Identity() # unconstrained
stddev_preconditioner = tfb.AffineScalar(scale=observed_stddev)
scaled_softplus = tfb.Chain([stddev_preconditioner, tfb.Softplus()])
super(SemiLocalLinearTrend, self).__init__(
parameters=[
Parameter('level_scale', level_scale_prior, scaled_softplus),
Parameter('slope_mean', slope_mean_prior, stddev_preconditioner),
Parameter('slope_scale', slope_scale_prior, scaled_softplus),
Parameter('autoregressive_coef',
autoregressive_coef_prior,
autoregressive_coef_bijector),
],
latent_size=2,
name=name)
def __init__(self,
num_timesteps,
num_seasons,
drift_scale,
initial_state_prior,
observation_noise_scale=0.,
num_steps_per_season=1,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a state space model implementing seasonal effects.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
num_seasons: Scalar Python `int` number of seasons.
drift_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
change in effect between consecutive occurrences of a given season.
This is assumed to be the same for all seasons.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states; must
have event shape `[num_seasons]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
Default value: 0.
num_steps_per_season: Python `int` number of steps in each
season. This may be either a scalar (shape `[]`), in which case all
seasons have the same length, or a NumPy array of shape `[num_seasons]`.
Default value: 1.
initial_step: Optional scalar `int` `Tensor` specifying the starting
timestep.
Default value: 0.
validate_args: Python `bool`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: "SeasonalStateSpaceModel".
Raises:
ValueError: if `num_steps_per_season` has invalid shape (neither
scalar nor `[num_seasons]`).
"""
with tf.name_scope(name, 'SeasonalStateSpaceModel',
values=[drift_scale, observation_noise_scale]) as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
drift_scale = tf.convert_to_tensor(
value=drift_scale, name='drift_scale', dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
# Coerce `num_steps_per_season` to a canonical form, an array of
# `num_seasons` integers.
num_steps_per_season = np.asarray(num_steps_per_season)
if not num_steps_per_season.shape:
num_steps_per_season = np.tile(num_steps_per_season, num_seasons)
elif num_steps_per_season.shape != (num_seasons,):
raise ValueError('num_steps_per_season must either be scalar (shape [])'
' or have length [num_seasons] = [{}] (saw: shape {})'.
format(num_seasons, num_steps_per_season.shape))
# Utility method to compute whether the season is changing.
num_steps_per_cycle = np.sum(num_steps_per_season)
changepoints = np.cumsum(num_steps_per_season) - 1
def is_last_day_of_season(t):
t_ = dist_util.maybe_get_static_value(t)
if t_ is not None: # static case
step_in_cycle = t_ % num_steps_per_cycle
return any(step_in_cycle == changepoints)
else:
step_in_cycle = tf.floormod(t, num_steps_per_cycle)
return tf.reduce_any(
input_tensor=tf.equal(step_in_cycle, changepoints))
# If the season is changing, the transition matrix rotates the latent
# state to shift all seasons up by a dimension, and sends the current
# season's effect to the bottom.
seasonal_permutation = np.concatenate(
[np.arange(1, num_seasons), [0]], axis=0)
seasonal_permutation_matrix = np.eye(num_seasons)[seasonal_permutation]
def seasonal_transition_matrix(t):
return tf.linalg.LinearOperatorFullMatrix(
matrix=dist_util.pick_scalar_condition(
is_last_day_of_season(t),
tf.constant(seasonal_permutation_matrix, dtype=dtype),
tf.eye(num_seasons, dtype=dtype)))
# If the season is changing, the transition noise model adds random drift
# to the effect of the outgoing season.
drift_scale_diag = tf.stack(
[tf.zeros_like(drift_scale)] * (num_seasons - 1) + [drift_scale],
axis=-1)
def seasonal_transition_noise(t):
noise_scale = dist_util.pick_scalar_condition(
is_last_day_of_season(t),
drift_scale_diag,
tf.zeros_like(drift_scale_diag, dtype=dtype))
return tfd.MultivariateNormalDiag(loc=tf.zeros(num_seasons,
dtype=dtype),
scale_diag=noise_scale)
self._drift_scale = drift_scale
self._observation_noise_scale = observation_noise_scale
self._num_seasons = num_seasons
self._num_steps_per_season = num_steps_per_season
super(SeasonalStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=seasonal_transition_matrix,
transition_noise=seasonal_transition_noise,
observation_matrix=tf.concat([tf.ones([1, 1], dtype=dtype),
tf.zeros([1, num_seasons-1],
dtype=dtype)],
axis=-1),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis]),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def make_weights_prior(self, dims, sigma):
return tfd.MultivariateNormalDiag(loc=tf.zeros([dims],
dtype=sigma.dtype),
scale_identity_multiplier=sigma)
def observation_noise_fn(t):
return tfd.MultivariateNormalDiag(
loc=(sum([ssm.get_observation_noise_for_timestep(t).mean()
for ssm in component_ssms]) + offset_at_step(t)),
scale_diag=observation_noise_scale[..., tf.newaxis])
def __init__(self,
level_scale_prior=None,
initial_level_prior=None,
observed_time_series=None,
name=None):
"""Specify a local level model.
Args:
level_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `level_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_level_prior: optional `tfd.Distribution` instance specifying a
prior on the initial level. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series). May
optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
a mask `Tensor` to specify timesteps with missing observations.
Default value: `None`.
name: the name of this model component.
Default value: 'LocalLevel'.
"""
with tf.compat.v1.name_scope(name,
'LocalLevel',
values=[observed_time_series]) as name:
dtype = dtype_util.common_dtype(
[level_scale_prior, initial_level_prior])
if level_scale_prior is None or initial_level_prior is None:
if observed_time_series is not None:
_, observed_stddev, observed_initial = (
sts_util.empirical_statistics(observed_time_series))
else:
observed_stddev, observed_initial = (tf.convert_to_tensor(
value=1., dtype=dtype),
tf.convert_to_tensor(
value=0.,
dtype=dtype))
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if level_scale_prior is None:
level_scale_prior = tfd.LogNormal(loc=tf.math.log(
.05 * observed_stddev),
scale=3.,
name='level_scale_prior')
if initial_level_prior is None:
self._initial_state_prior = tfd.MultivariateNormalDiag(
loc=observed_initial[..., tf.newaxis],
scale_diag=(tf.abs(observed_initial) +
observed_stddev)[..., tf.newaxis],
name='initial_level_prior')
else:
self._initial_state_prior = tfd.MultivariateNormalDiag(
loc=initial_level_prior.mean()[..., tf.newaxis],
scale_diag=initial_level_prior.stddev()[..., tf.newaxis])
super(LocalLevel, self).__init__(parameters=[
Parameter(
'level_scale', level_scale_prior,
tfb.Chain([
tfb.AffineScalar(scale=observed_stddev),
tfb.Softplus()
])),
],
latent_size=1,
name=name)
def target_log_prob_fn(event):
with tf.name_scope('nuts_test_target_log_prob'):
return tfd.MultivariateNormalDiag(
tf.zeros(event_size),
scale_identity_multiplier=1.).log_prob(event)
def test_batch_shape(self):
batch_shape = [3, 2]
partial_batch_shape = [2]
num_seasons = 24
initial_state_prior = tfd.MultivariateNormalDiag(
scale_diag=self._build_placeholder(
np.exp(
np.random.randn(*(partial_batch_shape +
[num_seasons - 1])))))
drift_scale = self._build_placeholder(
np.exp(np.random.randn(*batch_shape)))
observation_noise_scale = self._build_placeholder(
np.exp(np.random.randn(*partial_batch_shape)))
ssm = ConstrainedSeasonalStateSpaceModel(
num_timesteps=9,
num_seasons=24,
num_steps_per_season=2,
drift_scale=drift_scale,
observation_noise_scale=observation_noise_scale,
initial_state_prior=initial_state_prior)
# First check that the model's batch shape is the broadcast batch shape
# of parameters, as expected.
self.assertAllEqual(self.evaluate(ssm.batch_shape_tensor()),
batch_shape)
y_ = self.evaluate(ssm.sample())
self.assertAllEqual(y_.shape[:-2], batch_shape)
# Next check that the broadcasting works as expected, and the batch log_prob
# actually matches the log probs of independent models.
individual_ssms = []
for i in range(batch_shape[0]):
for j in range(batch_shape[1]):
individual_ssms.append(
ConstrainedSeasonalStateSpaceModel(
num_timesteps=9,
num_seasons=num_seasons,
num_steps_per_season=2,
drift_scale=drift_scale[i, j, ...],
observation_noise_scale=observation_noise_scale[j,
...],
initial_state_prior=tfd.MultivariateNormalDiag(
scale_diag=initial_state_prior.scale.diag[j,
...])))
batch_lps_ = self.evaluate(ssm.log_prob(y_)).flatten()
individual_ys = [
y_[i, j, ...] # pylint: disable=g-complex-comprehension
for i in range(batch_shape[0]) for j in range(batch_shape[1])
]
individual_lps_ = self.evaluate([
individual_ssm.log_prob(individual_y)
for (individual_ssm,
individual_y) in zip(individual_ssms, individual_ys)
])
self.assertAllClose(individual_lps_, batch_lps_)
def __init__(self,
level_scale_prior=None,
slope_scale_prior=None,
initial_level_prior=None,
initial_slope_prior=None,
observed_time_series=None,
name=None):
"""Specify a local linear trend model.
Args:
level_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `level_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `slope_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_level_prior: optional `tfd.Distribution` instance specifying a
prior on the initial level. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_slope_prior: optional `tfd.Distribution` instance specifying a
prior on the initial slope. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series).
Default value: `None`.
name: the name of this model component.
Default value: 'LocalLinearTrend'.
"""
with tf.name_scope(
name, 'LocalLinearTrend', values=[observed_time_series]) as name:
observed_stddev, observed_initial = (
sts_util.empirical_statistics(observed_time_series)
if observed_time_series is not None else (1., 0.))
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if level_scale_prior is None:
level_scale_prior = tfd.LogNormal(
loc=tf.log(.05 * observed_stddev),
scale=3.,
name='level_scale_prior')
if slope_scale_prior is None:
slope_scale_prior = tfd.LogNormal(
loc=tf.log(.05 * observed_stddev),
scale=3.,
name='slope_scale_prior')
if initial_level_prior is None:
initial_level_prior = tfd.Normal(
loc=observed_initial,
scale=tf.abs(observed_initial) + observed_stddev,
name='initial_level_prior')
if initial_slope_prior is None:
initial_slope_prior = tfd.Normal(
loc=0., scale=observed_stddev, name='initial_slope_prior')
tf.assert_same_float_dtype([
level_scale_prior, slope_scale_prior, initial_level_prior,
initial_slope_prior
])
self._initial_state_prior = tfd.MultivariateNormalDiag(
loc=tf.stack(
[initial_level_prior.mean(),
initial_slope_prior.mean()
], axis=-1),
scale_diag=tf.stack([
initial_level_prior.stddev(),
initial_slope_prior.stddev()
], axis=-1))
super(LocalLinearTrend, self).__init__(
parameters=[
Parameter('level_scale', level_scale_prior, tfb.Softplus()),
Parameter('slope_scale', slope_scale_prior, tfb.Softplus())
],
latent_size=2,
name=name)
def test_month_of_year_example(self):
num_days_per_month = np.array(
[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31])
# put wildly different near-deterministic priors on the effect for each
# month, so we can easily distinguish the months and diagnose off-by-one
# errors.
monthly_effect_prior_means = np.linspace(-1000, 1000, 12)
monthly_effect_prior_scales = 0.1 * np.ones(12)
drift_scale = 0.3
observation_noise_scale = 0.1
num_timesteps = 365 * 2 # 2 years.
initial_step = 22
month_of_year = SeasonalStateSpaceModel(
num_timesteps=num_timesteps,
num_seasons=12,
num_steps_per_season=num_days_per_month,
drift_scale=self._build_placeholder(drift_scale),
observation_noise_scale=observation_noise_scale,
initial_state_prior=tfd.MultivariateNormalDiag(
loc=self._build_placeholder(monthly_effect_prior_means),
scale_diag=self._build_placeholder(
monthly_effect_prior_scales)),
initial_step=initial_step)
sampled_series_, prior_mean_, prior_variance_ = self.evaluate(
(month_of_year.sample()[..., 0], month_of_year.mean()[..., 0],
month_of_year.variance()[..., 0]))
# For each month, ensure the mean (and samples) of each day matches the
# expected effect for that month, and that the variances all match
# including expected drift from year to year
current_idx = -initial_step # start at beginning of year
current_drift_variance = 0.
for _ in range(
3): # loop over three years to include entire 2-year period
for (num_days_in_month, prior_effect_mean,
prior_effect_scale) in zip(num_days_per_month,
monthly_effect_prior_means,
monthly_effect_prior_scales):
month_indices = range(current_idx,
current_idx + num_days_in_month)
current_idx += num_days_in_month
# Trim any indices outside of the observed window.
month_indices = [
idx for idx in month_indices if 0 <= idx < num_timesteps
]
if month_indices:
ones = np.ones(len(month_indices))
self.assertAllClose(sampled_series_[month_indices],
prior_effect_mean * ones,
atol=20.)
self.assertAllClose(prior_mean_[month_indices],
prior_effect_mean * ones,
atol=1e-4)
self.assertAllClose(
prior_variance_[month_indices],
(prior_effect_scale**2 + current_drift_variance +
observation_noise_scale**2) * ones,
atol=1e-4)
# At the end of each year, allow the effects to drift.
current_drift_variance += drift_scale**2
def __init__(self,
period,
frequency_multipliers,
allow_drift=True,
drift_scale_prior=None,
initial_state_prior=None,
observed_time_series=None,
name=None):
"""Specify a smooth seasonal effects model.
Args:
period: positive scalar `float` `Tensor` giving the number of timesteps
required for the longest cyclic effect to repeat.
frequency_multipliers: One-dimensional `float` `Tensor` listing the
frequencies (cyclic components) included in the model, as multipliers of
the base/fundamental frequency `2. * pi / period`. Each component is
specified by the number of times it repeats per period, and adds two
latent dimensions to the model. A smooth seasonal model that can
represent any periodic function is given by `frequency_multipliers = [1,
2, ..., floor(period / 2)]`. However, it is often desirable to enforce a
smoothness assumption (and reduce the computational burden) by dropping
some of the higher frequencies.
allow_drift: optional Python `bool` specifying whether the seasonal
effects can drift over time. Setting this to `False`
removes the `drift_scale` parameter from the model. This is
mathematically equivalent to
`drift_scale_prior = tfd.Deterministic(0.)`, but removing drift
directly is preferred because it avoids the use of a degenerate prior.
Default value: `True`.
drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
on the `drift_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_state_prior: instance of `tfd.MultivariateNormal` representing
the prior distribution on the latent states. Must have event shape
`[2 * len(frequency_multipliers)]`. If `None`, a heuristic default prior
is constructed based on the provided `observed_time_series`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series). May
optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
a mask `Tensor` to specify timesteps with missing observations.
Default value: `None`.
name: the name of this model component.
Default value: 'SmoothSeasonal'.
"""
with tf.name_scope(name or 'SmoothSeasonal') as name:
_, observed_stddev, observed_initial = (
sts_util.empirical_statistics(observed_time_series)
if observed_time_series is not None else (0., 1., 0.))
latent_size = 2 * static_num_frequencies(frequency_multipliers)
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if drift_scale_prior is None:
drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
.01 * observed_stddev),
scale=3.)
if initial_state_prior is None:
initial_state_scale = (tf.abs(observed_initial) +
observed_stddev)[..., tf.newaxis]
ones = tf.ones([latent_size], dtype=drift_scale_prior.dtype)
initial_state_prior = tfd.MultivariateNormalDiag(
scale_diag=initial_state_scale * ones)
self._initial_state_prior = initial_state_prior
self._period = period
self._frequency_multipliers = frequency_multipliers
parameters = []
if allow_drift:
parameters.append(
Parameter(
'drift_scale', drift_scale_prior,
tfb.Chain([
tfb.AffineScalar(scale=observed_stddev),
tfb.Softplus()
])))
self._allow_drift = allow_drift
super(SmoothSeasonal, self).__init__(parameters=parameters,
latent_size=latent_size,
name=name)
def __init__(self,
num_timesteps,
design_matrix,
drift_scale,
initial_state_prior,
observation_noise_scale=0.,
name=None,
**linear_gaussian_ssm_kwargs):
"""State space model for a dynamic linear regression.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
design_matrix: float `Tensor` of shape `concat([batch_shape,
[num_timesteps, num_features]])`.
drift_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
latent state transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states. Must have
event shape `[num_features]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
Default value: `0.`.
name: Python `str` name prefixed to ops created by this class.
Default value: 'DynamicLinearRegressionStateSpaceModel'.
**linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
to the base `tfd.LinearGaussianStateSpaceModel` constructor.
"""
parameters = dict(locals())
parameters.update(linear_gaussian_ssm_kwargs)
del parameters['linear_gaussian_ssm_kwargs']
with tf.name_scope(
name or 'DynamicLinearRegressionStateSpaceModel') as name:
dtype = dtype_util.common_dtype(
[design_matrix, drift_scale, initial_state_prior])
design_matrix = tf.convert_to_tensor(value=design_matrix,
name='design_matrix',
dtype=dtype)
design_matrix_with_time_in_first_dim = distribution_util.move_dimension(
design_matrix, -2, 0)
drift_scale = tf.convert_to_tensor(value=drift_scale,
name='drift_scale',
dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
num_features = prefer_static.shape(design_matrix)[-1]
def observation_matrix_fn(t):
observation_matrix = tf.linalg.LinearOperatorFullMatrix(
tf.gather(design_matrix_with_time_in_first_dim,
t)[..., tf.newaxis, :],
name='observation_matrix')
return observation_matrix
self._drift_scale = drift_scale
self._observation_noise_scale = observation_noise_scale
super(DynamicLinearRegressionStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.linalg.LinearOperatorIdentity(
num_rows=num_features,
dtype=dtype,
name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=(drift_scale[..., tf.newaxis] *
tf.ones([num_features], dtype=dtype)),
name='transition_noise'),
observation_matrix=observation_matrix_fn,
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
name=name,
**linear_gaussian_ssm_kwargs)
self._parameters = parameters
def observation_noise_fn(t):
predicted_slice = predicted_timeseries[..., t, :]
return tfd.MultivariateNormalDiag(
loc=predicted_slice, scale_diag=tf.zeros_like(predicted_slice))
def __init__(self,
design_matrix,
drift_scale_prior=None,
initial_weights_prior=None,
observed_time_series=None,
name=None):
"""Specify a dynamic linear regression.
Args:
design_matrix: float `Tensor` of shape `concat([batch_shape,
[num_timesteps, num_features]])`.
drift_scale_prior: instance of `tfd.Distribution` specifying a prior on
the `drift_scale` parameter. If `None`, a heuristic default prior is
constructed based on the provided `observed_time_series`.
Default value: `None`.
initial_weights_prior: instance of `tfd.MultivariateNormal` representing
the prior distribution on the latent states (the regression weights).
Must have event shape `[num_features]`. If `None`, a weakly-informative
Normal(0., 10.) prior is used.
Default value: `None`.
observed_time_series: optional `float` `Tensor` of shape
`batch_shape + [T, 1]` (omitting the trailing unit dimension is also
supported when `T > 1`), specifying an observed time series. Any `NaN`s
are interpreted as missing observations; missingness may be also be
explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
Any priors not explicitly set will be given default values according to
the scale of the observed time series (or batch of time series).
Default value: `None`.
name: Python `str` for the name of this component.
Default value: 'DynamicLinearRegression'.
"""
init_parameters = dict(locals())
with tf.name_scope(name or 'DynamicLinearRegression') as name:
dtype = dtype_util.common_dtype(
[design_matrix, drift_scale_prior, initial_weights_prior])
num_features = prefer_static.shape(design_matrix)[-1]
# Default to a weakly-informative Normal(0., 10.) for the initital state
if initial_weights_prior is None:
initial_weights_prior = tfd.MultivariateNormalDiag(
scale_diag=10. * tf.ones([num_features], dtype=dtype))
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if drift_scale_prior is None:
if observed_time_series is None:
observed_stddev = tf.constant(1.0, dtype=dtype)
else:
_, observed_stddev, _ = sts_util.empirical_statistics(
observed_time_series)
drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
.05 * observed_stddev),
scale=3.,
name='drift_scale_prior')
self._initial_state_prior = initial_weights_prior
self._design_matrix = design_matrix
super(DynamicLinearRegression,
self).__init__(parameters=[
Parameter(
'drift_scale', drift_scale_prior,
tfb.Chain([
tfb.Scale(scale=observed_stddev),
tfb.Softplus()
]))
],
latent_size=num_features,
init_parameters=init_parameters,
name=name)
def test_posterior_mode_high_rank(self, rank_o, rank_t, rank_i, rank_s):
def increase_rank(n, x):
# By choosing prime number dimensions we make it less
# likely that a test will pass for accidental reasons.
primes = [3, 5, 7]
for i in range(n):
x = primes[i] * [x]
return x
observation_locs_data = tf.constant(increase_rank(rank_o,
[[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]]),
dtype=self.dtype)
observation_scales_data = tf.constant(
[0.25, 0.25, 0.25, 0.25],
dtype=self.dtype)
transition_matrix_data = tf.constant(
increase_rank(rank_t, [[0.8, 0.1, 0.1, 0.0],
[0.1, 0.8, 0.0, 0.1],
[0.1, 0.0, 0.8, 0.1],
[0.0, 0.1, 0.1, 0.8]]),
dtype=self.dtype)
initial_prob_data = tf.constant(
increase_rank(rank_i, [0.25, 0.25, 0.25, 0.25]),
dtype=self.dtype)
(initial_prob, transition_matrix,
observation_locs, observation_scales) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_locs_data, observation_scales_data])
observations = tf.constant(
increase_rank(rank_s,
[[[0.91, 0.11], [0.21, 0.09]],
[[0.11, 0.97], [0.12, 0.08]],
[[0.01, 0.12], [0.92, 0.11]],
[[0.02, 0.11], [0.77, 0.11]],
[[0.81, 0.15], [0.21, 0.03]],
[[0.01, 0.13], [0.23, 0.91]],
[[0.11, 0.12], [0.23, 0.79]],
[[0.13, 0.11], [0.91, 0.29]]]),
dtype=self.dtype)
observation_distribution = tfp.distributions.TransformedDistribution(
tfd.MultivariateNormalDiag(observation_locs,
scale_diag=observation_scales),
tfp.bijectors.Reshape((2, 2)))
[num_steps] = self.make_placeholders([8])
model = tfd.HiddenMarkovModel(
tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
observation_distribution,
num_steps=num_steps,
validate_args=True)
inferred_states = model.posterior_mode(observations)
rank_e = max(rank_o, rank_t, rank_i, rank_s)
expected_states = increase_rank(rank_e, [0, 2, 2, 2, 2, 3, 3, 2])
self.assertAllEqual(inferred_states, expected_states)
def create_component():
loc = tf.random.normal(batch_and_event_shape)
scale_diag = 10 * tf.random.uniform(batch_and_event_shape)
tensorshape_util.set_shape(loc, static_batch_and_event_shape)
tensorshape_util.set_shape(scale_diag, static_batch_and_event_shape)
return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)
