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 Double-sided Maxwell Distribution."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  
  import numpy as np
  from scipy import stats
  
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  
  from tensorflow_probability.python.internal import test_util
  
  tfd = tfp.distributions
  
  
  @test_util.test_all_tf_execution_regimes
  class DoublesidedMaxwellTest(test_util.TestCase):
  
    def setUp(self):
      self._rng = np.random.RandomState(123)
  
    def _testParamShapes(self, sample_shape, expected):
      param_shapes = tfd.DoublesidedMaxwell.param_shapes(sample_shape)
      mu_shape, sigma_shape = param_shapes['loc'], param_shapes['scale']
  
      self.assertAllEqual(expected, self.evaluate(mu_shape))
      self.assertAllEqual(expected, self.evaluate(sigma_shape))
      mu = tf.zeros(mu_shape)
      sigma = tf.ones(sigma_shape)
      dsmaxwell = tfd.DoublesidedMaxwell(mu, sigma, validate_args=True)
      self.assertAllEqual(expected, self.evaluate(tf.shape(dsmaxwell.sample())))
  
    def _testParamStaticShapes(self, sample_shape, expected):
      param_shapes = tfd.DoublesidedMaxwell.param_static_shapes(sample_shape)
      mu_shape, sigma_shape = param_shapes['loc'], param_shapes['scale']
      self.assertEqual(expected, mu_shape)
      self.assertEqual(expected, sigma_shape)
  
    def testParamShapes(self):
      sample_shape = [10, 3, 4]
      self._testParamShapes(sample_shape, sample_shape)
      self._testParamShapes(tf.constant(sample_shape), sample_shape)
      self._testParamStaticShapes(sample_shape, sample_shape)
  
    def testDoublesidedMaxwellPDF(self):
      # Use only positive values since we will use the one sided maxwell
      # as a test.
      x = np.arange(1, 5, dtype=np.float64)
      loc = 1.
      scale = 1.
  
      # Test the pdf value by using its relationship to the 1-sided Maxwell.
      dsmaxwell = tfd.DoublesidedMaxwell(loc=loc, scale=scale, validate_args=True)
      log_prob = dsmaxwell.log_prob(x)
      expected_log_prob = tf.identity(
          stats.maxwell.logpdf(np.abs(x), loc, scale) - np.log(2))
      self.assertAllClose(expected_log_prob, log_prob)
  
    def testInvalidScale(self):
      scale = [-.01, 0., 2.]
      with self.assertRaisesOpError('Argument `scale` must be positive.'):
        dsmaxwell = tfd.DoublesidedMaxwell(
            loc=0., scale=scale, validate_args=True)
        self.evaluate(dsmaxwell.scale)
  
    @parameterized.named_parameters(
        ('test1', 4.0, 1.0),
        ('test2', 2.0, 3.0))
    def testDoublesidedMaxwellSample(self, loc, scale):
      n = int(100e3)
      dsmaxwell = tfd.DoublesidedMaxwell(loc=loc, scale=scale, validate_args=True)
      samples = dsmaxwell.sample(n, seed=100)
      mean = dsmaxwell.mean()
      variance = dsmaxwell.variance()
  
      [samples_, mean_, variance_] = self.evaluate([samples, mean, variance])
  
      # Check first and second moments.
      self.assertEqual((n,), samples_.shape)
      self.assertAllClose(np.mean(samples_), mean_, atol=0., rtol=0.1)
      self.assertAllClose(np.var(samples_), variance_, atol=0., rtol=0.1)
  
    def testDoublesidedMaxwellMean(self):
      # loc will be broadcast to [7, 7, 7]
      loc = [7.]
      sigma = [1., 2., 3.]
  
      dsmaxwell = tfd.DoublesidedMaxwell(loc=loc, scale=sigma)
  
      self.assertAllEqual((3,), dsmaxwell.mean().shape)
      self.assertAllEqual([7., 7., 7.], self.evaluate(dsmaxwell.mean()))
  
    def testDoublesidedMaxwellVariance(self):
      # sigma will be broadcast to [7, 7, 7]
      loc = [1., 2., 3.]
      sigma = [7.]
  
      dsmaxwell = tfd.DoublesidedMaxwell(loc=loc, scale=sigma)
  
      self.assertAllEqual((3,), dsmaxwell.variance().shape)
      self.assertAllClose([147., 147, 147], self.evaluate(dsmaxwell.variance()))
  
    def testDoublesidedMaxwellStandardDeviation(self):
      # sigma will be broadcast to [7, 7, 7]
      loc = [1., 2., 3.]
      sigma = [7.]
  
      dsmaxwell = tfd.DoublesidedMaxwell(loc=loc, scale=sigma)
  
      self.assertAllEqual((3,), dsmaxwell.stddev().shape)
      std = np.sqrt(147)
      self.assertAllClose(
          [std, std, std], self.evaluate(dsmaxwell.stddev()))
  
    def testDoublesidedMaxwellStandardSampleShape(self):
      # sigma will be broadcast to [7, 7, 7]
      loc = [1., 2., 3.]
      sigma = [7.]
  
      dsmaxwell = tfd.DoublesidedMaxwell(loc=loc, scale=sigma)
  
      n = 10
      samples = dsmaxwell.sample(n, seed=100)
      self.assertAllEqual((n, 3,), samples.shape)
  
  if __name__ == '__main__':
    tf.test.main()
  