# Lint as: python3
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Join a sequence of Ito Processes with specified correlations."""

import tensorflow.compat.v2 as tf

from tf_quant_finance.models import euler_sampling
from tf_quant_finance.models import generic_ito_process
from tf_quant_finance.models import ito_process
from tf_quant_finance.models import utils


class JoinedItoProcess(generic_ito_process.GenericItoProcess):
  """Join of Ito Processes with specified time dependent correlations.

  For a sequence of Ito processes `I_1, .., I_n` of dimensions `d_1,.., d_n`,
  the class initializes a process `I` of dimension `d_1 + .. + d_n` with
  marginal proceses `I_i` and a correlation function `Corr(t)`. That is, let the
  Ito Process `I_i` describe an SDE

  ```None
  dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dW^i_t
  ```

  where `a_i(t, x)` is a function taking values in `R^{d_i}`, `b_i(t, X_t)` is a
  function taking values in `d_i x d_i` matrices, `W_i` is a `d_i`-dimensional
  Brownian motion.

  Then `I` describes an SDE for the joint process `(X_1,..., X_n)` of dimension
  `d:= d_1 + ... + d_n`

  ```None
  dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dB^i_t,
  ```

  where `(B_1, ..., B_n) = chol(t) * (W_1, ..., W_n)` for a Cholesky
  decomposition `chol(t)` of the correlation matrix `Corr(t)` at time `t`.
  Here `(W_1, ..., W_n)` is `d`-dimensional vector and `Corr(t)` is a `d x d`
  correlation matrix.

  `Corr(t)` is represented as a block-diagonal formed of a list of matrices
  `[m_1(t), m_2(t), ..., m_k(t)]` with `sum(rank(m_i)) = d`.

  # Example. # Black-scholes and Heston model join.
  ```python
  import numpy as np
  import tensorflow as tf
  import tf_quant_finance as tff

  dtype = tf.float64
  # Define Black scholes model with zero rate and volatility `0.1`
  sigma = 0.1
  def drift_fn(t , x):
    return -sigma**2 / 2
  def vol_fn(t , x):
    return sigma * tf.ones([1, 1], dtype=x.dtype)
  black_scholes_process = tff.models.GenericItoProcess(
      dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)
  # Define Heston Model
  epsilon = tff.math.piecewise.PiecewiseConstantFunc(
      jump_locations=[0.5], values=[0.01, 0.02], dtype=np.float64)
  heston_process = tff.models.heston_model.HestonModel(
      kappa=1.0, theta=0.04, epsilon=epsilon, rho=0.2, dtype=dtype)
  # Define join process where correlation between `black_scholes_process` and
  # log-asset price of the `heston_process` is 0.5.
  corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]
  # `corr_structure` corresponds to a `3 x 3` correlation matrix. Here Brownian
  # motion of `black_scholes_process` is correlated only with the 1st dimension
  # of `heston_process` but not with the second one.
  join_process = JoinedItoProcess(
      processes=[black_scholes_process, heston_process],
      corr_structure=corr_structure)
  # Sample 100,000 sample paths at times [0.1, 1.0] from the join process using
  # Sobol random sequence
  times = [0.1, 1.0]
  # Wrap sample_paths method with a tf.function
  sample_paths_fn = tf.function(process.sample_paths)
  samples = sample_paths_fn(
      times=times, time_step=0.01, num_samples=100000,
      initial_state=np.array([0.0, 1.0, 0.04]),
      random_type=random.RandomType.SOBOL)
  # Estimated correlations.
  np.corrcoef(samples[:, -1, :], rowvar=False)
  # Expected result:
  # [[1.        , 0.49567078, 0.08128067],
  #  [0.49567078, 1.        , 0.16580689],
  #  [0.08128067, 0.16580689, 1.        ]]
  ```
  """

  def __init__(self, processes, corr_structure, dtype=None, name=None):
    """Initializes a JoinedItoProcess.

    Takes a list of `processes` which are instances of `tff.models.ItoProcess`
    and a list `corr_structure` of correlation matrices and creates an Ito
    process that joins `processes` using the correlation structure.
    `corr_structure` describes block-diagonal structure of correlations for
    the Brownian motions in `processes`. For example, if the dimension of the
    JoinedItoProcess is `3` and
    `corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]`, then the introduced
    correlation is
    `Corr(t) = [[1.0, 0.5, 0.0], [0.5, 1.0, 0.0], [0.0, 0.0, 1.0]]`,
    where `Corr(t)` is the same as in the `JoinedItoProcess` docstring.

    Args:
      processes: A sequence of instances of `tff.models.ItoProcess`. All
        processes should have the same `dtype.`
      corr_structure: A list of correlation matrices. Each correlation matrix
        is either a `Tensor` of the same `dtype` as the `processes` and
        square shape  (i.e., `[d_i, d_i]` for some `d_i`) or a callable. The
        callables should accept a scalar (stands for time `t`) and return a
        square `Tensor`. The total dimension
        `sum([m.shape[-1] for m in corr_structure]` of correlation
        structure should be the same as the dimension of the `JoinedItoProcess`
        `sum([p.dim() for p in processes])`.
      dtype: The default `dtype` of the `processes`.
        Default value: None which means that default dtypes inferred by
          TensorFlow are used.
      name: Python string. The name scope under which ops created by the methods
        of this class are nested.
        Default value: `None` which maps to the default name
          `join_ito_process`.

    Raises:
      ValueError:
        (a) If any of the `processes` is not an `ItoProcess`.
        (b) If `processes` do not have the same `dtype`.
    """
    self._name = name or "join_ito_process"
    with tf.name_scope(self._name):
      self._processes = []  # Input processes
      dim = 0  # Dimension of the process
      for process in processes:
        if not isinstance(process, ito_process.ItoProcess):
          raise ValueError(
              "All input process of JoinedItoProcess must be instances "
              "of the ItoProcess class.")
        self._processes.append(process)
        d = process.dim()
        dim += d  # Total dimension of the process
        if dtype is None:
          dtype = process.dtype()
        elif dtype != process.dtype():
          raise ValueError("All processes should have the same `dtype`")
      self._corr_structure = [
          corr if callable(corr) else tf.convert_to_tensor(
              corr, dtype=dtype, name="corr")
          for corr in corr_structure]
      self._dim = dim

      def _drift_fn(t, x):
        """Drift function of the JoinedItoProcess."""
        drifts = []
        i1 = 0
        i2 = 0
        for p in self._processes:
          dim = p.dim()
          i2 += dim
          position = x[..., i1:i2]
          drift = tf.convert_to_tensor(p.drift_fn()(t, position),
                                       dtype=dtype,
                                       name="drift")
          drift = tf.broadcast_to(drift, position.shape)
          drifts.append(drift)
          i1 += dim
        return tf.concat(drifts, -1)

      def _vol_fn(t, x):
        """Volatility function of the JoinedItoProcess."""
        vols = []
        i1 = 0
        i2 = 0
        for p in self._processes:
          dim = p.dim()
          i2 += dim
          position = x[..., i1:i2]
          vol = tf.convert_to_tensor(p.volatility_fn()(t, position),
                                     dtype=dtype,
                                     name="volatility")
          vol = tf.broadcast_to(vol, position.shape + [dim])
          vols.append(vol)
          i1 += dim
        # Convert block diagonal volatilities to a dense correlation matrix
        vol = utils.block_diagonal_to_dense(*vols)
        # Compute Cholesky decomposition of the correlation structure
        corr_structure = _get_parameters(
            tf.expand_dims(t, -1), *self._corr_structure)
        cholesky_decomp = [tf.linalg.cholesky(m) for m in corr_structure]
        cholesky_decomp = utils.block_diagonal_to_dense(*cholesky_decomp)
        return tf.linalg.matmul(vol, cholesky_decomp)
      # The following will initialize the Generic Ito Process that has
      # a sampling and PDE solving methods
      super().__init__(dim, _drift_fn, _vol_fn, dtype, name)

  def sample_paths(self,
                   times,
                   num_samples=1,
                   initial_state=None,
                   random_type=None,
                   seed=None,
                   time_step=None,
                   swap_memory=True,
                   skip=0,
                   name=None):
    """Returns a sample of paths from the process using Euler sampling.

    Args:
      times: Rank 1 `Tensor` of increasing positive real values. The times at
        which the path points are to be evaluated.
      num_samples: Positive scalar `int`. The number of paths to draw.
        Default value: 1.
      initial_state: `Tensor` of shape `[self._dim]`. The initial state of the
        process.
        Default value: None which maps to a zero initial state.
      random_type: Enum value of `RandomType`. The type of (quasi)-random number
        generator to use to generate the paths.
        Default value: None which maps to the standard pseudo-random numbers.
      seed: Seed for the random number generator. The seed is
        only relevant if `random_type` is one of
        `[STATELESS, PSEUDO, HALTON_RANDOMIZED, PSEUDO_ANTITHETIC,
          STATELESS_ANTITHETIC]`. For `PSEUDO`, `PSEUDO_ANTITHETIC` and
        `HALTON_RANDOMIZED` the seed should be an integer scalar `Tensor`. For
        `STATELESS` and  `STATELESS_ANTITHETIC `must be supplied as an integer
        `Tensor` of shape `[2]`.
        Default value: `None` which means no seed is set.
      time_step: Real scalar `Tensor`. The maximal distance between time points
        in grid in Euler scheme.
      swap_memory: A Python bool. Whether GPU-CPU memory swap is enabled for
        this op. See an equivalent flag in `tf.while_loop` documentation for
        more details. Useful when computing a gradient of the op since
        `tf.while_loop` is used to propagate stochastic process in time.
        Default value: True.
      skip: `int32` 0-d `Tensor`. The number of initial points of the Sobol or
        Halton sequence to skip. Used only when `random_type` is 'SOBOL',
        'HALTON', or 'HALTON_RANDOMIZED', otherwise ignored.
        Default value: `0`.
      name: Python string. The name to give this op.
        Default value: `None` which maps to `sample_paths` is used.

    Returns:
     A real `Tensor` of shape `[num_samples, k, n]` where `k` is the size of the
     `times`, and `n` is the dimension of the process.

    Raises:
      ValueError: If `time_step` is not supplied.
    """
    if time_step is None:
      raise ValueError("`time_step` has to be supplied for JoinedItoProcess "
                       "`sample_paths` method.")
    name = name or self._name + "sample_paths"
    with tf.name_scope(name):
      if initial_state is None:
        initial_state = tf.zeros(self._dim, dtype=self.dtype(),
                                 name="initial_state")
      else:
        if isinstance(initial_state, (tuple, list)):
          initial_state = [tf.convert_to_tensor(state, dtype=self.dtype(),
                                                name="initial_state")
                           for state in initial_state]
          initial_state = tf.stack(initial_state)
        else:
          initial_state = tf.convert_to_tensor(initial_state,
                                               dtype=self.dtype(),
                                               name="initial_state")
      samples = euler_sampling.sample(self.dim(),
                                      drift_fn=self.drift_fn(),
                                      volatility_fn=self.volatility_fn(),
                                      times=times,
                                      time_step=time_step,
                                      num_samples=num_samples,
                                      initial_state=initial_state,
                                      random_type=random_type,
                                      seed=seed,
                                      swap_memory=swap_memory,
                                      skip=skip,
                                      dtype=self.dtype())
      return samples


def _get_parameters(times, *params):
  """Gets parameter values at at specified `times`."""
  res = []
  for param in params:
    if callable(param):
      # Used only in drift and volatility computation.
      # Here `times` is of shape [1]
      t = tf.squeeze(times)
      # The result has to have shape [1] + param.shape
      param_value = tf.convert_to_tensor(param(t), dtype=times.dtype,
                                         name="param_value")
      res.append(tf.expand_dims(param_value, 0))
    else:
      res.append(param + tf.zeros(times.shape + param.shape, dtype=times.dtype))
  return res

# Lint as: python3
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Join a sequence of Ito Processes with specified correlations."""

