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)