def _grad_volatility_fn(current_time, current_state,
input_gradients):
return gradient.fwd_gradient(
functools.partial(vol_fn, current_time),
current_state,
input_gradients=input_gradients,
unconnected_gradients=tf.UnconnectedGradients.ZERO)
def vol_fn(time):
# We should consider changing the start time to be some small dt behind
# the time. In case the total covariance is being computed by a numerical
# integration, this will mean that we spend less time iterating.
start_time = tf.zeros_like(time)
total_covar_fn = lambda t: total_covariance_fn(start_time, t)
vol_sq = gradient.fwd_gradient(total_covar_fn, time)
return tf.cholesky(vol_sq, name='volatility')
def _instant_forward_rate_fn(t):
t = tf.convert_to_tensor(t, dtype=self._dtype)
def _log_zero_coupon_bond(x):
r = tf.convert_to_tensor(
initial_discount_rate_fn(x), dtype=self._dtype)
return -r * x
rate = -gradient.fwd_gradient(
_log_zero_coupon_bond, t, use_gradient_tape=True,
unconnected_gradients=tf.UnconnectedGradients.ZERO)
return rate
def _drift_fn(t, x):
"""Drift function of correlated Hull-White."""
# Get parameter values at time `t`
mean_reversion, volatility = _get_parameters( # pylint: disable=unbalanced-tuple-unpacking
tf.expand_dims(t, -1), self._mean_reversion, self._volatility)
fwd_rates = self._instant_forward_rate_fn(t)
fwd_rates_grad = gradient.fwd_gradient(
self._instant_forward_rate_fn, t)
drift = fwd_rates_grad + mean_reversion * fwd_rates
drift += (volatility**2 / 2 / mean_reversion *
(1 - tf.math.exp(-2 * mean_reversion * t)) -
mean_reversion * x)
return drift
def _instant_forward_rate_fn(t):
t = tf.convert_to_tensor(t, dtype=self._dtype)
def _log_zero_coupon_bond(x):
r = tf.convert_to_tensor(initial_discount_rate_fn(x),
dtype=self._dtype)
# If initial_discount_rate_fn returns a Tensor of the same input,
# expand the output dimension to `input_shape + [1]`. Otherwise,
# it is expected that `r` is of shape `input_shape + [dim]`.
if r.shape.rank == x.shape.rank:
r = tf.expand_dims(r, axis=-1)
# Shape `x.shape + [dim]`
return -r * tf.expand_dims(x, axis=-1)
rate = -gradient.fwd_gradient(
_log_zero_coupon_bond,
t,
use_gradient_tape=True,
unconnected_gradients=tf.UnconnectedGradients.ZERO)
return rate
def drift_from_total_drift(t):
start_time = tf.zeros_like(t)
return gradient.fwd_gradient(
lambda x: total_drift_fn(start_time, x), t)
