# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Affine Scalar Tests."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency
from tensorflow.python.platform import test


class AffineScalarBijectorTest(test.TestCase):
  """Tests correctness of the Y = scale @ x + shift transformation."""

  def testProperties(self):
    with self.test_session():
      mu = -1.
      # scale corresponds to 1.
      bijector = AffineScalar(shift=mu)
      self.assertEqual("affine_scalar", bijector.name)

  def testNoBatchScalar(self):
    with self.test_session() as sess:

      def static_run(fun, x):
        return fun(x).eval()

      def dynamic_run(fun, x_value):
        x_value = np.array(x_value)
        x = array_ops.placeholder(dtypes.float32, name="x")
        return sess.run(fun(x), feed_dict={x: x_value})

      for run in (static_run, dynamic_run):
        mu = -1.
        # Corresponds to scale = 2
        bijector = AffineScalar(shift=mu, scale=2.)
        x = [1., 2, 3]  # Three scalar samples (no batches).
        self.assertAllClose([1., 3, 5], run(bijector.forward, x))
        self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x))
        self.assertAllClose([-np.log(2.)] * 3,
                            run(bijector.inverse_log_det_jacobian, x))

  def testOneBatchScalarViaIdentityIn64BitUserProvidesShiftOnly(self):
    with self.test_session() as sess:

      def static_run(fun, x):
        return fun(x).eval()

      def dynamic_run(fun, x_value):
        x_value = np.array(x_value).astype(np.float64)
        x = array_ops.placeholder(dtypes.float64, name="x")
        return sess.run(fun(x), feed_dict={x: x_value})

      for run in (static_run, dynamic_run):
        mu = np.float64([1.])
        # One batch, scalar.
        # Corresponds to scale = 1.
        bijector = AffineScalar(shift=mu)
        x = np.float64([1.])  # One sample from one batches.
        self.assertAllClose([2.], run(bijector.forward, x))
        self.assertAllClose([0.], run(bijector.inverse, x))
        self.assertAllClose([0.], run(bijector.inverse_log_det_jacobian, x))

  def testOneBatchScalarViaIdentityIn64BitUserProvidesScaleOnly(self):
    with self.test_session() as sess:

      def static_run(fun, x):
        return fun(x).eval()

      def dynamic_run(fun, x_value):
        x_value = np.array(x_value).astype(np.float64)
        x = array_ops.placeholder(dtypes.float64, name="x")
        return sess.run(fun(x), feed_dict={x: x_value})

      for run in (static_run, dynamic_run):
        multiplier = np.float64([2.])
        # One batch, scalar.
        # Corresponds to scale = 2, shift = 0.
        bijector = AffineScalar(scale=multiplier)
        x = np.float64([1.])  # One sample from one batches.
        self.assertAllClose([2.], run(bijector.forward, x))
        self.assertAllClose([0.5], run(bijector.inverse, x))
        self.assertAllClose([np.log(0.5)],
                            run(bijector.inverse_log_det_jacobian, x))

  def testTwoBatchScalarIdentityViaIdentity(self):
    with self.test_session() as sess:

      def static_run(fun, x):
        return fun(x).eval()

      def dynamic_run(fun, x_value):
        x_value = np.array(x_value)
        x = array_ops.placeholder(dtypes.float32, name="x")
        return sess.run(fun(x), feed_dict={x: x_value})

      for run in (static_run, dynamic_run):
        mu = [1., -1]
        # Univariate, two batches.
        # Corresponds to scale = 1.
        bijector = AffineScalar(shift=mu)
        x = [1., 1]  # One sample from each of two batches.
        self.assertAllClose([2., 0], run(bijector.forward, x))
        self.assertAllClose([0., 2], run(bijector.inverse, x))
        self.assertAllClose([0., 0.], run(bijector.inverse_log_det_jacobian, x))

  def testTwoBatchScalarIdentityViaScale(self):
    with self.test_session() as sess:

      def static_run(fun, x):
        return fun(x).eval()

      def dynamic_run(fun, x_value):
        x_value = np.array(x_value)
        x = array_ops.placeholder(dtypes.float32, name="x")
        return sess.run(fun(x), feed_dict={x: x_value})

      for run in (static_run, dynamic_run):
        mu = [1., -1]
        # Univariate, two batches.
        # Corresponds to scale = 1.
        bijector = AffineScalar(shift=mu, scale=[2., 1])
        x = [1., 1]  # One sample from each of two batches.
        self.assertAllClose([3., 0], run(bijector.forward, x))
        self.assertAllClose([0., 2], run(bijector.inverse, x))
        self.assertAllClose(
            [-np.log(2), 0.], run(bijector.inverse_log_det_jacobian, x))

  def testScalarCongruency(self):
    with self.test_session():
      bijector = AffineScalar(shift=3.6, scale=0.42)
      assert_scalar_congruency(bijector, lower_x=-2., upper_x=2.)

