def testBijector(self):
with self.test_session():
self.assertEqual("tanh", tfb.Tanh().name)
x = np.linspace(-3., 3., 100).reshape([2, 5,
10]).astype(np.float32)
y = np.tanh(x)
ildj = -np.log1p(-np.square(np.tanh(x)))
bijector = tfb.Tanh()
self.assertAllClose(y,
bijector.forward(x).eval(),
atol=0.,
rtol=1e-2)
self.assertAllClose(x,
bijector.inverse(y).eval(),
atol=0.,
rtol=1e-4)
self.assertAllClose(ildj,
bijector.inverse_log_det_jacobian(
y, event_ndims=0).eval(),
atol=0.,
rtol=1e-6)
self.assertAllClose(-ildj,
bijector.forward_log_det_jacobian(
x, event_ndims=0).eval(),
atol=0.,
rtol=1e-4)
def testMatchWithAffineTransform(self):
direct_bj = tfb.Tanh()
indirect_bj = tfb.Chain([
tfb.Shift(tf.cast(-1.0, dtype=tf.float64)),
tfb.Scale(tf.cast(2.0, dtype=tf.float64)),
tfb.Sigmoid(),
tfb.Scale(tf.cast(2.0, dtype=tf.float64))
])
x = np.linspace(-3.0, 3.0, 100)
y = np.tanh(x)
self.assertAllClose(self.evaluate(direct_bj.forward(x)),
self.evaluate(indirect_bj.forward(x)))
self.assertAllClose(self.evaluate(direct_bj.inverse(y)),
self.evaluate(indirect_bj.inverse(y)))
self.assertAllClose(
self.evaluate(direct_bj.inverse_log_det_jacobian(y,
event_ndims=0)),
self.evaluate(
indirect_bj.inverse_log_det_jacobian(y, event_ndims=0)))
self.assertAllClose(
self.evaluate(direct_bj.forward_log_det_jacobian(x,
event_ndims=0)),
self.evaluate(
indirect_bj.forward_log_det_jacobian(x, event_ndims=0)))
def testScalarCongruency(self):
bijector_test_util.assert_scalar_congruency(tfb.Tanh(),
lower_x=-7.,
upper_x=7.,
eval_func=self.evaluate,
n=int(10e4),
rtol=.5)
def testMatchWithAffineTransform(self):
direct_bj = tfb.Tanh()
indirect_bj = tfb.Chain([
tfb.AffineScalar(shift=tf.to_double(-1.0),
scale=tf.to_double(2.0)),
tfb.Sigmoid(),
tfb.AffineScalar(scale=tf.to_double(2.0))
])
x = np.linspace(-3.0, 3.0, 100)
y = np.tanh(x)
self.assertAllClose(self.evaluate(direct_bj.forward(x)),
self.evaluate(indirect_bj.forward(x)))
self.assertAllClose(self.evaluate(direct_bj.inverse(y)),
self.evaluate(indirect_bj.inverse(y)))
self.assertAllClose(
self.evaluate(direct_bj.inverse_log_det_jacobian(y,
event_ndims=0)),
self.evaluate(
indirect_bj.inverse_log_det_jacobian(y, event_ndims=0)))
self.assertAllClose(
self.evaluate(direct_bj.forward_log_det_jacobian(x,
event_ndims=0)),
self.evaluate(
indirect_bj.forward_log_det_jacobian(x, event_ndims=0)))
def testBijectiveAndFinite(self):
x = np.linspace(-100., 100., 100).astype(np.float64)
eps = 1e-3
y = np.linspace(-1. + eps, 1. - eps, 100).astype(np.float64)
bijector_test_util.assert_bijective_and_finite(
tfb.Tanh(), x, y, eval_func=self.evaluate, event_ndims=0, atol=0.,
rtol=1e-4)
def testBijector(self):
self.assertStartsWith(tfb.Tanh().name, "tanh")
x = np.linspace(-3., 3., 100).reshape([2, 5, 10]).astype(np.float64)
y = np.tanh(x)
ildj = -np.log1p(-np.square(np.tanh(x)))
bijector = tfb.Tanh()
self.assertAllClose(
y, self.evaluate(bijector.forward(x)), atol=0., rtol=1e-2)
self.assertAllClose(
x, self.evaluate(bijector.inverse(y)), atol=0., rtol=1e-4)
self.assertAllClose(
ildj,
self.evaluate(bijector.inverse_log_det_jacobian(
y, event_ndims=0)), atol=0., rtol=1e-6)
self.assertAllClose(
-ildj,
self.evaluate(bijector.forward_log_det_jacobian(
x, event_ndims=0)), atol=0., rtol=1e-4)
def testBijectiveAndFinite(self):
with self.test_session():
x = np.linspace(-5., 5., 100).astype(np.float32)
eps = 1e-3
y = np.linspace(eps, 1. - eps, 100).astype(np.float32)
assert_bijective_and_finite(tfb.Tanh(),
x,
y,
event_ndims=0,
atol=0.,
rtol=1e-4)
def testTransformedNormalNormalKL(self):
batch_size = 6
mu_a = np.array([3.0] * batch_size).astype(np.float32)
sigma_a = np.array([1.0, 2.0, 3.0, 1.5, 2.5, 3.5]).astype(np.float32)
mu_b = np.array([-3.0] * batch_size).astype(np.float32)
sigma_b = np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0]).astype(np.float32)
n_a = tfd.Normal(loc=mu_a, scale=sigma_a, validate_args=True)
n_b = tfd.Normal(loc=mu_b, scale=sigma_b, validate_args=True)
kl_expected = ((mu_a - mu_b)**2 / (2 * sigma_b**2) + 0.5 * (
(sigma_a**2 / sigma_b**2) - 1 - 2 * np.log(sigma_a / sigma_b)))
bij1 = tfb.Shift(shift=1.)(tfb.Scale(scale=2.))
bij2 = (tfb.Shift(shift=np.array(2., dtype=np.float32))
(tfb.Scale(scale=np.array(3., dtype=np.float32))))
bij3 = tfb.Tanh()
for chain in bij2(bij1), bij3(bij2(bij1)):
td_a = tfd.TransformedDistribution(
distribution=n_a,
bijector=chain,
validate_args=True)
td_b = tfd.TransformedDistribution(
distribution=n_b,
bijector=copy.copy(chain),
validate_args=True)
kl = tfd.kl_divergence(td_a, td_b)
kl_val = self.evaluate(kl)
x = td_a.sample(int(1e5), seed=test_util.test_seed())
kl_sample = tf.reduce_mean(td_a.log_prob(x) - td_b.log_prob(x), axis=0)
kl_sample_ = self.evaluate(kl_sample)
self.assertEqual(kl.shape, (batch_size,))
self.assertAllClose(kl_val, kl_expected)
self.assertAllClose(kl_expected, kl_sample_, atol=0.0, rtol=1e-2)
def __init__(self,
order,
coefficients_prior=None,
level_scale_prior=None,
initial_state_prior=None,
coefficient_constraining_bijector=None,
observed_time_series=None,
name=None):
"""Specify an autoregressive model.
Args:
order: scalar Python positive `int` specifying the number of past
timesteps to regress on.
coefficients_prior: optional `tfd.Distribution` instance specifying a
prior on the `coefficients` parameter. If `None`, a default standard
normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([order]))`) prior
is used.
Default value: `None`.
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_state_prior: optional `tfd.Distribution` instance specifying a
prior on the initial state, corresponding to the values of the process
at a set of size `order` of imagined timesteps before the initial step.
If `None`, a heuristic default prior is constructed based on the
provided `observed_time_series`.
Default value: `None`.
coefficient_constraining_bijector: optional `tfb.Bijector` instance
representing a constraining mapping for the autoregressive coefficients.
For example, `tfb.Tanh()` constrains the coefficients to lie in
`(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
`tfb.Identity()` implies no constraint. If `None`, the default behavior
constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
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: 'Autoregressive'.
"""
with tf.name_scope(name or 'Autoregressive') as name:
masked_time_series = None
if observed_time_series is not None:
masked_time_series = (
sts_util.canonicalize_observed_time_series_with_mask(
observed_time_series))
dtype = dtype_util.common_dtype(
[(masked_time_series.time_series
if masked_time_series is not None else None),
coefficients_prior,
level_scale_prior,
initial_state_prior], dtype_hint=tf.float32)
if observed_time_series is not None:
_, observed_stddev, observed_initial = sts_util.empirical_statistics(
masked_time_series)
else:
observed_stddev, observed_initial = (
tf.convert_to_tensor(value=1., dtype=dtype),
tf.convert_to_tensor(value=0., dtype=dtype))
batch_ones = tf.ones(tf.concat([
tf.shape(observed_initial), # Batch shape
[order]], axis=0), dtype=dtype)
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if coefficients_prior is None:
coefficients_prior = tfd.MultivariateNormalDiag(
scale_diag=batch_ones)
if level_scale_prior is None:
level_scale_prior = tfd.LogNormal(
loc=tf.math.log(0.05 * observed_stddev), scale=3.)
if (coefficients_prior.event_shape.is_fully_defined() and
order != coefficients_prior.event_shape[0]):
raise ValueError("Prior dimension {} doesn't match order {}.".format(
coefficients_prior.event_shape[0], order))
if initial_state_prior is None:
initial_state_prior = tfd.MultivariateNormalDiag(
loc=observed_initial[..., tf.newaxis] * batch_ones,
scale_diag=(tf.abs(observed_initial) +
observed_stddev)[..., tf.newaxis] * batch_ones)
self._order = order
self._coefficients_prior = coefficients_prior
self._level_scale_prior = level_scale_prior
self._initial_state_prior = initial_state_prior
if coefficient_constraining_bijector is None:
coefficient_constraining_bijector = tfb.Tanh()
super(Autoregressive, self).__init__(
parameters=[
Parameter('coefficients',
coefficients_prior,
coefficient_constraining_bijector),
Parameter('level_scale', level_scale_prior,
tfb.Chain([tfb.AffineScalar(scale=observed_stddev),
tfb.Softplus()]))
],
latent_size=order,
name=name)
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.Scale(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 testScalarCongruency(self):
with self.test_session():
assert_scalar_congruency(tfb.Tanh(),
lower_x=-9.,
upper_x=9.,
n=int(10e4))
def __init__(self,
ar_order,
ma_order,
integration_degree=0,
ar_coefficients_prior=None,
ma_coefficients_prior=None,
level_drift_prior=None,
level_scale_prior=None,
initial_state_prior=None,
ar_coefficient_constraining_bijector=None,
ma_coefficient_constraining_bijector=None,
observed_time_series=None,
name=None):
"""Specifies an ARIMA(p=ar_order, d=integration_degree, q=ma_order) model.
Args:
ar_order: scalar Python positive `int` specifying the order of the
autoregressive process (`p` in `ARIMA(p, d, q)`).
ma_order: scalar Python positive `int` specifying the order of the
moving-average process (`q` in `ARIMA(p, d, q)`).
integration_degree: scalar Python positive `int` specifying the number
of times to integrate an ARMA process. (`d` in `ARIMA(p, d, q)`).
Default value: `0`.
ar_coefficients_prior: optional `tfd.Distribution` instance specifying a
prior on the `ar_coefficients` parameter. If `None`, a default standard
normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ar_order]))`)
prior is used.
Default value: `None`.
ma_coefficients_prior: optional `tfd.Distribution` instance specifying a
prior on the `ma_coefficients` parameter. If `None`, a default standard
normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ma_order]))`)
prior is used.
Default value: `None`.
level_drift_prior: optional `tfd.Distribution` instance specifying a prior
on the `level_drift` parameter. If `None`, the parameter is not inferred
and is instead fixed to zero.
Default value: `None`.
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_state_prior: optional `tfd.Distribution` instance specifying a
prior on the initial state, corresponding to the values of the process
at a set of size `order` of imagined timesteps before the initial step.
If `None`, a heuristic default prior is constructed based on the
provided `observed_time_series`.
Default value: `None`.
ar_coefficient_constraining_bijector: optional `tfb.Bijector` instance
representing a constraining mapping for the autoregressive coefficients.
For example, `tfb.Tanh()` constrains the coefficients to lie in
`(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
`tfb.Identity()` implies no constraint. If `None`, the default behavior
constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
Default value: `None`.
ma_coefficient_constraining_bijector: optional `tfb.Bijector` instance
representing a constraining mapping for the moving average coefficients.
For example, `tfb.Tanh()` constrains the coefficients to lie in
`(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
`tfb.Identity()` implies no constraint. If `None`, the default behavior
is to apply no constraint.
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: the name of this model component.
Default value: 'ARIMA'.
"""
init_parameters = dict(locals())
with tf.name_scope(name or 'ARIMA') as name:
masked_time_series = None
if observed_time_series is not None:
masked_time_series = (
sts_util.canonicalize_observed_time_series_with_mask(
observed_time_series))
dtype = dtype_util.common_dtype(
[(masked_time_series.time_series
if masked_time_series is not None else None),
ar_coefficients_prior,
ma_coefficients_prior,
level_scale_prior,
initial_state_prior], dtype_hint=tf.float32)
if observed_time_series is not None:
for _ in range(integration_degree):
# Compute statistics using `integration_order`-order differences.
masked_time_series = (
missing_values_util.differentiate_masked_time_series(
masked_time_series))
_, observed_stddev, observed_initial = sts_util.empirical_statistics(
masked_time_series)
else:
observed_stddev, observed_initial = (
tf.convert_to_tensor(value=1., dtype=dtype),
tf.convert_to_tensor(value=0., dtype=dtype))
batch_ones = ps.ones(ps.concat([
ps.shape(observed_initial), # Batch shape
[1]], axis=0), dtype=dtype)
# Heuristic default priors. Overriding these may dramatically
# change inference performance and results.
if ar_coefficients_prior is None:
ar_coefficients_prior = tfd.MultivariateNormalDiag(
scale_diag=batch_ones * ps.ones([ar_order]))
if ma_coefficients_prior is None:
ma_coefficients_prior = tfd.MultivariateNormalDiag(
scale_diag=batch_ones * ps.ones([ma_order]))
if level_scale_prior is None:
level_scale_prior = tfd.LogNormal(
loc=tf.math.log(0.05 * observed_stddev), scale=3.)
if (ar_coefficients_prior.event_shape.is_fully_defined() and
ar_order != ar_coefficients_prior.event_shape[0]):
raise ValueError(
"Autoregressive prior dimension {} doesn't match order {}.".format(
ar_coefficients_prior.event_shape[0], ar_order))
if (ma_coefficients_prior.event_shape.is_fully_defined() and
ma_order != ma_coefficients_prior.event_shape[0]):
raise ValueError(
"Moving average prior dimension {} doesn't match order {}.".format(
ma_coefficients_prior.event_shape[0], ma_order))
latent_size = ps.maximum(ar_order, ma_order + 1) + integration_degree
if initial_state_prior is None:
initial_state_prior = tfd.MultivariateNormalDiag(
loc=sts_util.pad_tensor_with_trailing_zeros(
observed_initial[..., tf.newaxis] * batch_ones,
num_zeros=latent_size - 1),
scale_diag=sts_util.pad_tensor_with_trailing_zeros(
(tf.abs(observed_initial) +
observed_stddev)[..., tf.newaxis] * batch_ones,
num_zeros=latent_size - 1))
self._ar_order = ar_order
self._ma_order = ma_order
self._integration_degree = integration_degree
self._ar_coefficients_prior = ar_coefficients_prior
self._ma_coefficients_prior = ma_coefficients_prior
self._level_scale_prior = level_scale_prior
self._initial_state_prior = initial_state_prior
parameters = []
if ar_order > 0:
parameters.append(
Parameter('ar_coefficients',
ar_coefficients_prior,
(ar_coefficient_constraining_bijector
if ar_coefficient_constraining_bijector
else tfb.Tanh())))
if ma_order > 0:
parameters.append(
Parameter('ma_coefficients',
ma_coefficients_prior,
(ma_coefficient_constraining_bijector
if ma_coefficient_constraining_bijector
else tfb.Identity())))
if level_drift_prior is not None:
parameters.append(
Parameter(
'level_drift',
level_drift_prior,
tfb.Chain([
tfb.Scale(scale=observed_stddev),
(level_drift_prior.
experimental_default_event_space_bijector())])))
super(AutoregressiveIntegratedMovingAverage, self).__init__(
parameters=parameters + [
Parameter('level_scale', level_scale_prior,
tfb.Chain([tfb.Scale(scale=observed_stddev),
tfb.Softplus(low=dtype_util.eps(dtype))]))
],
latent_size=latent_size,
init_parameters=init_parameters,
name=name)
