code: def _simplified_douglas_step_3d(
      self, u, dx, dy, dz, dxy, dyz, dxz, bx, by, bz, 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 + dt * A) u + b * dt
      def tridiag_contrib(tridiag_term, dim):
          return (
              tridiag_term[0] * _np_shift(u, dim, -1)
              + tridiag_term[1] * u
              + tridiag_term[2] * _np_shift(u, dim, 1)
          ) * dt
  
      def mixed_term_contrib(mixed_term, dim1, dim2):
          return (
              mixed_term
              * (
                  _np_shift(_np_shift(u, dim1, 1), dim2, 1)
                  - _np_shift(_np_shift(u, dim1, 1), dim2, -1)
                  - _np_shift(_np_shift(u, dim1, -1), dim2, 1)
                  + _np_shift(_np_shift(u, dim1, -1), dim2, -1)
              )
              * dt
          )
  
      u0 = (
          u
          + tridiag_contrib(dx, 2)
          + tridiag_contrib(dy, 1)
          + tridiag_contrib(dz, 0)
          + mixed_term_contrib(dxy, 2, 1)
          + mixed_term_contrib(dyz, 1, 0)
          + mixed_term_contrib(dxz, 2, 0)
      )
      u0 += (bx + by + bz) * dt
  
      # u1 = (1 - theta * dt * A_z)^(-1) (u0 - theta * dt * A_z u)
      theta *= dt  # Theta is always multiplied by dt.
      rhs = u0 - theta * (
          dz[0] * _np_shift(u, 0, -1)
          + dz[1] * u
          + dz[2] * _np_shift(u, 0, 1)
      )
  
      u1 = np.zeros_like(u0)
      for y in range(u0.shape[1]):
          for x in range(u0.shape[2]):
              diags = np.array(
                  [
                      -theta * dz[0, :, y, x],
                      1 - theta * dz[1, :, y, x],
                      -theta * dz[2, :, y, x],
                  ]
              )
              u1[:, y, x] = self._np_tridiagonal_solve(
                  diags, rhs[:, y, x]
              )
  
      # u2 = (1 - theta * dt * A_y)^(-1) (u1 - theta * dt * A_y u)
      rhs = u1 - theta * (
          dy[0] * _np_shift(u, 1, -1)
          + dy[1] * u
          + dy[2] * _np_shift(u, 1, 1)
      )
      u2 = np.zeros_like(u0)
      for z in range(u0.shape[0]):
          for x in range(u0.shape[2]):
              diags = np.array(
                  [
                      -theta * dy[0, z, :, x],
                      1 - theta * dy[1, z, :, x],
                      -theta * dy[2, z, :, x],
                  ]
              )
              u2[z, :, x] = self._np_tridiagonal_solve(
                  diags, rhs[z, :, x]
              )
  
      # u3 = (1 - theta * dt * A_x)^(-1) (u2 - theta * dt * A_x u)
      rhs = u2 - theta * (
          dx[0] * _np_shift(u, 2, -1)
          + dx[1] * u
          + dx[2] * _np_shift(u, 2, 1)
      )
      u3 = np.zeros_like(u0)
      for z in range(u0.shape[0]):
          for y in range(u0.shape[1]):
              diags = np.array(
                  [
                      -theta * dx[0, z, y, :],
                      1 - theta * dx[1, z, y, :],
                      -theta * dx[2, z, y, :],
                  ]
              )
              u3[z, y, :] = self._np_tridiagonal_solve(
                  diags, rhs[z, y, :]
              )
  
      return u3
  
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 5 for calling the DouglasAdiSchemeTest obj's underlying member method code with example usage: obj._simplified_douglas_step_3d(u, dx, dy, dz, dxy, dyz, dxz, bx, by, bz, t1, t2, theta) and returns: u3, unknown, unknown