import tensorflow.compat.v2 as tf

from tf_quant_finance.models import euler_sampling
from tf_quant_finance.models import generic_ito_process
from tf_quant_finance.models import ito_process
from tf_quant_finance.models import utils


class JoinedItoProcess(generic_ito_process.GenericItoProcess):
  """Join of Ito Processes with specified time dependent correlations.

  For a sequence of Ito processes `I_1, .., I_n` of dimensions `d_1,.., d_n`,
  the class initializes a process `I` of dimension `d_1 + .. + d_n` with
  marginal proceses `I_i` and a correlation function `Corr(t)`. That is, let the
  Ito Process `I_i` describe an SDE

  ```None
  dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dW^i_t
  ```

  where `a_i(t, x)` is a function taking values in `R^{d_i}`, `b_i(t, X_t)` is a
  function taking values in `d_i x d_i` matrices, `W_i` is a `d_i`-dimensional
  Brownian motion.

  Then `I` describes an SDE for the joint process `(X_1,..., X_n)` of dimension
  `d:= d_1 + ... + d_n`

  ```None
  dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dB^i_t,
  ```

  where `(B_1, ..., B_n) = chol(t) * (W_1, ..., W_n)` for a Cholesky
  decomposition `chol(t)` of the correlation matrix `Corr(t)` at time `t`.
  Here `(W_1, ..., W_n)` is `d`-dimensional vector and `Corr(t)` is a `d x d`
  correlation matrix.

  `Corr(t)` is represented as a block-diagonal formed of a list of matrices
  `[m_1(t), m_2(t), ..., m_k(t)]` with `sum(rank(m_i)) = d`.

  # Example. # Black-scholes and Heston model join.
  ```python
  import numpy as np
  import tensorflow as tf
  import tf_quant_finance as tff

  dtype = tf.float64
  # Define Black scholes model with zero rate and volatility `0.1`
  sigma = 0.1
  def drift_fn(t , x):
    return -sigma**2 / 2
  def vol_fn(t , x):
    return sigma * tf.ones([1, 1], dtype=x.dtype)
  black_scholes_process = tff.models.GenericItoProcess(
      dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)
  # Define Heston Model
  epsilon = tff.math.piecewise.PiecewiseConstantFunc(
      jump_locations=[0.5], values=[0.01, 0.02], dtype=np.float64)
  heston_process = tff.models.heston_model.HestonModel(
      kappa=1.0, theta=0.04, epsilon=epsilon, rho=0.2, dtype=dtype)
  # Define join process where correlation between `black_scholes_process` and
  # log-asset price of the `heston_process` is 0.5.
  corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]
  # `corr_structure` corresponds to a `3 x 3` correlation matrix. Here Brownian
  # motion of `black_scholes_process` is correlated only with the 1st dimension
  # of `heston_process` but not with the second one.
  join_process = JoinedItoProcess(
      processes=[black_scholes_process, heston_process],
      corr_structure=corr_structure)
  # Sample 100,000 sample paths at times [0.1, 1.0] from the join process using
  # Sobol random sequence
  times = [0.1, 1.0]
  # Wrap sample_paths method with a tf.function
  sample_paths_fn = tf.function(process.sample_paths)
  samples = sample_paths_fn(
      times=times, time_step=0.01, num_samples=100000,
      initial_state=np.array([0.0, 1.0, 0.04]),
      random_type=random.RandomType.SOBOL)
  # Estimated correlations.
  np.corrcoef(samples[:, -1, :], rowvar=False)
  # Expected result:
  # [[1.        , 0.49567078, 0.08128067],
  #  [0.49567078, 1.        , 0.16580689],
  #  [0.08128067, 0.16580689, 1.        ]]
  ```
  """

  def __init__(self, processes, corr_structure, dtype=None, name=None):
    """Initializes a JoinedItoProcess.

    Takes a list of `processes` which are instances of `tff.models.ItoProcess`
    and a list `corr_structure` of correlation matrices and creates an Ito
    process that joins `processes` using the correlation structure.
    `corr_structure` describes block-diagonal structure of correlations for
    the Brownian motions in `processes`. For example, if the dimension of the
    JoinedItoProcess is `3` and
    `corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]`, then the introduced
    correlation is
    `Corr(t) = [[1.0, 0.5, 0.0], [0.5, 1.0, 0.0], [0.0, 0.0, 1.0]]`,
    where `Corr(t)` is the same as in the `JoinedItoProcess` docstring.

    Args:
      processes: A sequence of instances of `tff.models.ItoProcess`. All
        processes should have the same `dtype.`
      corr_structure: A list of correlation matrices. Each correlation matrix
        is either a `Tensor` of the same `dtype` as the `processes` and
        square shape  (i.e., `[d_i, d_i]` for some `d_i`) or a callable. The
        callables should accept a scalar (stands for time `t`) and return a
        square `Tensor`. The total dimension
        `sum([m.shape[-1] for m in corr_structure]` of correlation
        structure should be the same as the dimension of the `JoinedItoProcess`
        `sum([p.dim() for p in processes])`.
      dtype: The default `dtype` of the `processes`.
        Default value: None which means that default dtypes inferred by
          TensorFlow are used.
      name: Python string. The name scope under which ops created by the methods
        of this class are nested.
        Default value: `None` which maps to the default name
          `join_ito_process`.

    Raises:
      ValueError:
        (a) If any of the `processes` is not an `ItoProcess`.
        (b) If `processes` do not have the same `dtype`.
    """
    self._name = name or "join_ito_process"
    with tf.name_scope(self._name):
      self._processes = []  # Input processes
      dim = 0  # Dimension of the process
      for process in processes:
        if not isinstance(process, ito_process.ItoProcess):
          raise ValueError(
              "All input process of JoinedItoProcess must be instances "
              "of the ItoProcess class.")
        self._processes.append(process)
        d = process.dim()
        dim += d  # Total dimension of the process
        if dtype is None:
          dtype = process.dtype()
        elif dtype != process.dtype():
          raise ValueError("All processes should have the same `dtype`")
      self._corr_structure = [
          corr if callable(corr) else tf.convert_to_tensor(
              corr, dtype=dtype, name="corr")
          for corr in corr_structure]
      self._dim = dim

      def _drift_fn(t, x):
        """Drift function of the JoinedItoProcess."""
        drifts = []
        i1 = 0
        i2 = 0
        for p in self._processes:
          dim = p.dim()
          i2 += dim
          position = x[..., i1:i2]
          drift = tf.convert_to_tensor(p.drift_fn()(t, position),
                                       dtype=dtype,
                                       name="drift")
          drift = tf.broadcast_to(drift, position.shape)
          drifts.append(drift)
          i1 += dim
        return tf.concat(drifts, -1)

      def _vol_fn(t, x):
        """Volatility function of the JoinedItoProcess."""
        vols = []
        i1 = 0
        i2 = 0
        for p in self._processes:
          dim = p.dim()
          i2 += dim
          position = x[..., i1:i2]
          vol = tf.convert_to_tensor(p.volatility_fn()(t, position),
                                     dtype=dtype,
                                     name="volatility")
          vol = tf.broadcast_to(vol, position.shape + [dim])
          vols.append(vol)
          i1 += dim
        # Convert block diagonal volatilities to a dense correlation matrix
        vol = utils.block_diagonal_to_dense(*vols)
        # Compute Cholesky decomposition of the correlation structure
        corr_structure = _get_parameters(
            tf.expand_dims(t, -1), *self._corr_structure)
        cholesky_decomp = [tf.linalg.cholesky(m) for m in corr_structure]
        cholesky_decomp = utils.block_diagonal_to_dense(*cholesky_decomp)
        return tf.linalg.matmul(vol, cholesky_decomp)
      # The following will initialize the Generic Ito Process that has
      # a sampling and PDE solving methods
      super().__init__(dim, _drift_fn, _vol_fn, dtype, name)

  def sample_paths(self,
                   times,
                   num_samples=1,
                   initial_state=None,
                   random_type=None,
                   seed=None,
                   time_step=None,
                   swap_memory=True,
                   skip=0,
                   name=None):
    """Returns a sample of paths from the process using Euler sampling.

    Args:
      times: Rank 1 `Tensor` of increasing positive real values. The times at
        which the path points are to be evaluated.
      num_samples: Positive scalar `int`. The number of paths to draw.
        Default value: 1.
      initial_state: `Tensor` of shape `[self._dim]`. The initial state of the
        process.
        Default value: None which maps to a zero initial state.
      random_type: Enum value of `RandomType`. The type of (quasi)-random number
        generator to use to generate the paths.
        Default value: None which maps to the standard pseudo-random numbers.
      seed: Seed for the random number generator. The seed is
        only relevant if `random_type` is one of
        `[STATELESS, PSEUDO, HALTON_RANDOMIZED, PSEUDO_ANTITHETIC,
          STATELESS_ANTITHETIC]`. For `PSEUDO`, `PSEUDO_ANTITHETIC` and
        `HALTON_RANDOMIZED` the seed should be an integer scalar `Tensor`. For
        `STATELESS` and  `STATELESS_ANTITHETIC `must be supplied as an integer
        `Tensor` of shape `[2]`.
        Default value: `None` which means no seed is set.
      time_step: Real scalar `Tensor`. The maximal distance between time points
        in grid in Euler scheme.
      swap_memory: A Python bool. Whether GPU-CPU memory swap is enabled for
        this op. See an equivalent flag in `tf.while_loop` documentation for
        more details. Useful when computing a gradient of the op since
        `tf.while_loop` is used to propagate stochastic process in time.
        Default value: True.
      skip: `int32` 0-d `Tensor`. The number of initial points of the Sobol or
        Halton sequence to skip. Used only when `random_type` is 'SOBOL',
        'HALTON', or 'HALTON_RANDOMIZED', otherwise ignored.
        Default value: `0`.
      name: Python string. The name to give this op.
        Default value: `None` which maps to `sample_paths` is used.

    Returns:
     A real `Tensor` of shape `[num_samples, k, n]` where `k` is the size of the
     `times`, and `n` is the dimension of the process.

    Raises:
      ValueError: If `time_step` is not supplied.
    """
    if time_step is None:
      raise ValueError("`time_step` has to be supplied for JoinedItoProcess "
                       "`sample_paths` method.")
    name = name or self._name + "sample_paths"
    with tf.name_scope(name):
      if initial_state is None:
        initial_state = tf.zeros(self._dim, dtype=self.dtype(),
                                 name="initial_state")
      else:
        if isinstance(initial_state, (tuple, list)):
          initial_state = [tf.convert_to_tensor(state, dtype=self.dtype(),
                                                name="initial_state")
                           for state in initial_state]
          initial_state = tf.stack(initial_state)
        else:
          initial_state = tf.convert_to_tensor(initial_state,
                                               dtype=self.dtype(),
                                               name="initial_state")
      samples = euler_sampling.sample(self.dim(),
                                      drift_fn=self.drift_fn(),
                                      volatility_fn=self.volatility_fn(),
                                      times=times,
                                      time_step=time_step,
                                      num_samples=num_samples,
                                      initial_state=initial_state,
                                      random_type=random_type,
                                      seed=seed,
                                      swap_memory=swap_memory,
                                      skip=skip,
                                      dtype=self.dtype())
      return samples


