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.
  # ============================================================================
  """Tests for TruncatedNormal distribution."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import itertools
  
  # Dependency imports
  from absl.testing import parameterized
  import numpy as np
  from scipy import stats as sp_stats
  
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  
  from tensorflow_probability.python.internal import test_case
  from tensorflow_probability.python.internal import test_util as tfp_test_util
  from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
  from tensorflow.python.ops import gradient_checker_v2  # pylint: disable=g-direct-tensorflow-import
  
  
  tfd = tfp.distributions
  
  EPSILON = 1e-5
  
  
  def scipy_trunc_norm_dist(loc, scale, low, high):
    """Construct a scipy.sp_stats.truncnorm for the (scalar) parameters given.
  
    Note: scipy's definition of the parameters is slightly different.
    https://github.com/scipy/scipy/issues/7591
  
    Args:
      loc: Params describing distribution (doesn't support batch)
      scale:
      low:
      high:
  
    Returns:
      scipy frozen distribution.
    """
    a = (low - loc) / scale
    b = (high - loc) / scale
    return sp_stats.truncnorm(a, b, loc=loc, scale=scale)
  
  
  class _TruncatedNormalTestCase(test_case.TestCase):
  
    def setUp(self):
      self._rng = np.random.RandomState(42)
  
    def assertAllFinite(self, a):
      is_finite = np.isfinite(a)
      all_true = np.ones_like(is_finite, dtype=np.bool)
      self.assertAllEqual(all_true, is_finite)
  
    def assertAllGreaterEqual(self, a, b):
      comparison = a >= b
      all_true = np.ones_like(comparison, dtype=np.bool)
      self.assertAllEqual(comparison, all_true)
  
    def assertAllLessEqual(self, a, b):
      comparison = a <= b
      all_true = np.ones_like(comparison, dtype=np.bool)
      self.assertAllEqual(comparison, all_true)
  
    def assertEmpiricalDistributionsEqual(self, sample_a, sample_b, rtol=1e-6,
                                          atol=1e-6):
      """Assert the empirical distribution of two set of samples is similar.
  
      Args:
        sample_a: Flat numpy array of samples from dist a.
        sample_b: Flat numpy array of samples from dist b.
        rtol: Relative tolerances in the histogram comparison.
        atol: Absolute tolerances in the histogram comparison.
      """
      self.assertAllFinite(sample_a)
      self.assertAllFinite(sample_b)
  
      lb = np.min([sample_a, sample_b])
      ub = np.max([sample_a, sample_b])
  
      hist_a = np.histogram(sample_a, range=(lb, ub), bins=30, density=True)[0]
      hist_b = np.histogram(sample_b, range=(lb, ub), bins=30, density=True)[0]
  
      self.assertAllClose(hist_a, hist_b, rtol=rtol, atol=atol)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class TruncatedNormalStandaloneTestCase(_TruncatedNormalTestCase,
                                          parameterized.TestCase):
  
    def _testParamShapes(self, desired_shape):
      tn_param_shapes = tfd.TruncatedNormal.param_shapes(desired_shape)
      # Check the shapes by comparison with the untruncated Normal.
      n_param_shapes = tfd.Normal.param_shapes(desired_shape)
      self.assertAllEqual(
          self.evaluate(tn_param_shapes["loc"]),
          self.evaluate(n_param_shapes["loc"]))
      self.assertAllEqual(
          self.evaluate(tn_param_shapes["scale"]),
          self.evaluate(n_param_shapes["scale"]))
      self.assertAllEqual(
          self.evaluate(tn_param_shapes["low"]),
          self.evaluate(n_param_shapes["loc"]))
      self.assertAllEqual(
          self.evaluate(tn_param_shapes["high"]),
          self.evaluate(n_param_shapes["loc"]))
  
      loc = tf.zeros(tn_param_shapes["loc"])
      scale = tf.ones(tn_param_shapes["scale"])
      high = tf.ones(tn_param_shapes["high"])
      low = tf.zeros(tn_param_shapes["low"])
      sample_shape = self.evaluate(
          tf.shape(
              input=tfd.TruncatedNormal(loc=loc, scale=scale, low=low,
                                        high=high).sample()))
      self.assertAllEqual(desired_shape, sample_shape)
  
    def testParamShapes(self):
      desired_shape = [10, 3, 4]
      self._testParamShapes(desired_shape)
      self._testParamShapes(tf.constant(desired_shape))
  
    def testParamStaticShapes(self):
      sample_shape = [7]
      self._testParamShapes(sample_shape)
      self._testParamShapes(tf.TensorShape(sample_shape))
  
    def testShapeWithPlaceholders(self):
      if tf.executing_eagerly():
        return
      loc = tf1.placeholder_with_default(input=5., shape=None)
      scale = tf1.placeholder_with_default(input=[1., 2], shape=None)
      ub = tf1.placeholder_with_default(input=[10., 11.], shape=None)
      lb = tf1.placeholder_with_default(input=[-1.], shape=None)
      dist = tfd.TruncatedNormal(loc, scale, lb, ub)
  
      self.assertEqual(dist.batch_shape, tf.TensorShape(None))
      self.assertEqual(dist.event_shape, ())
      self.assertAllEqual(self.evaluate(dist.event_shape_tensor()), [])
      self.assertAllEqual(self.evaluate(dist.batch_shape_tensor()), [2])
      self.assertAllEqual(self.evaluate(dist.sample(5)).shape, [5, 2])
  
      ub = tf1.placeholder_with_default(input=[[5., 11.]], shape=None)
      dist = tfd.TruncatedNormal(loc, scale, lb, ub)
      self.assertAllEqual(self.evaluate(dist.sample(5)).shape, [5, 1, 2])
  
    def testBatchSampling(self):
      """Check (empirically) the different parameters in a batch are respected.
      """
      n = int(1e5)
      lb = [[-1.0, 9.0], [0., 8.]]
      ub = [[1.0, 11.0], [5., 20.]]
      dist = tfd.TruncatedNormal(
          loc=[[0., 10.], [0., 10.]], scale=[[1., 1.], [5., 5.]], low=lb, high=ub)
      x = self.evaluate(dist.sample(n, seed=tfp_test_util.test_seed()))
      self.assertEqual(x.shape, (n, 2, 2))
  
      means = np.mean(x, axis=0)
      var = np.var(x, axis=0)
      self.assertAllClose(
          means, [[0., 10.], [2.299, 12.48]], rtol=1e-2, atol=1e-2)
      self.assertAllClose(var, [[0.29, 0.29], [1.99, 8.74]], rtol=1e-2, atol=1e-2)
  
      empirical_lb = np.min(x, axis=0)
      self.assertAllClose(empirical_lb, lb, atol=0.1)
      empirical_ub = np.max(x, axis=0)
      self.assertAllClose(empirical_ub, ub, atol=0.1)
  
    @parameterized.parameters(
        (0., 1., -1., 1.),
        (1., 1., 0., 2.),
        (-0.5, 0.5, -0.9, -0.4),
        (10., 3.0, 9.9, 25.),
        (2., 1.5, 0.1, 1.9))
    def testMomentsEmpirically(self, loc, scale, low, high):
      n = int(2e5)
      dist = tfd.TruncatedNormal(loc=loc, scale=scale, low=low, high=high)
      x = self.evaluate(dist.sample(n, seed=tfp_test_util.test_seed()))
      empirical_mean = np.mean(x)
      empirical_var = np.var(x)
      expected_mean = self.evaluate(dist.mean())
      expected_var = self.evaluate(dist.variance())
      self.assertAlmostEqual(expected_mean, empirical_mean, places=1)
      self.assertAlmostEqual(expected_var, empirical_var, places=1)
  
    def testNegativeSigmaFails(self):
      with self.assertRaisesOpError("Condition x > 0 did not hold"):
        dist = tfd.TruncatedNormal(
            loc=0., scale=-0.1, low=-1.0, high=1.0, validate_args=True)
        self.evaluate(dist.mean())
  
    def testIncorrectBoundsFails(self):
      with self.assertRaisesOpError("Condition x > 0 did not hold"):
        dist = tfd.TruncatedNormal(
            loc=0., scale=-0.1, low=1.0, high=-1.0, validate_args=True)
        self.evaluate(dist.mean())
  
      with self.assertRaisesOpError("Condition x > 0 did not hold"):
        dist = tfd.TruncatedNormal(
            loc=0., scale=-0.1, low=1.0, high=1.0, validate_args=True)
        self.evaluate(dist.mean())
  
    @parameterized.parameters(
        (0., 1., -1., 1.),
        (1., 1., 0., 2.),
        (-0.5, 0.5, -0.9, -0.4),
        (10., 3.0, 9.9, 25.),
        (2., 1.5, 0.1, 1.9),
        (-2., 0.2, -1.5, -0.5))
    def testMode(self, loc, scale, low, high):
      dist = tfd.TruncatedNormal(loc=loc, scale=scale, low=low, high=high)
      mode = np.asscalar(self.evaluate(dist.mode()))
      if loc < low:
        expected_mode = low
      elif loc > high:
        expected_mode = high
      else:
        expected_mode = loc
      self.assertAlmostEqual(mode, expected_mode)
  
    @parameterized.parameters((np.float32), (np.float64))
    def testReparametrizable(self, dtype=np.float32):
      loc = tf.Variable(dtype(0.1))
      scale = tf.Variable(dtype(1.1))
      low = tf.Variable(dtype(-10.0))
      high = tf.Variable(dtype(5.0))
      self.evaluate([v.initializer for v in (loc, scale, low, high)])
  
      def f(loc, scale, low, high):
        dist = tfd.TruncatedNormal(loc=loc, scale=scale, low=low, high=high)
  
        n = int(2e5)
        return tf.reduce_mean(
            input_tensor=tf.abs(dist.sample(n, seed=tfp_test_util.test_seed())))
  
      err = gradient_checker_v2.max_error(
          *gradient_checker_v2.compute_gradient(
              f, [loc, scale, low, high], delta=0.1))
  
      # These gradients are noisy due to sampling.
      self.assertLess(err, 0.05)
  
    def testReparametrizableBatch(self):
      def samples_sum(loc):
        dist = tfp.distributions.TruncatedNormal(
            loc=loc, scale=1., low=-1., high=1.)
        return tf.reduce_sum(input_tensor=dist.sample(100))
  
      loc = tf.constant([0., 1.])
      _, dy_loc = self.evaluate(tfp.math.value_and_gradient(samples_sum, loc))
      self.assertAllGreaterEqual(dy_loc, 0.)
  
    @parameterized.parameters(
        itertools.product((np.float32, np.float64),
                          ("prob", "log_prob", "cdf", "log_cdf",
                           "survival_function", "log_survival_function"))
    )
    def testGradientsFx(self, dtype, fn_name):
      loc = tf.Variable(dtype(0.1))
      scale = tf.Variable(dtype(3.0))
      low = tf.Variable(dtype(-10.0))
      high = tf.Variable(dtype(5.0))
  
      x = np.array([-1.0, 0.01, 0.1, 1., 4.9]).astype(dtype)
      self.evaluate([v.initializer for v in (loc, scale, low, high)])
  
      def f(loc, scale):
        dist = tfd.TruncatedNormal(loc=loc, scale=scale, low=low, high=high)
        func = getattr(dist, fn_name)
        return tf.reduce_mean(input_tensor=func(x))
  
      err = gradient_checker_v2.max_error(
          *gradient_checker_v2.compute_gradient(f, [loc, scale]))
      self.assertLess(err, 1e-2)
  
    @parameterized.parameters(
        itertools.product((np.float32, np.float64),
                          ("entropy", "mean", "variance", "mode"))
    )
    def testGradientsNx(self, dtype, fn_name):
      loc = tf.Variable(dtype(0.1))
      scale = tf.Variable(dtype(3.0))
      low = tf.Variable(dtype(-10.0))
      high = tf.Variable(dtype(5.0))
  
      self.evaluate([v.initializer for v in (loc, scale, low, high)])
  
      def f(loc, scale):
        dist = tfd.TruncatedNormal(loc=loc, scale=scale, low=low, high=high)
        func = getattr(dist, fn_name)
        return func()
  
      if fn_name not in ["mode"]:
        err = gradient_checker_v2.max_error(
            *gradient_checker_v2.compute_gradient(f, [loc, scale]))
        self.assertLess(err, 0.005)
      else:
        err = gradient_checker_v2.max_error(
            *gradient_checker_v2.compute_gradient(
                lambda x: f(x, scale), [loc]))
        self.assertLess(err, 0.005)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  @parameterized.parameters(
      (0.0, 1.0),
      (10.0, 1.0),
      (-0.3, 2.0),
      (100., 5.0),
      )
  class TruncatedNormalTestCompareWithNormal(_TruncatedNormalTestCase,
                                             parameterized.TestCase):
    """Test by comparing TruncatedNormals with wide bounds and unbounded Normal.
    """
  
    def constructDists(self, loc, scale):
      truncated_dist = tfd.TruncatedNormal(loc=loc, scale=scale,
                                           low=loc - (10. * scale),
                                           high=loc + (10. * scale))
      normal_dist = tfd.Normal(loc=loc, scale=scale)
      return truncated_dist, normal_dist
  
    def testEntropy(self, loc, scale):
      truncated_dist, normal_dist = self.constructDists(loc, scale)
      self.assertEqual(
          self.evaluate(truncated_dist.entropy()),
          self.evaluate(normal_dist.entropy()))
  
    def testSampling(self, loc, scale):
      n = 1000000
      truncated_dist, normal_dist = self.constructDists(loc, scale)
      truncated_samples = self.evaluate(truncated_dist.sample(n)).flatten()
      lb = self.evaluate(truncated_dist.low)
      ub = self.evaluate(truncated_dist.high)
      self.assertAllGreaterEqual(truncated_samples, lb)
      self.assertAllLessEqual(truncated_samples, ub)
  
      normal_samples = self.evaluate(normal_dist.sample(
          n, seed=tfp_test_util.test_seed())).flatten()
      # Rejection sample the normal distribution
      rejection_samples = normal_samples[normal_samples >= lb]
      rejection_samples = rejection_samples[rejection_samples <= ub]
  
      self.assertEmpiricalDistributionsEqual(
          truncated_samples, rejection_samples, rtol=1e-2, atol=1e-1)
  
    def testLogProb(self, loc, scale):
      truncated_dist, normal_dist = self.constructDists(loc, scale)
      low = self.evaluate(truncated_dist.low)
      high = self.evaluate(truncated_dist.high)
      test_x = list(np.float32(
          np.random.uniform(low, high, 10)))
      test_x += [low, high, low + EPSILON, high - EPSILON]
      tr_log_prob = self.evaluate(truncated_dist.log_prob(test_x))
      n_log_prob = self.evaluate(normal_dist.log_prob(test_x))
      self.assertAllClose(tr_log_prob, n_log_prob, rtol=1e-4, atol=1e-4)
  
      no_support_log_prob = self.evaluate(
          truncated_dist.log_prob(
              np.float32(
                  [low - EPSILON, high + EPSILON, low - 100., high + 100.]
              )))
      self.assertAllEqual(no_support_log_prob,
                          [np.log(0.)] * len(no_support_log_prob))
  
    def testCDF(self, loc, scale):
      truncated_dist, normal_dist = self.constructDists(loc, scale)
      low = self.evaluate(truncated_dist.low)
      high = self.evaluate(truncated_dist.high)
      test_x = list(
          np.float32(np.random.uniform(low, high, 10)))
      test_x += [low, high, low + EPSILON, high - EPSILON]
      tr_cdf = self.evaluate(truncated_dist.cdf(test_x))
      n_cdf = self.evaluate(normal_dist.cdf(test_x))
      self.assertAllClose(tr_cdf, n_cdf, rtol=1e-4, atol=1e-4)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  @parameterized.parameters(
      (0., 1., -1., 1.),
      (1., 1., 0., 2.),
      (-0.5, 0.5, -0.9, -0.4),
      (10., 3.0, 9.9, 25.),
      (2., 1.5, 0.1, 1.9),
      (-2., 0.2, -1.5, -0.5))
  class TruncatedNormalTestCompareWithScipy(_TruncatedNormalTestCase,
                                            parameterized.TestCase):
  
    def constructDists(self, loc, scale, low, high):
      tf_dist = tfd.TruncatedNormal(loc=loc, scale=scale,
                                    low=low,
                                    high=high)
      sp_dist = scipy_trunc_norm_dist(loc, scale, low, high)
      return tf_dist, sp_dist
  
    def testSampling(self, loc, scale, low, high):
      n = int(1000000)
      tf_dist, sp_dist = self.constructDists(loc, scale, low, high)
      tf_samples = self.evaluate(tf_dist.sample(
          n, seed=tfp_test_util.test_seed())).flatten()
      self.assertAllGreaterEqual(tf_samples, low)
      self.assertAllLessEqual(tf_samples, high)
  
      sp_samples = sp_dist.rvs(size=n)
      self.assertEmpiricalDistributionsEqual(
          tf_samples, sp_samples, atol=0.05, rtol=0.05)
  
    def testEntropy(self, loc, scale, low, high):
      tf_dist, sp_dist = self.constructDists(loc, scale, low, high)
      self.assertAlmostEqual(
          self.evaluate(tf_dist.entropy()), sp_dist.entropy(), places=2)
  
    def testLogProb(self, loc, scale, low, high):
      test_x = list(np.float32(np.random.uniform(low, high, 10)))
      test_x += [
          low, high, low + EPSILON, low - EPSILON, high + EPSILON,
          high - EPSILON
      ]
  
      tf_dist, sp_dist = self.constructDists(loc, scale, low, high)
  
      tf_log_prob = self.evaluate(tf_dist.log_prob(test_x))
      sp_log_prob = sp_dist.logpdf(test_x)
      self.assertAllClose(tf_log_prob, sp_log_prob, rtol=1e-4, atol=1e-4)
  
    def testCDF(self, loc, scale, low, high):
      test_x = list(np.float32(np.random.uniform(low, high, 10)))
      test_x += [
          low, high, low + EPSILON, low - EPSILON, high + EPSILON,
          high - EPSILON, low - 100., high + 100.
      ]
  
      tf_dist, sp_dist = self.constructDists(loc, scale, low, high)
  
      tf_cdf = self.evaluate(tf_dist.cdf(test_x))
      sp_cdf = sp_dist.cdf(test_x)
      self.assertAllClose(tf_cdf, sp_cdf, rtol=1e-4, atol=1e-4)
  
    def testMoments(self, loc, scale, low, high):
      tf_dist, sp_dist = self.constructDists(loc, scale, low, high)
      self.assertAlmostEqual(
          self.evaluate(tf_dist.mean()), sp_dist.mean(), places=3)
      self.assertAlmostEqual(
          self.evaluate(tf_dist.variance()), sp_dist.var(), places=3)
  
  
  if __name__ == "__main__":
    tf.test.main()
  