def test_generator2d():
grid = (5, 6)
size = grid[0] * grid[1]
x_min, x_max = 0.0, 1.0
y_min, y_max = -1.0, 0.0
x_std, y_std = 0.05, 0.06
generator = Generator2D(grid=grid,
xy_min=(x_min, y_min),
xy_max=(x_max, y_max),
method='equally-spaced-noisy')
x, y = generator.getter()
assert _check_shape_and_grad(generator, size, x, y)
generator = Generator2D(grid=grid,
xy_min=(x_min, y_min),
xy_max=(x_max, y_max),
method='equally-spaced-noisy',
xy_noise_std=(x_std, y_std))
x, y = generator.getter()
assert _check_shape_and_grad(generator, size, x, y)
generator = Generator2D(grid=grid,
xy_min=(x_min, y_min),
xy_max=(x_max, y_max),
method='equally-spaced')
x, y = generator.getter()
assert _check_shape_and_grad(generator, size, x, y)
assert _check_boundary((x, y), (x_min, y_min), (x_max, y_max))
str(generator)
repr(generator)
def __init__(
self,
pde_system,
conditions,
xy_min,
xy_max,
nets=None,
train_generator=None,
valid_generator=None,
analytic_solutions=None,
optimizer=None,
criterion=None,
n_batches_train=1,
n_batches_valid=4,
metrics=None,
n_output_units=1,
# deprecated arguments are listed below
batch_size=None,
shuffle=None):
if train_generator is None or valid_generator is None:
if xy_min is None or xy_max is None:
raise ValueError(
f"Either generator is not provided, xy_min and xy_max should be both provided: \n"
f"got xy_min={xy_min}, xy_max={xy_max}, "
f"train_generator={train_generator}, valid_generator={valid_generator}"
)
if train_generator is None:
train_generator = Generator2D((32, 32),
xy_min=xy_min,
xy_max=xy_max,
method='equally-spaced-noisy')
if valid_generator is None:
valid_generator = Generator2D((32, 32),
xy_min=xy_min,
xy_max=xy_max,
method='equally-spaced')
self.xy_min, self.xy_max = xy_min, xy_max
super(Solver2D, self).__init__(
diff_eqs=pde_system,
conditions=conditions,
nets=nets,
train_generator=train_generator,
valid_generator=valid_generator,
analytic_solutions=analytic_solutions,
optimizer=optimizer,
criterion=criterion,
n_batches_train=n_batches_train,
n_batches_valid=n_batches_valid,
metrics=metrics,
n_input_units=2,
n_output_units=n_output_units,
shuffle=shuffle,
batch_size=batch_size,
)
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, hidden_units=(32, 32))
solution_neural_net_laplace, _ = solve2D(
pde=laplace,
condition=bc,
xy_min=(0, 0),
xy_max=(1, 1),
net=net,
max_epochs=3,
train_generator=Generator2D((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 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, hidden_units=(32, 32))
solution_neural_net_laplace, _ = solve2D(pde=laplace,
condition=bc,
xy_min=(0, 0),
xy_max=(1, 1),
net=net,
max_epochs=300,
train_generator=Generator2D(
(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_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, hidden_units=(32, 32))
solution_neural_net_laplace, loss_history = solve2D(
pde=laplace,
condition=bc,
xy_min=(0, 0),
xy_max=(1, 1),
net=net,
max_epochs=3,
train_generator=Generator2D((32, 32), (0, 0), (1, 1),
method='equally-spaced-noisy',
xy_noise_std=(0.01, 0.01)),
batch_size=64)
assert isinstance(solution_neural_net_laplace, Solution2D)
assert isinstance(loss_history, dict)
keys = ['train_loss', 'valid_loss']
for key in keys:
assert key in loss_history
assert isinstance(loss_history[key], list)
assert len(loss_history[keys[0]]) == len(loss_history[keys[1]])
def test_filter_generator():
grid = (10, 10)
size = 100
x = [i * 1.0 for i in range(size)]
filter_fn = lambda a: (a[0] % 2 == 0)
filter_fn_2 = lambda a: (a % 2 == 0)
x_expected = filter(filter_fn_2, x)
generator = PredefinedGenerator(x)
filter_generator = FilterGenerator(generator,
filter_fn=filter_fn,
update_size=True)
x = filter_generator.get_examples()
assert _check_shape_and_grad(filter_generator, size // 2, x)
assert _check_iterable_equal(x_expected, x)
x = [i * 1.0 for i in range(size)]
y = [-i * 1.0 for i in range(size)]
filter_fn = lambda ab: (ab[0] % 2 == 0) & (ab[1] > -size / 2)
x_expected, y_expected = list(zip(*filter(filter_fn, zip(x, y))))
generator = PredefinedGenerator(x, y)
filter_generator = FilterGenerator(generator, filter_fn)
x, y = filter_generator.get_examples()
assert _check_shape_and_grad(filter_generator, size // 4, x, y)
assert _check_iterable_equal(x_expected, x)
assert _check_iterable_equal(y_expected, y)
generator = Generator2D(grid)
filter_fn = lambda ab: (ab[0] > 0.5) & (ab[1] < 0.5)
filter_generator = FilterGenerator(generator, filter_fn)
for _ in range(5):
x, y = filter_generator.get_examples()
assert _check_shape_and_grad(filter_generator, None, x, y)
fixed_size = 42
filter_generator = FilterGenerator(generator,
filter_fn,
size=fixed_size,
update_size=False)
for _ in range(5):
assert _check_shape_and_grad(filter_generator, fixed_size)
filter_generator.get_examples()
str(filter_generator)
repr(filter_generator)
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, hidden_units=(32, 32))
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=Generator2D(
(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_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, hidden_units=(32, 32))
solution_neural_net_heat, _ = solve2D(pde=heat,
condition=ibvp,
xy_min=(0, 0),
xy_max=(L, T),
net=net,
max_epochs=300,
train_generator=Generator2D(
(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_train_generator():
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, hidden_units=(32, 32))
solution_neural_net_laplace, _ = solve2D(
pde=laplace, condition=bc, xy_min=(0, 0), xy_max=(1, 1),
net=net, max_epochs=3,
train_generator=Generator2D((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))
)
train_gen = Generator2D((32, 32), (0, 0), (1, 1), method='equally-spaced')
solution_neural_net_laplace, _ = solve2D(
pde=laplace, condition=bc, xy_min=(0, 0), xy_max=(1, 1),
net=net, max_epochs=3, train_generator=train_gen, batch_size=64
)
train_gen = Generator2D((32, 32), (0, 0), (1, 1), method='equally-spaced-noisy')
solution_neural_net_laplace, _ = solve2D(
pde=laplace, condition=bc, xy_min=(0, 0), xy_max=(1, 1),
net=net, max_epochs=3, train_generator=train_gen, batch_size=64
)
with raises(ValueError):
train_gen = Generator2D((32, 32), (0, 0), (1, 1), method='magic')
valid_gen = Generator2D((32, 32), (0, 0), (1, 1), method='equally-spaced-noisy')
train_gen = Generator2D((32, 32), (0, 0), (1, 1), method='equally-spaced')
solution_neural_net_laplace, _ = solve2D(
pde=laplace, condition=bc,
net=net, max_epochs=3, train_generator=train_gen, valid_generator=valid_gen, batch_size=64
)
with raises(RuntimeError):
solution_neural_net_laplace, _ = solve2D(
pde=laplace, condition=bc,
net=net, max_epochs=3, batch_size=64
)
def test_solution():
x_grids = 7
y_grids = 11
xy_grids = (x_grids, y_grids)
N_SAMPLES = x_grids * y_grids
x0, y0 = random.random(), random.random()
x1, y1 = random.random() + 1, random.random() + 1
generator = Generator2D(xy_grids,
xy_min=(x0, y0),
xy_max=(x1, y1),
method='equally-spaced')
u00, u01, u10, u11 = random.random(), random.random(), random.random(
), random.random()
v00, v01, v10, v11 = random.random(), random.random(), random.random(
), random.random()
def get_single_boundary_func(z0, w0, z1, w1):
net = FCNN(1, 1)
condition = DirichletBVP(z0, w0, z1, w1)
def boundary_func(z):
return condition.enforce(net, z)
return boundary_func
def get_all_boundary_funcs(w00, w01, w10, w11):
wx0 = get_single_boundary_func(y0, w00, y1, w01)
wx1 = get_single_boundary_func(y0, w10, y1, w11)
wy0 = get_single_boundary_func(x0, w00, x1, w10)
wy1 = get_single_boundary_func(x0, w01, x1, w11)
return wx0, wx1, wy0, wy1
ux0, ux1, uy0, uy1 = get_all_boundary_funcs(u00, u01, u10, u11)
vx0, vx1, vy0, vy1 = get_all_boundary_funcs(v00, v01, v10, v11)
def get_solution(use_single: bool) -> Solution2D:
conditions = [
DirichletBVP2D(x0, ux0, x1, ux1, y0, uy0, y1, uy1),
DirichletBVP2D(x0, vx0, x1, vx1, y0, vy0, y1, vy1),
]
if use_single:
net = FCNN(2, 2)
for i, cond in enumerate(conditions):
cond.set_impose_on(i)
return Solution2D(net, conditions)
else:
nets = [FCNN(2, 1), FCNN(2, 1)]
return Solution2D(nets, conditions)
def check_output(uv, shape, type, msg=""):
msg += " "
assert isinstance(uv,
(list, tuple)), msg + "returned type is not a list"
assert len(uv) == 2, msg + "returned length is not 2"
assert isinstance(uv[0], type) and isinstance(
uv[1], type), msg + f"returned element is not {type}"
u, v = uv
assert u.shape == shape and v.shape == shape, msg + f"returned element shape is not {shape}"
u, v = u.reshape(*xy_grids), v.reshape(*xy_grids)
x = torch.linspace(x0, x1, steps=x_grids,
requires_grad=True).reshape(-1, 1)
y = torch.linspace(y0, y1, steps=y_grids,
requires_grad=True).reshape(-1, 1)
if type == torch.Tensor:
check_close = lambda a, b: torch.isclose(a, b).all()
elif type == np.ndarray:
check_close = lambda a, b: np.isclose(a,
b.cpu().detach().numpy()).all
else:
raise ValueError(f"Unrecognized type={type}")
assert check_close(u[0, :], torch.flatten(
ux0(y))), msg + "u on x0 not satisfied"
assert check_close(u[-1, :], torch.flatten(
ux1(y))), msg + "u on x1 not satisfied"
assert check_close(u[:, 0], torch.flatten(
uy0(x))), msg + "u on y0 not satisfied"
assert check_close(u[:, -1], torch.flatten(
uy1(x))), msg + "u on y1 not satisfied"
assert check_close(v[0, :], torch.flatten(
vx0(y))), msg + "v on x0 not satisfied"
assert check_close(v[-1, :], torch.flatten(
vx1(y))), msg + "v on x1 not satisfied"
assert check_close(v[:, 0], torch.flatten(
vy0(x))), msg + "v on y0 not satisfied"
assert check_close(v[:, -1], torch.flatten(
vy1(x))), msg + "v on y1 not satisfied"
for use_single in [True, False]:
solution = get_solution(use_single=use_single)
xs, ys = generator.get_examples()
us = solution(xs, ys)
check_output(us,
shape=(N_SAMPLES, ),
type=torch.Tensor,
msg=f"[use_single={use_single}]")
us = solution(xs, ys, as_type='np')
check_output(us,
shape=(N_SAMPLES, ),
type=np.ndarray,
msg=f"[use_single={use_single}]")
xs, ys = xs.reshape(-1, 1), ys.reshape(-1, 1)
us = solution(xs, ys)
check_output(us,
shape=(N_SAMPLES, 1),
type=torch.Tensor,
msg=f"[use_single={use_single}]")
us = solution(xs, ys, as_type='np')
check_output(us,
shape=(N_SAMPLES, 1),
type=np.ndarray,
msg=f"[use_single={use_single}]")
