code: def _simplified_douglas_step_2d(
      self, u, dx, dy, dxy, bx, by, t1, t2, theta
  ):
      # Simplified version of the step to test against: fixed number of
      # dimensions, np instead of tf, for loops, etc.
      dt = t2 - t1
  
      # u0 = (1 + A * dt) u + b * dt
      dx_contrib = (
          dx[0] * _np_shift(u, 1, -1)
          + dx[1] * u
          + dx[2] * _np_shift(u, 1, 1)
      )
      dy_contrib = (
          dy[0] * _np_shift(u, 0, -1)
          + dy[1] * u
          + dy[2] * _np_shift(u, 0, 1)
      )
      dxy_contrib = dxy * (
          _np_shift(_np_shift(u, 0, 1), 1, 1)
          - _np_shift(_np_shift(u, 0, 1), 1, -1)
          - _np_shift(_np_shift(u, 0, -1), 1, 1)
          + _np_shift(_np_shift(u, 0, -1), 1, -1)
      )
      u0 = u + (dx_contrib + dy_contrib + dxy_contrib) * dt
      u0 += (bx + by) * dt
  
      # u1 = (1 - theta * dt * A_y)^(-1) (u0 - theta * dt * A_y u)
      theta *= dt  # Theta is always multiplied by dt.
      rhs = u0 - theta * (
          dy[0] * _np_shift(u, 0, -1)
          + dy[1] * u
          + dy[2] * _np_shift(u, 0, 1)
      )
  
      u1 = np.zeros_like(u0)
      for i in range(u0.shape[1]):
          diags = np.array(
              [
                  -theta * dy[0, :, i],
                  1 - theta * dy[1, :, i],
                  -theta * dy[2, :, i],
              ]
          )
          u1[:, i] = self._np_tridiagonal_solve(diags, rhs[:, i])
  
      # u2 = (1 - theta * dt * A_x)^(-1) (u1 - theta * dt * A_x u)
      rhs = u1 - theta * (
          dx[0] * _np_shift(u, 1, -1)
          + dx[1] * u
          + dx[2] * _np_shift(u, 1, 1)
      )
      u2 = np.zeros_like(u0)
      for i in range(u0.shape[0]):
          diags = np.array(
              [
                  -theta * dx[0, i],
                  1 - theta * dx[1, i],
                  -theta * dx[2, i],
              ]
          )
          u2[i] = self._np_tridiagonal_solve(diags, rhs[i])
      return u2
  
file_path: tf-quant-finance/tf_quant_finance/math/pde/steppers/douglas_adi_scheme_test.py
imports: import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 3 for calling the DouglasAdiSchemeTest obj's underlying member method code with example usage: obj._simplified_douglas_step_2d(u, dx, dy, dxy, bx, by, t1, t2, theta) and returns: u2