# 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()

[((153, 2, 153, 13), 'tensorflow.python.platform.test.main', 'test.main', ({}, {}), '()', False, 'from tensorflow.python.platform import test\n'), ((37, 17, 37, 39), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((149, 17, 149, 52), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((150, 6, 150, 65), 'tensorflow.python.ops.distributions.bijector_test_util.assert_scalar_congruency', 'assert_scalar_congruency', (), '', False, 'from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency\n'), ((47, 18, 47, 35), 'numpy.array', 'np.array', ({(47, 27, 47, 34): 'x_value'}, {}), '(x_value)', True, 'import numpy as np\n'), ((48, 12, 48, 59), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (), '', False, 'from tensorflow.python.ops import array_ops\n'), ((54, 19, 54, 51), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((69, 12, 69, 59), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (), '', False, 'from tensorflow.python.ops import array_ops\n'), ((73, 13, 73, 29), 'numpy.float64', 'np.float64', ({(73, 24, 73, 28): '[1.0]'}, {}), '([1.0])', True, 'import numpy as np\n'), ((76, 19, 76, 41), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((77, 12, 77, 28), 'numpy.float64', 'np.float64', ({(77, 23, 77, 27): '[1.0]'}, {}), '([1.0])', True, 'import numpy as np\n'), ((90, 12, 90, 59), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (), '', False, 'from tensorflow.python.ops import array_ops\n'), ((94, 21, 94, 37), 'numpy.float64', 'np.float64', ({(94, 32, 94, 36): '[2.0]'}, {}), '([2.0])', True, 'import numpy as np\n'), ((97, 19, 97, 49), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((98, 12, 98, 28), 'numpy.float64', 'np.float64', ({(98, 23, 98, 27): '[1.0]'}, {}), '([1.0])', True, 'import numpy as np\n'), ((111, 18, 111, 35), 'numpy.array', 'np.array', ({(111, 27, 111, 34): 'x_value'}, {}), '(x_value)', True, 'import numpy as np\n'), ((112, 12, 112, 59), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (), '', False, 'from tensorflow.python.ops import array_ops\n'), ((119, 19, 119, 41), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((132, 18, 132, 35), 'numpy.array', 'np.array', ({(132, 27, 132, 34): 'x_value'}, {}), '(x_value)', True, 'import numpy as np\n'), ((133, 12, 133, 59), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', (), '', False, 'from tensorflow.python.ops import array_ops\n'), ((140, 19, 140, 56), 'tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar.AffineScalar', 'AffineScalar', (), '', False, 'from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar\n'), ((68, 18, 68, 35), 'numpy.array', 'np.array', ({(68, 27, 68, 34): 'x_value'}, {}), '(x_value)', True, 'import numpy as np\n'), ((89, 18, 89, 35), 'numpy.array', 'np.array', ({(89, 27, 89, 34): 'x_value'}, {}), '(x_value)', True, 'import numpy as np\n'), ((101, 29, 101, 40), 'numpy.log', 'np.log', ({(101, 36, 101, 39): '(0.5)'}, {}), '(0.5)', True, 'import numpy as np\n'), ((145, 14, 145, 23), 'numpy.log', 'np.log', ({(145, 21, 145, 22): '(2)'}, {}), '(2)', True, 'import numpy as np\n'), ((58, 30, 58, 40), 'numpy.log', 'np.log', ({(58, 37, 58, 39): '(2.0)'}, {}), '(2.0)', True, 'import numpy as np\n')]