Python tensorflow.compat.v2.float64() Examples

def test_sample_paths_dtypes(self):
    """Sampled paths have the expected dtypes."""
    for dtype in [np.float32, np.float64]:
      drift_fn = lambda t, x: tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)
      vol_fn = lambda t, x: t * tf.ones([1, 1], dtype=t.dtype)

      paths = self.evaluate(
          euler_sampling.sample(
              dim=1,
              drift_fn=drift_fn, volatility_fn=vol_fn,
              times=[0.1, 0.2],
              num_samples=10,
              initial_state=[0.1],
              time_step=0.01,
              seed=123,
              dtype=dtype))

      self.assertEqual(paths.dtype, dtype) 

def testMultiDimensionalShape(self):
    """Check that stateless_random_shuffle works with multi-dim shapes."""
    for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
      input_permutation = tf.constant([[[1], [2], [3]], [[4], [5], [6]]],
                                      dtype=dtype)
      random_shuffle = tff_rnd.stateless_random_shuffle(
          input_permutation, seed=(1, 42))
      random_permutation_first_call = self.evaluate(random_shuffle)
      random_permutation_next_call = self.evaluate(random_shuffle)
      input_permutation = self.evaluate(input_permutation)
      # Check that the dtype is correct
      np.testing.assert_equal(random_permutation_first_call.dtype,
                              dtype.as_numpy_dtype)
      # Check that the shuffles are the same
      np.testing.assert_array_equal(random_permutation_first_call,
                                    random_permutation_next_call)
      # Check that the output shape is correct
      np.testing.assert_equal(random_permutation_first_call.shape,
                              input_permutation.shape) 

def testOutputIsIndependentOfInputValues(self):
    """stateless_random_shuffle output is independent of input_tensor values."""
    # Generate sorted array of random numbers to control that the result
    # is independent of `input_tesnor` values
    np.random.seed(25)
    random_input = np.random.normal(size=[10])
    random_input.sort()
    for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
      # Permutation of a sequence [0, 1, .., 9]
      random_permutation = tff_rnd.stateless_random_shuffle(
          tf.range(10, dtype=dtype), seed=(100, 42))
      random_permutation = self.evaluate(random_permutation)
      # Shuffle `random_input` with the same seed
      random_shuffle_control = tff_rnd.stateless_random_shuffle(
          random_input, seed=(100, 42))
      random_shuffle_control = self.evaluate(random_shuffle_control)
      # Checks that the generated permutation does not depend on the underlying
      # values
      np.testing.assert_array_equal(
          np.argsort(random_permutation), np.argsort(random_shuffle_control)) 

def test_invalid_batch_size_piecewise(self):
    """Tests that the batch dimension should be [2] if it is not empty."""
    dtype = tf.float64
    # Batch shape is [1]. Should be [] or [2]
    mean_reversion = tff.math.piecewise.PiecewiseConstantFunc(
        [[0.1, 2.0]], values=[3 * [self.mean_reversion]], dtype=dtype)
    # Volatility with batch dimesnion
    volatility = self.volatility
    with self.assertRaises(ValueError):
      tff.models.hull_white.VectorHullWhiteModel(
          dim=2,
          mean_reversion=mean_reversion,
          volatility=volatility,
          corr_matrix=None,
          initial_discount_rate_fn=self.instant_forward_rate_2d_fn,
          dtype=dtype) 

def test_time_step_not_supplied(self):
    """Tests that the `time_step` should be supplied if Euler scheme is used."""
    dtype = tf.float64
    def volatility_fn(t):
      del t
      return self.volatility
    process = tff.models.hull_white.VectorHullWhiteModel(
        dim=2,
        mean_reversion=self.mean_reversion,
        volatility=volatility_fn,
        initial_discount_rate_fn=self.instant_forward_rate_2d_fn,
        dtype=dtype)
    with self.assertRaises(ValueError):
      process.sample_paths(
          [0.1, 2.0],
          num_samples=100) 

def true_divide(x1, x2):
  def _avoid_float64(x1, x2):
    if x1.dtype == x2.dtype and x1.dtype in (tf.int32, tf.int64):
      x1 = tf.cast(x1, dtype=tf.float32)
      x2 = tf.cast(x2, dtype=tf.float32)
    return x1, x2

  def f(x1, x2):
    if x1.dtype == tf.bool:
      assert x2.dtype == tf.bool
      float_ = dtypes.default_float_type()
      x1 = tf.cast(x1, float_)
      x2 = tf.cast(x2, float_)
    if not dtypes.is_allow_float64():
      # tf.math.truediv in Python3 produces float64 when both inputs are int32
      # or int64. We want to avoid that when is_allow_float64() is False.
      x1, x2 = _avoid_float64(x1, x2)
    return tf.math.truediv(x1, x2)
  return _bin_op(f, x1, x2) 