if __name__ == "__main__":
  test.main()

[((5795, 5806), 'tensorflow.python.platform.test.main', 'test.main', ([], {}), '()\n', (5804, 5806), False, 'from tensorflow.python.platform import test\n'), ((1410, 1432), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'shift': 'mu'}), '(shift=mu)\n', (1422, 1432), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((5663, 5698), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'shift': '(3.6)', 'scale': '(0.42)'}), '(shift=3.6, scale=0.42)\n', (5675, 5698), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((5705, 5766), 'tensorflow.python.ops.distributions.bijector_test_util.assert_scalar_congruency', 'assert_scalar_congruency', (['bijector'], {'lower_x': '(-2.0)', 'upper_x': '(2.0)'}), '(bijector, lower_x=-2.0, upper_x=2.0)\n', (5729, 5766), False, 'from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency\n'), ((1674, 1691), 'numpy.array', 'np.array', (['x_value'], {}), '(x_value)\n', (1682, 1691), True, 'import numpy as np\n'), ((1704, 1751), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (['dtypes.float32'], {'name': '"""x"""'}), "(dtypes.float32, name='x')\n", (1725, 1751), False, 'from tensorflow.python.ops import array_ops\n'), ((1924, 1957), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'shift': 'mu', 'scale': '(2.0)'}), '(shift=mu, scale=2.0)\n', (1936, 1957), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((2546, 2593), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (['dtypes.float64'], {'name': '"""x"""'}), "(dtypes.float64, name='x')\n", (2567, 2593), False, 'from tensorflow.python.ops import array_ops\n'), ((2708, 2725), 'numpy.float64', 'np.float64', (['[1.0]'], {}), '([1.0])\n', (2718, 2725), True, 'import numpy as np\n'), ((2809, 2831), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'shift': 'mu'}), '(shift=mu)\n', (2821, 2831), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((2844, 2861), 'numpy.float64', 'np.float64', (['[1.0]'], {}), '([1.0])\n', (2854, 2861), True, 'import numpy as np\n'), ((3365, 3412), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (['dtypes.float64'], {'name': '"""x"""'}), "(dtypes.float64, name='x')\n", (3386, 3412), False, 'from tensorflow.python.ops import array_ops\n'), ((3535, 3552), 'numpy.float64', 'np.float64', (['[2.0]'], {}), '([2.0])\n', (3545, 3552), True, 'import numpy as np\n'), ((3647, 3677), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'scale': 'multiplier'}), '(scale=multiplier)\n', (3659, 3677), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((3690, 3707), 'numpy.float64', 'np.float64', (['[1.0]'], {}), '([1.0])\n', (3700, 3707), True, 'import numpy as np\n'), ((4180, 4197), 'numpy.array', 'np.array', (['x_value'], {}), '(x_value)\n', (4188, 4197), True, 'import numpy as np\n'), ((4210, 4257), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (['dtypes.float32'], {'name': '"""x"""'}), "(dtypes.float32, name='x')\n", (4231, 4257), False, 'from tensorflow.python.ops import array_ops\n'), ((4471, 4493), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'shift': 'mu'}), '(shift=mu)\n', (4483, 4493), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((4964, 4981), 'numpy.array', 'np.array', (['x_value'], {}), '(x_value)\n', (4972, 4981), True, 'import numpy as np\n'), ((4994, 5041), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (['dtypes.float32'], {'name': '"""x"""'}), "(dtypes.float32, name='x')\n", (5015, 5041), False, 'from tensorflow.python.ops import array_ops\n'), ((5255, 5293), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', ([], {'shift': 'mu', 'scale': '[2.0, 1]'}), '(shift=mu, scale=[2.0, 1])\n', (5267, 5293), False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((2497, 2514), 'numpy.array', 'np.array', (['x_value'], {}), '(x_value)\n', (2505, 2514), True, 'import numpy as np\n'), ((3316, 3333), 'numpy.array', 'np.array', (['x_value'], {}), '(x_value)\n', (3324, 3333), True, 'import numpy as np\n'), ((3889, 3900), 'numpy.log', 'np.log', (['(0.5)'], {}), '(0.5)\n', (3895, 3900), True, 'import numpy as np\n'), ((5522, 5531), 'numpy.log', 'np.log', (['(2)'], {}), '(2)\n', (5528, 5531), True, 'import numpy as np\n'), ((2183, 2194), 'numpy.log', 'np.log', (['(2.0)'], {}), '(2.0)\n', (2189, 2194), True, 'import numpy as np\n')]