From 104146af261b43928224aec2e32a4eb72828debb Mon Sep 17 00:00:00 2001 From: Samuel Burbulla Date: Fri, 19 Apr 2024 15:24:54 +0200 Subject: [PATCH 1/3] Add 1D Poisson PI-DeepONet example as in DeepXDE. --- pyproject.toml | 1 + src/continuiti/__init__.py | 21 ++++++ tests/pde/test_deepxde.py | 99 ++++++++++++++++++++++++++++ tests/pde/test_pideeponet.py | 122 +++++++++++++++++++++++++++++++++++ 4 files changed, 243 insertions(+) create mode 100644 tests/pde/test_deepxde.py create mode 100644 tests/pde/test_pideeponet.py diff --git a/pyproject.toml b/pyproject.toml index a415d87c..8346d831 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,7 @@ dev = [ "neoteroi-mkdocs", "pygments", "gmsh", + "deepxde", ] [tool.setuptools.dynamic] diff --git a/src/continuiti/__init__.py b/src/continuiti/__init__.py index 3d390b10..f3e51017 100644 --- a/src/continuiti/__init__.py +++ b/src/continuiti/__init__.py @@ -37,3 +37,24 @@ """ __version__ = "0.0.0" + +__all__ = [ + "benchmarks", + "data", + "discrete", + "operators", + "pde", + "trainer", + "transforms", + "Trainer", +] + +from . import benchmarks +from . import data +from . import discrete +from . import operators +from . import pde +from . import trainer +from . import transforms + +from .trainer import Trainer diff --git a/tests/pde/test_deepxde.py b/tests/pde/test_deepxde.py new file mode 100644 index 00000000..b349dc94 --- /dev/null +++ b/tests/pde/test_deepxde.py @@ -0,0 +1,99 @@ +import torch +import deepxde as dde +import matplotlib.pyplot as plt +import numpy as np + +torch.manual_seed(0) + + +def test_deepxde(): + """Physics-informed DeepONet for Poisson equation in 1D. + Example from DeepXDE. + https://deepxde.readthedocs.io/en/latest/demos/operator/poisson.1d.pideeponet.html + """ + + # Poisson equation: -u_xx = f + def equation(x, y, f): + dy_xx = dde.grad.hessian(y, x) + return -dy_xx - f + + # Domain is interval [0, 1] + geom = dde.geometry.Interval(0, 1) + + # Zero Dirichlet BC + def u_boundary(_): + return 0 + + def boundary(_, on_boundary): + return on_boundary + + bc = dde.icbc.DirichletBC(geom, u_boundary, boundary) + + # Define PDE + pde = dde.data.PDE(geom, equation, bc, num_domain=100, num_boundary=2) + + # Function space for f(x) are polynomials + degree = 3 + space = dde.data.PowerSeries(N=degree + 1) + + # Choose evaluation points + num_eval_points = 10 + evaluation_points = geom.uniform_points(num_eval_points, boundary=True) + + # Define PDE operator + pde_op = dde.data.PDEOperatorCartesianProd( + pde, + space, + evaluation_points, + num_function=100, + ) + + # Setup DeepONet + dim_x = 1 + p = 32 + net = dde.nn.DeepONetCartesianProd( + [num_eval_points, 32, p], + [dim_x, 32, p], + activation="tanh", + kernel_initializer="Glorot normal", + ) + print("Params:", sum(p.numel() for p in net.parameters())) + + # Define and train model + model = dde.Model(pde_op, net) + model.compile("adam", lr=0.001) + model.train(epochs=1000) + + # Plot realizations of f(x) + n = 3 + features = space.random(n) + fx = space.eval_batch(features, evaluation_points) + + x = geom.uniform_points(100, boundary=True) + y = model.predict((fx, x)) + + # Setup figure + fig = plt.figure(figsize=(7, 8)) + plt.subplot(2, 1, 1) + plt.title("Poisson equation: Source term f(x) and solution u(x)") + plt.ylabel("f(x)") + z = np.zeros_like(x) + plt.plot(x, z, "k-", alpha=0.1) + + # Plot source term f(x) + for i in range(n): + plt.plot(evaluation_points, fx[i], ".") + + # Plot solution u(x) + plt.subplot(2, 1, 2) + plt.ylabel("u(x)") + plt.plot(x, z, "k-", alpha=0.1) + for i in range(n): + plt.plot(x, y[i], "-") + plt.xlabel("x") + + plt.show() + + +if __name__ == "__main__": + test_deepxde() diff --git a/tests/pde/test_pideeponet.py b/tests/pde/test_pideeponet.py new file mode 100644 index 00000000..7b90e2ba --- /dev/null +++ b/tests/pde/test_pideeponet.py @@ -0,0 +1,122 @@ +import torch +import deepxde as dde +import matplotlib.pyplot as plt +import numpy as np +import continuiti as cti + +torch.manual_seed(0) + + +def test_pideeponet(): + """Physics-informed DeepONet for Poisson equation in 1D. + Example from DeepXDE in *continuiti*. + https://deepxde.readthedocs.io/en/latest/demos/operator/poisson.1d.pideeponet.html + """ + + # Poisson equation: -v_xx = f + mse = torch.nn.MSELoss() + + def equation(_, f, y, v): + # PDE + dy_xx = dde.grad.hessian(v, y) + inner_loss = mse(-dy_xx, f) + + # BC + y_bnd, v_bnd = y[:, :, bnd_indices], v[:, :, bnd_indices] + boundary_loss = mse(v_bnd, v_boundary(y_bnd)) + + return inner_loss + boundary_loss + + # Domain is interval [0, 1] + geom = dde.geometry.Interval(0, 1) + + # Zero Dirichlet BC + def v_boundary(y): + return torch.zeros_like(y) + + # Sample domain and boundary points + num_domain = 100 + num_boundary = 2 + x_domain = geom.uniform_points(num_domain) + x_bnd = geom.uniform_boundary_points(num_boundary) + + x = np.concatenate([x_domain, x_bnd]) + num_points = len(x) + bnd_indices = range(num_domain, num_points) + + # Function space for f(x) are polynomials + degree = 3 + space = dde.data.PowerSeries(N=degree + 1) + + num_functions = 100 + coeffs = space.random(num_functions) + fx = space.eval_batch(coeffs, x) + + # Specify dataset + xt = torch.tensor(x.T).requires_grad_(True) + x_all = xt.expand(num_functions, -1, -1) # (num_functions, x_dim, num_domain) + u_all = torch.tensor(fx).unsqueeze(1) # (num_functions, u_dim, num_domain) + y_all = x_all # (num_functions, y_dim, num_domain) + v_all = torch.zeros_like(y_all) # (num_functions, v_dim, num_domain) + + dataset = cti.data.OperatorDataset( + x=x_all, + u=u_all, + y=y_all, # same as x_all + v=v_all, # only for shapes + ) + + # Define operator + operator = cti.operators.DeepONet( + dataset.shapes, + trunk_depth=1, + branch_depth=1, + basis_functions=32, + ) + # or any other operator, e.g.: + # operator = cti.operators.DeepNeuralOperator(dataset.shapes) + + # Define and train model + trainer = cti.Trainer( + operator, + loss_fn=cti.pde.PhysicsInformedLoss(equation), + ) + trainer.fit(dataset) + + # Plot realizations of f(x) + n = 3 + features = space.random(n) + fx = space.eval_batch(features, x) + + y = geom.uniform_points(200, boundary=True) + + x_plot = torch.tensor(x_domain.T).expand(n, -1, -1) + u_plot = torch.tensor(fx).unsqueeze(1) + y_plot = torch.tensor(y.T).expand(n, -1, -1) + v = operator(x_plot, u_plot, y_plot) + v = v.detach().numpy() + + fig = plt.figure(figsize=(7, 8)) + plt.subplot(2, 1, 1) + plt.title("Poisson equation: Source term f(x) and solution v(x)") + plt.ylabel("f(x)") + z = np.zeros_like(y) + plt.plot(y, z, "k-", alpha=0.1) + + # Plot source term f(x) + for i in range(n): + plt.plot(x, fx[i], ".") + + # Plot solution v(x) + plt.subplot(2, 1, 2) + plt.ylabel("v(x)") + plt.plot(y, z, "k-", alpha=0.1) + for i in range(n): + plt.plot(y, v[i].T, "-") + plt.xlabel("x") + + plt.show() + + +if __name__ == "__main__": + test_pideeponet() From 2968f1ffff71b79e3204eec3af054e700bc8d8c9 Mon Sep 17 00:00:00 2001 From: Samuel Burbulla Date: Mon, 22 Apr 2024 11:16:23 +0200 Subject: [PATCH 2/3] Remove deepxde example from tests. --- tests/pde/{test_deepxde.py => deepxde_example.py} | 4 ++-- tests/pde/test_pideeponet.py | 10 ++++------ 2 files changed, 6 insertions(+), 8 deletions(-) rename tests/pde/{test_deepxde.py => deepxde_example.py} (98%) diff --git a/tests/pde/test_deepxde.py b/tests/pde/deepxde_example.py similarity index 98% rename from tests/pde/test_deepxde.py rename to tests/pde/deepxde_example.py index b349dc94..fd0179d3 100644 --- a/tests/pde/test_deepxde.py +++ b/tests/pde/deepxde_example.py @@ -6,7 +6,7 @@ torch.manual_seed(0) -def test_deepxde(): +def deepxde_example(): """Physics-informed DeepONet for Poisson equation in 1D. Example from DeepXDE. https://deepxde.readthedocs.io/en/latest/demos/operator/poisson.1d.pideeponet.html @@ -96,4 +96,4 @@ def boundary(_, on_boundary): if __name__ == "__main__": - test_deepxde() + deepxde_example() diff --git a/tests/pde/test_pideeponet.py b/tests/pde/test_pideeponet.py index 7b90e2ba..c1104ad1 100644 --- a/tests/pde/test_pideeponet.py +++ b/tests/pde/test_pideeponet.py @@ -77,11 +77,9 @@ def v_boundary(y): # operator = cti.operators.DeepNeuralOperator(dataset.shapes) # Define and train model - trainer = cti.Trainer( - operator, - loss_fn=cti.pde.PhysicsInformedLoss(equation), - ) - trainer.fit(dataset) + loss_fn = cti.pde.PhysicsInformedLoss(equation) + trainer = cti.Trainer(operator, loss_fn=loss_fn) + trainer.fit(dataset, epochs=100) # Plot realizations of f(x) n = 3 @@ -115,7 +113,7 @@ def v_boundary(y): plt.plot(y, v[i].T, "-") plt.xlabel("x") - plt.show() + plt.savefig("pideeponet.png", dpi=500) if __name__ == "__main__": From 4e98a260c9d00c073121c8d7d1b091e6363de739 Mon Sep 17 00:00:00 2001 From: Samuel Burbulla Date: Mon, 22 Apr 2024 12:28:36 +0200 Subject: [PATCH 3/3] Accelerate test execution. --- tests/benchmarks/run/test_runner.py | 1 + tests/benchmarks/test_navierstokes.py | 32 ++++++++++++++------------- tests/pde/test_pideeponet.py | 6 +++-- tests/trainer/test_trainer.py | 14 ++++++------ 4 files changed, 29 insertions(+), 24 deletions(-) diff --git a/tests/benchmarks/run/test_runner.py b/tests/benchmarks/run/test_runner.py index 9b1c7e65..94f9fecd 100644 --- a/tests/benchmarks/run/test_runner.py +++ b/tests/benchmarks/run/test_runner.py @@ -10,5 +10,6 @@ def test_runner(): benchmark_factory=SineRegular, operator_factory=DeepNeuralOperator, max_epochs=2, + batch_size=128, ) BenchmarkRunner.run(config) diff --git a/tests/benchmarks/test_navierstokes.py b/tests/benchmarks/test_navierstokes.py index b7cf5a56..0ec041b9 100644 --- a/tests/benchmarks/test_navierstokes.py +++ b/tests/benchmarks/test_navierstokes.py @@ -35,18 +35,20 @@ def test_navierstokes_shapes_and_plot(): assert y.shape == (3, 64, 64, 10) assert v.shape == (1, 64, 64, 10) - fig, axs = plt.subplots(1, 2, subplot_kw={"projection": "3d"}, figsize=(10, 5)) - x, u, y, v = benchmark.test_dataset[0] - axs[0].scatter(x[2], x[0], x[1], s=1, c=u, cmap="jet", alpha=0.7) - axs[1].scatter(y[2], y[0], y[1], s=1, c=v, cmap="jet", alpha=0.7) - for i in range(2): - axs[i].set_xlabel("t") - axs[i].set_ylabel("x") - axs[i].set_zlabel("y") - axs[0].set_title("Input") - axs[1].set_title("Output") - - try: - fig.savefig("docs/benchmarks/img/navierstokes.png", dpi=500) - except FileNotFoundError: - pass + plot = False + if plot: + fig, axs = plt.subplots(1, 2, subplot_kw={"projection": "3d"}, figsize=(10, 5)) + x, u, y, v = benchmark.test_dataset[0] + axs[0].scatter(x[2], x[0], x[1], s=1, c=u, cmap="jet", alpha=0.7) + axs[1].scatter(y[2], y[0], y[1], s=1, c=v, cmap="jet", alpha=0.7) + for i in range(2): + axs[i].set_xlabel("t") + axs[i].set_ylabel("x") + axs[i].set_zlabel("y") + axs[0].set_title("Input") + axs[1].set_title("Output") + + try: + fig.savefig("docs/benchmarks/img/navierstokes.png", dpi=500) + except FileNotFoundError: + pass diff --git a/tests/pde/test_pideeponet.py b/tests/pde/test_pideeponet.py index c1104ad1..94cedc35 100644 --- a/tests/pde/test_pideeponet.py +++ b/tests/pde/test_pideeponet.py @@ -1,3 +1,4 @@ +import pytest import torch import deepxde as dde import matplotlib.pyplot as plt @@ -7,6 +8,7 @@ torch.manual_seed(0) +@pytest.mark.slow def test_pideeponet(): """Physics-informed DeepONet for Poisson equation in 1D. Example from DeepXDE in *continuiti*. @@ -35,7 +37,7 @@ def v_boundary(y): return torch.zeros_like(y) # Sample domain and boundary points - num_domain = 100 + num_domain = 32 num_boundary = 2 x_domain = geom.uniform_points(num_domain) x_bnd = geom.uniform_boundary_points(num_boundary) @@ -48,7 +50,7 @@ def v_boundary(y): degree = 3 space = dde.data.PowerSeries(N=degree + 1) - num_functions = 100 + num_functions = 10 coeffs = space.random(num_functions) fx = space.eval_batch(coeffs, x) diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py index ff1fa171..eb261037 100644 --- a/tests/trainer/test_trainer.py +++ b/tests/trainer/test_trainer.py @@ -7,15 +7,15 @@ def train(): - dataset = SineBenchmark().train_dataset + dataset = SineBenchmark(n_train=32).train_dataset operator = DeepONet(dataset.shapes, trunk_depth=16) - Trainer(operator).fit(dataset, tol=1e-3) + Trainer(operator).fit(dataset, tol=1e-2) # Make sure we can use operator output on cpu again x, u, y, v = dataset.x, dataset.u, dataset.y, dataset.v v_pred = operator(x, u, y) - assert ((v_pred - v.to("cpu")) ** 2).mean() < 1e-3 + assert ((v_pred - v.to("cpu")) ** 2).mean() < 1e-2 @pytest.mark.slow @@ -42,7 +42,7 @@ def f(x): input_size=1, output_size=1, width=32, - depth=3, + depth=8, ) # Define loss function (in continuiti style) @@ -56,14 +56,14 @@ def loss_fn(op, x, y): trainer = Trainer(model, loss_fn=loss_fn) logs = trainer.fit( train_dataset, - tol=1e-3, + tol=1e-2, test_dataset=test_dataset, ) # Test the model - assert logs.loss_test < 1e-3 + assert logs.loss_test < 1e-2 # Use ./run_parallel.sh to run test with CUDA if __name__ == "__main__": - train() + test_trainer_with_torch_model()