docstr: Tests solving 2d diffusion equation involving mixed terms.
  
  The equation is `u_{t} + D u_{xx} / 2 +  D u_{yy} / 2 + D u_{xy} = 0`.
  The final condition is a gaussian centered at (0, 0) with variance sigma.
  
  The equation can be rewritten as `u_{t} + D u_{zz} = 0`, where
  `z = (x + y) / sqrt(2)`.
  
  Thus variance should evolve as `sigma + 2D(t_final - t)` along z dimension
  and stay unchanged in the orthogonal dimension:
  `u(x, y, t) = gaussian((x + y)/sqrt(2), sigma) + 2D * (t_final - t)) *
  gaussian((x - y)/sqrt(2), sigma)`.
docstr_tok: ["Tests","solving","2d","diffusion","equation","involving","mixed","terms",".","The","equation","is","`","u_","{","t","}","+","D","u_","{","xx","}","\/","2","+","D","u_","{","yy","}","\/","2","+","D","u_","{","xy","}","=","0","`",".","The","final","condition","is","a","gaussian","centered","at","(","0",",","0",")","with","variance","sigma",".","The","equation","can","be","rewritten","as","`","u_","{","t","}","+","D","u_","{","zz","}","=","0","`",",","where","`","z","=","(","x","+","y",")","\/","sqrt","(","2",")","`",".","Thus","variance","should","evolve","as","`","sigma","+","2D","(","t_final","-","t",")","`","along","z","dimension","and","stay","unchanged","in","the","orthogonal","dimension",":","`","u","(","x",",","y",",","t",")","=","gaussian","(","(","x","+","y",")","\/sqrt","(","2",")",",","sigma",")","+","2D","*","(","t_final","-","t",")",")","*","gaussian","(","(","x","-","y",")","\/sqrt","(","2",")",",","sigma",")","`","."]
code: def _testDiffusionInDiagonalDirection(
      self, pack_second_order_coeff_fn, boundary_condition="dirichlet"
  ):
      """Tests solving 2d diffusion equation involving mixed terms.
  
      The equation is `u_{t} + D u_{xx} / 2 +  D u_{yy} / 2 + D u_{xy} = 0`.
      The final condition is a gaussian centered at (0, 0) with variance sigma.
  
      The equation can be rewritten as `u_{t} + D u_{zz} = 0`, where
      `z = (x + y) / sqrt(2)`.
  
      Thus variance should evolve as `sigma + 2D(t_final - t)` along z dimension
      and stay unchanged in the orthogonal dimension:
      `u(x, y, t) = gaussian((x + y)/sqrt(2), sigma) + 2D * (t_final - t)) *
      gaussian((x - y)/sqrt(2), sigma)`.
      """
      dtype = tf.float32
  
      grid = grids.uniform_grid(
          minimums=[-10, -20],
          maximums=[10, 20],
          sizes=[201, 301],
          dtype=dtype,
      )
      ys = self.evaluate(grid[0])
      xs = self.evaluate(grid[1])
  
      diff_coeff = 1  # D
      time_step = 0.1
      final_t = 3
      final_variance = 1
  
      def second_order_coeff_fn(t, location_grid):
          del t, location_grid
          return pack_second_order_coeff_fn(
              diff_coeff / 2, diff_coeff / 2, diff_coeff / 2
          )
  
      variance_along_diagonal = (
          final_variance + 2 * diff_coeff * final_t
      )
  
      def expected_fn(x, y):
          return _gaussian(
              (x + y) / _SQRT2, variance_along_diagonal
          ) * _gaussian((x - y) / _SQRT2, final_variance)
  
      expected = np.array([[expected_fn(x, y) for x in xs] for y in ys])
  
      final_values = tf.expand_dims(
          tf.constant(
              np.outer(
                  _gaussian(ys, final_variance),
                  _gaussian(xs, final_variance),
              ),
              dtype=dtype,
          ),
          axis=0,
      )
      if boundary_condition == "dirichlet":
          bound_cond = [
              (_zero_boundary, _zero_boundary),
              (_zero_boundary, _zero_boundary),
          ]
      elif boundary_condition == "mixed":
          bound_cond = [
              (_zero_boundary, None),
              (None, _zero_grad_boundary),
          ]
      elif boundary_condition == "default":
          bound_cond = [(None, None), (None, None)]
      else:
          raise ValueError(
              "`boundary_cond` should be either `dirichlet`, "
              "`mixed` or `default`."
          )
      step_fn = douglas_adi_step(theta=0.5)
      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=step_fn,
          boundary_conditions=bound_cond,
          second_order_coeff_fn=second_order_coeff_fn,
          dtype=grid[0].dtype,
      )
  
      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 9 for calling the MultidimParabolicEquationStepperTest obj's underlying member method code with example usage: obj._testDiffusionInDiagonalDirection(pack_second_order_coeff_fn, boundary_condition) and returns: pack_second_order_coeff_fn, unknown