def test_explicit_midpoint_method_with_constraints():
midpoint = ExplicitMidpointMethod()
y_shape = 5, 5, 1
y_0 = np.full(y_shape, 2.)
t_0 = 0.
d_t = .5
def d_y_over_d_t(t, y): return 2. * y - 4. * t
value = np.full(y_shape[:-1] + (1,), np.nan)
value[0, :] = value[-1, :] = 1.
value[:, 0] = value[:, -1] = 2.
mask = ~np.isnan(value)
value = value[mask]
y_constraint = Constraint(value, mask)
expected_y_next = np.full(y_shape, 4.5)
y_constraint.apply(expected_y_next[..., :1])
actual_y_next = midpoint.integral(
y_0,
t_0,
d_t,
d_y_over_d_t,
lambda _: [y_constraint])
assert np.allclose(actual_y_next, expected_y_next)
def create_4_by_1_test_constraint() -> Constraint:
array = np.full((4, 1), np.nan)
array[0, 0] = 1.
array[2, 0] = 3.
mask = ~np.isnan(array)
value = array[mask]
return Constraint(value, mask)
def create_3_by_3_test_constraint() -> Constraint:
array = np.full((3, 3), np.nan)
array[0, 0] = 1.
array[1, 1] = 2.
array[2, 2] = 3.
mask = ~np.isnan(array)
value = array[mask]
return Constraint(value, mask)
def test_crank_nicolson_method_with_constraints():
crank_nicolson = CrankNicolsonMethod()
y_0 = np.full((5, 1), -20.)
t_0 = 1.
d_t = .5
def d_y_over_d_t(t, y): return 5. * y
value = np.full(y_0.shape[:-1] + (1,), np.nan)
value[0] = value[-1] = 13.
mask = ~np.isnan(value)
value = value[mask]
y_constraint = Constraint(value, mask)
expected_y_next = (y_0 + .5 * 5. * d_t * y_0) / (1. - .5 * 5. * d_t)
y_constraint.apply(expected_y_next[..., :1])
actual_y_next = crank_nicolson.integral(
y_0, t_0, d_t, d_y_over_d_t, lambda _: [y_constraint])
assert np.allclose(actual_y_next, expected_y_next)
def test_backward_euler_method_with_constraints():
euler = BackwardEulerMethod()
y_0 = np.ones((5, 1))
t_0 = 1.
d_t = .5
def d_y_over_d_t(t, y): return 5. * y + t ** 2
value = np.full(y_0.shape[:-1] + (1,), np.nan)
value[0] = value[-1] = 999.
mask = ~np.isnan(value)
value = value[mask]
y_constraint = Constraint(value, mask)
expected_y_next = (y_0 + d_t * (t_0 + d_t) ** 2) / (1. - 5 * d_t)
y_constraint.apply(expected_y_next[..., :1])
actual_y_next = euler.integral(
y_0, t_0, d_t, d_y_over_d_t, lambda _: [y_constraint])
assert np.allclose(actual_y_next, expected_y_next)
def test_rk4_with_constraints():
rk4 = RK4()
y_shape = 8, 1
y_0 = np.zeros(y_shape)
t_0 = 1.
d_t = 1.
def d_y_over_d_t(_, y): return 2 * y + 1
value = np.full(y_shape[:-1] + (1,), np.nan)
value[0] = value[-1] = 0.
mask = ~np.isnan(value)
value = value[mask]
y_constraint = Constraint(value, mask)
expected_y_next = np.full(y_shape, 3.)
y_constraint.apply(expected_y_next[..., :1])
actual_y_next = rk4.integral(
y_0, t_0, d_t, d_y_over_d_t, lambda _: [y_constraint])
assert np.allclose(actual_y_next, expected_y_next)
def test_forward_euler_method_with_constraints():
euler = ForwardEulerMethod()
y_0 = np.ones((10, 2))
t_0 = 1.
d_t = .5
def d_y_over_d_t(t, y): return 5. * y + t ** 2
y_constraint_0 = Constraint(np.zeros(10), np.ones((10, 1), dtype=bool))
y_constraints = [y_constraint_0, None]
expected_y_next = np.concatenate(
(np.full((10, 1), 0.), np.full((10, 1), 4.)), axis=-1)
actual_y_next = euler.integral(
y_0, t_0, d_t, d_y_over_d_t, lambda _: y_constraints)
assert np.allclose(actual_y_next, expected_y_next)
def _create_boundary_constraints_for_all_y(
self, has_condition: bool,
condition_function: VectorizedBoundaryConditionFunction,
boundary_index_coordinates: np.ndarray,
t: Optional[float]) -> Sequence[Optional[Constraint]]:
"""
Creates a sequence of boundary constraints representing the boundary
condition, defined by the condition function, evaluated on a single
boundary for each element of y.
:param has_condition: whether there is a boundary condition specified
:param condition_function: the boundary condition function
:param boundary_index_coordinates: the coordinates of all the boundary
points
:param t: the time value
:return: a sequence of boundary constraints
"""
x_dimension = self._diff_eq.x_dimension
y_dimension = self._diff_eq.y_dimension
if not has_condition:
return [None] * y_dimension
x = boundary_index_coordinates.reshape((-1, x_dimension))
boundary_values = condition_function(x, t)
if boundary_values.shape != (len(x), y_dimension):
raise ValueError(
'expected boundary condition function output shape to be '
f'{(len(x), y_dimension)} but got {boundary_values.shape}')
boundary = boundary_values.reshape(
boundary_index_coordinates.shape[:-1] + (y_dimension, ))
boundary_constraints = []
for i in range(y_dimension):
boundary_i = boundary[..., i:i + 1]
mask = ~np.isnan(boundary_i)
value = boundary_i[mask]
boundary_constraints.append(Constraint(value, mask))
return boundary_constraints
def create_y_vertex_constraints(
self, y_boundary_vertex_constraints: Optional[np.ndarray]
) -> Optional[np.ndarray]:
"""
Creates a 1D array of solution value constraints evaluated on all
vertices of the mesh.
:param y_boundary_vertex_constraints: a 2D array (x dimension,
y dimension) of boundary value constraint pairs
:return: a 1D array (y dimension) of solution constraints
"""
diff_eq = self._diff_eq
if not diff_eq.x_dimension or y_boundary_vertex_constraints is None:
return None
slicer: List[Union[int, slice]] = \
[slice(None)] * len(self._y_vertices_shape)
y_constraints = np.empty(diff_eq.y_dimension, dtype=object)
y_element = np.empty(self._y_vertices_shape[:-1] + (1, ))
for y_ind in range(diff_eq.y_dimension):
y_element.fill(np.nan)
for axis in range(diff_eq.x_dimension):
for bc_ind, bc in \
enumerate(y_boundary_vertex_constraints[axis, y_ind]):
if bc is None:
continue
slicer[axis] = slice(-1, None) if bc_ind else slice(0, 1)
bc.apply(y_element[tuple(slicer)])
slicer[axis] = slice(None)
mask = ~np.isnan(y_element)
value = y_element[mask]
y_constraints[y_ind] = Constraint(value, mask)
return y_constraints