class_docstr: Heston Model with piecewise constant parameters.
  
  Represents the Ito process:
  
  ```None
    dX(t) = -V(t) / 2 * dt + sqrt(V(t)) * dW_{X}(t),
    dV(t) = mean_reversion(t) * (theta(t) - V(t)) * dt
            + volvol(t) * sqrt(V(t)) * dW_{V}(t)
  ```
  
  where `W_{X}` and `W_{V}` are 1D Brownian motions with a correlation
  `rho(t)`. `mean_reversion`, `theta`, `volvol`, and `rho` are positive
  piecewise constant functions of time. Here `V(t)` represents the process
  variance at time `t` and `X` represents logarithm of the spot price at time
  `t`.
  
  `mean_reversion` corresponds to the mean reversion rate, `theta` is the long
  run price variance, and `volvol` is the volatility of the volatility.
  
  See [1] and [2] for details.
  
  #### Example
  
  ```python
  import tf_quant_finance as tff
  import numpy as np
  volvol = tff.math.piecewise.PiecewiseConstantFunc(
      jump_locations=[0.5], values=[1, 1.1], dtype=np.float64)
  process = tff.models.HestonModel(
      mean_reversion=0.5, theta=0.04, volvol=volvol, rho=0.1, dtype=np.float64)
  times = np.linspace(0.0, 1.0, 1000)
  num_samples = 10000  # number of trajectories
  sample_paths = process.sample_paths(
      times,
      time_step=0.01,
      num_samples=num_samples,
      initial_state=np.array([1.0, 0.04]),
      random_type=random.RandomType.SOBOL)
  ```
  
  #### References:
    [1]: Cristian Homescu. Implied volatility surface: construction
      methodologies and characteristics.
      arXiv: https://arxiv.org/pdf/1107.1834.pdf
    [2]: Leif Andersen. Efficient Simulation of the Heston Stochastic
      Volatility Models. 2006.
      Link:
      http://www.ressources-actuarielles.net/ext/isfa/1226.nsf/d512ad5b22d73cc1c1257052003f1aed/1826b88b152e65a7c12574b000347c74/$FILE/LeifAndersenHeston.pdf