def setUp(self):
    super(LogicTest, self).setUp()
    self.array_transforms = [
        lambda x: x,  # Identity,
        tf.convert_to_tensor,
        np.array,
        lambda x: np.array(x, dtype=np.int32),
        lambda x: np.array(x, dtype=np.int64),
        lambda x: np.array(x, dtype=np.float32),
        lambda x: np.array(x, dtype=np.float64),
        array_ops.array,
        lambda x: array_ops.array(x, dtype=tf.int32),
        lambda x: array_ops.array(x, dtype=tf.int64),
        lambda x: array_ops.array(x, dtype=tf.float32),
        lambda x: array_ops.array(x, dtype=tf.float64),
    ] 

def testOutputIsPermutation(self):
    """Checks that stateless_random_shuffle outputs a permutation."""
    for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
      identity_permutation = tf.range(10, dtype=dtype)
      random_shuffle_seed_1 = tff_rnd.stateless_random_shuffle(
          identity_permutation, seed=tf.constant((1, 42), tf.int64))
      random_shuffle_seed_2 = tff_rnd.stateless_random_shuffle(
          identity_permutation, seed=tf.constant((2, 42), tf.int64))
      # Check that the shuffles are of the correct dtype
      for shuffle in (random_shuffle_seed_1, random_shuffle_seed_2):
        np.testing.assert_equal(shuffle.dtype, dtype.as_numpy_dtype)
      random_shuffle_seed_1 = self.evaluate(random_shuffle_seed_1)
      random_shuffle_seed_2 = self.evaluate(random_shuffle_seed_2)
      identity_permutation = self.evaluate(identity_permutation)
      # Check that the shuffles are different
      self.assertTrue(
          np.abs(random_shuffle_seed_1 - random_shuffle_seed_2).max())
      # Check that the shuffles are indeed permutations
      for shuffle in (random_shuffle_seed_1, random_shuffle_seed_2):
        self.assertAllEqual(set(shuffle), set(identity_permutation)) 

