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)
