def testInnerFirstAndSecondOrderCoeff(self): """Tests handling both inner_first_order_coeff and inner_second_order_coeff. We saw previously that the solution of `u_{t} - u_{xx} - u_{yy} - 2u_{x} - 4u_{y} - 5u = 0` is `u = exp(-x-2y) v`, where `v` solves the diffusion equation. Substitute now `u = exp(-x-2y) v` without expanding the derivatives: `v_{t} - exp(x)[exp(-x)v]_{xx} - exp(2y)[exp(-2y)v]_{yy} - 2exp(x)[exp(-x)v]_{x} - 4exp(2y)[exp(-2y)v]_{y} - 5v = 0`. Solve this equation and expect the solution of the diffusion equation. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 251], dtype=tf.float32) ys, xs = grid final_t = 0.1 time_step = 0.002 def second_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [[-tf.exp(2 * y), None], [None, -tf.exp(x)]] def inner_second_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [[tf.exp(-2 * y), None], [None, tf.exp(-x)]] def first_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [-4 * tf.exp(2 * y), -2 * tf.exp(x)] def inner_first_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [tf.exp(-2 * y), tf.exp(-x)] def zeroth_order_coeff_fn(t, coord_grid): del t, coord_grid return -5 initial = _reference_2d_pde_initial_cond(xs, ys) expected = _reference_2d_pde_solution(xs, ys, final_t) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, first_order_coeff_fn=first_order_coeff_fn, zeroth_order_coeff_fn=zeroth_order_coeff_fn, inner_second_order_coeff_fn=inner_second_order_coeff_fn, inner_first_order_coeff_fn=inner_first_order_coeff_fn)[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testNoTimeDependence(self): """Test for the case where all terms (quadratic, linear, shift) are null.""" 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.1 final_t = 1 variance = 1 final_cond = np.outer(_gaussian(ys, variance), _gaussian(xs, variance)) final_values = tf.expand_dims(tf.constant(final_cond, dtype=tf.float32), axis=0) bound_cond = [(_zero_boundary, _zero_boundary), (_zero_boundary, _zero_boundary)] 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, dtype=grid[0].dtype) expected = final_cond # No time dependence. self._assertClose(expected, result)
def _testDiffusionInDiagonalDirection(self, pack_second_order_coeff_fn): """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) bound_cond = [(_zero_boundary, _zero_boundary), (_zero_boundary, _zero_boundary)] step_fn = douglas_adi_step(theta=0.5) result = fd_solvers.step_back( 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)
def testInnerFirstAndSecondOrderCoeff(self): """Tests handling both inner_first_order_coeff and inner_second_order_coeff. We saw previously that the solution of `u_{t} - u_{xx} - 2u_{x} - u = 0` is `u = exp(-x) v`, where v solves the diffusion equation. Substitute now `u = exp(-x) v` without expanding the derivatives: `v_{t} - exp(x)[exp(-x)v]_{xx} - 2exp(x)[exp(-x)v]_{x} - v = 0`. Solve this equation and expect the solution of the diffusion equation. """ grid = grids.uniform_grid(minimums=[0], maximums=[1], sizes=[501], dtype=tf.float32) xs = grid[0] final_t = 0.1 time_step = 0.001 def second_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [[-tf.exp(x)]] def inner_second_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [[tf.exp(-x)]] def first_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [-2 * tf.exp(x)] def inner_first_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [tf.exp(-x)] def zeroth_order_coeff_fn(t, coord_grid): del t, coord_grid return -1 initial = _reference_pde_initial_cond(xs) expected = _reference_pde_solution(xs, final_t) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, first_order_coeff_fn=first_order_coeff_fn, zeroth_order_coeff_fn=zeroth_order_coeff_fn, inner_second_order_coeff_fn=inner_second_order_coeff_fn, inner_first_order_coeff_fn=inner_first_order_coeff_fn)[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testHeatEquation_InForwardDirection(self): """Test solving heat equation with various time marching schemes. Tests solving heat equation with the boundary conditions `u(x, t=1) = e * sin(x)`, `u(-2 pi n - pi / 2, t) = -e^t`, and `u(2 pi n + pi / 2, t) = -e^t` with some integer `n` for `u(x, t=0)`. The exact solution is `u(x, t=0) = sin(x)`. All time marching schemes should yield reasonable results given small enough time steps. First-order accurate schemes (explicit, implicit, weighted with theta != 0.5) require smaller time step than second-order accurate ones (Crank-Nicolson, Extrapolation). """ final_time = 1.0 def initial_cond_fn(x): return tf.sin(x) def expected_result_fn(x): return np.exp(-final_time) * tf.sin(x) @dirichlet def lower_boundary_fn(t, x): del x return -tf.exp(-t) @dirichlet def upper_boundary_fn(t, x): del x return tf.exp(-t) grid = grids.uniform_grid(minimums=[-10.5 * math.pi], maximums=[10.5 * math.pi], sizes=[1000], dtype=np.float32) def second_order_coeff_fn(t, x): del t, x return [[-1]] final_values = initial_cond_fn(grid[0]) result = fd_solvers.solve_forward( start_time=0.0, end_time=final_time, coord_grid=grid, values_grid=final_values, time_step=0.01, boundary_conditions=[(lower_boundary_fn, upper_boundary_fn)], second_order_coeff_fn=second_order_coeff_fn)[0] actual = self.evaluate(result) expected = self.evaluate(expected_result_fn(grid[0])) self.assertLess(np.max(np.abs(actual - expected)), 1e-3)
def testAnisotropicDiffusion(self): """Tests solving 2d diffusion equation. The equation is `u_{t} + Dx u_{xx} + Dy u_{yy} = 0`. The final condition is a gaussian centered at (0, 0) with variance sigma. The variance along each dimension should evolve as `sigma + 2 Dx t` and `sigma + 2 Dy (t_final - t)`. """ 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]) diff_coeff_x = 0.4 # Dx diff_coeff_y = 0.25 # Dy time_step = 0.1 final_t = 1 final_variance = 1 def quadratic_coeff_fn(t, location_grid): del t, location_grid u_xx = diff_coeff_x u_yy = diff_coeff_y u_xy = None return [[u_yy, u_xy], [u_xy, u_xx]] final_values = tf.expand_dims( tf.constant( np.outer( _gaussian(ys, final_variance), _gaussian(xs, final_variance)), dtype=tf.float32), axis=0) bound_cond = [(_zero_boundary, _zero_boundary), (_zero_boundary, _zero_boundary)] step_fn = douglas_adi_step(theta=0.5) result = fd_solvers.step_back( 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=quadratic_coeff_fn, dtype=grid[0].dtype) variance_x = final_variance + 2 * diff_coeff_x * final_t variance_y = final_variance + 2 * diff_coeff_y * final_t expected = np.outer(_gaussian(ys, variance_y), _gaussian(xs, variance_x)) self._assertClose(expected, result)
def testCrankNicolsonOscillationDamping(self): """Tests the Crank-Nicolson oscillation damping. Oscillations arise in Crank-Nicolson scheme when the initial (or final) conditions have discontinuities. We use Heaviside step function as initial conditions. The exact solution of the heat equation with unbounded x is ```None u(x, t) = (1 + erf(x/2sqrt(t))/2 ``` We take large enough x_min, x_max to be able to use this as a reference solution. CrankNicolsonWithOscillationDamping produces much smaller error than the usual crank_nicolson_scheme. """ final_t = 1 x_min = -10 x_max = 10 dtype = np.float32 def final_cond_fn(x): return 0.0 if x < 0 else 1.0 def expected_result_fn(x): return 1 / 2 + tf.math.erf(x / (2 * tf.sqrt(dtype(final_t)))) / 2 @dirichlet def lower_boundary_fn(t, x): del t, x return 0.0 @dirichlet def upper_boundary_fn(t, x): del t, x return 1.0 grid = grids.uniform_grid(minimums=[x_min], maximums=[x_max], sizes=[1000], dtype=dtype) self._testHeatEquation( grid=grid, final_t=final_t, time_step=0.01, final_cond_fn=final_cond_fn, expected_result_fn=expected_result_fn, one_step_fn=crank_nicolson_with_oscillation_damping_step(), lower_boundary_fn=lower_boundary_fn, upper_boundary_fn=upper_boundary_fn, error_tolerance=1e-3)
def testReferenceEquation_WithTransformationYieldingMixedTerm(self): """Tests an equation with mixed terms against exact solution. Take the reference equation `v_{t} = v_{xx} + v_{yy}` and substitute `v(x, y, t) = u(x, 2y - x, t)`. This yields `u_{t} = u_{xx} + 5u_{zz} - 2u_{xz}`, where `z = 2y - x`. Having `u(x, z, t) = v(x, (x+z)/2, t)` where `v(x, y, t)` is the known solution of the reference equation, we derive the boundary conditions and the expected solution for `u(x, y, t)`. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 251], dtype=tf.float32) final_t = 0.1 time_step = 0.002 def second_order_coeff_fn(t, coord_grid): del t, coord_grid return [[-5, 1], [None, -1]] @dirichlet def boundary_lower_z(t, coord_grid): x = coord_grid[1] return _reference_pde_solution(x, t) * _reference_pde_solution( x / 2, t) @dirichlet def boundary_upper_z(t, coord_grid): x = coord_grid[1] return _reference_pde_solution(x, t) * _reference_pde_solution( (x + 1) / 2, t) z_mesh, x_mesh = tf.meshgrid(grid[0], grid[1], indexing='ij') initial = (_reference_pde_initial_cond(x_mesh) * _reference_pde_initial_cond((x_mesh + z_mesh) / 2)) expected = (_reference_pde_solution(x_mesh, final_t) * _reference_pde_solution((x_mesh + z_mesh) / 2, final_t)) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, boundary_conditions=[(boundary_lower_z, boundary_upper_z), (_zero_boundary, _zero_boundary)])[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
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.step_back( 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)
def testShiftTerm(self): """Simple test for the shift term. The equation is `u_{t} + a u = 0`, the solution is `u(x, y, t) = exp(-a(t - t_final)) u(x, y, t_final)` """ 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.1 final_t = 1 variance = 1 a = 2 def zeroth_order_coeff_fn(t, location_grid): del t, location_grid return a expected = ( np.outer(_gaussian(ys, variance), _gaussian(xs, variance)) * np.exp(a * final_t)) 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)] step_fn = douglas_adi_step(theta=0.5) result = fd_solvers.step_back( 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, zeroth_order_coeff_fn=zeroth_order_coeff_fn, dtype=grid[0].dtype) self._assertClose(expected, result)
def test_compare_monte_carlo_to_backward_pde(self): dtype = tf.float64 kappa = 0.3 theta = 0.05 epsilon = 0.02 rho = 0.1 maturity_time = 1.0 initial_log_spot = 3.0 initial_vol = 0.05 strike = 15 discounting = 0.5 heston = heston_model.HestonModel(kappa=kappa, theta=theta, epsilon=epsilon, rho=rho, dtype=dtype) initial_state = np.array([initial_log_spot, initial_vol]) samples = heston.sample_paths(times=[maturity_time], initial_state=initial_state, time_step=0.01, num_samples=1000, seed=42) log_spots = samples[..., 0] monte_carlo_price = ( tf.constant(np.exp(-discounting * maturity_time), dtype=dtype) * tf.math.reduce_mean(tf.nn.relu(tf.math.exp(log_spots) - strike))) s_min, s_max = 2, 4 v_min, v_max = 0.03, 0.07 grid_size_s, grid_size_v = 101, 101 time_step = 0.01 grid = grids.uniform_grid(minimums=[s_min, v_min], maximums=[s_max, v_max], sizes=[grid_size_s, grid_size_v], dtype=dtype) s_mesh, _ = tf.meshgrid(grid[0], grid[1], indexing="ij") final_value_grid = tf.nn.relu(tf.math.exp(s_mesh) - strike) value_grid = heston.fd_solver_backward( start_time=1.0, end_time=0.0, coord_grid=grid, values_grid=final_value_grid, time_step=time_step, discounting=lambda *args: discounting)[0] pde_price = value_grid[int(grid_size_s / 2), int(grid_size_v / 2)] self.assertAllClose(monte_carlo_price, pde_price, atol=0.1, rtol=0.1)
def testHeatEquationWithVariousSchemes(self, one_step_fn, time_step): """Test solving heat equation with various time marching schemes. Tests solving heat equation with the boundary conditions `u(x, t=1) = e * sin(x)`, `u(-2 pi n - pi / 2, t) = -e^t`, and `u(2 pi n + pi / 2, t) = -e^t` with some integer `n` for `u(x, t=0)`. The exact solution is `u(x, t=0) = sin(x)`. All time marching schemes should yield reasonable results given small enough time steps. First-order accurate schemes (explicit, implicit, weighted with theta != 0.5) require smaller time step than second-order accurate ones (Crank-Nicolson, Extrapolation). Args: one_step_fn: one_step_fn representing a time marching scheme to use. time_step: time step for given scheme. """ def final_cond_fn(x): return math.e * math.sin(x) def expected_result_fn(x): return tf.sin(x) @dirichlet def lower_boundary_fn(t, x): del x return -tf.exp(t) @dirichlet def upper_boundary_fn(t, x): del x return tf.exp(t) grid = grids.uniform_grid(minimums=[-10.5 * math.pi], maximums=[10.5 * math.pi], sizes=[1000], dtype=np.float32) self._testHeatEquation(grid=grid, final_t=1, time_step=time_step, final_cond_fn=final_cond_fn, expected_result_fn=expected_result_fn, one_step_fn=one_step_fn, lower_boundary_fn=lower_boundary_fn, upper_boundary_fn=upper_boundary_fn, error_tolerance=1e-3)
def testReference_WithExponentMultiplier(self): """Tests solving diffusion equation with an exponent multiplier. Take the heat equation `v_{t} - v_{xx} - v_{yy} = 0` and substitute `v = exp(x + 2y) u`. This yields `u_{t} - u_{xx} - u_{yy} - 2u_{x} - 4u_{y} - 5u = 0`. The test compares numerical solution of this equation to the exact one, which is the diffusion equation solution times `exp(-x-2y)`. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 301], dtype=tf.float32) ys, xs = grid final_t = 0.1 time_step = 0.002 def second_order_coeff_fn(t, coord_grid): del t, coord_grid return [[-1, None], [None, -1]] def first_order_coeff_fn(t, coord_grid): del t, coord_grid return [-4, -2] def zeroth_order_coeff_fn(t, coord_grid): del t, coord_grid return -5 exp = _dir_prod(tf.exp(-2 * ys), tf.exp(-xs)) initial = exp * _reference_2d_pde_initial_cond(xs, ys) expected = exp * _reference_2d_pde_solution(xs, ys, final_t) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, first_order_coeff_fn=first_order_coeff_fn, zeroth_order_coeff_fn=zeroth_order_coeff_fn)[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testReferenceEquation(self): """Tests the equation used as reference for a few further tests. We solve the heat equation `u_t = u_xx + u_yy` on x = [0...1], y = [0...1] with boundary conditions `u(x, y, t=0) = (1/2 - |x-1/2|)(1/2-|y-1/2|), and zero Dirichlet on all spatial boundaries. The exact solution of the diffusion equation with zero-Dirichlet rectangular boundaries is `u(x, y, t) = u(x, t) * u(y, t)`, `u(z, t) = sum_{n=1..inf} b_n sin(pi n z) exp(-n^2 pi^2 t)`, `b_n = 2 integral_{0..1} sin(pi n z) u(z, t=0) dz.` The initial conditions are taken so that the integral easily calculates, and the sum can be approximated by a few first terms (given large enough `t`). See the result in _reference_heat_equation_solution. Using this solution helps to simplify the tests, as we don't have to maintain complicated boundary conditions in each test or tweak the parameters to keep the "support" of the function far from boundaries. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 301], dtype=tf.float32) ys, xs = grid final_t = 0.1 time_step = 0.002 def second_order_coeff_fn(t, coord_grid): del t, coord_grid return [[-1, None], [None, -1]] initial = _reference_2d_pde_initial_cond(xs, ys) expected = _reference_2d_pde_solution(xs, ys, final_t) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn)[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testInnerSecondOrderCoeff(self): """Tests handling inner_second_order_coeff. As in previous test, take the diffusion equation `v_{t} - v_{xx} - v_{yy} = 0` and substitute `v = exp(x + 2y) u`, but this time keep exponent under the derivative: `u_{t} - exp(-x)[exp(x)u]_{xx} - exp(-2y)[exp(2y)u]_{yy} = 0`. Expect the same solution as in previous test. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 251], dtype=tf.float32) ys, xs = grid final_t = 0.1 time_step = 0.002 def second_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [[-tf.exp(-2 * y), None], [None, -tf.exp(-x)]] def inner_second_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [[tf.exp(2 * y), None], [None, tf.exp(x)]] exp = _dir_prod(tf.exp(-2 * ys), tf.exp(-xs)) initial = exp * _reference_2d_pde_initial_cond(xs, ys) expected = exp * _reference_2d_pde_solution(xs, ys, final_t) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, inner_second_order_coeff_fn=inner_second_order_coeff_fn)[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testInnerSecondOrderCoeff(self): """Tests handling inner_second_order_coeff. As in previous test, take the diffusion equation `v_{t} - v_{xx} = 0` and substitute `v = exp(x) u`, but this time keep exponent under the derivative: `u_{t} - exp(-x)[exp(x)u]_{xx} = 0`. Expect the same solution as in previous test. """ grid = grids.uniform_grid(minimums=[0], maximums=[1], sizes=[501], dtype=tf.float32) xs = grid[0] final_t = 0.1 time_step = 0.001 def second_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [[-tf.exp(-x)]] def inner_second_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [[tf.exp(x)]] initial = tf.exp(-xs) * _reference_pde_initial_cond(xs) expected = tf.exp(-xs) * _reference_pde_solution(xs, final_t) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, inner_second_order_coeff_fn=inner_second_order_coeff_fn)[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testHeatEquation_WithMixedBoundaryConditions(self): """Test for mixed boundary conditions. Tests solving heat equation with the following boundary conditions: `u(x, t=1) = e * sin(x)`, `u_x(0, t) = e^t`, and `u(2 pi n + pi/2, t) = e^t`, where `n` is some integer. The exact solution `u(x, t=0) = e^t sin(x)`. """ def final_cond_fn(x): return math.e * math.sin(x) def expected_result_fn(x): return tf.sin(x) @neumann def lower_boundary_fn(t, x): del x return -tf.exp(t) @dirichlet def upper_boundary_fn(t, x): del x return tf.exp(t) grid = grids.uniform_grid(minimums=[0], maximums=[10.5 * math.pi], sizes=[1000], dtype=np.float32) self._testHeatEquation(grid, final_t=1, time_step=0.01, final_cond_fn=final_cond_fn, expected_result_fn=expected_result_fn, one_step_fn=crank_nicolson_step, lower_boundary_fn=lower_boundary_fn, upper_boundary_fn=upper_boundary_fn, error_tolerance=1e-3)
def testHeatEquation_WithRobinBoundaryConditions(self): """Test for Robin boundary conditions. Tests solving heat equation with the following boundary conditions: `u(x, t=1) = e * sin(x)`, `u_x(0, t) + 2u(0, t) = e^t`, and `2u(x_max, t) + u_x(x_max, t) = 2*e^t`, where `x_max = 2 pi n + pi/2` with some integer `n`. The exact solution `u(x, t=0) = e^t sin(x)`. """ def final_cond_fn(x): return math.e * math.sin(x) def expected_result_fn(x): return tf.sin(x) def lower_boundary_fn(t, x): del x return 2, -1, tf.exp(t) def upper_boundary_fn(t, x): del x return 2, 1, 2 * tf.exp(t) grid = grids.uniform_grid(minimums=[0], maximums=[4.5 * math.pi], sizes=[1000], dtype=np.float64) self._testHeatEquation(grid, final_t=1, time_step=0.01, final_cond_fn=final_cond_fn, expected_result_fn=expected_result_fn, one_step_fn=crank_nicolson_step, lower_boundary_fn=lower_boundary_fn, upper_boundary_fn=upper_boundary_fn, error_tolerance=1e-2)
def test_solving_backward_pde_for_sde_with_const_coeffs(self): # Integration test for converting 2d SDE with constant coeffs to a # backward Kolmogorov PDE and solving it. # The SDE is: # dS_x = (dW_1 + dW_2) / sqrt(2) # dS_y = (dW_1 + dW_2) / sqrt(2) # It is of course trivial, but we'll solve it the hard way for the sake of # testing. # The Kolmogorov backwards PDE is: # u_{t} + D u_{xx} / 2 + D u_{yy} / 2 + D u_{xy} = 0 # The equation can be rewritten as `u_{t} + D u_{zz} = 0`, where # z = (x + y) / sqrt(2). # If the final condition is a gaussian centered at (0, 0) with variance # sigma, then the solution is: # `u(x, y, t) = gaussian((x + y)/sqrt(2), sigma + 2D(t_final - t)) * # gaussian((x - y)/sqrt(2), sigma)`. def vol_fn(t, grid): del t xs = grid[..., 1] vol_elem = tf.ones_like(xs) / np.sqrt( 2) # all 4 elements are equal. return tf.stack( (tf.stack((vol_elem, vol_elem), axis=-1), tf.stack((vol_elem, vol_elem), axis=-1)), axis=-1) drift_fn = lambda t, grid: tf.zeros(grid.shape) process = generic_ito_process.GenericItoProcess(dim=2, volatility_fn=vol_fn, drift_fn=drift_fn, dtype=tf.float32) 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]) diff_coeff = 1 time_step = 0.1 final_t = 3 final_variance = 1 variance_along_diagonal = final_variance + 2 * diff_coeff * final_t def expected_fn(x, y): return (_gaussian( (x + y) / np.sqrt(2), variance_along_diagonal) * _gaussian( (x - y) / np.sqrt(2), 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=tf.float32), axis=0) result = self.evaluate( process.fd_solver_backward(start_time=final_t, end_time=0, coord_grid=grid, values_grid=final_values, time_step=time_step, dtype=tf.float32)[0]) self.assertLess( np.max(np.abs(result - expected)) / np.max(expected), 0.01)
def testCompareExpandedAndNotExpandedPdes(self): """Tests comparing PDEs with expanded derivatives and without. Take equation `u_{t} - [x^2 u]_{xx} + [x u]_{x} = 0`. Expanding the derivatives yields `u_{t} - x^2 u_{xx} - 3x u_{x} - u = 0`. Solve both equations and expect the results to be equal. """ grid = grids.uniform_grid(minimums=[0], maximums=[1], sizes=[501], dtype=tf.float32) xs = grid[0] final_t = 0.1 time_step = 0.001 initial = _reference_pde_initial_cond(xs) # arbitrary def inner_second_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [[-tf.square(x)]] def inner_first_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [x] result_not_expanded = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, inner_second_order_coeff_fn=inner_second_order_coeff_fn, inner_first_order_coeff_fn=inner_first_order_coeff_fn)[0] def second_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [[-tf.square(x)]] def first_order_coeff_fn(t, coord_grid): del t x = coord_grid[0] return [-3 * x] def zeroth_order_coeff_fn(t, coord_grid): del t, coord_grid return -1 result_expanded = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, first_order_coeff_fn=first_order_coeff_fn, zeroth_order_coeff_fn=zeroth_order_coeff_fn)[0] self.assertAllClose(result_not_expanded, result_expanded, atol=1e-3, rtol=1e-3)
def testCompareExpandedAndNotExpandedPdes(self): """Tests comparing PDEs with expanded derivatives and without. The equation is `u_{t} + [x u]_{x} + [y^2 u]_{y} - [sin(x) u]_{xx} - [cos(y) u]_yy + [x^3 y^2 u]_{xy} = 0`. Solve the equation, expand the derivatives and solve the equation again. Expect the results to be equal. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 251], dtype=tf.float32) final_t = 0.1 time_step = 0.002 y, x = grid initial = _reference_2d_pde_initial_cond(x, y) # arbitrary def inner_second_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [[-tf.math.cos(y), x**3 * y**2 / 2], [None, -tf.math.sin(x)]] def inner_first_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [y**2, x] result_not_expanded = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, inner_second_order_coeff_fn=inner_second_order_coeff_fn, inner_first_order_coeff_fn=inner_first_order_coeff_fn)[0] def second_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [[-tf.math.cos(y), x**3 * y**2 / 2], [None, -tf.math.sin(x)]] def first_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return [ y**2 * (1 + 3 * x**2) + 2 * tf.math.sin(y), x * (1 + 2 * x**2 * y) - 2 * tf.math.cos(x) ] def zeroth_order_coeff_fn(t, coord_grid): del t y, x = tf.meshgrid(*coord_grid, indexing='ij') return 1 + 2 * y + tf.math.sin(x) + tf.math.cos(x) + 6 * x**2 * y result_expanded = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, first_order_coeff_fn=first_order_coeff_fn, zeroth_order_coeff_fn=zeroth_order_coeff_fn)[0] self.assertAllClose(result_not_expanded, result_expanded, atol=1e-3, rtol=1e-3)
def testInnerMixedSecondOrderCoeffs(self): """Tests handling coefficients under the mixed second derivative. Take the equation from the previous test, `u_{t} = u_{xx} + 5u_{zz} - 2u_{xz}` and substitute `u = exp(xz) w`, leaving the exponent under the derivatives: `w_{t} = exp(-xz) [exp(xz) u]_{xx} + 5 exp(-xz) [exp(xz) u]_{zz} - 2 exp(-xz) [exp(xz) u]_{xz}`. We now have a coefficient under the mixed derivative. Test that the solution is `w = exp(-xz) u`, where u is from the previous test. """ grid = grids.uniform_grid(minimums=[0, 0], maximums=[1, 1], sizes=[201, 251], dtype=tf.float32) final_t = 0.1 time_step = 0.002 def second_order_coeff_fn(t, coord_grid): del t, z, x = tf.meshgrid(*coord_grid, indexing='ij') exp = tf.math.exp(-z * x) return [[-5 * exp, exp], [None, -exp]] def inner_second_order_coeff_fn(t, coord_grid): del t, z, x = tf.meshgrid(*coord_grid, indexing='ij') exp = tf.math.exp(z * x) return [[exp, exp], [None, exp]] @dirichlet def boundary_lower_z(t, coord_grid): x = coord_grid[1] return _reference_pde_solution(x, t) * _reference_pde_solution( x / 2, t) @dirichlet def boundary_upper_z(t, coord_grid): x = coord_grid[1] return tf.exp(-x) * _reference_pde_solution( x, t) * _reference_pde_solution((x + 1) / 2, t) z, x = tf.meshgrid(*grid, indexing='ij') exp = tf.math.exp(-z * x) initial = exp * (_reference_pde_initial_cond(x) * _reference_pde_initial_cond((x + z) / 2)) expected = exp * (_reference_pde_solution(x, final_t) * _reference_pde_solution((x + z) / 2, final_t)) actual = fd_solvers.solve_forward( start_time=0, end_time=final_t, coord_grid=grid, values_grid=initial, time_step=time_step, second_order_coeff_fn=second_order_coeff_fn, inner_second_order_coeff_fn=inner_second_order_coeff_fn, boundary_conditions=[(boundary_lower_z, boundary_upper_z), (_zero_boundary, _zero_boundary)])[0] self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testAnisotropicDiffusion_WithDirichletBoundaries(self): """Tests solving 2d diffusion equation with Dirichlet boundary conditions. The equation is `u_{t} + u_{xx} + 2 u_{yy} = 0`. The final condition is `u(t=1, x, y) = e * sin(x/sqrt(2)) * cos(y / 2)`. The following function satisfies this PDE and final condition: `u(t, x, y) = exp(t) * sin(x / sqrt(2)) * cos(y / 2)`. We impose Dirichlet boundary conditions using this function: `u(t, x_min, y) = exp(t) * sin(x_min / sqrt(2)) * cos(y / 2)`, etc. The other tests below are similar, but with other types of boundary conditions. """ time_step = 0.01 final_t = 1 x_min = -20 x_max = 20 y_min = -10 y_max = 10 grid = grids.uniform_grid(minimums=[y_min, x_min], maximums=[y_max, x_max], sizes=[201, 301], dtype=tf.float32) ys = self.evaluate(grid[0]) xs = self.evaluate(grid[1]) def second_order_coeff_fn(t, location_grid): del t, location_grid return [[2, None], [None, 1]] @dirichlet def lower_bound_x(t, location_grid): del location_grid f = tf.exp(t) * np.sin(x_min / _SQRT2) * tf.sin(ys / 2) return tf.expand_dims(f, 0) @dirichlet def upper_bound_x(t, location_grid): del location_grid f = tf.exp(t) * np.sin(x_max / _SQRT2) * tf.sin(ys / 2) return tf.expand_dims(f, 0) @dirichlet def lower_bound_y(t, location_grid): del location_grid f = tf.exp(t) * tf.sin(xs / _SQRT2) * np.sin(y_min / 2) return tf.expand_dims(f, 0) @dirichlet def upper_bound_y(t, location_grid): del location_grid f = tf.exp(t) * tf.sin(xs / _SQRT2) * np.sin(y_max / 2) return tf.expand_dims(f, 0) expected = np.outer(np.sin(ys / 2), np.sin(xs / _SQRT2)) final_values = tf.expand_dims(tf.constant( np.outer(np.sin(ys / 2), np.sin(xs / _SQRT2)) * np.exp(final_t), dtype=tf.float32), axis=0) bound_cond = [(lower_bound_y, upper_bound_y), (lower_bound_x, upper_bound_x)] 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)
def testAnisotropicDiffusion_WithRobinBoundaries(self): """Tests solving 2d diffusion equation with Robin boundary conditions.""" time_step = 0.01 final_t = 1 x_min = -20 x_max = 20 y_min = -10 y_max = 10 grid = grids.uniform_grid(minimums=[y_min, x_min], maximums=[y_max, x_max], sizes=[201, 301], dtype=tf.float32) ys = self.evaluate(grid[0]) xs = self.evaluate(grid[1]) def second_order_coeff_fn(t, location_grid): del t, location_grid return [[2, None], [None, 1]] def lower_bound_x(t, location_grid): del location_grid f = tf.exp(t) * tf.sin(ys / 2) * (np.sin(x_min / _SQRT2) - np.cos(x_min / _SQRT2) / _SQRT2) return 1, 1, tf.expand_dims(f, 0) def upper_bound_x(t, location_grid): del location_grid f = tf.exp(t) * tf.sin(ys / 2) * ( np.sin(x_max / _SQRT2) + 2 * np.cos(x_max / _SQRT2) / _SQRT2) return 1, 2, tf.expand_dims(f, 0) def lower_bound_y(t, location_grid): del location_grid f = tf.exp(t) * tf.sin( xs / _SQRT2) * (np.sin(y_min / 2) - 3 * np.cos(y_min / 2) / 2) return 1, 3, tf.expand_dims(f, 0) def upper_bound_y(t, location_grid): del location_grid f = tf.exp(t) * tf.sin(xs / _SQRT2) * (2 * np.sin(y_max / 2) + 3 * np.cos(y_max / 2) / 2) return 2, 3, tf.expand_dims(f, 0) expected = np.outer(np.sin(ys / 2), np.sin(xs / _SQRT2)) final_values = tf.expand_dims(tf.constant( np.outer(np.sin(ys / 2), np.sin(xs / _SQRT2)) * np.exp(final_t), dtype=tf.float32), axis=0) bound_cond = [(lower_bound_y, upper_bound_y), (lower_bound_x, upper_bound_x)] 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)