def testHomogeneous(self, scheme, accuracy_order):
    # Tests solving du/dt = At for a time step.
    # Compares with exact solution u(t) = exp(At) u(0).

    # Time step should be small enough to "resolve" different orders of accuracy
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(
        scheme(u, 0, time_step, lambda t: (tridiag_form, None)))
    expected = self.evaluate(
        tf.squeeze(
            tf.matmul(tf.linalg.expm(matrix * time_step), tf.expand_dims(u,
                                                                         1))))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def testHomogeneousBackwards(self, scheme, accuracy_order):
    # Tests solving du/dt = At for a backward time step.
    # Compares with exact solution u(0) = exp(-At) u(t).
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(
        scheme(u, time_step, 0, lambda t: (tridiag_form, None)))

    expected = self.evaluate(
        tf.squeeze(
            tf.matmul(
                tf.linalg.expm(-matrix * time_step), tf.expand_dims(u, 1))))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def test_correctness_1d(self, use_analytic_pricing, error_tol):
    """Tests model with constant parameters in 1 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0])
    maturities = np.array([5.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=self.volatility_1d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 1])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.02817777]], rtol=error_tol, atol=error_tol) 

def testFindsRootForFlatFunction(self):
    # Flat in the [-0.5, 0.5] range.
    objective_fn = lambda x: 0 if x == 0 else x * exp(-1 / x**2)

    left_bracket = [-10]
    right_bracket = [1]
    expected_num_iterations = [13]

    expected_num_iterations, result = self.evaluate([
        tf.constant(expected_num_iterations, dtype=tf.int32),
        root_search.brentq(objective_fn,
                           tf.constant(left_bracket, dtype=tf.float64),
                           tf.constant(right_bracket, dtype=tf.float64))
    ])

    _, value_at_roots, num_iterations, _ = result

    # Simply check that the objective function is close to the root for the
    # returned estimate. Do not check the estimate itself.
    # Unlike Brent's original algorithm (and the SciPy implementation), this
    # implementation stops the search as soon as a good enough root estimate is
    # found. As a result, the estimate may significantly differ from the one
    # returned by SciPy for functions which are extremely flat around the root.
    self.assertAllClose(value_at_roots, [0.])
    self.assertAllEqual(num_iterations, expected_num_iterations) 

def testWithNoIteration(self):
    left_bracket = [-10, 1]
    right_bracket = [10, -1]

    first_guess = tf.constant(left_bracket, dtype=tf.float64)
    second_guess = tf.constant(right_bracket, dtype=tf.float64)

    # Skip iteration entirely.
    # Should return a Tensor built from the best guesses in input positions.
    guess, result = self.evaluate([
        tf.constant([-10, -1], dtype=tf.float64),
        root_search.brentq(
            polynomial5, first_guess, second_guess, max_iterations=0)
    ])

    self.assertAllEqual(result.estimated_root, guess) 

def test_multiivariate_xla_compatible(self):
    """Tests that multiivariate GBM sampling is XLA-compatible."""
    corr_matrix = [[1, 0.1], [0.1, 1]]
    process = tff.models.MultivariateGeometricBrownianMotion(
        dim=2, means=0.05, volatilities=[0.1, 0.2], corr_matrix=corr_matrix,
        dtype=tf.float64)
    times = [0.1, 0.5, 1.0]
    initial_state = [1.0, 2.0]
    @tf.function
    def sample_fn():
      return process.sample_paths(
          times=times, initial_state=initial_state, num_samples=10000)
    samples = tf.xla.experimental.compile(sample_fn)[0]
    log_s = tf.math.log(samples)
    mean = tf.reduce_mean(log_s, axis=0)
    expected_mean = ((process._means - process._vols**2 / 2)
                     * np.array(np.expand_dims(times, -1))
                     + np.log(initial_state))
    self.assertAllClose(mean, expected_mean, atol=1e-2, rtol=1e-2) 

def testInhomogeneousBackwards(self, scheme, accuracy_order):
    # Tests solving du/dt = At + b for a backward time step.
    # Compares with exact solution u(0) = exp(-At) u(t)
    # + (exp(-At) - 1) A^(-1) b.
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)
    b = tf.constant([1, -1, -2, 2], dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(scheme(u, time_step, 0, lambda t: (tridiag_form, b)))

    exponent = tf.linalg.expm(-matrix * time_step)
    eye = tf.eye(4, 4, dtype=tf.float64)
    u = tf.expand_dims(u, 1)
    b = tf.expand_dims(b, 1)
    expected = (
        tf.matmul(exponent, u) +
        tf.matmul(exponent - eye, tf.matmul(tf.linalg.inv(matrix), b)))
    expected = self.evaluate(tf.squeeze(expected))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def test_construct_vol_covar_and_vol_callables(self):
    dtype = np.float64
    vol_matrix = np.array([[1.0, 0.21, -0.33], [0.61, 1.5, 1.77],
                           [-0.3, 1.19, -0.55]]).astype(dtype)
    covar_matrix = np.matmul(vol_matrix, vol_matrix.transpose())
    vol_fn = lambda time: bm_utils.outer_multiply(time, vol_matrix)

    def tc_fn(t1, t2):
      return bm_utils.outer_multiply((t2**2 - t1**2) / 2, covar_matrix)

    times = np.array([[0.12, 0.44], [0.48, 1.698]]).astype(dtype)
    actual_vol_fn, actual_tc_fn = bm_utils.construct_vol_data(
        vol_fn, tc_fn, 3, dtype)
    actual_vols = self.evaluate(actual_vol_fn(times))
    np.testing.assert_array_equal(actual_vols.shape, [2, 2, 3, 3])
    np.testing.assert_allclose(actual_vols, self.evaluate(vol_fn(times)))
    times2 = times + np.array([[0.12, 0.34], [0.56, 0.78]]).astype(dtype)
    actual_tc = self.evaluate(actual_tc_fn(times, times2))
    np.testing.assert_array_equal(actual_tc.shape, [2, 2, 3, 3])
    np.testing.assert_allclose(actual_tc,
                               self.evaluate(actual_tc_fn(times, times2))) 

def test_construct_vol_covar_and_vector_vol(self):
    dtype = np.float64
    vol = np.array([0.94, 1.1, 0.42], dtype=dtype)
    # Note that the total covariance function we supply is deliberately not
    # the one that is implied by the volatility function.
    np.random.seed(5321)
    dim = 3
    vol_matrix = np.random.randn(dim, dim)
    covar_matrix = np.matmul(vol_matrix, vol_matrix.transpose())

    def tc_fn(t1, t2):
      return bm_utils.outer_multiply(t2 - t1, covar_matrix)

    times = np.array([[0.12], [0.48]]).astype(dtype)
    actual_vol_fn, actual_tc_fn = bm_utils.construct_vol_data(
        tf.constant(vol), tc_fn, dim, dtype)
    actual_vols = self.evaluate(actual_vol_fn(times))
    np.testing.assert_array_equal(actual_vols.shape, [2, 1, dim, dim])
    for i in range(2):
      np.testing.assert_allclose(actual_vols[i, 0], np.diag(vol))
    actual_tc = self.evaluate(actual_tc_fn(times, times + 0.22))
    np.testing.assert_array_equal(actual_tc.shape, [2, 1, dim, dim])
    for i in range(2):
      np.testing.assert_allclose(actual_tc[i, 0], covar_matrix * 0.22) 

def test_spline_batch(self, optimize_for_tpu):
    """Tests batching of four splines."""
    for dtype in (np.float32, np.float64):
      x_data = np.linspace(-11, 12, 24)
      x_data = np.reshape(x_data, [2, 2, 6])
      y_data = 1.0 / (1.0 + x_data * x_data)
      search_args = np.array([[[-10.5, -5.], [-4.5, 1]],
                              [[1.5, 2.], [7.5, 12.]]])

      spline = tff.math.interpolation.cubic.build_spline(
          x_data, y_data, dtype=dtype)
      result = tff.math.interpolation.cubic.interpolate(
          search_args, spline,
          optimize_for_tpu=optimize_for_tpu, dtype=dtype)

      expected = np.array([[[0.00900778, 0.02702703],
                            [0.04705774, 1.]],
                           [[0.33135411, 0.2],
                            [0.01756963, 0.00689655]]],
                          dtype=dtype)
      self.assertEqual(result.dtype.as_numpy_dtype, dtype)
      result = self.evaluate(result)
      np.testing.assert_almost_equal(expected, result) 

def test_error_calc(self, optimize_for_tpu):
    """Test the deviation of the interpolated values from the actual."""
    sampling_points = 1000
    spline_x = np.linspace(0.0, 10.0, num=11, dtype=np.float64)
    spline_y = [1.0 / (1.0 + x * x) for x in spline_x]
    x_series = np.array([spline_x])
    y_series = np.array([spline_y])
    spline = tff.math.interpolation.cubic.build_spline(x_series, y_series)

    # There is an error if we go to 10.0
    test_range_x = np.linspace(0.0, 9.99, num=sampling_points, dtype=np.float64)
    search_args = tf.constant(np.array([test_range_x]), dtype=tf.float64)
    projected_y = tff.math.interpolation.cubic.interpolate(
        search_args, spline, optimize_for_tpu=optimize_for_tpu)
    expected_y = tf.constant([[1.0 / (1.0 + x * x) for x in test_range_x]],
                             dtype=tf.float64)
    errors = expected_y - projected_y
    deviation = self.evaluate(tfp.stats.stddev(errors[0], sample_axis=0))
    limit = 0.02
    self.assertLess(deviation, limit) 

def test_construct_vol_covar_and_vol_matrix(self):
    dtype = np.float64
    vol_matrix = np.array([[1.0, 0.21, -0.33], [0.61, 1.5, 1.77],
                           [-0.3, 1.19, -0.55]]).astype(dtype)
    covar_matrix = np.matmul(vol_matrix, vol_matrix.transpose())

    def tc_fn(t1, t2):
      return bm_utils.outer_multiply((t2 - t1), covar_matrix)

    times = np.array([[0.12, 0.44], [0.48, 1.698]]).astype(dtype)
    actual_vol_fn, actual_tc_fn = bm_utils.construct_vol_data(
        vol_matrix, tc_fn, 3, dtype)
    actual_vols = self.evaluate(actual_vol_fn(times))
    np.testing.assert_array_equal(actual_vols.shape, [2, 2, 3, 3])
    for i in range(2):
      for j in range(2):
        np.testing.assert_allclose(actual_vols[i, j], vol_matrix)

    times2 = times + np.array([[0.12, 0.34], [0.56, 0.78]]).astype(dtype)
    actual_tc = self.evaluate(actual_tc_fn(times, times2))
    np.testing.assert_array_equal(actual_tc.shape, [2, 2, 3, 3])
    np.testing.assert_allclose(actual_tc,
                               self.evaluate(actual_tc_fn(times, times2))) 

def test_construct_vol_no_covar_and_scalar_vol(self):
    dtype = np.float64
    vol = tf.constant(0.94, dtype=dtype)
    np.random.seed(1235)
    dim = 5
    covar_matrix = np.eye(dim) * 0.94 * 0.94
    times = np.array([[0.12, 0.44], [0.48, 1.698]]).astype(dtype)
    actual_vol_fn, actual_tc_fn = bm_utils.construct_vol_data(
        vol, None, dim, dtype)
    actual_vols = self.evaluate(actual_vol_fn(times))
    np.testing.assert_array_equal(actual_vols.shape, [2, 2, dim, dim])
    for i in range(2):
      for j in range(2):
        np.testing.assert_allclose(actual_vols[i, j],
                                   np.eye(dim).astype(dtype) * 0.94)

    actual_tc = self.evaluate(actual_tc_fn(times, times + 1.0))
    np.testing.assert_array_equal(actual_tc.shape, [2, 2, dim, dim])
    for i in range(2):
      for j in range(2):
        np.testing.assert_allclose(actual_tc[i, j], covar_matrix) 

def test_construct_vol_no_covar_and_vector_vol(self):
    dtype = np.float64
    vol = np.array([0.94, 1.1, 0.42], dtype=dtype)
    np.random.seed(5321)
    dim = 3
    vol_matrix = np.diag(vol)
    covar_matrix = np.matmul(vol_matrix, vol_matrix.transpose())

    times = np.array([[0.12], [0.48]]).astype(dtype)
    actual_vol_fn, actual_tc_fn = bm_utils.construct_vol_data(
        tf.constant(vol), None, dim, dtype)
    actual_vols = self.evaluate(actual_vol_fn(times))
    np.testing.assert_array_equal(actual_vols.shape, [2, 1, dim, dim])
    for i in range(2):
      np.testing.assert_allclose(actual_vols[i, 0], np.diag(vol))
    actual_tc = self.evaluate(actual_tc_fn(times, times + 0.22))
    np.testing.assert_array_equal(actual_tc.shape, [2, 1, dim, dim])
    for i in range(2):
      np.testing.assert_allclose(actual_tc[i, 0], covar_matrix * 0.22) 

def test_construct_vol_no_covar_and_vol_matrix(self):
    dtype = np.float64
    vol_matrix = np.array([[1.0, 0.21, -0.33], [0.61, 1.5, 1.77],
                           [-0.3, 1.19, -0.55]]).astype(dtype)
    covar_matrix = np.matmul(vol_matrix, vol_matrix.transpose())
    times = np.array([[0.12, 0.44], [0.48, 1.698]]).astype(dtype)
    actual_vol_fn, actual_tc_fn = bm_utils.construct_vol_data(
        vol_matrix, None, 3, dtype)
    actual_vols = self.evaluate(actual_vol_fn(times))
    np.testing.assert_array_equal(actual_vols.shape, [2, 2, 3, 3])
    for i in range(2):
      for j in range(2):
        np.testing.assert_allclose(actual_vols[i, j], vol_matrix)

    dt = np.array([[0.12, 0.34], [0.56, 0.78]]).astype(dtype)
    times2 = times + dt
    actual_tc = self.evaluate(actual_tc_fn(times, times2))
    np.testing.assert_array_equal(actual_tc.shape, [2, 2, 3, 3])
    for i in range(2):
      for j in range(2):
        np.testing.assert_allclose(actual_tc[i, j], covar_matrix * dt[i, j]) 

def test_multivariate_drift_and_volatility_no_corr(self, dtype):
    """Tests multivariate GBM drift and volatility functions."""
    means = [0.05, 0.02]
    volatilities = [0.1, 0.2]
    corr_matrix = [[1, 0.0], [0.0, 1]]
    process = tff.models.MultivariateGeometricBrownianMotion(
        dim=2, means=means, volatilities=volatilities,
        dtype=tf.float64)
    drift_fn = process.drift_fn()
    volatility_fn = process.volatility_fn()
    state = np.array([[1., 2.], [3., 4.], [5., 6.]], dtype=dtype)
    with self.subTest("Drift"):
      drift = drift_fn(0.2, state)
      expected_drift = np.array(means) * state
      self.assertAllClose(drift, expected_drift, atol=1e-8, rtol=1e-8)
    with self.subTest("Volatility"):
      vol = volatility_fn(0.2, state)
      expected_vol = np.expand_dims(
          np.array(volatilities) * state,
          axis=-1) * np.linalg.cholesky(corr_matrix)
      self.assertAllClose(vol, expected_vol, atol=1e-8, rtol=1e-8) 

def test_volatility(self):
    """Tests volatility stays close to its mean for small vol of vol."""
    theta = 0.05
    process = HestonModel(
        kappa=1.0, theta=theta, epsilon=0.00001,
        rho=-0.0, dtype=np.float64)
    years = 1.0
    times = np.linspace(0.0, years, int(365 * years))
    num_samples = 2
    paths = process.sample_paths(
        times,
        time_step=0.01,
        num_samples=num_samples,
        initial_state=np.array([np.log(100), 0.045]),
        seed=None)
    # For small values of epsilon, volatility should stay close to theta
    volatility_trace = self.evaluate(paths)[..., 1]
    max_deviation = np.max(abs(volatility_trace[:, 50:] - theta))
    self.assertAlmostEqual(
        max_deviation, 0.0, places=2) 

def test_correctness_2d(self, use_analytic_pricing, error_tol):
    """Tests model with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0])
    maturities = np.array([5.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=2,
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 2])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.02817777, 0.01309971]],
                        rtol=error_tol, atol=error_tol) 

