def testBijector(self):
low = np.array([[-3.], [0.], [5.]]).astype(np.float32)
high = 12.
bijector = tfb.Sigmoid(low=low, high=high, validate_args=True)
equivalent_bijector = tfb.Chain(
[tfb.Shift(shift=low),
tfb.Scale(scale=high - low),
tfb.Sigmoid()])
x = [[[1., 2., -5., -0.3]]]
y = self.evaluate(equivalent_bijector.forward(x))
self.assertAllClose(y, self.evaluate(bijector.forward(x)))
self.assertAllClose(x,
self.evaluate(bijector.inverse(y)[..., :1, :]),
rtol=1e-5)
self.assertAllClose(
self.evaluate(
equivalent_bijector.inverse_log_det_jacobian(y,
event_ndims=1)),
self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1)),
rtol=1e-5)
self.assertAllClose(
self.evaluate(
equivalent_bijector.forward_log_det_jacobian(x,
event_ndims=1)),
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)))
def testBijector(self):
with self.cached_session():
self.assertEqual("sigmoid", tfb.Sigmoid().name)
x = np.linspace(-10., 10., 100).reshape([2, 5,
10]).astype(np.float32)
y = special.expit(x)
ildj = -np.log(y) - np.log1p(-y)
bijector = tfb.Sigmoid()
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 make_asvi_trainable_variables(prior):
"""Generates parameter dictionaries given a prior distribution and list."""
with tf.name_scope('make_asvi_trainable_variables'):
param_dicts = []
prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access
for dist in prior_dists:
actual_dist = dist.distribution if isinstance(dist, Root) else dist
dist_params = actual_dist.parameters
new_params_dict = {}
# Build trainable ASVI representation for each distribution's parameters.
for param, value in dist_params.items():
if param in _NON_STATISTICAL_PARAMS or value is None:
continue
new_params_dict[param] = ASVIParameters(
prior_weight=tfp.util.TransformedVariable(
0.5,
bijector=tfb.Sigmoid(),
name='prior_weight/{}/{}'.format(dist.name, param)),
mean_field_parameter=tfp.util.TransformedVariable(
0.5,
bijector=parameter_constraints.constraint_for(param),
name='mean_field_parameter/{}/{}'.format(dist.name, param)))
param_dicts.append(new_params_dict)
return param_dicts
def _make_asvi_trainable_variables(prior,
mean_field=False,
initial_prior_weight=0.5):
"""Generates parameter dictionaries given a prior distribution and list."""
with tf.name_scope('make_asvi_trainable_variables'):
param_dicts = []
prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access
for dist in prior_dists:
original_dist = dist.distribution if isinstance(dist, Root) else dist
substituted_dist = _as_trainable_family(original_dist)
# Grab the base distribution if it exists
try:
actual_dist = substituted_dist.distribution
except AttributeError:
actual_dist = substituted_dist
new_params_dict = {}
# Build trainable ASVI representation for each distribution's parameters.
parameter_properties = actual_dist.parameter_properties(
dtype=actual_dist.dtype)
sample_shape = tf.concat(
[dist.batch_shape_tensor(),
dist.event_shape_tensor()], axis=0)
for param, value in actual_dist.parameters.items():
if param in (_NON_STATISTICAL_PARAMS +
_NON_TRAINABLE_PARAMS) or value is None:
continue
try:
bijector = parameter_properties[
param].default_constraining_bijector_fn()
except NotImplementedError:
bijector = tfb.Identity()
unconstrained_ones = tf.ones(
shape=bijector.inverse_event_shape_tensor(
parameter_properties[param].shape_fn(
sample_shape=sample_shape)),
dtype=actual_dist.dtype)
if mean_field:
new_params_dict[param] = ASVIParameters(
prior_weight=None,
mean_field_parameter=tfp_util.TransformedVariable(
value,
bijector=bijector,
name='mean_field_parameter/{}/{}'.format(dist.name, param)))
else:
new_params_dict[param] = ASVIParameters(
prior_weight=tfp_util.TransformedVariable(
initial_prior_weight * unconstrained_ones,
bijector=tfb.Sigmoid(),
name='prior_weight/{}/{}'.format(dist.name, param)),
mean_field_parameter=tfp_util.TransformedVariable(
value,
bijector=bijector,
name='mean_field_parameter/{}/{}'.format(dist.name, param)))
param_dicts.append(new_params_dict)
return param_dicts
def testAssertHighGtLow(self):
low = np.array([1., 1., 1.], dtype=np.float32)
high = np.array([1., 2., 3.], dtype=np.float32)
with self.assertRaisesOpError('not defined when `low` >= `high`'):
bijector = tfb.Sigmoid(low=low, high=high, validate_args=True)
self.evaluate(bijector.forward(3.))
def testSigmoidBetaTargetConservation(self):
# Not inverting the sigmoid bijector makes a kooky distribution, but nuts
# should still conserve it (with a smaller step size).
sigmoid_beta_dist = tfb.Identity(tfb.Sigmoid())(
tfd.Beta(concentration0=1., concentration1=2.))
self.evaluate(assert_univariate_target_conservation(
self, sigmoid_beta_dist, step_size=0.02))
def testLogitBetaTargetConservation(self):
logit_beta_dist = tfb.Invert(tfb.Sigmoid())(tfd.Beta(
concentration0=1., concentration1=2.))
self.evaluate(
assert_univariate_target_conservation(self,
logit_beta_dist,
step_size=0.2))
def make_asvi_trainable_variables(prior):
"""Generates parameter dictionaries given a prior distribution and list."""
param_dicts = []
prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access
d = 0
for dist in prior_dists:
actual_dist = dist.distribution if isinstance(dist, Root) else dist
dist_params = actual_dist.parameters
new_params_dict = {}
# make and store mean_field_parameters & prior_weights in new_params_dict
for param, value in dist_params.items():
if param in _NON_STATISTICAL_PARAMS or value is None:
new_params_dict[param] = value
else:
# TODO(kateslin): Change dictionary naming scheme to use namedtuples.
prior_weight_name = 'prior_weight_{}'.format(d)
mean_field_parameter_name = 'mean_field_parameter_{}'.format(d)
prior_weight = tfp.util.TransformedVariable(
0.5, bijector=tfb.Sigmoid(), name=prior_weight_name)
new_params_dict['{}_prior_weight'.format(param)] = prior_weight
mean_field_parameter_constraint_bijector = (
parameter_constraints.constraint_for(param))
mean_field_parameter = tfp.util.TransformedVariable(
0.5,
mean_field_parameter_constraint_bijector,
name=mean_field_parameter_name)
new_params_dict['{}_mean_field_parameter'.format(
param)] = mean_field_parameter
d += 1
param_dicts.append(new_params_dict)
return param_dicts
def testBijectiveAndFinite(self):
x = np.linspace(-100., 100., 100).astype(np.float32)
eps = 1e-3
y = np.linspace(eps, 1. - eps, 100).astype(np.float32)
bijector_test_util.assert_bijective_and_finite(
tfb.Sigmoid(), x, y, eval_func=self.evaluate, event_ndims=0, atol=0.,
rtol=1e-4)
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 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):
low = -2.
high = 5.
bijector = tfb.Sigmoid(low=low, high=high, validate_args=True)
bijector_test_util.assert_scalar_congruency(bijector,
lower_x=-5.,
upper_x=3.5,
eval_func=self.evaluate,
rtol=0.05)
def testNumericallySuperiorToEquivalentChain(self):
x = np.array([-5., 3., 17., 23.]).astype(np.float32)
low = -0.08587775
high = 0.12498104
bijector = tfb.Sigmoid(low=low, high=high, validate_args=True)
equivalent_bijector = tfb.Chain(
[tfb.Shift(shift=low),
tfb.Scale(scale=high - low),
tfb.Sigmoid()])
self.assertAllLessEqual(self.evaluate(bijector.forward(x)), high)
# The mathematically equivalent `Chain` bijector can return values greater
# than the intended upper bound of `high`.
self.assertTrue(
(self.evaluate(equivalent_bijector.forward(x)) > high).any())
def testQuantile(self):
logit_normal = self._cls()(distribution=tfd.Normal(loc=0., scale=1.),
bijector=tfb.Sigmoid(),
validate_args=True)
grid = [0., 0.25, 0.5, 0.75, 1.]
q = logit_normal.quantile(grid)
cdf = logit_normal.cdf(q)
cdf_ = self.evaluate(cdf)
self.assertAllClose(grid, cdf_, rtol=1e-6, atol=0.)
def testCopyOverride(self):
sigmoid = tfb.Sigmoid(low=-1., high=2., validate_args=True)
self.assertEqual(sigmoid.parameters, sigmoid.copy().parameters)
unused_sigmoid_copy = sigmoid.copy(validate_args=False)
base_params = sigmoid.parameters.copy()
copy_params = sigmoid.copy(validate_args=False).parameters.copy()
self.assertNotEqual(
base_params.pop('validate_args'), copy_params.pop('validate_args'))
self.assertEqual(base_params, copy_params)
def __init__(self,
temperature,
logits=None,
probs=None,
validate_args=False,
allow_nan_stats=True,
name="RelaxedBernoulli"):
"""Construct RelaxedBernoulli distributions.
Args:
temperature: An 0-D `Tensor`, representing the temperature
of a set of RelaxedBernoulli distributions. The temperature should be
positive.
logits: An N-D `Tensor` representing the log-odds
of a positive event. Each entry in the `Tensor` parametrizes
an independent RelaxedBernoulli distribution where the probability of an
event is sigmoid(logits). Only one of `logits` or `probs` should be
passed in.
probs: An N-D `Tensor` representing the probability of a positive event.
Each entry in the `Tensor` parameterizes an independent Bernoulli
distribution. Only one of `logits` or `probs` should be passed in.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
(e.g., mean, mode, variance) use the value "`NaN`" to indicate the
result is undefined. When `False`, an exception is raised if one or
more of the statistic's batch members are undefined.
name: Python `str` name prefixed to Ops created by this class.
Raises:
ValueError: If both `probs` and `logits` are passed, or if neither.
"""
parameters = dict(locals())
with tf.name_scope(name, values=[logits, probs, temperature]) as name:
dtype = dtype_util.common_dtype([logits, probs, temperature], tf.float32)
self._temperature = tf.convert_to_tensor(
temperature, name="temperature", dtype=dtype)
if validate_args:
with tf.control_dependencies([tf.assert_positive(temperature)]):
self._temperature = tf.identity(self._temperature)
self._logits, self._probs = distribution_util.get_logits_and_probs(
logits=logits, probs=probs, validate_args=validate_args, dtype=dtype)
super(RelaxedBernoulli, self).__init__(
distribution=logistic.Logistic(
self._logits / self._temperature,
1. / self._temperature,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name + "/Logistic"),
bijector=bijectors.Sigmoid(validate_args=validate_args),
validate_args=validate_args,
name=name)
self._parameters = parameters
def __init__(self,
centered_returns,
name='stochastic_volatility',
pretty_name='Stochastic Volatility'):
"""Construct the stochastic volatility model.
Args:
centered_returns: Float `Tensor` of shape `[num_timesteps]` giving the
mean-adjusted return (change in asset price, minus the average change)
observed at each step.
name: Python `str` name prefixed to Ops created by this class.
pretty_name: A Python `str`. The pretty name of this model.
"""
with tf.name_scope(name):
num_timesteps = ps.size0(centered_returns)
if tf.is_tensor(num_timesteps):
raise ValueError(
'Returns series length must be static, but saw '
'shape {}.'.format(centered_returns.shape))
self._prior_dist = tfd.JointDistributionCoroutine(
functools.partial(stochastic_volatility_prior_fn,
num_timesteps=num_timesteps))
self._log_likelihood_fn = functools.partial(
stochastic_volatility_log_likelihood_fn,
centered_returns=centered_returns)
def _ext_identity(params):
res = collections.OrderedDict()
res['persistence_of_volatility'] = params[0]
res['mean_log_volatility'] = params[1]
res['white_noise_shock_scale'] = params[2]
res['log_volatility'] = tf.stack(params[3], axis=-1)
return res
sample_transformations = {
'identity':
model.Model.SampleTransformation(
fn=_ext_identity,
pretty_name='Identity',
)
}
super(StochasticVolatility, self).__init__(
default_event_space_bijector=(tfb.Sigmoid(
-1., 1.), tfb.Identity(), tfb.Softplus()) +
((tfb.Identity(), ) * num_timesteps, ),
event_shape=self._prior_dist.event_shape,
dtype=self._prior_dist.dtype,
name=name,
pretty_name=pretty_name,
sample_transformations=sample_transformations,
)
def testCopyExtraArgs(self):
# Note: we cannot easily test all bijectors since each requires
# different initialization arguments. We therefore spot test a few.
sigmoid = tfb.Sigmoid(low=-1., high=2., validate_args=True)
self.assertEqual(sigmoid.parameters, sigmoid.copy().parameters)
chain = tfb.Chain(
[
tfb.Softplus(hinge_softness=[1., 2.], validate_args=True),
tfb.MatrixInverseTriL(validate_args=True)
], validate_args=True)
self.assertEqual(chain.parameters, chain.copy().parameters)
def testBijectiveAndFinite(self):
with self.test_session():
x = np.linspace(-7., 7., 100).astype(np.float32)
eps = 1e-3
y = np.linspace(eps, 1. - eps, 100).astype(np.float32)
assert_bijective_and_finite(tfb.Sigmoid(),
x,
y,
event_ndims=0,
atol=0.,
rtol=1e-4)
def testBijectiveAndFinite(self):
low = -5.
high = 8.
bijector = tfb.Sigmoid(low=low, high=high, validate_args=True)
x = np.linspace(-10, 10, num=100).astype(np.float32)
eps = 1e-6
y = np.linspace(low + eps, high - eps, num=100).astype(np.float32)
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=0)
def testEdgeCaseRequiringClipping(self):
np.set_printoptions(floatmode='unique', precision=None)
lo = np.float32(0.010489981)
hi = test_util.floats_near(0.010499111, 100,
dtype=np.float32)[:, np.newaxis]
self.assertAllEqual([100, 1], hi.shape)
xs = test_util.floats_near(9.814646, 100, dtype=np.float32)
bijector = tfb.Sigmoid(low=lo, high=hi, validate_args=True)
answers = bijector.forward(xs)
self.assertAllEqual([100, 100], answers.shape)
for ans1, hi1 in zip(self.evaluate(answers), hi):
self.assertAllLessEqual(ans1, hi1)
def test_support_works_correctly_with_HMC(self):
num_results = 2000
target = tfd.Beta(concentration1=self.dtype(1.),
concentration0=self.dtype(10.))
transformed_hmc = tfp.mcmc.TransformedTransitionKernel(
inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=tf.function(target.log_prob,
autograph=False),
step_size=1.64,
num_leapfrog_steps=2,
seed=_maybe_seed(55)),
bijector=tfb.Sigmoid())
# Recall, tfp.mcmc.sample_chain calls
# transformed_hmc.bootstrap_results too.
states, kernel_results = tfp.mcmc.sample_chain(
num_results=num_results,
# The initial state is used by inner_kernel.bootstrap_results.
# Note the input is *after* bijector.forward.
current_state=self.dtype(0.25),
kernel=transformed_hmc,
num_burnin_steps=200,
num_steps_between_results=1,
parallel_iterations=1)
self.assertEqual(num_results,
tf.compat.dimension_value(states.shape[0]))
sample_mean = tf.reduce_mean(states, axis=0)
sample_var = tf.reduce_mean(tf.math.squared_difference(
states, sample_mean),
axis=0)
[
sample_mean_,
sample_var_,
is_accepted_,
true_mean_,
true_var_,
] = self.evaluate([
sample_mean,
sample_var,
kernel_results.inner_results.is_accepted,
target.mean(),
target.variance(),
])
self.assertAllClose(true_mean_, sample_mean_, atol=0.06, rtol=0.)
self.assertAllClose(true_var_, sample_var_, atol=0.01, rtol=0.1)
self.assertNear(0.6, is_accepted_.mean(), err=0.05)
def test_support_works_correctly_with_MALA(self):
num_results = 2000
target = tfd.Beta(concentration1=self.dtype(1.),
concentration0=self.dtype(10.))
transformed_mala = tfp.mcmc.TransformedTransitionKernel(
inner_kernel=tfp.mcmc.MetropolisAdjustedLangevinAlgorithm(
target_log_prob_fn=tf.function(target.log_prob,
autograph=False),
step_size=1.,
seed=_maybe_seed(test_util.test_seed())),
bijector=tfb.Sigmoid())
# Recall, tfp.mcmc.sample_chain calls
# transformed_hmc.bootstrap_results too.
states, _ = tfp.mcmc.sample_chain(
num_results=num_results,
# The initial state is used by inner_kernel.bootstrap_results.
# Note the input is *after* bijector.forward.
current_state=self.dtype(0.25),
kernel=transformed_mala,
num_burnin_steps=200,
num_steps_between_results=1,
parallel_iterations=1)
self.assertEqual(num_results,
tf.compat.dimension_value(states.shape[0]))
sample_mean = tf.reduce_mean(states, axis=0)
sample_var = tf.reduce_mean(tf.math.squared_difference(
states, sample_mean),
axis=0)
[
sample_mean_,
sample_var_,
true_mean_,
true_var_,
] = self.evaluate([
sample_mean,
sample_var,
target.mean(),
target.variance(),
])
self.assertAllClose(true_mean_, sample_mean_, atol=0.06, rtol=0.)
self.assertAllClose(true_var_, sample_var_, atol=0.01, rtol=0.1)
def test_support_works_correctly_with_rwm(self):
num_results = 500
target = tfd.Beta(
concentration1=self.dtype(1.),
concentration0=self.dtype(10.))
transformed_rwm = tfp.mcmc.TransformedTransitionKernel(
inner_kernel=tfp.mcmc.RandomWalkMetropolis(
target_log_prob_fn=tf.function(target.log_prob, autograph=False),
new_state_fn=tfp.mcmc.random_walk_normal_fn(scale=1.5)),
bijector=tfb.Sigmoid())
# Recall, tfp.mcmc.sample_chain calls
# transformed_hmc.bootstrap_results too.
states = tfp.mcmc.sample_chain(
num_results=num_results,
# The initial state is used by inner_kernel.bootstrap_results.
# Note the input is *after* bijector.forward.
current_state=self.dtype(0.25),
kernel=transformed_rwm,
num_burnin_steps=200,
num_steps_between_results=1,
trace_fn=None,
seed=test_util.test_seed())
self.assertEqual(num_results, tf.compat.dimension_value(states.shape[0]))
sample_mean = tf.reduce_mean(states, axis=0)
sample_var = tf.reduce_mean(
tf.math.squared_difference(states, sample_mean), axis=0)
[
sample_mean_,
sample_var_,
true_mean_,
true_var_,
] = self.evaluate([
sample_mean,
sample_var,
target.mean(),
target.variance(),
])
self.assertAllClose(true_mean_, sample_mean_,
atol=0.15, rtol=0.)
self.assertAllClose(true_var_, sample_var_,
atol=0.03, rtol=0.2)
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):
bijector_test_util.assert_scalar_congruency(tfb.Sigmoid(),
lower_x=-7.,
upper_x=7.,
eval_func=self.evaluate)
def mk_sigmoid_beta():
beta = tfd.Beta(concentration0=1., concentration1=2.)
# Not inverting the sigmoid bijector makes a kooky distribution, but
# nuts should still conserve it (with a smaller step size).
return tfb.Sigmoid()(beta)
def mk_logit_beta():
beta = tfd.Beta(concentration0=1., concentration1=2.)
return tfb.Invert(tfb.Sigmoid())(beta)
def _make_asvi_trainable_variables(prior,
mean_field=False,
initial_prior_weight=0.5):
"""Generates parameter dictionaries given a prior distribution and list."""
with tf.name_scope('make_asvi_trainable_variables'):
param_dicts = []
prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access
for dist in prior_dists:
original_dist = dist.distribution if isinstance(dist,
Root) else dist
substituted_dist = _as_trainable_family(original_dist)
# Grab the base distribution if it exists
try:
actual_dist = substituted_dist.distribution
except AttributeError:
actual_dist = substituted_dist
new_params_dict = {}
# Build trainable ASVI representation for each distribution's parameters.
parameter_properties = actual_dist.parameter_properties(
dtype=actual_dist.dtype)
if isinstance(original_dist, sample.Sample):
posterior_batch_shape = ps.concat([
actual_dist.batch_shape_tensor(),
original_dist.sample_shape
],
axis=0)
else:
posterior_batch_shape = actual_dist.batch_shape_tensor()
for param, value in actual_dist.parameters.items():
if param in (_NON_STATISTICAL_PARAMS +
_NON_TRAINABLE_PARAMS) or value is None:
continue
actual_event_shape = parameter_properties[param].shape_fn(
actual_dist.event_shape_tensor())
try:
bijector = parameter_properties[
param].default_constraining_bijector_fn()
except NotImplementedError:
bijector = tfb.Identity()
if mean_field:
prior_weight = None
else:
unconstrained_ones = tf.ones(shape=ps.concat([
posterior_batch_shape,
bijector.inverse_event_shape_tensor(actual_event_shape)
],
axis=0),
dtype=actual_dist.dtype)
prior_weight = tfp_util.TransformedVariable(
initial_prior_weight * unconstrained_ones,
bijector=tfb.Sigmoid(),
name='prior_weight/{}/{}'.format(dist.name, param))
# If the prior distribution was a tfd.Sample wrapping a base
# distribution, we want to give every single sample in the prior its
# own lambda and alpha value (rather than having a single lambda and
# alpha).
if isinstance(original_dist, sample.Sample):
value = tf.reshape(
value,
ps.concat([
actual_dist.batch_shape_tensor(),
ps.ones(
ps.rank_from_shape(
original_dist.sample_shape)),
actual_event_shape
],
axis=0))
value = tf.broadcast_to(
value,
ps.concat([posterior_batch_shape, actual_event_shape],
axis=0))
new_params_dict[param] = ASVIParameters(
prior_weight=prior_weight,
mean_field_parameter=tfp_util.TransformedVariable(
value,
bijector=bijector,
name='mean_field_parameter/{}/{}'.format(
dist.name, param)))
param_dicts.append(new_params_dict)
return param_dicts
def fn(x):
return tf.math.log(tfb.Sigmoid().forward(x))