class_docstr_tok: ["Heston","Model","with","piecewise","constant","parameters",".","Represents","the","Ito","process",":","``","`","None","dX","(","t",")","=","-V","(","t",")","\/","2","*","dt","+","sqrt","(","V","(","t",")",")","*","dW_","{","X","}","(","t",")",",","dV","(","t",")","=","mean_reversion","(","t",")","*","(","theta","(","t",")","-","V","(","t",")",")","*","dt","+","volvol","(","t",")","*","sqrt","(","V","(","t",")",")","*","dW_","{","V","}","(","t",")","``","`","where","`","W_","{","X","}","`","and","`","W_","{","V","}","`","are","1D","Brownian","motions","with","a","correlation","`","rho","(","t",")","`",".","`","mean_reversion","`",",","`","theta","`",",","`","volvol","`",",","and","`","rho","`","are","positive","piecewise","constant","functions","of","time",".","Here","`","V","(","t",")","`","represents","the","process","variance","at","time","`","t","`","and","`","X","`","represents","logarithm","of","the","spot","price","at","time","`","t","`",".","`","mean_reversion","`","corresponds","to","the","mean","reversion","rate",",","`","theta","`","is","the","long","run","price","variance",",","and","`","volvol","`","is","the","volatility","of","the","volatility",".","See","[","1","]","and","[","2","]","for","details",".","#","#","#","#","Example","``","`","python","import","tf_quant_finance","as","tff","import","numpy","as","np","volvol","=","tff.math.piecewise.PiecewiseConstantFunc","(","jump_locations=","[","0.5","]",",","values=","[","1",",","1.1","]",",","dtype=np.float64",")","process","=","tff.models.HestonModel","(","mean_reversion=0.5",",","theta=0.04",",","volvol=volvol",",","rho=0.1",",","dtype=np.float64",")","times","=","np.linspace","(","0.0",",","1.0",",","1000",")","num_samples","=","10000","#","number","of","trajectories","sample_paths","=","process.sample_paths","(","times",",","time_step=0.01",",","num_samples=num_samples",",","initial_state=np.array","(","[","1.0",",","0.04","]",")",",","random_type=random.RandomType.SOBOL",")","``","`","#","#","#","#","References",":","[","1","]",":","Cristian","Homescu",".","Implied","volatility","surface",":","construction","methodologies","and","characteristics",".","arXiv",":","https",":","\/\/arxiv.org\/pdf\/1107.1834.pdf","[","2","]",":","Leif","Andersen",".","Efficient","Simulation","of","the","Heston","Stochastic","Volatility","Models",".","2006",".","Link",":","http",":","\/\/www.ressources-actuarielles.net\/ext\/isfa\/1226.nsf\/d512ad5b22d73cc1c1257052003f1aed\/1826b88b152e65a7c12574b000347c74\/","$","FILE\/LeifAndersenHeston.pdf"]
docstr: Returns a sample of paths from the process.
docstr_tok: ["Returns","a","sample","of","paths","from","the","process","."]
code: def _sample_paths(
      self,
      times,
      num_requested_times,
      current_log_spot,
      current_vol,
      num_samples,
      random_type,
      keep_mask,
      seed,
      skip,
      tolerance,
      precompute_normal_draws,
      normal_draws,
  ):
      """Returns a sample of paths from the process."""
      # Note: all the notations below are the same as in [1].
      dt = times[1:] - times[:-1]
      # Compute the parameters at `times`. Here + tf.reduce_min(dt) / 2 ensures
      # that the value is constant between `times`.
      (
          mean_reversion,
          theta,
          volvol,
          rho,
      ) = _get_parameters(  # pylint: disable=unbalanced-tuple-unpacking
          times + tf.reduce_min(dt) / 2,
          self._mean_reversion,
          self._theta,
          self._volvol,
          self._rho,
      )
      # In order random_type which is not PSEUDO,  sequence of independent random
      # normals should be generated upfront.
      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:
          if precompute_normal_draws or 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=2,
                  num_time_steps=steps_num,
                  num_sample_paths=num_samples,
                  random_type=random_type,
                  dtype=self.dtype(),
                  seed=seed,
                  skip=skip,
              )
          else:
              # If pseudo or anthithetic sampling is used, proceed with random
              # sampling at each step.
              normal_draws = None
      # Prepare results format
      written_count = 0
      if (
          isinstance(num_requested_times, int)
          and num_requested_times == 1
      ):
          record_samples = False
          log_spot_paths = current_log_spot
          vol_paths = current_vol
      else:
          # If more than one sample has to be recorded, create a TensorArray
          record_samples = True
          element_shape = current_log_spot.shape
          log_spot_paths = tf.TensorArray(
              dtype=times.dtype,
              size=num_requested_times,
              element_shape=element_shape,
              clear_after_read=False,
          )
          vol_paths = tf.TensorArray(
              dtype=times.dtype,
              size=num_requested_times,
              element_shape=element_shape,
              clear_after_read=False,
          )
          # Include initial state, if necessary
          log_spot_paths = log_spot_paths.write(
              written_count, current_log_spot
          )
          vol_paths = vol_paths.write(written_count, current_vol)
      written_count += tf.cast(keep_mask[0], dtype=tf.int32)
  
      # Define sampling while_loop body function
      def cond_fn(i, written_count, *args):
          # It can happen that `times_grid[-1] > times[-1]` in which case we have
          # to terminate when `written_count` reaches `num_requested_times`
          del args
          return tf.math.logical_and(
              i < steps_num, written_count < num_requested_times
          )
  
      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,
          )
  
          if record_samples:
              # Update volatility paths
              vol_paths = vol_paths.write(written_count, next_vol)
              # Update log-spot paths
              log_spot_paths = log_spot_paths.write(
                  written_count, next_log_spot
              )
          else:
              vol_paths = next_vol
              log_spot_paths = next_log_spot
          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,
          )
  
      # Sample paths
      _, _, _, _, vol_paths, log_spot_paths = tf.while_loop(
          cond_fn,
          body_fn,
          (
              0,
              0,
              current_vol,
              current_log_spot,
              vol_paths,
              log_spot_paths,
          ),
          maximum_iterations=steps_num,
      )
      if not record_samples:
          # shape [num_samples, 1]
          vol_paths = tf.expand_dims(vol_paths, axis=-1)
          log_spot_paths = tf.expand_dims(log_spot_paths, axis=-1)
          # shape [num_samples, 1, 1]
          return tf.stack([log_spot_paths, vol_paths], -1)
      # Shape [num_time_points] + [num_samples]
      vol_paths = vol_paths.stack()
      log_spot_paths = log_spot_paths.stack()
      # transpose to shape [num_samples, num_time_points]
      vol_paths = tf.transpose(vol_paths)
      log_spot_paths = tf.transpose(log_spot_paths)
      # Shape [num_samples, num_time_points, 2]
      return tf.stack([log_spot_paths, vol_paths], -1)
  
file_path: tf-quant-finance/tf_quant_finance/models/heston/heston_model.py
imports: from typing import Union, Callable, Optional
  import numpy
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
  from tf_quant_finance.math import piecewise
  from tf_quant_finance.math import random_ops
  from tf_quant_finance.models import generic_ito_process
  from tf_quant_finance.models import utils
label_desc: Use image node_id 3 for calling the HestonModel obj's underlying member method code with example usage: obj._sample_paths(times, num_requested_times, current_log_spot, current_vol, num_samples, random_type, keep_mask, seed, skip, tolerance, precompute_normal_draws, normal_draws) and returns: tf, tf, unknown,  written_count,  next_vol,  next_log_spot,  vol_paths,  log_spot_paths, tf, _update_variance, _update_log_spot