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 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_InForwardDirection(self):
"""Tests solving 2d diffusion equation in forward direction.
The equation is `u_{t} - Dx u_{xx} - Dy u_{yy} = 0`.
The initial condition is a gaussian centered at (0, 0) with variance sigma.
The variance along each dimension should evolve as `sigma + 2 Dx (t - t_0)`
and `sigma + 2 Dy (t - t_0)`.
"""
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.0
initial_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, initial_variance), _gaussian(xs, initial_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.solve_forward(
start_time=0.0,
end_time=final_t,
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 = initial_variance + 2 * diff_coeff_x * final_t
variance_y = initial_variance + 2 * diff_coeff_y * final_t
expected = np.outer(_gaussian(ys, variance_y),
_gaussian(xs, variance_x))
self._assertClose(expected, result)
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 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 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 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)