def _euler_step(*, i, written_count, current_state, result, drift_fn,
volatility_fn, wiener_mean, num_samples, times, dt, sqrt_dt,
keep_mask, random_type, seed, normal_draws):
"""Performs one step of Euler scheme."""
current_time = times[i + 1]
written_count = tf.cast(written_count, tf.int32)
if normal_draws is not None:
dw = normal_draws[i]
else:
dw = random.mv_normal_sample((num_samples, ),
mean=wiener_mean,
random_type=random_type,
seed=seed)
dw = dw * sqrt_dt[i]
dt_inc = dt[i] * drift_fn(current_time, current_state) # pylint: disable=not-callable
dw_inc = tf.linalg.matvec(volatility_fn(current_time, current_state), dw) # pylint: disable=not-callable
next_state = current_state + dt_inc + dw_inc
result = utils.maybe_update_along_axis(tensor=result,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(
next_state, axis=1))
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return i + 1, written_count, next_state, result
def body_fn(i, written_count, current_x, rate_paths):
"""Simulate hull-white process to the next time point."""
if normal_draws is None:
normals = random.mv_normal_sample(
(num_samples, ),
mean=tf.zeros((self._dim, ), dtype=mean_reversion.dtype),
random_type=random_type,
seed=seed)
else:
normals = normal_draws[i]
if corr_matrix_root is not None:
normals = tf.linalg.matvec(corr_matrix_root[i], normals)
next_x = (
tf.math.exp(-mean_reversion[:, i + 1] * dt[i]) * current_x +
exp_x_t[:, i] + tf.math.sqrt(var_x_t[:, i]) * normals)
f_0_t = self._instant_forward_rate_fn(times[i + 1])
# Update `rate_paths`
rate_paths = utils.maybe_update_along_axis(
tensor=rate_paths,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_x, axis=1) + f_0_t)
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return (i + 1, written_count, next_x, rate_paths)
def body_fn(i, written_count, current_rates,
current_instant_forward_rates, rate_paths):
"""Simulate Heston process to the next time point."""
current_time = times[i]
next_time = times[i + 1]
if normal_draws is None:
normals = random.mv_normal_sample(
(num_samples, ),
mean=tf.zeros((self._dim, ), dtype=mean_reversion.dtype),
random_type=random_type,
seed=seed)
else:
normals = normal_draws[i]
next_rates, next_instant_forward_rates = _sample_at_next_time(
i, next_time, current_time, mean_reversion[i], volatility[i],
self._instant_forward_rate_fn, current_instant_forward_rates,
current_rates, corr_matrix_root, normals)
# Update `rate_paths`
rate_paths = utils.maybe_update_along_axis(
tensor=rate_paths,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_rates, axis=1))
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return (i + 1, written_count, next_rates,
next_instant_forward_rates, rate_paths)
def body_fn(i, written_count, current_x, rate_paths):
"""Simulate hull-white process to the next time point."""
if normal_draws is None:
normals = random.mv_normal_sample(
(num_samples, ),
mean=tf.zeros((self._dim, ), dtype=mean_reversion.dtype),
random_type=random_type,
seed=seed)
else:
normals = normal_draws[i]
if corr_matrix_root is not None:
normals = tf.linalg.matvec(corr_matrix_root[i], normals)
vol_x_t = tf.math.sqrt(tf.nn.relu(tf.transpose(var_x_t)[i]))
# If numerically `vol_x_t == 0`, the gradient of `vol_x_t` becomes `NaN`.
# To prevent this, we explicitly set `vol_x_t` to zero tensor at zero
# values so that the gradient is set to zero at this values.
vol_x_t = tf.where(vol_x_t > 0.0, vol_x_t, 0.0)
next_x = (
tf.math.exp(-tf.transpose(mean_reversion)[i + 1] * dt[i]) *
current_x + tf.transpose(exp_x_t)[i] + vol_x_t * normals)
f_0_t = self._instant_forward_rate_fn(times[i + 1])
# Update `rate_paths`
rate_paths = utils.maybe_update_along_axis(
tensor=rate_paths,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_x, axis=1) + f_0_t)
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return (i + 1, written_count, next_x, rate_paths)
def body_fn(i, written_count, current_vol, current_log_spot, vol_paths,
log_spot_paths):
"""Simulate Heston process to the next time point."""
time_step = dt[i]
if normal_draws is None:
normals = random.mv_normal_sample(
(num_samples, ),
mean=tf.zeros([2], dtype=mean_reversion.dtype),
seed=seed)
else:
normals = normal_draws[i]
def _next_vol_fn():
return _update_variance(mean_reversion[i], theta[i], volvol[i],
rho[i], current_vol, time_step,
normals[..., 0])
# Do not update variance if `time_step > tolerance`
next_vol = tf.cond(time_step > tolerance, _next_vol_fn,
lambda: current_vol)
def _next_log_spot_fn():
return _update_log_spot(mean_reversion[i], theta[i], volvol[i],
rho[i], current_vol, next_vol,
current_log_spot, time_step,
normals[..., 1])
# Do not update state if `time_step > tolerance`
next_log_spot = tf.cond(time_step > tolerance, _next_log_spot_fn,
lambda: current_log_spot)
# Update volatility paths
vol_paths = utils.maybe_update_along_axis(
tensor=vol_paths,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_vol, axis=1))
# Update log-spot paths
log_spot_paths = utils.maybe_update_along_axis(
tensor=log_spot_paths,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_log_spot, axis=1))
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return (i + 1, written_count, next_vol, next_log_spot, vol_paths,
log_spot_paths)
def body_fn(i, written_count,
current_x,
current_y,
x_paths,
y_paths):
"""Simulate qG-HJM process to the next time point."""
if normal_draws is None:
normals = random.mv_normal_sample(
(num_samples,),
mean=tf.zeros((self._dim,), dtype=self._dtype),
random_type=random_type, seed=seed)
else:
normals = normal_draws[i]
if self._sqrt_rho is not None:
normals = tf.linalg.matvec(self._sqrt_rho, normals)
vol = self._volatility(times[i + 1], current_x)
next_x = (current_x
+ (current_y - self._mean_reversion * current_x) * dt[i]
+ vol * normals * tf.math.sqrt(dt[i]))
next_y = current_y + (vol**2 -
2.0 * self._mean_reversion * current_y) * dt[i]
# Update `x_paths` and `y_paths`
x_paths = utils.maybe_update_along_axis(
tensor=x_paths,
do_update=True,
ind=written_count + 1,
axis=1,
new_tensor=tf.expand_dims(next_x, axis=1))
y_paths = utils.maybe_update_along_axis(
tensor=y_paths,
do_update=True,
ind=written_count + 1,
axis=1,
new_tensor=tf.expand_dims(next_y, axis=1))
written_count += 1
return (i + 1, written_count, next_x, next_y, x_paths, y_paths)
def body_fn(index, current_time, forward, vol, forward_paths, vol_paths,
normal_draws_index):
"""Simulate Sabr process to the next time point."""
forward, vol, normal_draws_index = self._propagate_to_time(
forward, vol, current_time, times[index], time_step, random_type,
seed, normal_draws, normal_draws_index)
# Always update paths in outer loop.
forward_paths = utils.maybe_update_along_axis(
tensor=forward_paths,
do_update=True,
ind=index,
axis=1,
new_tensor=tf.expand_dims(forward, axis=1))
vol_paths = utils.maybe_update_along_axis(
tensor=vol_paths,
do_update=True,
ind=index,
axis=1,
new_tensor=tf.expand_dims(vol, axis=1))
return index + 1, times[
index], forward, vol, forward_paths, vol_paths, normal_draws_index
def _while_loop(*, dim, steps_num, current_state, drift_fn, volatility_fn,
grad_volatility_fn, wiener_mean, num_samples, times, dt,
sqrt_dt, time_step, num_requested_times, keep_mask,
swap_memory, random_type, seed, normal_draws, input_gradients,
stratonovich_order, aux_normal_draws):
"""Smaple paths using tf.while_loop."""
cond_fn = lambda i, *args: i < steps_num
def step_fn(i, written_count, current_state, result):
return _milstein_step(dim=dim,
i=i,
written_count=written_count,
current_state=current_state,
result=result,
drift_fn=drift_fn,
volatility_fn=volatility_fn,
grad_volatility_fn=grad_volatility_fn,
wiener_mean=wiener_mean,
num_samples=num_samples,
times=times,
dt=dt,
sqrt_dt=sqrt_dt,
keep_mask=keep_mask,
random_type=random_type,
seed=seed,
normal_draws=normal_draws,
input_gradients=input_gradients,
stratonovich_order=stratonovich_order,
aux_normal_draws=aux_normal_draws)
maximum_iterations = (tf.cast(1. / time_step, dtype=tf.int32) +
tf.size(times))
# Include initial state, if necessary
result = tf.zeros((num_samples, num_requested_times, dim),
dtype=current_state.dtype)
result = utils.maybe_update_along_axis(tensor=result,
do_update=keep_mask[0],
ind=0,
axis=1,
new_tensor=tf.expand_dims(
current_state, axis=1))
written_count = tf.cast(keep_mask[0], dtype=tf.int32)
# Sample paths
_, _, _, result = tf.while_loop(cond_fn,
step_fn,
(0, written_count, current_state, result),
maximum_iterations=maximum_iterations,
swap_memory=swap_memory)
return result
def _while_loop(*, dim, batch_shape, steps_num, current_state, drift_fn,
volatility_fn, wiener_mean, num_samples, times, dt, sqrt_dt,
num_requested_times, keep_mask, swap_memory, random_type, seed,
normal_draws):
"""Smaple paths using tf.while_loop."""
cond_fn = lambda i, *args: i < steps_num
def step_fn(i, written_count, current_state, result):
return _euler_step(i=i,
written_count=written_count,
current_state=current_state,
result=result,
drift_fn=drift_fn,
volatility_fn=volatility_fn,
wiener_mean=wiener_mean,
num_samples=num_samples,
times=times,
dt=dt,
sqrt_dt=sqrt_dt,
keep_mask=keep_mask,
random_type=random_type,
seed=seed,
normal_draws=normal_draws)
# Include initial state, if necessary
result_shape = tf.concat(
[batch_shape, [num_samples, num_requested_times, dim]], axis=0)
result = tf.zeros(result_shape, dtype=current_state.dtype)
result = utils.maybe_update_along_axis(tensor=result,
do_update=keep_mask[0],
ind=0,
axis=result.shape.rank - 2,
new_tensor=tf.expand_dims(
current_state, axis=-2))
written_count = tf.cast(keep_mask[0], dtype=tf.int32)
# Sample paths
_, _, _, result = tf.while_loop(cond_fn,
step_fn,
(0, written_count, current_state, result),
maximum_iterations=steps_num,
swap_memory=swap_memory)
return result
def _milstein_step(*, i, written_count, current_state, result, drift_fn,
volatility_fn, grad_volatility_fn, wiener_mean, num_samples,
times, dt, sqrt_dt, keep_mask, random_type, seed,
normal_draws):
"""Performs one step of Milstein scheme."""
current_time = times[i + 1]
written_count = tf.cast(written_count, tf.int32)
if normal_draws is not None:
dw = normal_draws[i]
else:
dw = random.mv_normal_sample((num_samples,),
mean=wiener_mean,
random_type=random_type,
seed=seed)
dw = dw * sqrt_dt[i]
dt_inc = dt[i] * drift_fn(current_time, current_state) # pylint: disable=not-callable
dw_inc = tf.linalg.matvec(volatility_fn(current_time, current_state), dw) # pylint: disable=not-callable
# Higher order terms. For dim 1, the product here is elementwise.
# Will need to adjust for higher dims.
hot_vol = tf.squeeze(
tf.multiply(
volatility_fn(current_time, current_state),
grad_volatility_fn(current_time, current_state)), -1)
hot_dw = dw * dw - dt[i]
hot_inc = tf.multiply(hot_vol, hot_dw) / 2
next_state = current_state + dt_inc + dw_inc + hot_inc
result = utils.maybe_update_along_axis(
tensor=result,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_state, axis=1))
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return i + 1, written_count, next_state, result
def _sample_paths(self,
times,
num_samples,
random_type,
skip,
seed,
normal_draws=None,
times_grid=None,
validate_args=False):
"""Returns a sample of paths from the process."""
# Note: all the notations below are the same as in [1].
num_requested_times = tf.shape(times)[0]
params = [self._mean_reversion, self._volatility]
if self._corr_matrix is not None:
params = params + [self._corr_matrix]
times, keep_mask = _prepare_grid(times, times_grid, *params)
# Add zeros as a starting location
dt = times[1:] - times[:-1]
if dt.shape.is_fully_defined():
steps_num = dt.shape.as_list()[-1]
else:
steps_num = tf.shape(dt)[-1]
# TODO(b/148133811): Re-enable Sobol test when TF 2.2 is released.
if random_type == random.RandomType.SOBOL:
raise ValueError(
'Sobol sequence for Euler sampling is temporarily '
'unsupported when `time_step` or `times` have a '
'non-constant value')
if normal_draws is None:
# In order to use low-discrepancy random_type we need to generate the
# sequence of independent random normals upfront. We also precompute
# random numbers for stateless random type in order to ensure independent
# samples for multiple function calls whith different seeds.
if random_type in (random.RandomType.SOBOL,
random.RandomType.HALTON,
random.RandomType.HALTON_RANDOMIZED,
random.RandomType.STATELESS,
random.RandomType.STATELESS_ANTITHETIC):
normal_draws = utils.generate_mc_normal_draws(
num_normal_draws=self._dim,
num_time_steps=steps_num,
num_sample_paths=num_samples,
random_type=random_type,
seed=seed,
dtype=self._dtype,
skip=skip)
else:
normal_draws = None
else:
if validate_args:
draws_times = tf.shape(normal_draws)[0]
asserts = tf.assert_equal(
draws_times,
tf.shape(times)[0] - 1, # We have added `0` to `times`
message='`tf.shape(normal_draws)[1]` should be equal to the '
'number of all `times` plus the number of all jumps of '
'the piecewise constant parameters.')
with tf.compat.v1.control_dependencies([asserts]):
normal_draws = tf.identity(normal_draws)
# The below is OK because we support exact discretization with piecewise
# constant mr and vol.
mean_reversion = self._mean_reversion(times)
volatility = self._volatility(times)
if self._corr_matrix is not None:
corr_matrix = _get_parameters(times + tf.math.reduce_min(dt) / 2,
self._corr_matrix)[0]
corr_matrix_root = tf.linalg.cholesky(corr_matrix)
else:
corr_matrix_root = None
exp_x_t = self._conditional_mean_x(times, mean_reversion, volatility)
var_x_t = self._conditional_variance_x(times, mean_reversion,
volatility)
if self._dim == 1:
mean_reversion = tf.expand_dims(mean_reversion, axis=0)
cond_fn = lambda i, *args: i < tf.size(dt)
def body_fn(i, written_count, current_x, rate_paths):
"""Simulate hull-white process to the next time point."""
if normal_draws is None:
normals = random.mv_normal_sample(
(num_samples, ),
mean=tf.zeros((self._dim, ), dtype=mean_reversion.dtype),
random_type=random_type,
seed=seed)
else:
normals = normal_draws[i]
if corr_matrix_root is not None:
normals = tf.linalg.matvec(corr_matrix_root[i], normals)
vol_x_t = tf.math.sqrt(tf.nn.relu(tf.transpose(var_x_t)[i]))
# If numerically `vol_x_t == 0`, the gradient of `vol_x_t` becomes `NaN`.
# To prevent this, we explicitly set `vol_x_t` to zero tensor at zero
# values so that the gradient is set to zero at this values.
vol_x_t = tf.where(vol_x_t > 0.0, vol_x_t, 0.0)
next_x = (
tf.math.exp(-tf.transpose(mean_reversion)[i + 1] * dt[i]) *
current_x + tf.transpose(exp_x_t)[i] + vol_x_t * normals)
f_0_t = self._instant_forward_rate_fn(times[i + 1])
# Update `rate_paths`
rate_paths = utils.maybe_update_along_axis(
tensor=rate_paths,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(next_x, axis=1) + f_0_t)
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return (i + 1, written_count, next_x, rate_paths)
rate_paths = tf.zeros((num_samples, num_requested_times, self._dim),
dtype=self._dtype)
# Include initial state, if necessary
f0_t = self._instant_forward_rate_fn(times[0])
rate_paths = utils.maybe_update_along_axis(tensor=rate_paths,
do_update=keep_mask[0],
ind=0,
axis=1,
new_tensor=f0_t)
written_count = tf.cast(keep_mask[0], dtype=tf.int32)
initial_x = tf.zeros((num_samples, self._dim), dtype=self._dtype)
# TODO(b/157232803): Use tf.cumsum instead?
_, _, _, rate_paths = tf.while_loop(
cond_fn, body_fn, (0, written_count, initial_x, rate_paths))
return rate_paths
def maybe_update_along_axis(do_update):
return utils.maybe_update_along_axis(
tensor=tensor, new_tensor=new_tensor, axis=1, ind=2,
do_update=do_update)
def _milstein_step(*, dim, i, written_count, current_state, result, drift_fn,
volatility_fn, grad_volatility_fn, wiener_mean, num_samples,
times, dt, sqrt_dt, keep_mask, random_type, seed,
normal_draws, input_gradients, stratonovich_order,
aux_normal_draws):
"""Performs one step of Milstein scheme."""
current_time = times[i + 1]
written_count = tf.cast(written_count, tf.int32)
if normal_draws is not None:
dw = normal_draws[i]
else:
dw = random.mv_normal_sample((num_samples, ),
mean=wiener_mean,
random_type=random_type,
seed=seed)
if aux_normal_draws is not None:
stratonovich_draws = []
for j in range(3):
stratonovich_draws.append(
tf.reshape(aux_normal_draws[j][i],
[num_samples, dim, stratonovich_order]))
else:
stratonovich_draws = []
# Three sets of normal draws for stratonovich integrals.
for j in range(3):
stratonovich_draws.append(
random.mv_normal_sample(
(num_samples, ),
mean=tf.zeros((dim, stratonovich_order),
dtype=current_state.dtype,
name='stratonovich_draws_{}'.format(j)),
random_type=random_type,
seed=seed))
if dim == 1:
drift = drift_fn(current_time, current_state)
vol = volatility_fn(current_time, current_state)
grad_vol = grad_volatility_fn(current_time, current_state,
tf.ones_like(current_state))
next_state = _milstein_1d(dw=dw,
dt=dt[i],
sqrt_dt=sqrt_dt[i],
current_state=current_state,
drift=drift,
vol=vol,
grad_vol=grad_vol)
else:
drift = drift_fn(current_time, current_state)
vol = volatility_fn(current_time, current_state)
# This is a list of size equal to the dimension of the state space `dim`.
# It contains tensors of shape [num_samples, dim, wiener_dim] representing
# the gradient of the volatility function. In our case, the dimension of the
# wiener process `wiener_dim` is equal to the state dimension `dim`.
grad_vol = [
grad_volatility_fn(current_time, current_state, start)
for start in input_gradients
]
next_state = _milstein_nd(dim=dim,
num_samples=num_samples,
dw=dw,
dt=dt[i],
sqrt_dt=sqrt_dt[i],
current_state=current_state,
drift=drift,
vol=vol,
grad_vol=grad_vol,
stratonovich_draws=stratonovich_draws,
stratonovich_order=stratonovich_order)
result = utils.maybe_update_along_axis(tensor=result,
do_update=keep_mask[i + 1],
ind=written_count,
axis=1,
new_tensor=tf.expand_dims(
next_state, axis=1))
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return i + 1, written_count, next_state, result
