def solve(self,
ivp: InitialValueProblem,
parallel_enabled: bool = True) -> Solution:
cp = ivp.constrained_problem
t = discretize_time_domain(ivp.t_interval, self._d_t)
y = np.empty((len(t) - 1, ) + cp.y_vertices_shape)
y_i = ivp.initial_condition.discrete_y_0(True)
if not cp.are_all_boundary_conditions_static:
init_boundary_constraints = cp.create_boundary_constraints(
True, t[0])
init_y_constraints = cp.create_y_vertex_constraints(
init_boundary_constraints[0])
apply_constraints_along_last_axis(init_y_constraints, y_i)
y_constraints_cache: YConstraintsCache = {}
boundary_constraints_cache: BoundaryConstraintsCache = {}
y_next = self._create_y_next_function(ivp, y_constraints_cache,
boundary_constraints_cache)
for i, t_i in enumerate(t[:-1]):
y[i] = y_i = y_next(t_i, y_i)
if not cp.are_all_boundary_conditions_static:
y_constraints_cache.clear()
boundary_constraints_cache.clear()
return Solution(ivp, t[1:], y, vertex_oriented=True, d_t=self._d_t)
def _create_discrete_y_0(self, vertex_oriented: bool) -> np.ndarray:
"""
Creates the discretized initial values of y evaluated at the vertices
or cell centers of the spatial mesh.
:param vertex_oriented: whether the initial conditions are to be
evaluated at the vertices or cell centers of the spatial mesh
:return: the discretized initial values
"""
diff_eq = self._cp.differential_equation
if not diff_eq.x_dimension:
y_0 = np.array(self._y_0_func(None))
if y_0.shape != self._cp.y_shape():
raise ValueError(
'expected initial condition function output shape to be '
f'{self._cp.y_shape()} but got {y_0.shape}')
return y_0
x = self._cp.mesh.all_index_coordinates(vertex_oriented, flatten=True)
y_0 = self._y_0_func(x)
if y_0.shape != (len(x), diff_eq.y_dimension):
raise ValueError(
'expected initial condition function output shape to be '
f'{(len(x), diff_eq.y_dimension)} but got {y_0.shape}')
y_0 = y_0.reshape(self._cp.y_shape(vertex_oriented))
if vertex_oriented:
apply_constraints_along_last_axis(
self._cp.static_y_vertex_constraints, y_0)
return y_0
def discrete_y(self,
vertex_oriented: Optional[bool] = None,
interpolation_method: str = 'linear') -> np.ndarray:
"""
Returns the discrete solution evaluated either at vertices or the cell
centers of the spatial mesh.
:param vertex_oriented: whether the solution returned should be
evaluated at the vertices or the cell centers of the spatial mesh;
only interpolation is supported, therefore, it is not possible to
evaluate the solution at the vertices based on a cell-oriented
solution
:param interpolation_method: the interpolation method to use
:return: the discrete solution
"""
if vertex_oriented is None:
vertex_oriented = self._vertex_oriented
cp = self._ivp.constrained_problem
if not cp.differential_equation.x_dimension \
or self._vertex_oriented == vertex_oriented:
return np.copy(self._discrete_y)
x = cp.mesh.all_index_coordinates(vertex_oriented)
discrete_y = self.y(x, interpolation_method)
if vertex_oriented:
apply_constraints_along_last_axis(cp.static_y_vertex_constraints,
discrete_y)
return discrete_y
def __init__(
self,
cp: ConstrainedProblem,
y_0: np.ndarray,
vertex_oriented: Optional[bool] = None,
interpolation_method: str = 'linear'):
"""
:param cp: the constrained problem to turn into an initial value
problem by providing the initial conditions for it
:param y_0: the array containing the initial values of y over a spatial
mesh (which may be 0 dimensional in case of an ODE)
:param vertex_oriented: whether the initial conditions are evaluated at
the vertices or cell centers of the spatial mesh; it the
constrained problem is an ODE, it can be None
:param interpolation_method: the interpolation method to use to
calculate values that do not exactly fall on points of the y_0
grid; if the constrained problem is based on an ODE, it can be None
"""
if cp.differential_equation.x_dimension and vertex_oriented is None:
raise ValueError('vertex orientation must be defined for PDEs')
if y_0.shape != cp.y_shape(vertex_oriented):
raise ValueError(
f'discrete initial value shape {y_0.shape} must match '
'constrained problem solution shape '
f'{cp.y_shape(vertex_oriented)}')
self._cp = cp
self._y_0 = np.copy(y_0)
self._vertex_oriented = vertex_oriented
self._interpolation_method = interpolation_method
if vertex_oriented:
apply_constraints_along_last_axis(
cp.static_y_vertex_constraints, self._y_0)
def integral(
self,
y: np.ndarray,
t: float,
d_t: float,
d_y_over_d_t: Callable[[float, np.ndarray], np.ndarray],
y_constraint_function: Callable[
[Optional[float]],
Optional[Union[Sequence[Constraint], np.ndarray]]
]) -> np.ndarray:
half_d_t = d_t / 2.
y_half_next_constraints = y_constraint_function(t + half_d_t)
y_next_constraints = y_constraint_function(t + d_t)
k1 = d_t * d_y_over_d_t(t, y)
k2 = d_t * d_y_over_d_t(
t + half_d_t,
apply_constraints_along_last_axis(
y_half_next_constraints,
y + k1 / 2.))
k3 = d_t * d_y_over_d_t(
t + half_d_t,
apply_constraints_along_last_axis(
y_half_next_constraints,
y + k2 / 2.))
k4 = d_t * d_y_over_d_t(
t + d_t,
apply_constraints_along_last_axis(
y_next_constraints,
y + k3))
return apply_constraints_along_last_axis(
y_next_constraints,
y + (k1 + 2. * k2 + 2. * k3 + k4) / 6.)
def test_apply_constraints_along_last_axis_with_one_dimensional_array():
constraints = [
create_4_by_1_test_constraint(), create_4_by_1_test_constraint()
]
array = np.zeros(1)
with pytest.raises(ValueError):
apply_constraints_along_last_axis(constraints, array)
def test_apply_constraints_along_last_axis_with_wrong_last_array_axis_size():
constraints = [
create_4_by_1_test_constraint(), create_4_by_1_test_constraint()
]
array = np.zeros((3, 3, 1))
with pytest.raises(ValueError):
apply_constraints_along_last_axis(constraints, array)
def anti_laplacian(self,
laplacian: np.ndarray,
mesh: Mesh,
y_constraints: Union[Sequence[Optional[Constraint]],
np.ndarray],
derivative_boundary_constraints: Optional[
np.ndarray] = None,
y_init: Optional[np.ndarray] = None) -> np.ndarray:
"""
Computes the inverse of the element-wise scalar Laplacian using the
Jacobi method.
:param laplacian: the right-hand side of the equation
:param mesh: the mesh representing the discretized spatial domain
:param y_constraints: a sequence of constraints on the values of the
solution containing a constraint for each element of y; each
constraint must constrain the boundary values of corresponding
element of y for the system to be solvable
:param derivative_boundary_constraints: an optional 2D array
(x dimension, y dimension) of boundary constraint pairs that
specify constraints on the first derivatives of the solution
:param y_init: an optional initial estimate of the solution; if it is
None, a random array is used
:return: the array representing the solution to Poisson's equation at
every point of the mesh
"""
self._verify_input_shape_matches_mesh(laplacian, mesh, 'Laplacian')
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
laplacian.shape[-1])
if y_init is None:
y = np.random.random(laplacian.shape)
else:
if y_init.shape != laplacian.shape:
raise ValueError
y = y_init
apply_constraints_along_last_axis(y_constraints, y)
diff = np.inf
while diff > self._tol:
y_old = y
y = self._next_anti_laplacian_estimate(
y_old, laplacian, mesh, derivative_boundary_constraints)
apply_constraints_along_last_axis(y_constraints, y)
diff = float(np.linalg.norm(y - y_old))
return y
def test_apply_constraints_along_last_axis():
constraints = [
create_4_by_1_test_constraint(), create_4_by_1_test_constraint()
]
array = np.zeros((1, 4, 2))
expected_array = [[
[1., 1.],
[0., 0.],
[3., 3.],
[0., 0.]
]]
apply_constraints_along_last_axis(constraints, array)
assert np.array_equal(array, expected_array)
def discrete_y_0(
self,
vertex_oriented: Optional[bool] = None) -> np.ndarray:
if vertex_oriented is None:
vertex_oriented = self._vertex_oriented
if not self._cp.differential_equation.x_dimension \
or vertex_oriented == self._vertex_oriented:
return np.copy(self._y_0)
y_0 = self.y_0(self._cp.mesh.all_index_coordinates(vertex_oriented))
if vertex_oriented:
apply_constraints_along_last_axis(
self._cp.static_y_vertex_constraints, y_0)
return y_0
def y_next_function(t: float, y: np.ndarray) -> np.ndarray:
y_next = self._integrator.integral(y, t, self._d_t,
d_y_over_d_t_function,
y_constraint_func)
if len(y_eq_indices):
y_constraint = y_constraint_func(t + self._d_t)
y_constraint = None if y_constraint is None \
else y_constraint[y_eq_indices]
y_rhs = symbol_mapper.map_concatenated(
FDMSymbolMapArg(t, y, d_y_constraint_func), Lhs.Y)
y_next[..., y_eq_indices] = \
apply_constraints_along_last_axis(y_constraint, y_rhs)
if len(y_laplacian_eq_indices):
y_constraint = y_constraint_func(t + self._d_t)
y_constraint = None if y_constraint is None \
else y_constraint[y_laplacian_eq_indices]
d_y_constraint = d_y_constraint_func(t + self._d_t)
d_y_constraint = None if d_y_constraint is None \
else d_y_constraint[:, y_laplacian_eq_indices]
y_laplacian_rhs = symbol_mapper.map_concatenated(
FDMSymbolMapArg(t, y, d_y_constraint_func),
Lhs.Y_LAPLACIAN)
y_next[..., y_laplacian_eq_indices] = \
self._differentiator.anti_laplacian(
y_laplacian_rhs,
cp.mesh,
y_constraint,
d_y_constraint)
return y_next
def integral(
self,
y: np.ndarray,
t: float,
d_t: float,
d_y_over_d_t: Callable[[float, np.ndarray], np.ndarray],
y_constraint_function: Callable[
[Optional[float]],
Optional[Union[Sequence[Constraint], np.ndarray]]
]) -> np.ndarray:
half_d_t = d_t / 2.
y_half_next_constraints = y_constraint_function(t + half_d_t)
y_next_constraints = y_constraint_function(t + d_t)
y_hat = apply_constraints_along_last_axis(
y_half_next_constraints,
y + half_d_t * d_y_over_d_t(t, y))
return apply_constraints_along_last_axis(
y_next_constraints,
y + d_t * d_y_over_d_t(t + half_d_t, y_hat))
def integral(
self,
y: np.ndarray,
t: float,
d_t: float,
d_y_over_d_t: Callable[[float, np.ndarray], np.ndarray],
y_constraint_function: Callable[
[Optional[float]],
Optional[Union[Sequence[Constraint], np.ndarray]]
]) -> np.ndarray:
t_next = t + d_t
y_next_constraints = y_constraint_function(t_next)
y_next_init = apply_constraints_along_last_axis(
y_next_constraints,
y + d_t * d_y_over_d_t(t, y))
def y_next_residual_function(y_next: np.ndarray) -> np.ndarray:
return y_next - apply_constraints_along_last_axis(
y_next_constraints,
y + d_t * d_y_over_d_t(t_next, y_next))
return self._solve(y_next_residual_function, y_next_init)
def test_cp_2d_pde():
diff_eq = WaveEquation(2)
mesh = Mesh([(2., 6.), (-3., 3.)], [.1, .2])
bcs = ((DirichletBoundaryCondition(
vectorize_bc_function(lambda x, t: (999., None)), is_static=True),
NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (100., -100.)),
is_static=True)),
(NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (-x[0], None)),
is_static=True),
DirichletBoundaryCondition(
vectorize_bc_function(lambda x, t: (x[0], x[1])),
is_static=True)))
cp = ConstrainedProblem(diff_eq, mesh, bcs)
assert cp.are_all_boundary_conditions_static
assert cp.are_there_boundary_conditions_on_y
y_vertices = np.full(cp.y_shape(True), 13.)
apply_constraints_along_last_axis(cp.static_y_vertex_constraints,
y_vertices)
assert np.all(y_vertices[0, :-1, 0] == 999.)
assert np.all(y_vertices[0, :-1, 1] == 13.)
assert np.all(y_vertices[-1, :-1, :] == 13.)
assert np.all(y_vertices[1:, 0, :] == 13.)
assert np.allclose(y_vertices[:, -1, 0],
np.linspace(2., 6., y_vertices.shape[0]))
assert np.all(y_vertices[:, -1, 1] == 3.)
y_vertices = np.zeros(cp.y_shape(True))
diff = ThreePointCentralDifferenceMethod()
d_y_boundary_constraints = cp.static_boundary_vertex_constraints[1]
d_y_0_over_d_x_0 = diff.gradient(y_vertices[..., :1], mesh, 0,
d_y_boundary_constraints[:, :1])
assert np.all(d_y_0_over_d_x_0[-1, :, :] == 100.)
assert np.all(d_y_0_over_d_x_0[:-1, :, :] == 0.)
d_y_0_over_d_x_1 = diff.gradient(y_vertices[..., :1], mesh, 1,
d_y_boundary_constraints[:, :1])
assert np.allclose(d_y_0_over_d_x_1[:, 0, 0],
np.linspace(-2., -6., y_vertices.shape[0]))
assert np.all(d_y_0_over_d_x_1[:, 1:, :] == 0.)
d_y_1_over_d_x_0 = diff.gradient(y_vertices[..., 1:], mesh, 0,
d_y_boundary_constraints[:, 1:])
assert np.all(d_y_1_over_d_x_0[-1, :, :] == -100.)
assert np.all(d_y_1_over_d_x_0[:-1, :, :] == 0.)
d_y_1_over_d_x_1 = diff.gradient(y_vertices[..., 1:], mesh, 1,
d_y_boundary_constraints[:, 1:])
assert np.all(d_y_1_over_d_x_1 == 0.)
y_boundary_cell_constraints = cp.static_boundary_cell_constraints[0]
assert np.all(y_boundary_cell_constraints[0, 0][0].mask == [True] *
cp.y_cells_shape[1])
assert np.all(y_boundary_cell_constraints[0, 0][0].values == 999.)
assert np.all(y_boundary_cell_constraints[0, 1][0].mask == [False] *
cp.y_cells_shape[1])
assert y_boundary_cell_constraints[0, 1][0].values.size == 0
assert np.all(y_boundary_cell_constraints[1, 0][1].mask == [True] *
cp.y_cells_shape[0])
assert np.allclose(y_boundary_cell_constraints[1, 0][1].values,
np.linspace(2.05, 5.95, cp.y_cells_shape[0]))
assert np.all(y_boundary_cell_constraints[1, 1][1].mask == [True] *
cp.y_cells_shape[0])
assert np.all(y_boundary_cell_constraints[1, 1][1].values == 3.)
assert y_boundary_cell_constraints[1, 0][0] is None
def y_next_residual_function(y_next: np.ndarray) -> np.ndarray:
return y_next - apply_constraints_along_last_axis(
y_next_constraints,
y +
self._a * d_t * d_y_over_d_t(t_next, y_next) +
self._b * forward_update)
def y_next_residual_function(y_next: np.ndarray) -> np.ndarray:
return y_next - apply_constraints_along_last_axis(
y_next_constraints,
y + d_t * d_y_over_d_t(t_next, y_next))