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 Proximal Hessian Optimizer."""
  
  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
  import tensorflow_probability as tfp
  
  from tensorflow_probability.python.internal import test_case
  from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
  
  
  class _ProximalHessianTest(object):
  
    def _make_placeholder(self, x):
      return tf1.placeholder_with_default(
          input=x, shape=(x.shape if self.use_static_shape else None))
  
    def _adjust_dtype_and_shape_hints(self, x):
      x_ = tf.cast(x, self.dtype)
  
      # Since there is no sparse_placeholder_with_default, we manually feed in the
      # constituent dense Tensors to create a defacto placeholder SparseTensor.
      if isinstance(x_, tf.SparseTensor):
        indices_placeholder = self._make_placeholder(x_.indices)
        values_placeholder = self._make_placeholder(x_.values)
  
        if self.use_static_shape:
          dense_shape_placeholder = x_.dense_shape
        else:
          dense_shape_placeholder = self._make_placeholder(x_.dense_shape)
  
        x_ = tf.SparseTensor(
            indices=indices_placeholder,
            values=values_placeholder,
            dense_shape=dense_shape_placeholder)
      else:
        x_ = self._make_placeholder(x_)
      return x_
  
    def _test_finding_sparse_solution(self, batch_shape=None):
      # Test that Proximal Hessian descent prefers sparse solutions when
      # l1_regularizer is large enough.
      #
      # Define
      #
      #     Loss(x) := (x[0] - a[0])**2 + epsilon * sum(
      #                    (x[i] - a[i])**2 for i in range(1, n))
      #
      # where `a` is a constant and epsilon is small.  Set
      # l2_regularizer = 0 and set l1_regularizer such that
      #
      #     epsilon << l1_regularizer << 1.
      #
      # L1 regularization should cause the computed optimum to have zeros in all
      # but the 0th coordinate: optimal_x ~= [a[0], 0, ..., 0].
      n = 10
      epsilon = 1e-6
      if batch_shape is None:
        batch_shape = []
      # Set a[0] explicitly to make sure it's not very close to zero
      a0 = 6.
      a_ = np.concatenate([
          np.full(batch_shape + [1], a0),
          np.random.random(size=batch_shape + [n - 1])
      ],
                          axis=-1)
      a = self._adjust_dtype_and_shape_hints(a_)
  
      def _grad_and_hessian_unregularized_loss_fn(x):
        diff = x - a
        grad = 2. * tf.concat([diff[..., :1], epsilon * diff[..., 1:]], axis=-1)
  
        hessian_outer = tf.SparseTensor(
            indices=[
                b + (i, i) for i in range(n) for b in np.ndindex(*batch_shape)
            ],
            values=tf.ones(shape=[np.product(batch_shape) * n], dtype=self.dtype),
            dense_shape=batch_shape + [n, n])
  
        hessian_middle_per_batch = 2 * tf.concat(
            [[1.], epsilon * tf.ones([n - 1], dtype=self.dtype)], axis=0)
        hessian_middle = tf.zeros(
            batch_shape + [n], dtype=self.dtype) + hessian_middle_per_batch
        return grad, hessian_outer, hessian_middle
  
      w, is_converged, num_iter = tfp.optimizer.proximal_hessian_sparse_minimize(
          _grad_and_hessian_unregularized_loss_fn,
          x_start=tf.zeros(batch_shape + [n], dtype=self.dtype),
          l1_regularizer=1e-2,
          l2_regularizer=None,
          maximum_iterations=10,
          maximum_full_sweeps_per_iteration=10,
          tolerance=1e-5,
          learning_rate=1.)
  
      w_, is_converged_, _ = self.evaluate([w, is_converged, num_iter])
  
      expected_w = tf.concat(
          [a[..., :1], tf.zeros(batch_shape + [n - 1], self.dtype)], axis=-1)
  
      # Using atol=0 ensures that w must be exactly zero in all coordinates
      # where expected_w is exactly zero.
      self.assertAllEqual(is_converged_, True)
      self.assertAllClose(w_, expected_w, atol=0., rtol=1e-3)
  
    def testFindingSparseSolution_SingleInstance(self):
      self._test_finding_sparse_solution()
  
    def testFindingSparseSolution_SingleBatch(self):
      self._test_finding_sparse_solution(batch_shape=[1])
  
    def testFindingSparseSolution_BatchOfRank2(self):
      self._test_finding_sparse_solution(batch_shape=[2, 3])
  
    def testL2Regularization(self):
      # Define Loss(x) := ||x - a||_2**2, where a is a constant.
      # Set l1_regularizer = 0 and l2_regularizer = 1.
      # Then the regularized loss is
      #
      #     ||x - a||_2**2 + ||x||_2**2
      #
      # And the true optimum is x = 0.5 * a.
      n = 100
      np.random.seed(42)
      a_ = np.random.random(size=(n,))
      a = self._adjust_dtype_and_shape_hints(a_)
  
      def _grad_and_hessian_unregularized_loss_fn(x):
        grad = 2 * (x - a)
        hessian_outer = tf.eye(n, dtype=a.dtype)
        hessian_middle = 2. * tf.ones_like(a)
        return grad, hessian_outer, hessian_middle
  
      w, is_converged, num_iter = tfp.optimizer.proximal_hessian_sparse_minimize(
          _grad_and_hessian_unregularized_loss_fn,
          x_start=tf.zeros_like(a_, dtype=self.dtype),
          l1_regularizer=0.,
          l2_regularizer=1.,
          maximum_iterations=4,
          maximum_full_sweeps_per_iteration=4,
          tolerance=1e-5,
          learning_rate=1.)
  
      w_, is_converged_, _ = self.evaluate([w, is_converged, num_iter])
  
      expected_w = 0.5 * a
      self.assertAllEqual(is_converged_, True)
      self.assertAllClose(w_, expected_w, atol=0., rtol=0.03)
  
    def testNumIter(self):
      # Same as testL2Regularization, except we set
      # maximum_full_sweeps_per_iteration = 1 and check that the number of sweeps
      # is equals what we expect it to (usually we don't know the exact number,
      # but in this simple case we do -- explanation below).
      #
      # Since l1_regularizer = 0, the soft threshold operator is actually the
      # identity operator, hence the `proximal_hessian_sparse_minimize` algorithm
      # becomes literally coordinatewise Newton's method being used to find the
      # zeros of grad Loss(x), which in this case is a linear function of x. Hence
      # Newton's method should find the exact correct answer in 1 sweep.  At the
      # end of the first sweep the algorithm does not yet know it has converged;
      # it takes a second sweep, when the algorithm notices that its answer hasn't
      # changed at all, to become aware that convergence has happened.  Hence we
      # expect two sweeps.  So with maximum_full_sweeps_per_iteration = 1, that
      # means we expect 2 iterations of the outer loop.
      n = 100
      np.random.seed(42)
      a_ = np.random.random(size=(n,))
      a = self._adjust_dtype_and_shape_hints(a_)
  
      def _grad_and_hessian_unregularized_loss_fn(x):
        grad = 2 * (x - a)
        hessian_outer = tf.linalg.tensor_diag(tf.ones_like(a))
        hessian_middle = 2. * tf.ones_like(a)
        return grad, hessian_outer, hessian_middle
  
      w, is_converged, num_iter = tfp.optimizer.proximal_hessian_sparse_minimize(
          _grad_and_hessian_unregularized_loss_fn,
          x_start=tf.zeros_like(a_, dtype=self.dtype),
          l1_regularizer=0.,
          l2_regularizer=1.,
          maximum_iterations=4,
          maximum_full_sweeps_per_iteration=1,
          tolerance=1e-5,
          learning_rate=1.)
  
      w_, is_converged_, num_iter_ = self.evaluate([w, is_converged, num_iter])
  
      expected_w = 0.5 * a
      self.assertAllEqual(is_converged_, True)
      self.assertAllEqual(num_iter_, 2)
      self.assertAllClose(w_, expected_w, atol=0., rtol=0.03)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class ProximalHessianTestStaticShapeFloat32(test_case.TestCase,
                                              _ProximalHessianTest):
    dtype = tf.float32
    use_static_shape = True
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class ProximalHessianTestDynamicShapeFloat32(test_case.TestCase,
                                               _ProximalHessianTest):
    dtype = tf.float32
    use_static_shape = False
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class ProximalHessianTestStaticShapeFloat64(test_case.TestCase,
                                              _ProximalHessianTest):
    dtype = tf.float64
    use_static_shape = True
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class ProximalHessianTestDynamicShapeFloat64(test_case.TestCase,
                                               _ProximalHessianTest):
    dtype = tf.float64
    use_static_shape = False
  
  
  if __name__ == '__main__':
    tf.test.main()
  