text: # Copyright 2021 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Tests for MarkovChain."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  from tensorflow_probability.python import bijectors as tfb
  from tensorflow_probability.python import distributions as tfd
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import prefer_static as ps
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  
  
  @test_util.test_graph_and_eager_modes
  class MarkovChainTest(test_util.TestCase):
  
    def test_error_when_transition_modifies_batch_shape(self):
      loses_batch_shape = tfd.MarkovChain(
          initial_state_prior=tfd.Normal(loc=0., scale=[1., 1.]),
          transition_fn=lambda _, x: tfd.Independent(  # pylint: disable=g-long-lambda
              tfd.Normal(loc=0., scale=tf.ones_like(x)),
              reinterpreted_batch_ndims=1),
          num_steps=5)
      x = self.evaluate(loses_batch_shape.sample([2], seed=test_util.test_seed()))
      with self.assertRaisesRegexp(ValueError, 'batch shape is incorrect'):
        loses_batch_shape.log_prob(x)
  
      gains_batch_shape = tfd.MarkovChain(
          initial_state_prior=tfd.Independent(
              tfd.Normal(loc=0., scale=[1., 1.]),
              reinterpreted_batch_ndims=1),
          transition_fn=lambda _, x: tfd.Normal(loc=0., scale=tf.ones_like(x)),
          num_steps=5)
      x = self.evaluate(gains_batch_shape.sample([2], seed=test_util.test_seed()))
      with self.assertRaisesRegexp(ValueError, 'batch shape is incorrect'):
        gains_batch_shape.log_prob(x)
  
    def test_log_prob_matches_linear_gaussian_ssm(self):
      dim = 2
      batch_shape = [3, 1]
      seed, *model_seeds = samplers.split_seed(test_util.test_seed(), n=6)
  
      # Sample a random linear Gaussian process.
      prior_loc = self.evaluate(
          tfd.Normal(0., 1.).sample(batch_shape + [dim], seed=model_seeds[0]))
      prior_scale = self.evaluate(
          tfd.InverseGamma(1., 1.).sample(batch_shape + [dim],
                                          seed=model_seeds[1]))
      transition_matrix = self.evaluate(
          tfd.Normal(0., 1.).sample([dim, dim], seed=model_seeds[2]))
      transition_bias = self.evaluate(
          tfd.Normal(0., 1.).sample(batch_shape + [dim], seed=model_seeds[3]))
      transition_scale_tril = self.evaluate(
          tf.linalg.cholesky(
              tfd.WishartTriL(df=dim, scale_tril=tf.eye(dim)).sample(
                  seed=model_seeds[4])))
  
      initial_state_prior = tfd.MultivariateNormalDiag(
          loc=prior_loc, scale_diag=prior_scale, name='initial_state_prior')
  
      lgssm = tfd.LinearGaussianStateSpaceModel(
          num_timesteps=7,
          transition_matrix=transition_matrix,
          transition_noise=tfd.MultivariateNormalTriL(
              loc=transition_bias, scale_tril=transition_scale_tril),
          # Trivial observation model to pass through the latent state.
          observation_matrix=tf.eye(dim),
          observation_noise=tfd.MultivariateNormalDiag(loc=tf.zeros(dim),
                                                       scale_diag=tf.zeros(dim)),
          initial_state_prior=initial_state_prior)
  
      markov_chain = tfd.MarkovChain(
          initial_state_prior=initial_state_prior,
          transition_fn=lambda _, x: tfd.MultivariateNormalTriL(  # pylint: disable=g-long-lambda
              loc=tf.linalg.matvec(transition_matrix, x) + transition_bias,
              scale_tril=transition_scale_tril),
          num_steps=7)
  
      x = markov_chain.sample(5, seed=seed)
      self.assertAllClose(lgssm.log_prob(x), markov_chain.log_prob(x), rtol=1e-5)
  
    @test_util.numpy_disable_test_missing_functionality(
        'JointDistributionNamedAutoBatched')
    def test_docstring_example_autoregressive_process(self):
  
      def transition_fn(_, previous_state):
        return tfd.JointDistributionNamedAutoBatched(
            # The previous state may include batch dimensions. Since the log scale
            # is a scalar quantity, its shape is the batch shape.
            batch_ndims=ps.rank(previous_state['log_scale']),
            model={
                # The autoregressive coefficients and the `log_scale` each follow
                # an independent slow-moving random walk.
                'coefs': tfd.Normal(loc=previous_state['coefs'], scale=0.01),
                'log_scale': tfd.Normal(loc=previous_state['log_scale'],
                                        scale=0.01),
                # The level is a linear combination of the previous *two* levels,
                # with additional noise of scale `exp(log_scale)`.
                'level': lambda coefs, log_scale: tfd.Normal(  # pylint: disable=g-long-lambda
                    loc=(coefs[..., 0] * previous_state['level'] +
                         coefs[..., 1] * previous_state['previous_level']),
                    scale=tf.exp(log_scale)),
                # Store the previous level to access at the next step.
                'previous_level': tfd.Deterministic(previous_state['level'])})
  
      process = tfd.MarkovChain(
          # For simplicity, define the prior as a 'transition' from fixed values.
          initial_state_prior=transition_fn(
              0, previous_state={
                  'coefs': [0.7, -0.2],
                  'log_scale': -1.,
                  'level': 0.,
                  'previous_level': 0.}),
          transition_fn=transition_fn,
          num_steps=100)
      self.assertAllEqualNested(process.event_shape,
                                {'coefs': [100, 2], 'log_scale': [100],
                                 'level': [100], 'previous_level': [100]})
      self.assertAllEqual(process.batch_shape, [])
  
      x = process.sample(5, seed=test_util.test_seed())
      self.assertAllEqual(x['coefs'].shape, [5, 100, 2])
      self.assertAllEqual(x['log_scale'].shape, [5, 100])
      self.assertAllEqual(x['level'].shape, [5, 100])
      self.assertAllEqual(x['previous_level'].shape, [5, 100])
  
      lp = process.log_prob(x)
      self.assertAllEqual(lp.shape, [5])
  
      x2, lp2 = process.experimental_sample_and_log_prob(
          2, seed=test_util.test_seed())
      self.assertAllClose(lp2, process.log_prob(x2))
  
    @parameterized.named_parameters(
        ('float32_dynamic', tf.float32, True),
        ('float64_static', tf.float64, False))
    def test_docstring_example_batch_gaussian_walk(self,
                                                   float_dtype,
                                                   use_dynamic_shapes):
      if tf.executing_eagerly() and use_dynamic_shapes:
        self.skipTest('No dynamic shapes in eager mode.')
      def _as_tensor(x, dtype=None):
        x = ps.cast(x, dtype=dtype if dtype else float_dtype)
        if use_dynamic_shapes:
          x = tf1.placeholder_with_default(x, shape=None)
        return x
  
      scales = _as_tensor([0.5, 0.3, 0.2, 0.2, 0.3, 0.2, 0.7])
      batch_gaussian_walk = tfd.MarkovChain(
          # The prior distribution determines the batch shape for the chain.
          # Transitions must respect this batch shape.
          initial_state_prior=tfd.Normal(loc=_as_tensor([-10., 0., 10.]),
                                         scale=_as_tensor([1., 1., 1.])),
          transition_fn=lambda t, x: tfd.Normal(  # pylint: disable=g-long-lambda
              loc=x,
              # The `num_steps` dimension will always be leftmost in `x`, so we
              # pad the scale to the same rank as `x` so that the shapes line up.
              scale=tf.reshape(
                  tf.gather(scales, t),
                  ps.concat([[-1],
                             ps.ones(ps.rank(x) - 1, dtype=tf.int32)], axis=0))),
          # Limit to eight steps since we only specified scales for seven
          # transitions.
          num_steps=8)
      self.assertAllEqual(batch_gaussian_walk.event_shape_tensor(), [8])
      self.assertAllEqual(batch_gaussian_walk.batch_shape_tensor(), [3])
  
      x = batch_gaussian_walk.sample(5, seed=test_util.test_seed())
      self.assertAllEqual(ps.shape(x), [5, 3, 8])
      lp = batch_gaussian_walk.log_prob(x)
      self.assertAllEqual(ps.shape(lp), [5, 3])
  
      x2, lp2 = batch_gaussian_walk.experimental_sample_and_log_prob(
          [2], seed=test_util.test_seed())
      self.assertAllClose(lp2, batch_gaussian_walk.log_prob(x2))
  
    def test_docstring_example_gaussian_walk(self):
      gaussian_walk = tfd.MarkovChain(
          initial_state_prior=tfd.Normal(loc=0., scale=1.),
          transition_fn=lambda _, x: tfd.Normal(loc=x, scale=1.),
          num_steps=100)
      self.assertAllEqual(gaussian_walk.event_shape, [100])
      self.assertAllEqual(gaussian_walk.batch_shape, [])
  
      x = gaussian_walk.sample(5, seed=test_util.test_seed())
      self.assertAllEqual(x.shape, [5, 100])
      lp = gaussian_walk.log_prob(x)
      self.assertAllEqual(lp.shape, [5])
  
      n = tfd.Normal(0., 1.)
      expected_lp = (n.log_prob(x[:, 0]) +
                     tf.reduce_sum(n.log_prob(x[:, 1:] - x[:, :-1]), axis=-1))
      self.assertAllClose(lp, expected_lp)
  
      x2, lp2 = gaussian_walk.experimental_sample_and_log_prob(
          [2], seed=test_util.test_seed())
      self.assertAllClose(lp2, gaussian_walk.log_prob(x2))
  
    def test_non_autobatched_joint_distribution(self):
  
      def transition_fn(_, previous_state):
        return tfd.JointDistributionNamed(
            {
                # The autoregressive coefficients and the `log_scale` each follow
                # an independent slow-moving random walk.
                'coefs': tfd.Independent(
                    tfd.Normal(loc=previous_state['coefs'], scale=0.01),
                    reinterpreted_batch_ndims=1),
                'log_scale': tfd.Normal(loc=previous_state['log_scale'],
                                        scale=0.01),
                # The level is a linear combination of the previous *two* levels,
                # with additional noise of scale `exp(log_scale)`.
                'level': lambda coefs, log_scale: tfd.Normal(  # pylint: disable=g-long-lambda
                    loc=(coefs[..., 0] * previous_state['level'] +
                         coefs[..., 1] * previous_state['previous_level']),
                    scale=tf.exp(log_scale)),
                # Store the previous level to access at the next step.
                'previous_level': tfd.Deterministic(previous_state['level'])})
  
      process = tfd.MarkovChain(
          # For simplicity, define the prior as a 'transition' from fixed values.
          initial_state_prior=transition_fn(
              0, previous_state={
                  'coefs': [0.7, -0.2],
                  'log_scale': -1.,
                  'level': 0.,
                  'previous_level': 0.}),
          transition_fn=transition_fn,
          num_steps=100)
      self.assertAllEqualNested(process.event_shape,
                                {'coefs': [100, 2], 'log_scale': [100],
                                 'level': [100], 'previous_level': [100]})
      self.assertAllEqual(process.batch_shape,
                          {'coefs': [], 'log_scale': [],
                           'level': [], 'previous_level': []})
  
      x = process.sample(5, seed=test_util.test_seed())
      self.assertAllEqual(x['coefs'].shape, [5, 100, 2])
      self.assertAllEqual(x['log_scale'].shape, [5, 100])
      self.assertAllEqual(x['level'].shape, [5, 100])
      self.assertAllEqual(x['previous_level'].shape, [5, 100])
  
      lp = process.log_prob(x)
      self.assertAllEqual(lp.shape, [5])
  
      x2, lp2 = process.experimental_sample_and_log_prob(
          2, seed=test_util.test_seed())
      self.assertAllClose(lp2, process.log_prob(x2))
  
    def test_log_prob_ratio(self):
      p = tfd.MarkovChain(
          initial_state_prior=tfd.Normal(0., 1.),
          transition_fn=lambda _, x: tfd.Normal(x, tf.nn.softplus(x)),
          num_steps=10)
      q = tfd.MarkovChain(
          initial_state_prior=tfd.Normal(-10, 3.),
          transition_fn=lambda _, x: tfd.Normal(x, tf.abs(x)),
          num_steps=10)
      x = self.evaluate(p.sample(4, seed=test_util.test_seed()))
      y = self.evaluate(q.sample(4, seed=test_util.test_seed()))
      self.assertAllClose(
          p.log_prob(x) - q.log_prob(y),
          log_prob_ratio.log_prob_ratio(p, x, q, y), atol=1e-5)
  
    def test_unexpected_num_steps_raises(self):
      p = tfd.MarkovChain(
          initial_state_prior=tfd.Normal(0., 1.),
          transition_fn=lambda _, x: tfd.Normal(x, tf.nn.softplus(x)),
          num_steps=10,
          validate_args=True)
      with self.assertRaisesRegex(
          (ValueError, tf.errors.InvalidArgumentError),
          'does not match the expected num_steps'):
        p.log_prob(tf.zeros([11]))
  
  
  @test_util.test_graph_and_eager_modes
  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())
  
  
  if __name__ == '__main__':
    tf.test.main()
  