docstr: Tests solving 2d drift equation.
  
  The equation is `u_{t} + vx u_{x} + vy u_{y} = 0`.
  The final condition is a gaussian centered at (0, 0) with variance sigma.
  The gaussian should drift with velocity `[vx, vy]`.
docstr_tok: ["Tests","solving","2d","drift","equation",".","The","equation","is","`","u_","{","t","}","+","vx","u_","{","x","}","+","vy","u_","{","y","}","=","0","`",".","The","final","condition","is","a","gaussian","centered","at","(","0",",","0",")","with","variance","sigma",".","The","gaussian","should","drift","with","velocity","`","[","vx",",","vy","]","`","."]
code: def testSimpleDrift(self):
      """Tests solving 2d drift equation.
  
      The equation is `u_{t} + vx u_{x} + vy u_{y} = 0`.
      The final condition is a gaussian centered at (0, 0) with variance sigma.
      The gaussian should drift with velocity `[vx, vy]`.
      """
      grid = grids.uniform_grid(
          minimums=[-10, -20],
          maximums=[10, 20],
          sizes=[201, 301],
          dtype=tf.float32,
      )
      ys = self.evaluate(grid[0])
      xs = self.evaluate(grid[1])
  
      time_step = 0.01
      final_t = 3
      variance = 1
      vx = 0.1
      vy = 0.3
  
      def first_order_coeff_fn(t, location_grid):
          del t, location_grid
          return [vy, vx]
  
      final_values = tf.expand_dims(
          tf.constant(
              np.outer(
                  _gaussian(ys, variance), _gaussian(xs, variance)
              ),
              dtype=tf.float32,
          ),
          axis=0,
      )
  
      bound_cond = [
          (_zero_boundary, _zero_boundary),
          (_zero_boundary, _zero_boundary),
      ]
  
      result = fd_solvers.solve_backward(
          start_time=final_t,
          end_time=0,
          coord_grid=grid,
          values_grid=final_values,
          time_step=time_step,
          one_step_fn=douglas_adi_step(theta=0.5),
          boundary_conditions=bound_cond,
          first_order_coeff_fn=first_order_coeff_fn,
          dtype=grid[0].dtype,
      )
  
      expected = np.outer(
          _gaussian(ys + vy * final_t, variance),
          _gaussian(xs + vx * final_t, variance),
      )
  
      self._assertClose(expected, result)
  
imports: from absl.testing import parameterized
  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 MultidimParabolicEquationStepperTest obj's underlying member method code with example usage: obj.testSimpleDrift() and returns: list