def test_fdm_operator_on_pde_with_t_and_x_dependent_rhs():
class TestDiffEq(DifferentialEquation):
def __init__(self):
super(TestDiffEq, self).__init__(2, 1)
@property
def symbolic_equation_system(self) -> SymbolicEquationSystem:
return SymbolicEquationSystem([
self.symbols.t / 100. *
(self.symbols.x[0] + self.symbols.x[1])**2
])
diff_eq = TestDiffEq()
mesh = Mesh([(-5., 5.), (0., 3.)], [2., 1.])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))))] * 2
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = ContinuousInitialCondition(cp, lambda x: np.zeros((len(x), 1)))
ivp = InitialValueProblem(cp, (0., 5.), ic)
op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .25)
solution = op.solve(ivp)
y = solution.discrete_y()
assert solution.vertex_oriented
assert solution.d_t == .25
assert y.shape == (20, 6, 4, 1)
def test_pidon_operator_on_spherical_pde():
set_random_seed(0)
diff_eq = DiffusionEquation(3)
mesh = Mesh(
[(1., 11.), (0., 2 * np.pi), (.25 * np.pi, .75 * np.pi)],
[2., np.pi / 5., np.pi / 4],
CoordinateSystem.SPHERICAL)
bcs = [
(DirichletBoundaryCondition(
lambda x, t: np.ones((len(x), 1)), is_static=True),
DirichletBoundaryCondition(
lambda x, t: np.full((len(x), 1), 1. / 11.), is_static=True)),
(NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 1)), is_static=True),
NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 1)), is_static=True)),
(NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 1)), is_static=True),
NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 1)), is_static=True))
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = ContinuousInitialCondition(cp, lambda x: 1. / x[:, :1])
t_interval = (0., .5)
ivp = InitialValueProblem(cp, t_interval, ic)
sampler = UniformRandomCollocationPointSampler()
pidon = PIDONOperator(sampler, .001, True)
training_loss_history, test_loss_history = pidon.train(
cp,
t_interval,
training_data_args=DataArgs(
y_0_functions=[ic.y_0],
n_domain_points=20,
n_boundary_points=10,
n_batches=1
),
model_args=ModelArgs(
latent_output_size=20,
branch_hidden_layer_sizes=[30, 30],
trunk_hidden_layer_sizes=[30, 30],
),
optimization_args=OptimizationArgs(
optimizer=optimizers.Adam(learning_rate=2e-5),
epochs=3,
verbose=False
)
)
assert len(training_loss_history) == 3
for i in range(2):
assert np.all(
training_loss_history[i + 1].weighted_total_loss.numpy() <
training_loss_history[i].weighted_total_loss.numpy())
solution = pidon.solve(ivp)
assert solution.d_t == .001
assert solution.discrete_y().shape == (500, 6, 11, 3, 1)
def test_solution_generate_plots_for_2d_pde_with_scalar_and_vector_fields():
diff_eq = ShallowWaterEquation(.5)
mesh = Mesh([(0., 5.), (0., 5.)], [1., 1.])
bcs = [(NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (.0, None, None)), is_static=True),
NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (.0, None, None)),
is_static=True))] * 2
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(
cp, [(np.array([2.5, 1.25]), np.array([[1., 0.], [0., 1.]]))] * 3,
[1., .0, .0])
ivp = InitialValueProblem(cp, (0., 20.), ic)
t_coordinates = np.array([10., 20.])
discrete_y = np.arange(216).reshape((2, 6, 6, 3))
solution = Solution(ivp, t_coordinates, discrete_y, vertex_oriented=True)
plots = list(solution.generate_plots())
try:
assert len(plots) == 4
assert isinstance(plots[0], QuiverPlot)
assert isinstance(plots[1], StreamPlot)
assert isinstance(plots[2], ContourPlot)
assert isinstance(plots[3], SurfacePlot)
finally:
for plot in plots:
plot.close()
def test_fdm_operator_on_pde_with_t_and_x_dependent_rhs():
class TestDiffEq(DifferentialEquation):
def __init__(self):
super(TestDiffEq, self).__init__(2, 1)
@property
def symbolic_equation_system(self) -> SymbolicEquationSystem:
return SymbolicEquationSystem([
self.symbols.t / 100. *
(self.symbols.x[0] + self.symbols.x[1]) ** 2
])
diff_eq = TestDiffEq()
mesh = Mesh([(-1., 1.), (0., 2.)], [2., 1.])
bcs = [
(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))))
] * 2
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = ContinuousInitialCondition(cp, lambda x: np.zeros((len(x), 1)))
t_interval = (0., 1.)
ivp = InitialValueProblem(cp, t_interval, ic)
sampler = UniformRandomCollocationPointSampler()
pidon = PIDONOperator(sampler, .05, True)
training_loss_history, test_loss_history = pidon.train(
cp,
t_interval,
training_data_args=DataArgs(
y_0_functions=[ic.y_0],
n_domain_points=20,
n_boundary_points=10,
n_batches=1
),
model_args=ModelArgs(
latent_output_size=20,
branch_hidden_layer_sizes=[30, 30],
trunk_hidden_layer_sizes=[30, 30],
),
optimization_args=OptimizationArgs(
optimizer=optimizers.Adam(learning_rate=2e-5),
epochs=3,
verbose=False
)
)
assert len(training_loss_history) == 3
for i in range(2):
assert np.all(
training_loss_history[i + 1].weighted_total_loss.numpy() <
training_loss_history[i].weighted_total_loss.numpy())
solution = pidon.solve(ivp)
assert solution.d_t == .05
assert solution.discrete_y().shape == (20, 2, 3, 1)
def test_neumann_boundary_condition():
bc = NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 2)), is_static=True)
assert bc.is_static
assert not bc.has_y_condition
assert bc.has_d_y_condition
assert np.allclose(
bc.d_y_condition(np.ones((7, 3)), 5.),
np.zeros((7, 2)))
with pytest.raises(RuntimeError):
bc.y_condition(np.ones((7, 1)), 5.)
def test_solution_pde_with_no_vertex_orientation_defined():
diff_eq = WaveEquation(1)
mesh = Mesh([(0., 2.)], [1.])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))))]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([1.]), np.array([[1.]]))] * 2)
ivp = InitialValueProblem(cp, (0., 2.), ic)
t_coordinates = np.array([1., 2.])
discrete_y = np.arange(12).reshape((2, 3, 2))
with pytest.raises(ValueError):
Solution(ivp, t_coordinates, discrete_y)
def test_pidon_operator_in_ar_mode_on_pde():
set_random_seed(0)
diff_eq = WaveEquation(1)
mesh = Mesh([(0., 1.)], (.2,))
bcs = [
(NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 2)), is_static=True),
NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 2)), is_static=True)),
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
t_interval = (0., 1.)
ic = BetaInitialCondition(cp, [(3.5, 3.5), (3.5, 3.5)])
ivp = InitialValueProblem(cp, t_interval, ic)
training_y_0_functions = [
BetaInitialCondition(cp, [(p, p), (p, p)]).y_0
for p in [2., 3., 4., 5.]
]
sampler = UniformRandomCollocationPointSampler()
pidon = PIDONOperator(sampler, .25, False, auto_regression_mode=True)
assert pidon.auto_regression_mode
pidon.train(
cp,
(0., .25),
training_data_args=DataArgs(
y_0_functions=training_y_0_functions,
n_domain_points=50,
n_boundary_points=20,
n_batches=2
),
model_args=ModelArgs(
latent_output_size=50,
branch_hidden_layer_sizes=[50, 50],
trunk_hidden_layer_sizes=[50, 50],
),
optimization_args=OptimizationArgs(
optimizer=optimizers.Adam(learning_rate=1e-4),
epochs=2,
ic_loss_weight=10.,
verbose=False
)
)
sol = pidon.solve(ivp)
assert np.allclose(sol.t_coordinates, [.25, .5, .75, 1.])
assert sol.discrete_y().shape == (4, 5, 2)
def test_beta_initial_condition_with_wrong_number_of_alpha_and_betas():
diff_eq = DiffusionEquation(1)
mesh = Mesh([(0., 1.)], [.1])
bcs = [(NeumannBoundaryCondition(lambda x: np.zeros((len(x), 1))), ) * 2]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
with pytest.raises(ValueError):
BetaInitialCondition(cp, [(1., 1.), (1., 1)])
def test_beta_initial_condition_with_more_than_1d_pde():
diff_eq = DiffusionEquation(2)
mesh = Mesh([(0., 1.), (0., 1.)], [.1, .1])
bcs = [(NeumannBoundaryCondition(lambda x: np.zeros((len(x), 1))), ) * 2
] * 2
cp = ConstrainedProblem(diff_eq, mesh, bcs)
with pytest.raises(ValueError):
BetaInitialCondition(cp, [(1., 1.), (1., 1)])
def test_fdm_operator_on_1d_pde():
diff_eq = BurgerEquation(1, 1000.)
mesh = Mesh([(0., 10.)], [.1])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)),
is_static=True),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)),
is_static=True))]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([2.5]), np.array([[1.]]))])
ivp = InitialValueProblem(cp, (0., 50.), ic)
op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .25)
solution = op.solve(ivp)
assert solution.vertex_oriented
assert solution.d_t == .25
assert solution.discrete_y().shape == (200, 101, 1)
assert solution.discrete_y(False).shape == (200, 100, 1)
def test_fdm_operator_on_3d_pde():
diff_eq = CahnHilliardEquation(3)
mesh = Mesh([(0., 5.), (0., 5.), (0., 10.)], [.5, 1., 2.])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2)),
is_static=True),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2)),
is_static=True))] * 3
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = DiscreteInitialCondition(
cp, .05 * np.random.uniform(-1., 1., cp.y_shape(True)), True)
ivp = InitialValueProblem(cp, (0., 5.), ic)
op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .05)
solution = op.solve(ivp)
assert solution.vertex_oriented
assert solution.d_t == .05
assert solution.discrete_y().shape == (100, 11, 6, 6, 2)
assert solution.discrete_y(False).shape == (100, 10, 5, 5, 2)
def test_fdm_operator_on_spherical_pde():
diff_eq = DiffusionEquation(3)
mesh = Mesh([(1., 11.), (0., 2. * np.pi), (.1 * np.pi, .9 * np.pi)],
[2., np.pi / 5., np.pi / 5], CoordinateSystem.SPHERICAL)
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)),
is_static=True),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)),
is_static=True))] * 3
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = ContinuousInitialCondition(cp, lambda x: 1. / x[:, :1])
ivp = InitialValueProblem(cp, (0., 5.), ic)
op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .1)
solution = op.solve(ivp)
assert solution.vertex_oriented
assert solution.d_t == .1
assert solution.discrete_y().shape == (50, 6, 11, 5, 1)
assert solution.discrete_y(False).shape == (50, 5, 10, 4, 1)
def test_fdm_operator_on_polar_pde():
diff_eq = ShallowWaterEquation(.5)
mesh = Mesh([(1., 11.), (0., 2 * np.pi)], [2., np.pi / 5.],
CoordinateSystem.POLAR)
bcs = [(NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (.0, None, None)), is_static=True),
NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (.0, None, None)),
is_static=True))] * 2
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(
cp, [(np.array([-6., 0.]), np.array([[.25, 0.], [0., .25]]))] * 3,
[1., .0, .0])
ivp = InitialValueProblem(cp, (0., 5.), ic)
op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .1)
solution = op.solve(ivp)
assert solution.vertex_oriented
assert solution.d_t == .1
assert solution.discrete_y().shape == (50, 6, 11, 3)
assert solution.discrete_y(False).shape == (50, 5, 10, 3)
def test_solution_pde():
diff_eq = WaveEquation(1)
mesh = Mesh([(0., 2.)], [1.])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))))]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([1.]), np.array([[1.]]))] * 2)
ivp = InitialValueProblem(cp, (0., 2.), ic)
t_coordinates = np.array([1., 2.])
discrete_y = np.arange(12).reshape((2, 3, 2))
solution = Solution(ivp, t_coordinates, discrete_y, vertex_oriented=True)
assert solution.initial_value_problem == ivp
assert np.array_equal(solution.t_coordinates, t_coordinates)
assert np.isclose(solution.d_t, 1.)
assert solution.vertex_oriented
x_coordinates = np.array([[.5], [1.]])
expected_y = [[[1., 2.], [2., 3.]], [[7., 8.], [8., 9.]]]
assert np.allclose(solution.y(x_coordinates), expected_y)
expected_cell_y = [[[1., 2.], [3., 4.]], [[7., 8.], [9., 10.]]]
assert np.allclose(solution.discrete_y(False), expected_cell_y)
assert np.allclose(solution.discrete_y(), discrete_y)
other_solutions = [
Solution(ivp,
np.linspace(.5, 2., 4),
np.arange(16).reshape((4, 2, 2)),
vertex_oriented=False)
]
expected_differences = [
[[3., 3.], [3., 3.], [3., 3.]],
[[5., 5.], [5., 5.], [5., 5.]],
]
diff = solution.diff(other_solutions)
assert np.allclose(diff.matching_time_points, [1., 2.])
assert np.allclose(diff.differences, expected_differences)
def test_solution_generate_plots_for_1d_pde_with_scalar_fields():
diff_eq = WaveEquation(1)
mesh = Mesh([(0., 2.)], [1.])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))))]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([1.]), np.array([[1.]]))] * 2)
ivp = InitialValueProblem(cp, (0., 2.), ic)
t_coordinates = np.array([1., 2.])
discrete_y = np.arange(12).reshape((2, 3, 2))
solution = Solution(ivp, t_coordinates, discrete_y, vertex_oriented=True)
plots = list(solution.generate_plots())
try:
assert len(plots) == 2
assert isinstance(plots[0], SpaceLinePlot)
assert isinstance(plots[1], SpaceLinePlot)
finally:
for plot in plots:
plot.close()
def test_beta_initial_condition():
diff_eq = WaveEquation(1)
mesh = Mesh([(0., 1.)], [.5])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2)),
is_static=True),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2)),
is_static=True))]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
initial_condition = BetaInitialCondition(cp, [(2., 3.), (3., 2.)])
x_coordinates = np.array([[.125], [.625]])
expected_y_0 = [[1.1484375, .1640625], [1.0546875, 1.7578125]]
actual_y_0 = initial_condition.y_0(x_coordinates)
assert np.allclose(actual_y_0, expected_y_0)
expected_vertex_discrete_y_0 = [[0., 0.], [1.5, 1.5], [0., 0.]]
actual_vertex_discrete_y_0 = initial_condition.discrete_y_0(True)
assert np.allclose(actual_vertex_discrete_y_0,
expected_vertex_discrete_y_0)
expected_cell_discrete_y_0 = [[1.6875, .5625], [.5625, 1.6875]]
actual_cell_discrete_y_0 = initial_condition.discrete_y_0(False)
assert np.allclose(actual_cell_discrete_y_0, expected_cell_discrete_y_0)
def test_parareal_operator_on_pde():
diff_eq = DiffusionEquation(2)
mesh = Mesh([(0., 5.), (0., 5.)], [1., 1.])
bcs = [(NeumannBoundaryCondition(lambda x, _: np.zeros((len(x), 1)),
is_static=True),
NeumannBoundaryCondition(lambda x, _: np.zeros((len(x), 1)),
is_static=True))] * 2
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(
cp, [(np.array([2.5, 2.5]), np.array([[1., 0.], [0., 1.]]))])
ivp = InitialValueProblem(cp, (0., 5.), ic)
f = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .05)
g = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .5)
p = PararealOperator(f, g, .005)
f_solution = f.solve(ivp)
p_solution = p.solve(ivp)
assert p_solution.vertex_oriented == f_solution.vertex_oriented
assert p_solution.d_t == f_solution.d_t
assert np.array_equal(p_solution.t_coordinates, f_solution.t_coordinates)
assert np.allclose(p_solution.discrete_y(), f_solution.discrete_y())
def test_ode_operator_on_pde():
diff_eq = DiffusionEquation(1, 1.5)
mesh = Mesh([(0., 10.)], [.1])
bcs = [
(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))),
DirichletBoundaryCondition(lambda x, t: np.zeros((len(x), 1)))),
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([5.]), np.array([[2.5]]))],
[20.])
ivp = InitialValueProblem(cp, (0., 10.), ic)
op = ODEOperator('RK23', 2.5e-3)
with pytest.raises(ValueError):
op.solve(ivp)
def test_solution_generate_plots_for_3d_pde_with_vector_field():
diff_eq = BurgerEquation(3)
mesh = Mesh([(0., 2.)] * 3, [1.] * 3)
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 2))))] * 3
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array(
[1., 1., 1.]), np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))] *
3)
ivp = InitialValueProblem(cp, (0., 2.), ic)
t_coordinates = np.array([1., 2.])
discrete_y = np.arange(162).reshape((2, 3, 3, 3, 3))
solution = Solution(ivp, t_coordinates, discrete_y, vertex_oriented=True)
plots = list(solution.generate_plots())
try:
assert len(plots) == 1
assert isinstance(plots[0], QuiverPlot)
finally:
for plot in plots:
plot.close()
def test_pidon_operator_in_ar_mode_training_with_dynamic_boundary_conditions():
diff_eq = DiffusionEquation(1, .25)
mesh = Mesh([(0., .5)], (.05,))
bcs = [
(NeumannBoundaryCondition(
lambda x, t: np.full((len(x), 1), t), is_static=False),
NeumannBoundaryCondition(
lambda x, t: np.zeros((len(x), 1)), is_static=True)),
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
t_interval = (0., .5)
y_0_functions = [lambda x: np.zeros((len(x), 1))]
pidon = PIDONOperator(UniformRandomCollocationPointSampler(), .001, True)
with pytest.raises(ValueError):
pidon.train(
cp,
t_interval,
training_data_args=DataArgs(
y_0_functions=y_0_functions,
n_domain_points=50,
n_boundary_points=20,
n_batches=2
),
model_args=ModelArgs(
latent_output_size=50,
branch_hidden_layer_sizes=[50, 50],
trunk_hidden_layer_sizes=[50, 50],
),
optimization_args=OptimizationArgs(
optimizer={'class_name': 'Adam'},
epochs=3,
ic_loss_weight=10.,
verbose=False
)
)
def test_fdm_operator_conserves_density_on_zero_flux_diffusion_equation():
diff_eq = DiffusionEquation(1, 5.)
mesh = Mesh([(0., 500.)], [.1])
bcs = [
(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)),
is_static=True),
NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1)),
is_static=True)),
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([250]), np.array([[250.]]))],
[1000.])
ivp = InitialValueProblem(cp, (0., 20.), ic)
y_0 = ic.discrete_y_0(True)
y_0_sum = np.sum(y_0)
fdm_op = FDMOperator(CrankNicolsonMethod(),
ThreePointCentralDifferenceMethod(), 1e-3)
solution = fdm_op.solve(ivp)
y = solution.discrete_y()
y_sums = np.sum(y, axis=tuple(range(1, y.ndim)))
assert np.allclose(y_sums, y_0_sum)
def test_cp_1d_pde():
diff_eq = DiffusionEquation(1)
mesh = Mesh([(0., 1.)], [.1])
bcs = [(NeumannBoundaryCondition(lambda x, t: np.zeros((1, 1)),
is_static=True), ) * 2]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
assert cp.are_all_boundary_conditions_static
assert not cp.are_there_boundary_conditions_on_y
assert cp.y_shape(True) == (11, 1)
assert cp.y_shape(False) == (10, 1)
assert cp.differential_equation == diff_eq
assert cp.mesh == mesh
assert np.array_equal(cp.boundary_conditions, bcs)
y_vertex_constraints = cp.static_y_vertex_constraints
assert y_vertex_constraints.shape == (1, )
assert np.all(y_vertex_constraints[0].mask == [False])
assert np.all(y_vertex_constraints[0].values == [])
vertex_boundary_constraints = cp.static_boundary_constraints(True)
y_vertex_boundary_constraints = vertex_boundary_constraints[0]
assert y_vertex_boundary_constraints.shape == (1, 1)
assert y_vertex_boundary_constraints[0, 0][0] is None
assert y_vertex_boundary_constraints[0, 0][1] is None
d_y_vertex_boundary_constraints = vertex_boundary_constraints[1]
assert d_y_vertex_boundary_constraints.shape == (1, 1)
assert np.all(d_y_vertex_boundary_constraints[0, 0][0].mask == [True])
assert np.all(d_y_vertex_boundary_constraints[0, 0][0].values == [0.])
assert np.all(d_y_vertex_boundary_constraints[0, 0][1].mask == [True])
assert np.all(d_y_vertex_boundary_constraints[0, 0][1].values == [0.])
cell_boundary_constraints = cp.static_boundary_constraints(False)
y_cell_boundary_constraints = cell_boundary_constraints[0]
assert y_cell_boundary_constraints.shape == (1, 1)
assert y_cell_boundary_constraints[0, 0][0] is None
assert y_cell_boundary_constraints[0, 0][1] is None
d_y_cell_boundary_constraints = cell_boundary_constraints[1]
assert d_y_cell_boundary_constraints.shape == (1, 1)
assert np.all(d_y_cell_boundary_constraints[0, 0][0].mask == [True])
assert np.all(d_y_cell_boundary_constraints[0, 0][0].values == [0.])
assert np.all(d_y_cell_boundary_constraints[0, 0][1].mask == [True])
assert np.all(d_y_cell_boundary_constraints[0, 0][1].values == [0.])
def test_fdm_operator_on_pde_with_dynamic_boundary_conditions():
diff_eq = DiffusionEquation(1, 1.5)
mesh = Mesh([(0., 10.)], [1.])
bcs = [
(NeumannBoundaryCondition(lambda x, t: np.zeros((len(x), 1))),
DirichletBoundaryCondition(lambda x, t: np.full((len(x), 1), t / 5.))
),
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
ic = GaussianInitialCondition(cp, [(np.array([5.]), np.array([[2.5]]))],
[20.])
ivp = InitialValueProblem(cp, (0., 10.), ic)
op = FDMOperator(RK4(), ThreePointCentralDifferenceMethod(), .5)
solution = op.solve(ivp)
y = solution.discrete_y()
assert solution.vertex_oriented
assert solution.d_t == .5
assert y.shape == (20, 11, 1)
assert solution.discrete_y(False).shape == (20, 10, 1)
assert np.isclose(y[0, -1, 0], .1)
assert np.isclose(y[-1, -1, 0], 2.)
def test_cp_3d_pde():
mesh = Mesh([(2., 6.), (-3., 3.), (10., 12.)], [.1, .2, .5])
assert mesh.shape(True) == (41, 31, 5)
assert mesh.shape(False) == (40, 30, 4)
diff_eq = WaveEquation(3)
cp = ConstrainedProblem(
diff_eq, mesh,
((DirichletBoundaryCondition(
vectorize_bc_function(lambda x, t: (999., None)), is_static=True),
NeumannBoundaryCondition(
vectorize_bc_function(lambda x, t: (None, None)),
is_static=True)),
(DirichletBoundaryCondition(
vectorize_bc_function(lambda x, t: (0., 0.)), is_static=True),
NeumannBoundaryCondition(lambda x, t: np.full((len(x), 2), t))),
(NeumannBoundaryCondition(lambda x, t: -x[:, :2] * x[:, 1:3],
is_static=True),
DirichletBoundaryCondition(
vectorize_bc_function(lambda x, t: (-999., None))))))
assert cp.y_shape(True) == (41, 31, 5, 2)
assert cp.y_shape(False) == (40, 30, 4, 2)
assert not cp.are_all_boundary_conditions_static
assert cp.are_there_boundary_conditions_on_y
assert cp.static_y_vertex_constraints.shape == (2, )
y = np.full(cp._y_vertices_shape, -1)
cp.static_y_vertex_constraints[0].apply(y[..., :1])
cp.static_y_vertex_constraints[1].apply(y[..., 1:])
assert np.all(y[0, 1:, :, 0] == 999.)
assert np.all(y[:, 0, :, 0] == 0.)
assert np.all(y[1:, 1:, :, 0] == -1.)
assert np.all(y[:, 0, :, 1] == 0.)
assert np.all(y[:, 1:, :, 1] == -1.)
vertex_boundary_constraints = cp.static_boundary_vertex_constraints
cell_boundary_constraints = cp.static_boundary_cell_constraints
for y_boundary_constraints in \
[vertex_boundary_constraints[0], cell_boundary_constraints[0]]:
assert y_boundary_constraints.shape == (3, 2)
assert y_boundary_constraints[0, 0][0] is not None
assert y_boundary_constraints[0, 1][0] is not None
assert y_boundary_constraints[0, 0][1] is None
assert y_boundary_constraints[0, 1][1] is None
assert y_boundary_constraints[1, 0][0] is not None
assert y_boundary_constraints[1, 1][0] is not None
assert y_boundary_constraints[1, 0][1] is None
assert y_boundary_constraints[1, 1][1] is None
assert y_boundary_constraints[2, 0][0] is None
assert y_boundary_constraints[2, 1][0] is None
assert y_boundary_constraints[2, 0][1] is None
assert y_boundary_constraints[2, 1][1] is None
for d_y_boundary_constraints in \
[vertex_boundary_constraints[1], cell_boundary_constraints[1]]:
assert d_y_boundary_constraints.shape == (3, 2)
assert d_y_boundary_constraints[0, 0][0] is None
assert d_y_boundary_constraints[0, 1][0] is None
assert d_y_boundary_constraints[0, 0][1] is not None
assert d_y_boundary_constraints[0, 1][1] is not None
assert d_y_boundary_constraints[1, 0][0] is None
assert d_y_boundary_constraints[1, 1][0] is None
assert d_y_boundary_constraints[1, 0][1] is None
assert d_y_boundary_constraints[1, 1][1] is None
assert d_y_boundary_constraints[2, 0][0] is not None
assert d_y_boundary_constraints[2, 1][0] is not None
assert d_y_boundary_constraints[2, 0][1] is None
assert d_y_boundary_constraints[2, 1][1] is None
new_vertex_boundary_constraints = cp.create_boundary_constraints(True, 1.)
new_y_boundary_constraints = new_vertex_boundary_constraints[0]
new_d_y_boundary_constraints = new_vertex_boundary_constraints[1]
assert new_y_boundary_constraints[2, 0][1] is not None
assert new_y_boundary_constraints[2, 1][1] is not None
assert new_d_y_boundary_constraints[1, 0][1] is not None
assert new_d_y_boundary_constraints[1, 1][1] is not None
d_y_boundary = np.full((41, 1, 5, 2), np.nan)
new_d_y_boundary_constraints[1, 0][1].apply(d_y_boundary[..., :1])
new_d_y_boundary_constraints[1, 1][1].apply(d_y_boundary[..., 1:])
assert np.all(d_y_boundary == 1.)
new_y_vertex_constraints = \
cp.create_y_vertex_constraints(new_y_boundary_constraints)
assert new_y_vertex_constraints.shape == (2, )
y = np.full(cp._y_vertices_shape, -1)
new_y_vertex_constraints[0].apply(y[..., :1])
new_y_vertex_constraints[1].apply(y[..., 1:])
assert np.all(y[0, 1:, :-1, 0] == 999.)
assert np.all(y[:, 0, :-1, 0] == 0.)
assert np.all(y[:, :, -1, 0] == -999.)
assert np.all(y[1:, 1:, :-1, 0] == -1.)
assert np.all(y[:, 0, :, 1] == 0.)
assert np.all(y[:, 1:, :, 1] == -1.)
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 test_pidon_operator_on_pde_with_dynamic_boundary_conditions():
set_random_seed(0)
diff_eq = DiffusionEquation(1, .25)
mesh = Mesh([(0., 1.)], (.1,))
bcs = [
(NeumannBoundaryCondition(lambda x, t: np.full((len(x), 1), t)),
NeumannBoundaryCondition(lambda x, t: np.full((len(x), 1), t))),
]
cp = ConstrainedProblem(diff_eq, mesh, bcs)
t_interval = (0., .5)
training_y_0_functions = [
BetaInitialCondition(cp, [(p, p)]).y_0 for p in [2., 3., 4., 5.]
]
test_y_0_functions = [
BetaInitialCondition(cp, [(p, p)]).y_0 for p in [2.5, 3.5, 4.5]
]
sampler = UniformRandomCollocationPointSampler()
pidon = PIDONOperator(sampler, .001, True)
training_loss_history, test_loss_history = pidon.train(
cp,
t_interval,
training_data_args=DataArgs(
y_0_functions=training_y_0_functions,
n_domain_points=50,
n_boundary_points=20,
n_batches=2
),
test_data_args=DataArgs(
y_0_functions=test_y_0_functions,
n_domain_points=25,
n_boundary_points=10,
n_batches=1
),
model_args=ModelArgs(
latent_output_size=50,
branch_hidden_layer_sizes=[50, 50],
trunk_hidden_layer_sizes=[50, 50],
),
optimization_args=OptimizationArgs(
optimizer={
'class_name': 'Adam',
'config': {
'learning_rate': 1e-4
}
},
epochs=3,
ic_loss_weight=10.,
verbose=False
)
)
assert len(training_loss_history) == 3
assert len(test_loss_history) == 3
for i in range(2):
assert training_loss_history[i + 1].weighted_total_loss.numpy() < \
training_loss_history[i].weighted_total_loss.numpy()
assert test_loss_history[i + 1].weighted_total_loss.numpy() < \
test_loss_history[i].weighted_total_loss.numpy()
ic = BetaInitialCondition(cp, [(3.5, 3.5)])
ivp = InitialValueProblem(cp, t_interval, ic)
solution = pidon.solve(ivp)
assert solution.d_t == .001
assert solution.discrete_y().shape == (500, 11, 1)