def test_laplace():
laplace = lambda u, x, y: diff(u, x, order=2) + diff(u, y, order=2)
bc = DirichletBVP2D(x_min=0,
x_min_val=lambda y: torch.sin(np.pi * y),
x_max=1,
x_max_val=lambda y: 0,
y_min=0,
y_min_val=lambda x: 0,
y_max=1,
y_max_val=lambda x: 0)
net = FCNN(n_input_units=2, n_hidden_units=32, n_hidden_layers=1)
solution_neural_net_laplace, _ = solve2D(
pde=laplace,
condition=bc,
xy_min=(0, 0),
xy_max=(1, 1),
net=net,
max_epochs=300,
train_generator=ExampleGenerator2D((32, 32), (0, 0), (1, 1),
method='equally-spaced-noisy',
xy_noise_std=(0.01, 0.01)),
batch_size=64)
solution_analytical_laplace = lambda x, y: np.sin(np.pi * y) * np.sinh(
np.pi * (1 - x)) / np.sinh(np.pi)
xs, ys = np.linspace(0, 1, 101), np.linspace(0, 1, 101)
xx, yy = np.meshgrid(xs, ys)
sol_net = solution_neural_net_laplace(xx, yy, as_type='np')
sol_ana = solution_analytical_laplace(xx, yy)
assert isclose(sol_net, sol_ana, atol=0.01).all()
print('Laplace test passed.')
def test_monitor():
laplace = lambda u, x, y: diff(u, x, order=2) + diff(u, y, order=2)
bc = DirichletBVP2D(x_min=0,
x_min_val=lambda y: torch.sin(np.pi * y),
x_max=1,
x_max_val=lambda y: 0,
y_min=0,
y_min_val=lambda x: 0,
y_max=1,
y_max_val=lambda x: 0)
net = FCNN(n_input_units=2, n_hidden_units=32, n_hidden_layers=1)
solution_neural_net_laplace, _ = solve2D(
pde=laplace,
condition=bc,
xy_min=(0, 0),
xy_max=(1, 1),
net=net,
max_epochs=3,
train_generator=ExampleGenerator2D((32, 32), (0, 0), (1, 1),
method='equally-spaced-noisy'),
batch_size=64,
monitor=Monitor2D(check_every=1, xy_min=(0, 0), xy_max=(1, 1)))
print('Monitor test passed.')
def laplacian_spherical(u, r, theta, phi):
"""a helper function that computes the Laplacian in spherical coordinates
"""
r_component = diff(u * r, r, order=2) / r
theta_component = diff(torch.sin(theta) * diff(u, theta),
theta) / (r**2 * torch.sin(theta))
phi_component = diff(u, phi, order=2) / (r**2 * torch.sin(theta)**2)
return r_component + theta_component + phi_component
def test_cylindrical_laplacian(u, x):
out = cylindrical_laplacian(u, *x)
rho, phi, z = x
assert torch.allclose(
out,
diff(rho * diff(u, rho), rho) / rho + diff(u, phi, order=2) / rho ** 2 + diff(u, z, order=2)
)
def test_robin_2d(ones):
x0, x1 = 0, 1
y0, y1 = 0, 1
fd_x = FCNNSplitInput(2, 1)
fd_y = FCNNSplitInput(2, 1)
fr_x = FCNNSplitInput(2, 1)
a_x, b_x = 1.0, 2.0
fr_y = FCNNSplitInput(2, 1)
a_y, b_y = 3.0, 4.0
comp_x1 = ConditionComponent(x1,
f_dirichlet=fd_x,
f_neumann=None,
coord_index=0)
comp_y0 = ConditionComponent(y0,
f_dirichlet=fd_y,
f_neumann=None,
coord_index=1)
comp_x0 = RobinConditionComponent(x0, a_x, b_x, fr_x, coord_index=0)
comp_y1 = RobinConditionComponent(y1, a_y, b_y, fr_y, coord_index=1)
condition = ComposedCondition([comp_x0, comp_x1, comp_y0, comp_y1])
net = FCNN(2, 1)
x = (x1 - EPS) * ones
y = linspace_without_endpoints(y0, y1, N,
requires_grad=True).reshape(-1, 1)
u = condition.enforce(net, x, y)
assert all_close(u, fd_x(x, y), rtol=1e-4)
x = linspace_without_endpoints(x0, x1, N,
requires_grad=True).reshape(-1, 1)
y = (y0 + EPS) * ones
u = condition.enforce(net, x, y)
assert all_close(u, fd_y(x, y), rtol=1e-4)
x = (x0 + EPS) * ones
y = linspace_without_endpoints(y0, y1, N,
requires_grad=True).reshape(-1, 1)
u = condition.enforce(net, x, y)
pred = a_x * u + b_x * diff(u, x)
true = fr_x(x, y)
assert all_close(pred, true, rtol=1e-3)
x = linspace_without_endpoints(x0, x1, N,
requires_grad=True).reshape(-1, 1)
y = (y1 - EPS) * ones
u = condition.enforce(net, x, y)
pred = a_y * u + b_y * diff(u, y)
true = fr_y(x, y)
assert all_close(pred, true, rtol=1e-3)
def solver():
return Solver1D(
ode_system=lambda u, t: [u + diff(u, t)],
conditions=[NoCondition()],
t_min=0.0,
t_max=1.0,
)
def test_set_criterion_callback_reset(solver):
coords = torch.rand(10, 1, requires_grad=True)
funcs = torch.exp(coords * 1.01)
residuals = diff(funcs, coords) - funcs
# w/o resetting
callback = SetCriterion(criterion='l1', reset=False)
callback(solver)
assert torch.allclose(solver.criterion(residuals, funcs, coords),
torch.abs(residuals).mean())
solver._set_criterion('l2')
callback(solver)
assert torch.allclose(solver.criterion(residuals, funcs, coords),
(residuals**2).mean())
# w/ resetting
callback = SetCriterion(criterion='l1', reset=True)
callback(solver)
assert torch.allclose(solver.criterion(residuals, funcs, coords),
torch.abs(residuals).mean())
solver._set_criterion('l2')
callback(solver)
assert torch.allclose(solver.criterion(residuals, funcs, coords),
torch.abs(residuals).mean())
def test_laplacian(n_dim):
xs = [torch.rand(N_SAMPLES, 1, requires_grad=True) for _ in range(n_dim)]
net = FCNN(n_dim, 1)
u = net(torch.cat(xs, dim=1))
ans = torch.cat([diff(u, x, order=2) for x in xs], dim=1).sum(dim=1,
keepdim=True)
output = laplacian(u, *xs)
assert (ans == output).all()
def test_set_criterion_callback_polymorphism(solver, criterion):
coords = torch.rand(10, 1, requires_grad=True)
funcs = torch.exp(coords * 1.01)
residuals = diff(funcs, coords) - funcs
callback = SetCriterion(criterion=criterion)
callback(solver)
assert torch.allclose(solver.criterion(residuals, funcs, coords),
torch.abs(residuals).mean())
def get_dn(self, net, *x):
projection = self._get_projection(*x)
p_tensor = torch.cat(projection, dim=1)
output_val = net(p_tensor)
normal_der = diff(output_val, projection[self.idx])
R = (self.f(*projection) - self.a * output_val -
self.b * normal_der) / 2
return R / self.a, R / self.b
def test_ode_system():
parametric_circle = lambda x1, x2, t: [diff(x1, t) - x2, diff(x2, t) + x1]
init_vals_pc = [IVP(t_0=0.0, x_0=0.0), IVP(t_0=0.0, x_0=1.0)]
solution_pc, _ = solve_system(
ode_system=parametric_circle,
conditions=init_vals_pc,
t_min=0.0,
t_max=2 * np.pi,
max_epochs=5000,
)
ts = np.linspace(0, 2 * np.pi, 100)
x1_net, x2_net = solution_pc(ts, as_type='np')
x1_ana, x2_ana = np.sin(ts), np.cos(ts)
assert isclose(x1_net, x1_ana, atol=0.1).all()
assert isclose(x2_net, x2_ana, atol=0.1).all()
print('solve_system basic test passed.')
def test_div(n_dim):
xs = [torch.rand(N_SAMPLES, 1, requires_grad=True) for _ in range(n_dim)]
nets = [FCNN(n_dim, 1) for _ in range(n_dim)]
us = [net(torch.cat(xs, dim=1)) for net in nets]
grads = [diff(u, x) for u, x in zip(us, xs)]
ans = torch.cat(grads, dim=1).sum(dim=1, keepdim=True)
output = div(*us, *xs)
assert (output == ans).all()
def test_monitor():
exponential = lambda x, t: diff(x, t) - x
init_val_ex = IVP(t_0=0.0, x_0=1.0)
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.0,
t_max=2.0,
max_epochs=3,
monitor=Monitor(t_min=0.0, t_max=2.0,
check_every=1))
print('Monitor test passed.')
def get_rfx(n_input, n_output, n_equation):
coords = [torch.rand((N, 1), requires_grad=True) for _ in range(n_input)]
coords_tensor = torch.cat(coords, dim=1)
funcs = [
torch.sigmoid(torch.sum(coords_tensor, dim=1, keepdim=True))
for _ in range(n_output)
]
residual = [
diff(funcs[0], coords[0]) + funcs[0] for _ in range(n_equation)
]
residual = torch.cat(residual, dim=1)
return residual, funcs, coords
def test_condition_component_3d_get_dn(w0, f, g, xyz, net_31, coord_index):
component = ConditionComponent3D(w0, f, g, coord_index=coord_index)
p = component._get_projection(*xyz)
D, N = component.get_dn(net_31, *xyz)
p_tensor = torch.cat(p, dim=1)
assert all_close(D, f(*p) - net_31(p_tensor))
assert all_close(N, g(*p) - diff(net_31(p_tensor), p[coord_index]))
component = ConditionComponent3D(w0, None, g, coord_index=coord_index)
p = component._get_projection(*xyz)
D, N = component.get_dn(net_31, *xyz)
p_tensor = torch.cat(p, dim=1)
assert all_close(D, 0.0)
assert all_close(N, g(*p) - diff(net_31(p_tensor), p[coord_index]))
component = ConditionComponent3D(w0, f, None, coord_index=coord_index)
p = component._get_projection(*xyz)
D, N = component.get_dn(net_31, *xyz)
p_tensor = torch.cat(p, dim=1)
assert all_close(D, f(*p) - net_31(p_tensor))
assert all_close(N, 0.0)
def test_heat():
k, L, T = 0.3, 2, 3
heat = lambda u, x, t: diff(u, t) - k * diff(u, x, order=2)
ibvp = IBVP1D(x_min=0,
x_min_val=lambda t: 0,
x_max=L,
x_max_val=lambda t: 0,
t_min=0,
t_min_val=lambda x: torch.sin(np.pi * x / L))
net = FCNN(n_input_units=2, n_hidden_units=32, n_hidden_layers=1)
def mse(u, x, y):
true_u = torch.sin(np.pi * y) * torch.sinh(np.pi *
(1 - x)) / np.sinh(np.pi)
return torch.mean((u - true_u)**2)
solution_neural_net_heat, _ = solve2D(pde=heat,
condition=ibvp,
xy_min=(0, 0),
xy_max=(L, T),
net=net,
max_epochs=300,
train_generator=ExampleGenerator2D(
(32, 32), (0, 0), (L, T),
method='equally-spaced-noisy'),
batch_size=64,
metrics={'mse': mse})
solution_analytical_heat = lambda x, t: np.sin(np.pi * x / L) * np.exp(
-k * np.pi**2 * t / L**2)
xs = np.linspace(0, L, 101)
ts = np.linspace(0, T, 101)
xx, tt = np.meshgrid(xs, ts)
make_animation(solution_neural_net_heat, xs, ts) # test animation
sol_ana = solution_analytical_heat(xx, tt)
sol_net = solution_neural_net_heat(xx, tt, as_type='np')
assert isclose(sol_net, sol_ana, atol=0.01).all()
print('Heat test passed.')
def test_ode_ivp():
oscillator = lambda x, t: diff(x, t, order=2) + x
init_val_ho = IVP(t_0=0.0, x_0=0.0, x_0_prime=1.0)
solution_ho, _ = solve(ode=oscillator,
condition=init_val_ho,
max_epochs=3000,
t_min=0.0,
t_max=2 * np.pi)
ts = np.linspace(0, 2 * np.pi, 100)
x_net = solution_ho(ts, as_type='np')
x_ana = np.sin(ts)
assert isclose(x_net, x_ana, atol=0.1).all()
print('IVP basic test passed.')
def test_ode_bvp():
oscillator = lambda x, t: diff(x, t, order=2) + x
bound_val_ho = DirichletBVP(t_0=0.0, x_0=0.0, t_1=1.5 * np.pi, x_1=-1.0)
solution_ho, _ = solve(ode=oscillator,
condition=bound_val_ho,
max_epochs=3000,
t_min=0.0,
t_max=1.5 * np.pi)
ts = np.linspace(0, 1.5 * np.pi, 100)
x_net = solution_ho(ts, as_type='np')
x_ana = np.sin(ts)
assert isclose(x_net, x_ana, atol=0.1).all()
print('BVP basic test passed.')
def test_composed_condition_3d_single_component(w0, f, g, xyz, net_31,
coord_index, ComponentClass,
ConditionClass, k, detach):
component = ComponentClass(w0, f, g, coord_index=coord_index)
condition = ConditionClass(components=[component], k=k, detach=detach)
xyz = list(xyz)
xyz[coord_index] = torch.ones_like(xyz[coord_index],
requires_grad=True) * (w0 + EPS)
xyz_tensor = torch.cat(xyz, dim=1)
u = condition.enforce(net_31, *xyz)
assert all_close(u, f(*xyz), atol=1e-4, rtol=1e-2)
assert all_close(diff(u, xyz[coord_index]), g(*xyz), atol=1e-4, rtol=1e-2)
def test_cylindrical_curl(U, x):
out_rho, out_phi, out_z = cylindrical_curl(*U, *x)
rho, phi, z = x
urho, uphi, uz = U
assert torch.allclose(out_rho, diff(uz, phi) / rho - diff(uphi, z))
assert torch.allclose(out_phi, diff(urho, z) - diff(uz, rho))
assert torch.allclose(out_z, (diff(rho * uphi, rho) - diff(urho, phi)) / rho)
def test_neumann_boundaries_2():
k, L, T = 0.3, 2, 3
heat = lambda u, x, t: diff(u, t) - k * diff(u, x, order=2)
solution_analytical_heat = lambda x, t: np.sin(np.pi * x / L) * np.exp(
-k * np.pi**2 * t / L**2)
# Neumann on the left Dirichlet on the right
ibvp = IBVP1D(
x_min=0,
x_min_prime=lambda t: np.pi / L * torch.exp(-k * np.pi**2 * t / L**2),
x_max=L,
x_max_val=lambda t: 0,
t_min=0,
t_min_val=lambda x: torch.sin(np.pi * x / L))
net = FCNN(n_input_units=2, n_hidden_units=32, n_hidden_layers=1)
solution_neural_net_heat, _ = solve2D(pde=heat,
condition=ibvp,
xy_min=(0, 0),
xy_max=(L, T),
net=net,
max_epochs=300,
train_generator=ExampleGenerator2D(
(32, 32), (0, 0), (L, T),
method='equally-spaced-noisy'),
batch_size=64)
xs = np.linspace(0, L, 101)
ts = np.linspace(0, T, 101)
xx, tt = np.meshgrid(xs, ts)
make_animation(solution_neural_net_heat, xs, ts) # test animation
sol_ana = solution_analytical_heat(xx, tt)
sol_net = solution_neural_net_heat(xx, tt, as_type='np')
assert isclose(sol_net, sol_ana, atol=0.1).all()
print('Neumann on the left Dirichlet on the right test passed.')
def test_lotka_volterra():
alpha, beta, delta, gamma = 1, 1, 1, 1
lotka_volterra = lambda x, y, t: [
diff(x, t) - (alpha * x - beta * x * y),
diff(y, t) - (delta * x * y - gamma * y)
]
init_vals_lv = [IVP(t_0=0.0, x_0=1.5), IVP(t_0=0.0, x_0=1.0)]
nets_lv = [
FCNN(n_hidden_units=32, n_hidden_layers=1, actv=SinActv),
FCNN(n_hidden_units=32, n_hidden_layers=1, actv=SinActv)
]
solution_lv, _ = solve_system(ode_system=lotka_volterra,
conditions=init_vals_lv,
t_min=0.0,
t_max=12,
nets=nets_lv,
max_epochs=12000,
monitor=Monitor(t_min=0.0,
t_max=12,
check_every=100))
ts = np.linspace(0, 12, 100)
prey_net, pred_net = solution_lv(ts, as_type='np')
def dPdt(P, t):
return [
P[0] * alpha - beta * P[0] * P[1],
delta * P[0] * P[1] - gamma * P[1]
]
P0 = [1.5, 1.0]
Ps = odeint(dPdt, P0, ts)
prey_num = Ps[:, 0]
pred_num = Ps[:, 1]
assert isclose(prey_net, prey_num, atol=0.1).all()
assert isclose(pred_net, pred_num, atol=0.1).all()
print('Lotka Volterra test passed.')
def get_dn(self, net, *x):
projection = self._get_projection(*x)
p_tensor = torch.cat(projection, dim=1)
if self.f_d is None:
D = self.bias(net)
else:
D = self.f_d(*projection) - net(p_tensor)
if self.f_n is None:
N = 0.0
else:
N = self.f_n(*projection) - diff(net(p_tensor),
projection[self.idx])
return D, N
def test_curl(x3):
nets = [FCNN(3, 1) for _ in range(3)]
x, y, z = x3
u, v, w = [net(torch.cat(x3, dim=1)) for net in nets]
ans = [
diff(w, y) - diff(v, z),
diff(u, z) - diff(w, x),
diff(v, x) - diff(u, y),
]
output = curl(u, v, w, x, y, z)
for a, o in zip(ans, output):
assert (a == o).all()
def test_ode():
def mse(x, t):
true_x = torch.sin(t)
return torch.mean((x - true_x)**2)
exponential = lambda x, t: diff(x, t) - x
init_val_ex = IVP(t_0=0.0, x_0=1.0)
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.0,
t_max=2.0,
shuffle=False,
max_epochs=2000,
return_best=True,
metrics={'mse': mse})
ts = np.linspace(0, 2.0, 100)
x_net = solution_ex(ts, as_type='np')
x_ana = np.exp(ts)
assert isclose(x_net, x_ana, atol=0.1).all()
print('solve basic test passed.')
def test_train_generator():
exponential = lambda x, t: diff(x, t) - x
init_val_ex = IVP(t_0=0.0, x_0=1.0)
train_gen = Generator1D(size=32, t_min=0.0, t_max=2.0, method='uniform')
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.0,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
train_gen = Generator1D(size=32,
t_min=0.0,
t_max=2.0,
method='equally-spaced')
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.0,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
train_gen = Generator1D(size=32,
t_min=0.0,
t_max=2.0,
method='equally-spaced-noisy')
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.0,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
train_gen = Generator1D(size=32,
t_min=0.0,
t_max=2.0,
method='equally-spaced-noisy',
noise_std=0.01)
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.0,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
train_gen = Generator1D(size=32,
t_min=np.log10(0.1),
t_max=np.log10(2.0),
method='log-spaced')
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.1,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
train_gen = Generator1D(size=32,
t_min=np.log10(0.1),
t_max=np.log10(2.0),
method='log-spaced-noisy')
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.1,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
train_gen = Generator1D(size=32,
t_min=np.log10(0.1),
t_max=np.log10(2.0),
method='log-spaced-noisy',
noise_std=0.01)
solution_ex, _ = solve(ode=exponential,
condition=init_val_ex,
t_min=0.1,
t_max=2.0,
train_generator=train_gen,
max_epochs=3)
with raises(ValueError):
train_gen = Generator1D(size=32, t_min=0.0, t_max=2.0, method='magic')
print('ExampleGenerator test passed.')
def particle_squarewell(y1, y2, t):
return [(-1 / 2) * diff(y1, t, order=2) - 3 - (y2) * (y1), diff(y2, t)]
def test_cylindrical_grad(u, x):
out_rho, out_phi, out_z = cylindrical_grad(u, *x)
rho, phi, z = x
assert torch.allclose(out_rho, diff(u, rho))
assert torch.allclose(out_phi, diff(u, phi) / rho)
assert torch.allclose(out_z, diff(u, z))
def test_cylindrical_div(U, x):
out = cylindrical_div(*U, *x)
rho, phi, z = x
urho, uphi, uz = U
assert torch.allclose(out, diff(rho * urho, rho) / rho + diff(uphi, phi) / rho + diff(uz, z))
def test_cylindrical_vector_laplacian(U, x):
out_rho, out_phi, out_z = cylindrical_vector_laplacian(*U, *x)
rho, phi, z = x
urho, uphi, uz = U
def scalar_lap(u):
return diff(rho * diff(u, rho), rho) / rho + diff(u, phi, order=2) / rho ** 2 + diff(u, z, order=2)
assert torch.allclose(out_rho, scalar_lap(urho) - urho / rho ** 2 - 2 / rho ** 2 * diff(uphi, phi))
assert torch.allclose(out_phi, scalar_lap(uphi) - uphi / rho ** 2 + 2 / rho ** 2 * diff(urho, phi))
assert torch.allclose(out_z, scalar_lap(uz))
