code: def douglas_adi_scheme(theta):
      """Applies Douglas time marching scheme (see [1] and Eq. 3.1 in [2]).
  
      Time marching schemes solve the space-discretized equation
      `du_inner/dt = A(t) u_inner(t) + A_mixed u(t) + b(t)`,
      where `u`, `u_inner` and `b` are vectors and `A`, `A_mixed` are matrices.
      `u_inner` is `u` with all boundaries having Robin boundary conditions
      trimmed and `A_mixed` are contributions of mixed derivative terms.
      See more details in multidim_parabolic_equation_stepper.py.
  
      In Douglas scheme (as well as other ADI schemes), the matrix `A` is
      represented as sum `A = sum_i A_i`. `A_i` is the contribution of
      terms with partial derivatives w.r.t. dimension `i`. The shift term is split
      evenly between `A_i`. Similarly, inhomogeneous term is represented as sum
      `b = sum_i b_i`, where `b_i` comes from boundary conditions on boundary
      orthogonal to dimension `i`.
  
      Given the current values vector u(t1), the step is defined as follows
      (using the notation of Eq. 3.1 in [2]):
      `Y_0 = (1 + (A(t1) + A_mixed(t1)) dt) U_{n-1} + b(t1) dt`,
      `Y_j = Y_{j-1} + theta dt (A_j(t2) Y_j - A_j(t1) U_{n-1} + b_j(t2) - b_j(t1))`
      for each spatial dimension `j`, and
      `U_n = Y_{n_dims-1}`.
  
      Here the parameter `theta` is a non-negative number, `U_{n-1} = u(t1)`,
      `U_n = u(t2)`, and `dt = t2 - t1`.
  
      Note: Douglas scheme is only first-order accurate if mixed terms are
      present. More advanced schemes, such as Craig-Sneyd scheme, are needed to
      achieve the second-order accuracy.
  
      #### References:
      [1] Douglas Jr., Jim (1962), "Alternating direction methods for three space
        variables", Numerische Mathematik, 4 (1): 41-63
      [2] Tinne Haentjens, Karek J. in't Hout. ADI finite difference schemes for
        the Heston-Hull-White PDE. https://arxiv.org/abs/1111.4087
  
      Args:
        theta: Number between 0 and 1 (see the step definition above). `theta = 0`
          corresponds to fully-explicit scheme.
  
      Returns:
        A callable consumes the following arguments by keyword:
          1. inner_value_grid: Grid of solution values at the current time of
            the same `dtype` as `value_grid` and shape of
            `batch_shape` + `[d_1 - 2 + n_def_i , ..., d_n -2 + n_def_i]`
            where `d_i` is the number of space discretization points along dimension
            `i` and `n_def_i` is the number of default boundaries along that
            dimension. `n_def_i` takes values 0, 1, 2 (default boundary),
          2. t1: Time before the step.
          3. t2: Time after the step.
          4. equation_params_fn: A callable that takes a scalar `Tensor` argument
            representing time, and returns a tuple of two objects:
              * First object is a nested list `L` such that `L[i][i]` is a tuple of
              three `Tensor`s, main, upper, and lower diagonal of the tridiagonal
              matrix `A` in a direction `i`. Each element `L[i][j]` corresponds
              to the mixed terms and is either None (meaning there are no mixed
              terms present) or a tuple of `Tensor`s representing contributions of
              mixed terms in directions (i + 1, j + 1), (i + 1, j - 1),
              (i - 1, j + 1), and (i - 1, j - 1).
              * The second object is a tuple of inhomogeneous terms for each
              dimension.
            All of the `Tensor`s are of the same `dtype` as `inner_value_grid` and
            of the shape broadcastable with the shape of `inner_value_grid`.
          5. A callable that accepts a `Tensor` of shape `inner_value_grid` and
            appends boundaries according to the boundary conditions, i.e. transforms
            `u_inner` to `u`.
          6. n_dims: A Python integer, the spatial dimension of the PDE.
          7. has_default_lower_boundary: A Python list of booleans of length
            `n_dims`. List indices enumerate the dimensions with `True` values
            marking default lower boundary condition along corresponding dimensions,
            and  `False` values indicating Robin boundary conditions.
          8. has_default_upper_boundary: Similar to has_default_lower_boundary, but
            for upper boundaries.
        The callable returns a `Tensor` of the same shape and `dtype` a
        `values_grid` and represents an approximate solution `u(t2)`.
      """
  
      if theta < 0 or theta > 1:
          raise ValueError("Theta should be in the interval [0, 1].")
  
      def _marching_scheme(
          value_grid,
          t1,
          t2,
          equation_params_fn,
          append_boundaries_fn,
          n_dims,
          has_default_lower_boundary,
          has_default_upper_boundary,
      ):
          """Constructs the Douglas ADI time marching scheme."""
          current_grid = value_grid
          matrix_params_t1, inhomog_terms_t1 = equation_params_fn(t1)
          matrix_params_t2, inhomog_terms_t2 = equation_params_fn(t2)
  
          # Explicit substep: Y_0 = (1 + A(t1) dt) U_{n-1} + b(t1) dt,
          # where dt = t2 - t1
          value_grid_with_boundaries = append_boundaries_fn(value_grid)
          for i in range(n_dims - 1):
              for j in range(i + 1, n_dims):
                  mixed_term = matrix_params_t1[i][j]
                  if mixed_term is not None:
                      current_grid += _apply_mixed_term_explicitly(
                          value_grid_with_boundaries,
                          mixed_term,
                          t2 - t1,
                          i,
                          j,
                          has_default_lower_boundary,
                          has_default_upper_boundary,
                          n_dims,
                      )
  
          # These are A_i(t1) * U_{n-1} * dt; caching them because they appear again
          # later in the correction substeps.
          explicit_contributions = []
  
          for i in range(n_dims):
              superdiag, diag, subdiag = (
                  matrix_params_t1[i][i][d] for d in range(3)
              )
              contribution = _apply_tridiag_matrix_explicitly(
                  value_grid, superdiag, diag, subdiag, i, n_dims
              ) * (t2 - t1)
              explicit_contributions.append(contribution)
              current_grid += contribution
  
          for inhomog_term in inhomog_terms_t1:
              current_grid += inhomog_term * (t2 - t1)
  
          # Correction substeps. For each dimension i:
          # Y_i = (1 - theta * A_i(t2) * dt)^(-1) *
          #      (Y_{i-1} - theta * dt * A_i(t1) * U_{n-1} + dt * (b_i(t2) - b_i(t1)))
          if theta == 0:
              return current_grid
  
          for i in range(n_dims):
              inhomog_term_delta = (
                  inhomog_terms_t2[i] - inhomog_terms_t1[i]
              )
              superdiag, diag, subdiag = (
                  matrix_params_t2[i][i][d] for d in range(3)
              )
              current_grid = _apply_correction(
                  theta,
                  current_grid,
                  explicit_contributions[i],
                  superdiag,
                  diag,
                  subdiag,
                  inhomog_term_delta,
                  t1,
                  t2,
                  i,
                  n_dims,
              )
  
          return current_grid
  
      return _marching_scheme
  
file_path: tf-quant-finance/tf_quant_finance/math/pde/steppers/douglas_adi.py
imports: import numpy
  import tensorflow.compat.v2
  from tf_quant_finance.math.pde.steppers.multidim_parabolic_equation_stepper import multidim_parabolic_equation_step
label_desc: Use image node_id 2 for calling a global function with example usage: douglas_adi_scheme(theta) and returns: _marching_scheme, current_grid, current_grid