text: # Copyright 2018 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.
  # ============================================================================
  """Dynamic Linear Regression State Space Model Tests."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # Dependency imports
  
  import numpy as np
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  from tensorflow_probability.python import distributions as tfd
  from tensorflow_probability.python.internal import test_case
  from tensorflow_probability.python.sts import DynamicLinearRegression
  from tensorflow_probability.python.sts import DynamicLinearRegressionStateSpaceModel
  
  from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
  
  
  class _DynamicLinearRegressionStateSpaceModelTest(object):
  
    def test_basic_statistics_no_latent_variance(self):
      batch_shape = [4, 3]
      num_timesteps = 10
      num_features = 2
      drift_scale = 0.
  
      design_matrix = self._build_placeholder(
          np.random.randn(*(batch_shape + [num_timesteps, num_features])))
  
      initial_state_loc = self._build_placeholder(
          np.random.randn(*(batch_shape + [num_features])))
      initial_state_scale = tf.zeros_like(initial_state_loc)
      initial_state_prior = tfd.MultivariateNormalDiag(
          loc=initial_state_loc, scale_diag=initial_state_scale)
  
      ssm = DynamicLinearRegressionStateSpaceModel(
          num_timesteps=num_timesteps,
          design_matrix=design_matrix,
          drift_scale=drift_scale,
          initial_state_prior=initial_state_prior)
  
      predicted_time_series = tf.linalg.matmul(
          design_matrix, initial_state_loc[..., tf.newaxis])
  
      self.assertAllEqual(self.evaluate(ssm.mean()), predicted_time_series)
      self.assertAllEqual(*self.evaluate((ssm.stddev(),
                                          tf.zeros_like(predicted_time_series))))
  
    def test_initial_state_broadcasts_over_batch(self):
      batch_shape = [4, 3]
      num_timesteps = 10
      num_features = 2
      drift_scale = 1.
  
      initial_state_loc = self._build_placeholder([0.1, -0.2])
      initial_state_scale = self._build_placeholder([0.42, 0.314])
  
      design_matrix = self._build_placeholder(
          np.random.randn(*(batch_shape + [num_timesteps, num_features])))
  
      initial_state_prior = tfd.MultivariateNormalDiag(
          loc=initial_state_loc, scale_diag=initial_state_scale)
  
      ssm = DynamicLinearRegressionStateSpaceModel(
          num_timesteps=num_timesteps,
          design_matrix=design_matrix,
          drift_scale=drift_scale,
          initial_state_prior=initial_state_prior)
  
      sample = ssm.sample()
  
      ll, means, _, _, _, _, _ = ssm.forward_filter(sample)
  
      self.assertAllEqual(batch_shape + [num_timesteps, 1],
                          self.evaluate(sample).shape)
  
      self.assertAllEqual(batch_shape + [num_timesteps],
                          self.evaluate(ll).shape)
  
      self.assertAllEqual(batch_shape + [num_timesteps, num_features],
                          self.evaluate(means).shape)
  
    def test_matrices_from_component(self):
      num_timesteps = 4
      num_features = 3
      drift_scale = 1.23
  
      design_matrix = self._build_placeholder(
          np.random.randn(num_timesteps, num_features))
  
      component = DynamicLinearRegression(design_matrix=design_matrix)
  
      ssm = component.make_state_space_model(num_timesteps, [drift_scale])
  
      for t in range(num_timesteps):
        observation_matrix = self.evaluate(
            ssm.get_observation_matrix_for_timestep(t).to_dense())
  
        self.assertAllClose(self.evaluate(design_matrix[tf.newaxis, t]),
                            observation_matrix)
  
        observation_noise_mean = self.evaluate(
            ssm.get_observation_noise_for_timestep(t).mean())
        observation_noise_covariance = self.evaluate(
            ssm.get_observation_noise_for_timestep(t).covariance())
  
        self.assertAllClose([0.0], observation_noise_mean)
        self.assertAllClose([[0.0]], observation_noise_covariance)
  
        transition_matrix = self.evaluate(
            ssm.get_transition_matrix_for_timestep(t).to_dense())
  
        self.assertAllClose(np.eye(num_features), transition_matrix)
  
        transition_noise_mean = self.evaluate(
            ssm.get_transition_noise_for_timestep(t).mean())
        transition_noise_covariance = self.evaluate(
            ssm.get_transition_noise_for_timestep(t).covariance())
  
        self.assertAllClose(np.zeros([num_features]),
                            transition_noise_mean)
        self.assertAllClose(np.square(drift_scale) * np.eye(num_features),
                            transition_noise_covariance)
  
    def _build_placeholder(self, ndarray):
      ndarray = np.asarray(ndarray).astype(self.dtype)
      return tf1.placeholder_with_default(
          input=ndarray, shape=ndarray.shape if self.use_static_shape else None)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class DynamicRegressionStateSpaceModelTestStaticShape32(
      test_case.TestCase, _DynamicLinearRegressionStateSpaceModelTest):
    dtype = np.float32
    use_static_shape = True
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class DynamicRegressionStateSpaceModelTestDynamicShape32(
      test_case.TestCase, _DynamicLinearRegressionStateSpaceModelTest):
    dtype = np.float32
    use_static_shape = False
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class DynamicRegressionStateSpaceModelTestStaticShape64(
      test_case.TestCase, _DynamicLinearRegressionStateSpaceModelTest):
    dtype = np.float64
    use_static_shape = True
  
  
  if __name__ == "__main__":
    tf.test.main()
  