Merge pull request #133 from aai-institute/feature/minimal-example

Add 1D Poisson PI-DeepONet example as in DeepXDE.

pyproject.toml

@@ -72,6 +72,7 @@ dev = [
     "neoteroi-mkdocs",
     "pygments",
     "gmsh",
+    "deepxde",
 ]
 
 [tool.setuptools.dynamic]

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

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)

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

tests/pde/deepxde_example.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 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
+    """
+
+    # 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__":
+    deepxde_example()

tests/pde/test_pideeponet.py

@@ -0,0 +1,122 @@
+import pytest
+import torch
+import deepxde as dde
+import matplotlib.pyplot as plt
+import numpy as np
+import continuiti as cti
+
+torch.manual_seed(0)
+
+
+@pytest.mark.slow
+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 = 32
+    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 = 10
+    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
+    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
+    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.savefig("pideeponet.png", dpi=500)
+
+
+if __name__ == "__main__":
+    test_pideeponet()

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()