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