def _get_parameters(times, *params):
  """Gets parameter values at at specified `times`."""
  res = []
  for param in params:
    if callable(param):
      # Used only in drift and volatility computation.
      # Here `times` is of shape [1]
      t = tf.squeeze(times)
      # The result has to have shape [1] + param.shape
      param_value = tf.convert_to_tensor(param(t), dtype=times.dtype,
                                         name="param_value")
      res.append(tf.expand_dims(param_value, 0))
    else:
      res.append(param + tf.zeros(times.shape + param.shape, dtype=times.dtype))
  return res

# Lint as: python3 # Copyright 2020 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. Join a sequence of Ito Processes with specified correlations. Join of Ito Processes with specified time dependent correlations.    For a sequence of Ito processes `I_1, .., I_n` of dimensions `d_1,.., d_n`,   the class initializes a process `I` of dimension `d_1 + .. + d_n` with   marginal proceses `I_i` and a correlation function `Corr(t)`. That is, let the   Ito Process `I_i` describe an SDE    ```None   dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dW^i_t   ```    where `a_i(t, x)` is a function taking values in `R^{d_i}`, `b_i(t, X_t)` is a   function taking values in `d_i x d_i` matrices, `W_i` is a `d_i`-dimensional   Brownian motion.    Then `I` describes an SDE for the joint process `(X_1,..., X_n)` of dimension   `d:= d_1 + ... + d_n`    ```None   dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dB^i_t,   ```    where `(B_1, ..., B_n) = chol(t) * (W_1, ..., W_n)` for a Cholesky   decomposition `chol(t)` of the correlation matrix `Corr(t)` at time `t`.   Here `(W_1, ..., W_n)` is `d`-dimensional vector and `Corr(t)` is a `d x d`   correlation matrix.    `Corr(t)` is represented as a block-diagonal formed of a list of matrices   `[m_1(t), m_2(t), ..., m_k(t)]` with `sum(rank(m_i)) = d`.    # Example. # Black-scholes and Heston model join.   ```python   import numpy as np   import tensorflow as tf   import tf_quant_finance as tff    dtype = tf.float64   # Define Black scholes model with zero rate and volatility `0.1`   sigma = 0.1   def drift_fn(t , x):     return -sigma**2 / 2   def vol_fn(t , x):     return sigma * tf.ones([1, 1], dtype=x.dtype)   black_scholes_process = tff.models.GenericItoProcess(       dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)   # Define Heston Model   epsilon = tff.math.piecewise.PiecewiseConstantFunc(       jump_locations=[0.5], values=[0.01, 0.02], dtype=np.float64)   heston_process = tff.models.heston_model.HestonModel(       kappa=1.0, theta=0.04, epsilon=epsilon, rho=0.2, dtype=dtype)   # Define join process where correlation between `black_scholes_process` and   # log-asset price of the `heston_process` is 0.5.   corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]   # `corr_structure` corresponds to a `3 x 3` correlation matrix. Here Brownian   # motion of `black_scholes_process` is correlated only with the 1st dimension   # of `heston_process` but not with the second one.   join_process = JoinedItoProcess(       processes=[black_scholes_process, heston_process],       corr_structure=corr_structure)   # Sample 100,000 sample paths at times [0.1, 1.0] from the join process using   # Sobol random sequence   times = [0.1, 1.0]   # Wrap sample_paths method with a tf.function   sample_paths_fn = tf.function(process.sample_paths)   samples = sample_paths_fn(       times=times, time_step=0.01, num_samples=100000,       initial_state=np.array([0.0, 1.0, 0.04]),       random_type=random.RandomType.SOBOL)   # Estimated correlations.   np.corrcoef(samples[:, -1, :], rowvar=False)   # Expected result:   # [[1.        , 0.49567078, 0.08128067],   #  [0.49567078, 1.        , 0.16580689],   #  [0.08128067, 0.16580689, 1.        ]]   ``` Initializes a JoinedItoProcess.      Takes a list of `processes` which are instances of `tff.models.ItoProcess`     and a list `corr_structure` of correlation matrices and creates an Ito     process that joins `processes` using the correlation structure.     `corr_structure` describes block-diagonal structure of correlations for     the Brownian motions in `processes`. For example, if the dimension of the     JoinedItoProcess is `3` and     `corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]`, then the introduced     correlation is     `Corr(t) = [[1.0, 0.5, 0.0], [0.5, 1.0, 0.0], [0.0, 0.0, 1.0]]`,     where `Corr(t)` is the same as in the `JoinedItoProcess` docstring.      Args:       processes: A sequence of instances of `tff.models.ItoProcess`. All         processes should have the same `dtype.`       corr_structure: A list of correlation matrices. Each correlation matrix         is either a `Tensor` of the same `dtype` as the `processes` and         square shape  (i.e., `[d_i, d_i]` for some `d_i`) or a callable. The         callables should accept a scalar (stands for time `t`) and return a         square `Tensor`. The total dimension         `sum([m.shape[-1] for m in corr_structure]` of correlation         structure should be the same as the dimension of the `JoinedItoProcess`         `sum([p.dim() for p in processes])`.       dtype: The default `dtype` of the `processes`.         Default value: None which means that default dtypes inferred by           TensorFlow are used.       name: Python string. The name scope under which ops created by the methods         of this class are nested.         Default value: `None` which maps to the default name           `join_ito_process`.      Raises:       ValueError:         (a) If any of the `processes` is not an `ItoProcess`.         (b) If `processes` do not have the same `dtype`. # Input processes # Dimension of the process # Total dimension of the process Drift function of the JoinedItoProcess. Volatility function of the JoinedItoProcess. # Convert block diagonal volatilities to a dense correlation matrix # Compute Cholesky decomposition of the correlation structure # The following will initialize the Generic Ito Process that has # a sampling and PDE solving methods Returns a sample of paths from the process using Euler sampling.      Args:       times: Rank 1 `Tensor` of increasing positive real values. The times at         which the path points are to be evaluated.       num_samples: Positive scalar `int`. The number of paths to draw.         Default value: 1.       initial_state: `Tensor` of shape `[self._dim]`. The initial state of the         process.         Default value: None which maps to a zero initial state.       random_type: Enum value of `RandomType`. The type of (quasi)-random number         generator to use to generate the paths.         Default value: None which maps to the standard pseudo-random numbers.       seed: Seed for the random number generator. The seed is         only relevant if `random_type` is one of         `[STATELESS, PSEUDO, HALTON_RANDOMIZED, PSEUDO_ANTITHETIC,           STATELESS_ANTITHETIC]`. For `PSEUDO`, `PSEUDO_ANTITHETIC` and         `HALTON_RANDOMIZED` the seed should be an integer scalar `Tensor`. For         `STATELESS` and  `STATELESS_ANTITHETIC `must be supplied as an integer         `Tensor` of shape `[2]`.         Default value: `None` which means no seed is set.       time_step: Real scalar `Tensor`. The maximal distance between time points         in grid in Euler scheme.       swap_memory: A Python bool. Whether GPU-CPU memory swap is enabled for         this op. See an equivalent flag in `tf.while_loop` documentation for         more details. Useful when computing a gradient of the op since         `tf.while_loop` is used to propagate stochastic process in time.         Default value: True.       skip: `int32` 0-d `Tensor`. The number of initial points of the Sobol or         Halton sequence to skip. Used only when `random_type` is 'SOBOL',         'HALTON', or 'HALTON_RANDOMIZED', otherwise ignored.         Default value: `0`.       name: Python string. The name to give this op.         Default value: `None` which maps to `sample_paths` is used.      Returns:      A real `Tensor` of shape `[num_samples, k, n]` where `k` is the size of the      `times`, and `n` is the dimension of the process.      Raises:       ValueError: If `time_step` is not supplied. Gets parameter values at at specified `times`. # Used only in drift and volatility computation. # Here `times` is of shape [1] # The result has to have shape [1] + param.shape