code: def multidim_parabolic_equation_step(
      time,
      next_time,
      coord_grid,
      value_grid,
      boundary_conditions,
      time_marching_scheme,
      second_order_coeff_fn=None,
      first_order_coeff_fn=None,
      zeroth_order_coeff_fn=None,
      inner_second_order_coeff_fn=None,
      inner_first_order_coeff_fn=None,
      dtype=None,
      name=None,
  ):
      """Performs one step in time to solve a multidimensional PDE.
  
      Typically one doesn't need to use this function directly, unless they have
      a custom time marching scheme. A simple stepper function for multidimensional
      PDEs can be found in `douglas_adi.py`.
  
      The PDE is of the form
  
      ```None
        dV/dt + Sum[a_ij d2(A_ij V)/dx_i dx_j, 1 <= i, j <=n] +
           Sum[b_i d(B_i V)/dx_i, 1 <= i <= n] + c V = 0.
      ```
      from time `t0` to time `t1`. The solver can go both forward and backward in
      time. Here `a_ij`, `A_ij`, `b_i`, `B_i` and `c` are coefficients that may
      depend on spatial variables `x` and time `t`.
  
      Here `V` is the unknown function, `V_{...}` denotes partial derivatives
      w.r.t. dimensions specified in curly brackets, `i` and `j` denote spatial
      dimensions, `r` is the spatial radius-vector.
  
      Args:
        time: Real scalar `Tensor`. The time before the step.
        next_time: Real scalar `Tensor`. The time after the step.
        coord_grid: List of `n` rank 1 real `Tensor`s. `n` is the dimension of the
          domain. The i-th `Tensor` has shape either `[d_i]` or `B + [d_i]` where
          `d_i` is the size of the grid along axis `i` and `B` is a batch shape. The
          coordinates of the grid points. Corresponds to the spatial grid `G` above.
        value_grid: Real `Tensor` containing the function values at time
          `time` which have to be evolved to time `next_time`. The shape of the
          `Tensor` must broadcast with `B + [d_1, d_2, ..., d_n]`. `B` is the batch
          dimensions (one or more), which allow multiple functions (with potentially
          different boundary/final conditions and PDE coefficients) to be evolved
          simultaneously.
        boundary_conditions: The boundary conditions. Only rectangular boundary
          conditions are supported. A list of tuples of size `n` (space dimension
          of the PDE). The elements of the Tuple can be either a Python Callable or
          `None` representing the boundary conditions at the minimum and maximum
          values of the spatial variable indexed by the position in the list. E.g.,
          for `n=2`, the length of `boundary_conditions` should be 2,
          `boundary_conditions[0][0]` describes the boundary `(y_min, x)`, and
          `boundary_conditions[1][0]`- the boundary `(y, x_min)`. `None` values mean
          that the second order terms for that dimension on the boundary are assumed
          to be zero, i.e., if `boundary_conditions[k][0]` is None,
          'dV/dt + Sum[a_ij d2(A_ij V)/dx_i dx_j, 1 <= i, j <=n, i!=k+1, j!=k+1] +
             Sum[b_i d(B_i V)/dx_i, 1 <= i <= n] + c V = 0.'
          For not `None` values, the boundary conditions are accepted in the form
          `alpha(t, x) V + beta(t, x) V_n = gamma(t, x)`, where `V_n` is the
          derivative with respect to the exterior normal to the boundary.
          Each callable receives the current time `t` and the `coord_grid` at the
          current time, and should return a tuple of `alpha`, `beta`, and `gamma`.
          Each can be a number, a zero-rank `Tensor` or a `Tensor` whose shape is
          the grid shape with the corresponding dimension removed.
          For example, for a two-dimensional grid of shape `(b, ny, nx)`, where `b`
          is the batch size, `boundary_conditions[0][i]` with `i = 0, 1` should
          return a tuple of either numbers, zero-rank tensors or tensors of shape
          `(b, nx)`. Similarly for `boundary_conditions[1][i]`, except the tensor
          shape should be `(b, ny)`. `alpha` and `beta` can also be `None` in case
          of Neumann and Dirichlet conditions, respectively.
          Default value: `None`. Unlike setting `None` to individual elements of
          `boundary_conditions`, setting the entire `boundary_conditions` object to
          `None` means Dirichlet conditions with zero value on all boundaries are
          applied.
        time_marching_scheme: A callable which represents the time marching scheme
          for solving the PDE equation. If `u(t)` is space-discretized vector of the
          solution of a PDE, a time marching scheme approximately solves the
          equation `du_inner/dt = A(t) u_inner(t) + A_mixed(t) u(t) + b(t)` for
          `u(t2)` given `u(t1)`, or vice versa if going backwards in time.
          Here `A` is a banded matrix containing contributions from the current and
          neighboring points in space, `A_mixed` are contributions of mixed terms,
          `b` is an arbitrary vector (inhomogeneous term), and `u_inner` is `u` with
          boundaries with Robin conditions trimmed.
          Multidimensional time marching schemes are usually based on the idea of
          ADI (alternating direction implicit) method: the time step is split into
          substeps, and in each substep only one dimension is treated "implicitly",
          while all the others are treated "explicitly". This way one has to solve
          only tridiagonal systems of equations, but not more complicated banded
          ones. A few examples of time marching schemes (Douglas, Craig-Sneyd, etc.)
          can be found in [1].
          The 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 `value_grid[..., 1:-1]`.
            2. t1: Lesser of the two times defining the step.
            3. t2: Greater of the two times defining the step.
            4. equation_params_fn: A callable that takes a scalar `Tensor` argument
              representing time and returns a tuple of two elements.
              The first one represents `A`. The length must be the number of
              dimensions (`n_dims`), and A[i] must have length `n_dims - i`.
              `A[i][0]` is a tridiagonal matrix representing influence of the
              neighboring points along the dimension `i`. It is a tuple of
              superdiagonal, diagonal, and subdiagonal parts of the tridiagonal
              matrix. The shape of these tensors must be same as of `value_grid`.
              superdiagonal[..., -1] and subdiagonal[..., 0] are ignored.
              `A[i][j]` with `i < j < n_dims` are tuples of four Tensors with same
              shape as `value_grid` representing the influence of four points placed
              diagonally from the given point in the plane of dimensions `i` and
              `j`. Denoting `k`, `l` the indices of a given grid point in the plane,
              the four Tensors represent contributions of points `(k+1, l+1)`,
              `(k+1, l-1)`, `(k-1, l+1)`, and `(k-1, l-1)`, in this order.
              The second element in the tuple is a list of contributions to `b(t)`
              associated with each dimension. E.g. if `b(t)` comes from boundary
              conditions, then it is split correspondingly. Each element in the list
              is a Tensor with the shape of `value_grid`.
              For example a 2D problem with `value_grid.shape = (b, ny, nx)`, where
              `b` is the batch size. The elements `Aij` are non-zero if `i = j` or
              `i` is a neighbor of `j` in the x-y plane. Depict these non-zero
              elements on the grid as follows:
              ```
              a_mm    a_y-   a_mp
              a_x-    a_0    a_x+
              a_pm   a_y+   a_pp
              ```
              The callable should return
              ```
              ([[(a_y-, a_0y, a_y+), (a_pp, a_pm, a_mp, a_pp)],
                [None, (a_x-, a_0x, a_x+)]],
              [b_y, b_x])
              ```
              where `a_0x + a_0y = a_0` (the splitting is arbitrary). Note that
              there is no need to repeat the non-diagonal term
              `(a_pp, a_pm, a_mp, a_pp)` for the second time: it's replaced with
              `None`.
              All the elements `a_...` may be different for each point in the grid,
              so they are `Tensors` of shape `(B, ny, nx)`. `b_y` and `b_x` are also
              `Tensors` of that shape.
            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 should return a `Tensor` of the same shape and `dtype` as
          `values_grid` that represents an approximate solution of the
          space-discretized PDE.
        second_order_coeff_fn: Callable returning the second order coefficient
          `a_{ij}(t, r)` evaluated at given time `t`.
          The callable accepts the following arguments:
            `t`: The time at which the coefficient should be evaluated.
            `locations_grid`: a `Tensor` representing a grid of locations `r` at
              which the coefficient should be evaluated.
          Returns an object `A` such that `A[i][j]` is defined and
          `A[i][j]=a_{ij}(r, t)`, where `0 <= i < n_dims` and `i <= j < n_dims`.
          For example, the object may be a list of lists or a rank 2 Tensor.
          Only the elements with `j >= i` will be used, and it is assumed that
          `a_{ji} = a_{ij}`, so `A[i][j] with `j < i` may return `None`.
          Each `A[i][j]` should be a Number, a `Tensor` broadcastable to the
          shape of the grid represented by `locations_grid`, or `None` if
          corresponding term is absent in the equation. Also, the callable itself
          may be None, meaning there are no second-order derivatives in the
          equation.
          For example, for `n_dims=2`, the callable may return either
          `[[a_yy, a_xy], [a_xy, a_xx]]` or `[[a_yy, a_xy], [None, a_xx]]`.
        first_order_coeff_fn: Callable returning the first order coefficients
          `b_{i}(t, r)` evaluated at given time `t`.
          The callable accepts the following arguments:
            `t`: The time at which the coefficient should be evaluated.
            `locations_grid`: a `Tensor` representing a grid of locations `r` at
              which the coefficient should be evaluated.
          Returns a list or an 1D `Tensor`, `i`-th element of which represents
          `b_{i}(t, r)`. Each element should be a Number, a `Tensor` broadcastable
           to the shape of of the grid represented by `locations_grid`, or None if
           corresponding term is absent in the equation. The callable itself may be
           None, meaning there are no first-order derivatives in the equation.
        zeroth_order_coeff_fn: Callable returning the zeroth order coefficient
          `c(t, r)` evaluated at given time `t`.
          The callable accepts the following arguments:
            `t`: The time at which the coefficient should be evaluated.
            `locations_grid`: a `Tensor` representing a grid of locations `r` at
              which the coefficient should be evaluated.
          Should return a Number or a `Tensor` broadcastable to the shape of
          the grid represented by `locations_grid`. May also return None or be None
          if the shift term is absent in the equation.
        inner_second_order_coeff_fn: Callable returning the coefficients under the
          second derivatives (i.e. `A_ij(t, x)` above) at given time `t`. The
          requirements are the same as for `second_order_coeff_fn`.
        inner_first_order_coeff_fn: Callable returning the coefficients under the
          first derivatives (i.e. `B_i(t, x)` above) at given time `t`. The
          requirements are the same as for `first_order_coeff_fn`.
        dtype: The dtype to use.
        name: The name to give to the ops.
          Default value: None which means `parabolic_equation_step` is used.
  
      Returns:
        A sequence of two `Tensor`s. The first one is a `Tensor` of the same
        `dtype` and `shape` as `coord_grid` and represents a new coordinate grid
        after one iteration. The second `Tensor` is of the same shape and `dtype`
        as`values_grid` and represents an approximate solution of the equation after
        one iteration.
  
      #### References:
      [1] Tinne Haentjens, Karek J. in't Hout. ADI finite difference schemes
      for the Heston-Hull-White PDE. https://arxiv.org/abs/1111.4087
      """
      with tf.compat.v1.name_scope(
          name,
          "multidim_parabolic_equation_step",
          values=[time, next_time, coord_grid, value_grid],
      ):
          time = tf.convert_to_tensor(time, dtype=dtype, name="time")
          next_time = tf.convert_to_tensor(
              next_time, dtype=dtype, name="next_time"
          )
          coord_grid = [
              tf.convert_to_tensor(
                  x, dtype=dtype, name="coord_grid_axis_{}".format(ind)
              )
              for ind, x in enumerate(coord_grid)
          ]
          # Try broadcasting batch_shapes of coord_grid
          coord_grid = list(
              utils.broadcast_common_batch_shape(*coord_grid)
          )
          value_grid = tf.convert_to_tensor(
              value_grid, dtype=dtype, name="value_grid"
          )
  
          n_dims = len(coord_grid)
  
          # Sanitize the coeff callables.
          second_order_coeff_fn = second_order_coeff_fn or (
              lambda *args: [[None] * n_dims] * n_dims
          )
          first_order_coeff_fn = first_order_coeff_fn or (
              lambda *args: [None] * n_dims
          )
          zeroth_order_coeff_fn = zeroth_order_coeff_fn or (
              lambda *args: None
          )
          inner_second_order_coeff_fn = inner_second_order_coeff_fn or (
              lambda *args: [[None] * n_dims] * n_dims
          )
          inner_first_order_coeff_fn = inner_first_order_coeff_fn or (
              lambda *args: [None] * n_dims
          )
  
          batch_rank = len(value_grid.shape.as_list()) - len(coord_grid)
  
          # Get information on default boundary conditions
          # For each dimension we verify if either of the upper of lower boundaries
          # is `default`. If it is, than there is no need to perform trimming for the
          # value grid.
          has_default_lower_boundary = []
          has_default_upper_boundary = []
          lower_trim_indices = []
          upper_trim_indices = []
          for d in range(n_dims):
              num_discretization_pts = utils.get_shape(value_grid)[
                  batch_rank + d
              ]
              if boundary_conditions[d][0] is None:
                  # lower_inds are used to build an inner grid on which boundary
                  # conditions are imposed. For the default BC, no need for the value grid
                  # truncation.
                  has_default_lower_boundary.append(True)
                  lower_trim_indices.append(0)
              else:
                  has_default_lower_boundary.append(False)
                  lower_trim_indices.append(1)
              if boundary_conditions[d][1] is None:
                  # For the default BC, no need for the the value grid truncation.
                  upper_trim_indices.append(num_discretization_pts - 1)
                  has_default_upper_boundary.append(True)
              else:
                  upper_trim_indices.append(num_discretization_pts - 2)
                  has_default_upper_boundary.append(False)
  
          def equation_params_fn(t):
              return _construct_discretized_equation_params(
                  coord_grid,
                  value_grid,
                  boundary_conditions,
                  has_default_lower_boundary,
                  has_default_upper_boundary,
                  lower_trim_indices,
                  upper_trim_indices,
                  second_order_coeff_fn,
                  first_order_coeff_fn,
                  zeroth_order_coeff_fn,
                  inner_second_order_coeff_fn,
                  inner_first_order_coeff_fn,
                  batch_rank,
                  t,
              )
  
          # Construct the inner grid
          inner_grid_in = _trim_boundaries(
              value_grid,
              batch_rank,
              lower_trim_indices=lower_trim_indices,
              upper_trim_indices=upper_trim_indices,
          )
  
          # Apply time marching scheme to the inner grid
          def _append_boundaries_fn(inner_value_grid):
              # Add boundaries to the inner grid. The result has the same
              # shape as value_grid
              value_grid_with_boundaries = _append_boundaries(
                  value_grid,
                  inner_value_grid,
                  coord_grid,
                  boundary_conditions,
                  has_default_lower_boundary,
                  has_default_upper_boundary,
                  lower_trim_indices,
                  upper_trim_indices,
                  batch_rank,
                  time,
              )
              return value_grid_with_boundaries
  
          inner_grid_out = time_marching_scheme(
              value_grid=inner_grid_in,
              t1=time,
              t2=next_time,
              equation_params_fn=equation_params_fn,
              append_boundaries_fn=_append_boundaries_fn,
              has_default_lower_boundary=has_default_lower_boundary,
              has_default_upper_boundary=has_default_upper_boundary,
              n_dims=n_dims,
          )
          # Apply boundary conditions to restore values on the original grid
          updated_value_grid = _append_boundaries(
              value_grid,
              inner_grid_out,
              coord_grid,
              boundary_conditions,
              has_default_lower_boundary,
              has_default_upper_boundary,
              lower_trim_indices,
              upper_trim_indices,
              batch_rank,
              next_time,
          )
  
          return coord_grid, updated_value_grid
  
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
label_desc: Use image node_id 1 for calling a global function with example usage: multidim_parabolic_equation_step(time, next_time, coord_grid, value_grid, boundary_conditions, time_marching_scheme, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, dtype, name) and returns: coord_grid,  updated_value_grid, _construct_discretized_equation_params, value_grid_with_boundaries