docstr: Tests model with piecewise constant or constant parameters in 2 dim.
docstr_tok: ["Tests","model","with","piecewise","constant","or","constant","parameters","in","2","dim","."]
code: def test_mean_variance_correlation_piecewise_constant_2d(
      self, supply_draws, supply_grid
  ):
      """Tests model with piecewise constant or constant parameters in 2 dim."""
      dtype = tf.float64
      mean_reversion = self.mean_reversion
      volatility = tff.math.piecewise.PiecewiseConstantFunc(
          2 * [[0.2, 1.0]],
          values=np.array(3 * [self.volatility]).transpose(),
          dtype=dtype,
      )
      expected_corr_matrix = [[1.0, 0.5], [0.5, 1.0]]
      corr_matrix = tff.math.piecewise.PiecewiseConstantFunc(
          [0.5, 2.0],
          values=[
              expected_corr_matrix,
              [[1.0, 0.6], [0.6, 1.0]],
              [[1.0, 0.9], [0.9, 1.0]],
          ],
          dtype=dtype,
      )
      process = tff.models.hull_white.VectorHullWhiteModel(
          dim=2,
          mean_reversion=mean_reversion,
          volatility=volatility,
          corr_matrix=corr_matrix,
          initial_discount_rate_fn=self.instant_forward_rate_2d_fn,
          dtype=dtype,
      )
      num_samples = 100000
      test_num_samples = 100000
      if supply_draws:
          normal_draws = tf.random.stateless_normal(
              [num_samples // 2, 9, 2], seed=[4, 2], dtype=dtype
          )
          normal_draws = tf.concat(
              [normal_draws, -normal_draws], axis=0
          )
          test_num_samples = 1
      else:
          normal_draws = None
      times_grid = None
      if supply_grid:
          times_grid = [0.0, 0.1, 0.2, 0.5, 1.0]
          if supply_draws:
              normal_draws = tf.random.stateless_normal(
                  [num_samples // 2, 4, 2], seed=[4, 2], dtype=dtype
              )
              normal_draws = tf.concat(
                  [normal_draws, -normal_draws], axis=0
              )
      paths = process.sample_paths(
          [0.1, 0.5, 1.0],
          num_samples=test_num_samples,
          random_type=tff.math.random.RandomType.STATELESS_ANTITHETIC,
          seed=[4, 2],
          normal_draws=normal_draws,
          times_grid=times_grid,
      )
      with self.subTest("Dtype"):
          self.assertEqual(paths.dtype, dtype)
      with self.subTest("Shape"):
          self.assertAllEqual(paths.shape, [num_samples, 3, 2])
      paths = self.evaluate(paths)
      estimated_corr_matrix = np.corrcoef(
          paths[:, 1, 0], paths[:, 1, 1]
      )
      paths = paths[
          :, -1, :
      ]  # Extract paths values for the terminal time
      mean = np.mean(paths, axis=0)
      variance = np.var(paths, axis=0)
      with self.subTest("Mean"):
          self.assertAllClose(
              mean, self.true_mean(1.0), rtol=1e-3, atol=1e-3
          )
      with self.subTest("Variance"):
          self.assertAllClose(
              variance, self.true_var(1.0), rtol=1e-3, atol=1e-3
          )
      with self.subTest("CorrMatrix"):
          self.assertAllClose(
              estimated_corr_matrix,
              expected_corr_matrix,
              rtol=1e-2,
              atol=1e-2,
          )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 6 for calling the HullWhiteTest obj's underlying member method code with example usage: obj.test_mean_variance_correlation_piecewise_constant_2d(supply_draws, supply_grid) without return types