def test_1d_batch(self, use_analytic_pricing, error_tol):
    """Tests model with 1d batch of options."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0, 1.0, 1.0])
    maturities = np.array([5.0, 5.0, 5.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=self.volatility_1d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [3, 1])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.02817777], [0.02817777], [0.02817777]],
                        rtol=error_tol, atol=error_tol) 

def test_times_wrong_rank(self):
    """Tests that the `times` should be a rank 1 `Tensor`."""
    dtype = tf.float64
    process = tff.models.hull_white.VectorHullWhiteModel(
        dim=2,
        mean_reversion=self.mean_reversion,
        volatility=self.volatility,
        initial_discount_rate_fn=self.instant_forward_rate_2d_fn,
        dtype=dtype)
    with self.assertRaises(ValueError):
      process.sample_paths(
          [[0.1, 2.0]],
          num_samples=100) 

def test_construct_vol_covar_and_no_vol(self):
    # Check the None, None
    dtype = np.float64
    covar_matrix = np.array([[0.15, 0.30], [0.30, 0.60]]).astype(dtype)

    def covar_fn(t1, t2):
      return bm_utils.outer_multiply(t2 - t1, covar_matrix)

    vol_fn, covar_fn = bm_utils.construct_vol_data(None, covar_fn, 2, dtype)

    times = tf.constant([0.5, 1.0, 2.0, 3.0], dtype=dtype)
    vols = self.evaluate(vol_fn(times))
    np.testing.assert_array_equal(vols.shape, [4, 2, 2])
    # Directly testing vols is not sensible because given a covariance matrix
    # there are many volatility matrices that would lead to the same covar.
    # So we instead check the implied covariance.
    for i in range(4):
      actual_covar = np.matmul(vols[i], vols[i].transpose())
      np.testing.assert_allclose(actual_covar, covar_matrix)
    times2 = times + 0.31415
    # Check that the total covariance is correct.
    tc = self.evaluate(covar_fn(times, times2))
    np.testing.assert_array_equal(tc.shape, [4, 2, 2])
    # Directly testing vols is not sensible because given a covariance matrix
    # there are many volatility matrices that would lead to the same covar.
    # So we instead check the implied covariance.
    for i in range(4):
      np.testing.assert_allclose(tc[i], covar_matrix * 0.31415) 

def test_time_dependent_mr(self):
    """Tests that time depemdent mr uses generic sampling."""
    dtype = tf.float64
    mean_reversion = tff.math.piecewise.PiecewiseConstantFunc(
        [1.0], values=[0.03, 0.04], dtype=dtype)

    process = tff.models.hull_white.VectorHullWhiteModel(
        dim=1,
        mean_reversion=mean_reversion,
        volatility=[0.015],
        corr_matrix=None,
        initial_discount_rate_fn=self.instant_forward_rate_1d_fn,
        dtype=dtype)
    self.assertTrue(process._sample_with_generic)