Returns a sample of paths from the process using the Milstein method.
(*, dim, drift_fn, volatility_fn, grad_volatility_fn, times,
time_step, keep_mask, num_requested_times, num_samples,
initial_state, random_type, seed, swap_memory, skip,
precompute_normal_draws, watch_params, time_indices,
input_gradients, stratonovich_order, dtype)
| 256 | |
| 257 | |
| 258 | def _sample(*, dim, drift_fn, volatility_fn, grad_volatility_fn, times, |
| 259 | time_step, keep_mask, num_requested_times, num_samples, |
| 260 | initial_state, random_type, seed, swap_memory, skip, |
| 261 | precompute_normal_draws, watch_params, time_indices, |
| 262 | input_gradients, stratonovich_order, dtype): |
| 263 | """Returns a sample of paths from the process using the Milstein method.""" |
| 264 | dt = times[1:] - times[:-1] |
| 265 | sqrt_dt = tf.sqrt(dt) |
| 266 | current_state = initial_state + tf.zeros([num_samples, dim], |
| 267 | dtype=initial_state.dtype) |
| 268 | if dt.shape.is_fully_defined(): |
| 269 | steps_num = dt.shape.as_list()[-1] |
| 270 | else: |
| 271 | steps_num = tf.shape(dt)[-1] |
| 272 | # In order to use low-discrepancy random_type we need to generate the sequence |
| 273 | # of independent random normals upfront. We also precompute random numbers |
| 274 | # for stateless random type in order to ensure independent samples for |
| 275 | # multiple function calls with different seeds. |
| 276 | if precompute_normal_draws or random_type in ( |
| 277 | random.RandomType.SOBOL, random.RandomType.HALTON, |
| 278 | random.RandomType.HALTON_RANDOMIZED, random.RandomType.STATELESS, |
| 279 | random.RandomType.STATELESS_ANTITHETIC): |
| 280 | # Process dimension plus auxiliary random variables for stratonovich |
| 281 | # integral computation. |
| 282 | all_normal_draws = utils.generate_mc_normal_draws( |
| 283 | num_normal_draws=dim + 3 * dim * stratonovich_order, |
| 284 | num_time_steps=steps_num, |
| 285 | num_sample_paths=num_samples, |
| 286 | random_type=random_type, |
| 287 | dtype=dtype, |
| 288 | seed=seed, |
| 289 | skip=skip) |
| 290 | normal_draws = all_normal_draws[:, :, :dim] |
| 291 | wiener_mean = None |
| 292 | # Auxiliary normal draws for use with the stratonovich integral |
| 293 | # approximation. |
| 294 | aux_normal_draws = [] |
| 295 | start = dim |
| 296 | for _ in range(3): |
| 297 | end = start + dim * stratonovich_order |
| 298 | aux_normal_draws.append(all_normal_draws[:, :, start:end]) |
| 299 | start = end |
| 300 | else: |
| 301 | # If pseudo or anthithetic sampling is used, proceed with random sampling |
| 302 | # at each step. |
| 303 | wiener_mean = tf.zeros((dim,), dtype=dtype, name='wiener_mean') |
| 304 | normal_draws = None |
| 305 | aux_normal_draws = None |
| 306 | if watch_params is None: |
| 307 | # Use while_loop if `watch_params` is not passed |
| 308 | return _while_loop( |
| 309 | dim=dim, |
| 310 | steps_num=steps_num, |
| 311 | current_state=current_state, |
| 312 | drift_fn=drift_fn, |
| 313 | volatility_fn=volatility_fn, |
| 314 | grad_volatility_fn=grad_volatility_fn, |
| 315 | wiener_mean=wiener_mean, |
no test coverage detected