Add experimental submodule(examples)

examples/experimental_examples/examples-function/dataset.py

@@ -0,0 +1,31 @@
+import os
+
+import brainstate as bst
+import numpy as np
+
+import deepxde.experimental as deepxde
+
+PATH = os.path.dirname(os.path.abspath(__file__))
+train_data = np.loadtxt(os.path.join(PATH, "../..", "dataset", "dataset.train"))
+test_data = np.loadtxt(os.path.join(PATH, "../..", "dataset", "dataset.test"))
+
+net = deepxde.nn.Model(
+    deepxde.nn.DictToArray(x=None),
+    deepxde.nn.FNN([1] + [50] * 3 + [1], "tanh", bst.init.KaimingUniform()),
+    deepxde.nn.ArrayToDict(y=None),
+)
+
+data = deepxde.problem.DataSet(
+    X_train={"x": train_data[:, 0]},
+    y_train={"y": train_data[:, 1]},
+    X_test={"x": test_data[:, 0]},
+    y_test={"y": test_data[:, 1]},
+    standardize=True,
+    approximator=net,
+)
+
+model = deepxde.Trainer(data)
+model.compile(bst.optim.Adam(0.001), metrics=["l2 relative error"]).train(
+    iterations=50000
+)
+model.saveplot(issave=True, isplot=True)

examples/experimental_examples/examples-function/func.py

@@ -0,0 +1,27 @@
+import brainstate as bst
+import brainunit as u
+
+import deepxde.experimental as deepxde
+
+
+def func(x):
+    return {"y": x["x"] * u.math.sin(5 * x["x"])}
+
+
+geom = deepxde.geometry.Interval(-1, 1).to_dict_point("x")
+num_train = 160
+num_test = 100
+
+net = deepxde.nn.Model(
+    deepxde.nn.DictToArray(x=None),
+    deepxde.nn.FNN([1] + [20] * 3 + [1], "tanh", bst.init.LecunUniform()),
+    deepxde.nn.ArrayToDict(y=None),
+)
+
+data = deepxde.problem.Function(geom, func, num_train, num_test, approximator=net)
+
+trainer = deepxde.Trainer(data)
+trainer.compile(bst.optim.Adam(0.001), metrics=["l2 relative error"]).train(
+    iterations=10000
+)
+trainer.saveplot(issave=False, isplot=True)

examples/experimental_examples/examples-function/func_uncertainty.py

@@ -0,0 +1,28 @@
+import brainstate as bst
+import brainunit as u
+
+import deepxde.experimental as deepxde
+
+
+def func(x):
+    return {"y": x["x"] * u.math.sin(5 * x["x"])}
+
+
+layer_size = [1] + [50] * 3 + [1]
+net = deepxde.nn.Model(
+    deepxde.nn.DictToArray(x=None),
+    deepxde.nn.FNN(layer_size, "tanh", bst.init.KaimingUniform()),
+    deepxde.nn.ArrayToDict(y=None),
+)
+
+geom = deepxde.geometry.Interval(-1, 1).to_dict_point("x")
+num_train = 100
+num_test = 1000
+data = deepxde.problem.Function(geom, func, num_train, num_test, approximator=net)
+
+trainer = deepxde.Trainer(data)
+uncertainty = deepxde.callbacks.DropoutUncertainty(period=1000)
+trainer.compile(bst.optim.Adam(0.001), metrics=["l2 relative error"]).train(
+    iterations=30000, callbacks=uncertainty
+)
+trainer.saveplot(issave=True, isplot=True)

examples/experimental_examples/examples-operator/ADR_solver.py

@@ -0,0 +1,73 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def solve_ADR(xmin, xmax, tmin, tmax, k, v, g, dg, f, u0, Nx, Nt):
+    """Solve 1D
+    u_t = (k(x) u_x)_x - v(x) u_x + g(u) + f(x, t)
+    with zero boundary condition.
+    """
+    x = np.linspace(xmin, xmax, Nx)
+    t = np.linspace(tmin, tmax, Nt)
+    h = x[1] - x[0]
+    dt = t[1] - t[0]
+    h2 = h**2
+
+    D1 = np.eye(Nx, k=1) - np.eye(Nx, k=-1)
+    D2 = -2 * np.eye(Nx) + np.eye(Nx, k=-1) + np.eye(Nx, k=1)
+    D3 = np.eye(Nx - 2)
+    k = k(x)
+    M = -np.diag(D1 @ k) @ D1 - 4 * np.diag(k) @ D2
+    m_bond = 8 * h2 / dt * D3 + M[1:-1, 1:-1]
+    v = v(x)
+    v_bond = 2 * h * np.diag(v[1:-1]) @ D1[1:-1, 1:-1] + 2 * h * np.diag(
+        v[2:] - v[: Nx - 2]
+    )
+    mv_bond = m_bond + v_bond
+    c = 8 * h2 / dt * D3 - M[1:-1, 1:-1] - v_bond
+    f = f(x[:, None], t)
+
+    u = np.zeros((Nx, Nt))
+    u[:, 0] = u0(x)
+    for i in range(Nt - 1):
+        gi = g(u[1:-1, i])
+        dgi = dg(u[1:-1, i])
+        h2dgi = np.diag(4 * h2 * dgi)
+        A = mv_bond - h2dgi
+        b1 = 8 * h2 * (0.5 * f[1:-1, i] + 0.5 * f[1:-1, i + 1] + gi)
+        b2 = (c - h2dgi) @ u[1:-1, i].T
+        u[1:-1, i + 1] = np.linalg.solve(A, b1 + b2)
+    return x, t, u
+
+
+def main():
+    xmin, xmax = -1, 1
+    tmin, tmax = 0, 1
+    k = lambda x: x**2 - x**2 + 1
+    v = lambda x: np.ones_like(x)
+    g = lambda u: u**3
+    dg = lambda u: 3 * u**2
+    f = lambda x, t: np.exp(-t) * (1 + x**2 - 2 * x) - (np.exp(-t) * (1 - x**2)) ** 3
+    u0 = lambda x: (x + 1) * (1 - x)
+    u_true = lambda x, t: np.exp(-t) * (1 - x**2)
+
+    # xmin, xmax = 0, 1
+    # tmin, tmax = 0, 1
+    # k = lambda x: np.ones_like(x)
+    # v = lambda x: np.zeros_like(x)
+    # g = lambda u: u ** 2
+    # dg = lambda u: 2 * u
+    # f = lambda x, t: x * (1 - x) + 2 * t - t ** 2 * (x - x ** 2) ** 2
+    # u0 = lambda x: np.zeros_like(x)
+    # u_true = lambda x, t: t * x * (1 - x)
+
+    Nx, Nt = 100, 100
+    x, t, u = solve_ADR(xmin, xmax, tmin, tmax, k, v, g, dg, f, u0, Nx, Nt)
+
+    print(np.max(abs(u - u_true(x[:, None], t))))
+    plt.plot(x, u)
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()

examples/experimental_examples/examples-operator/advection_aligned_pideeponet.py

@@ -0,0 +1,93 @@
+import brainstate as bst
+import brainunit as u
+import jax
+import matplotlib.pyplot as plt
+import numpy as np
+
+import deepxde
+import deepxde.experimental as deepxde_exp
+
+geom = deepxde_exp.geometry.Interval(0, 1)
+timedomain = deepxde_exp.geometry.TimeDomain(0, 1)
+geomtime = deepxde_exp.geometry.GeometryXTime(geom, timedomain)
+geomtime = geomtime.to_dict_point("x", "t")
+
+
+# PDE
+def pde_eqs(x, y, aux):
+    def solve_jac(inp1):
+        f1 = lambda i: deepxde_exp.grad.jacobian(
+            lambda inp: net((x[0], inp))["y"][i], inp1, vmap=False
+        )
+        return jax.vmap(f1)(np.arange(x[0].shape[0]))
+
+    jacobian = jax.vmap(solve_jac, out_axes=1)(jax.numpy.expand_dims(x[1], 1))
+    dy_x = u.math.squeeze(jacobian[..., 0])
+    dy_t = u.math.squeeze(jacobian[..., 1])
+    return dy_t + dy_x
+
+
+# Net
+def periodic(x):
+    x, t = x[..., :1], x[..., 1:]
+    x = x * 2 * u.math.pi
+    return u.math.concatenate(
+        [u.math.cos(x), u.math.sin(x), u.math.cos(2 * x), u.math.sin(2 * x), t], axis=-1
+    )
+
+
+dim_x = 5
+net = bst.nn.Sequential(
+    deepxde_exp.nn.DeepONetCartesianProd(
+        [50, 128, 128, 128],
+        [dim_x, 128, 128, 128],
+        "tanh",
+        input_transform=periodic,
+    ),
+    deepxde_exp.nn.ArrayToDict(y=None),
+)
+
+ic = deepxde_exp.icbc.IC(lambda x, aux: {"y": aux})
+
+# Function space
+func_space = deepxde.data.GRF(kernel="ExpSineSquared", length_scale=1)
+
+# Problem
+eval_pts = np.linspace(0, 1, num=50)[:, None]
+problem = deepxde_exp.problem.PDEOperatorCartesianProd(
+    geomtime,
+    pde_eqs,
+    ic,
+    func_space,
+    eval_pts,
+    approximator=net,
+    num_function=1000,
+    function_variables=[0],
+    num_fn_test=100,
+    batch_size=32,
+    num_domain=250,
+    num_initial=50,
+    num_test=500,
+)
+
+model = deepxde_exp.Trainer(problem)
+model.compile(bst.optim.Adam(0.0005)).train(iterations=50000)
+model.saveplot(issave=True, isplot=True)
+
+x = np.linspace(0, 1, num=100)
+t = np.linspace(0, 1, num=100)
+u_true = np.sin(2 * np.pi * (x - t[:, None]))
+plt.figure()
+plt.imshow(u_true)
+plt.colorbar()
+
+v_branch = np.sin(2 * np.pi * eval_pts).T
+xv, tv = np.meshgrid(x, t)
+x_trunk = np.vstack((np.ravel(xv), np.ravel(tv))).T
+u_pred = model.predict((v_branch, x_trunk))["y"]
+u_pred = u_pred.reshape((100, 100))
+plt.figure()
+plt.imshow(u_pred)
+plt.colorbar()
+plt.show()
+print(deepxde_exp.metrics.l2_relative_error(u_true, u_pred))

examples/experimental_examples/examples-operator/advection_aligned_pideeponet_2d.py

@@ -0,0 +1,100 @@
+import brainstate as bst
+import brainunit as u
+import jax
+import matplotlib.pyplot as plt
+import numpy as np
+
+import deepxde
+import deepxde.experimental as deepxde_exp
+
+dim_x = 5
+
+
+# PDE
+def pde(x, y, aux):
+    def solve_jac(inp1):
+        f1 = lambda i: deepxde_exp.grad.jacobian(
+            lambda inp: net((x[0], inp))["u"][i], inp1, vmap=False
+        )
+        return jax.vmap(f1)(np.arange(x[0].shape[0]))
+
+    jacobian = jax.vmap(solve_jac, out_axes=1)(jax.numpy.expand_dims(x[1], 1))
+    dy_x = u.math.squeeze(jacobian[..., 0])
+    dy_t = u.math.squeeze(jacobian[..., 1])
+    return dy_t + dy_x
+
+
+# The same problem as advection_aligned_pideeponet.py
+# But consider time as the 2nd space coordinate
+# to demonstrate the implementation of 2D problems
+geom = deepxde_exp.geometry.Rectangle([0, 0], [1, 1])
+geom = geom.to_dict_point("x", "y")
+
+
+def periodic(x):
+    x, t = x[..., :1], x[..., 1:]
+    x = x * 2 * np.pi
+    return u.math.concatenate(
+        [u.math.cos(x), u.math.sin(x), u.math.cos(2 * x), u.math.sin(2 * x), t], axis=-1
+    )
+
+
+# Net
+net = bst.nn.Sequential(
+    deepxde_exp.nn.DeepONetCartesianProd(
+        [50, 128, 128, 128],
+        [dim_x, 128, 128, 128],
+        "tanh",
+        input_transform=periodic,
+    ),
+    deepxde_exp.nn.ArrayToDict(u=None),
+)
+
+
+def boundary(x, on_boundary):
+    return u.math.logical_and(on_boundary, u.math.isclose(x["y"], 0))
+
+
+ic = deepxde_exp.icbc.DirichletBC(lambda x, aux: {"u": aux}, boundary)
+
+# Function space
+func_space = deepxde.data.GRF(kernel="ExpSineSquared", length_scale=1)
+
+# Problem
+eval_pts = np.linspace(0, 1, num=50)[:, None]
+data = deepxde_exp.problem.PDEOperatorCartesianProd(
+    geom,
+    pde,
+    ic,
+    func_space,
+    eval_pts,
+    approximator=net,
+    num_function=1000,
+    function_variables=[0],
+    num_fn_test=100,
+    batch_size=32,
+    num_domain=200,
+    num_boundary=200,
+)
+
+trainer = deepxde_exp.Trainer(data)
+trainer.compile(bst.optim.Adam(0.0005)).train(iterations=30000)
+trainer.saveplot()
+
+x = np.linspace(0, 1, num=100)
+t = np.linspace(0, 1, num=100)
+u_true = np.sin(2 * np.pi * (x - t[:, None]))
+plt.figure()
+plt.imshow(u_true)
+plt.colorbar()
+
+v_branch = np.sin(2 * np.pi * eval_pts).T
+xv, tv = np.meshgrid(x, t)
+x_trunk = np.vstack((np.ravel(xv), np.ravel(tv))).T
+u_pred = trainer.predict((v_branch, x_trunk))["u"]
+u_pred = u_pred.reshape((100, 100))
+plt.figure()
+plt.imshow(u_pred)
+plt.colorbar()
+plt.show()
+print(deepxde_exp.metrics.l2_relative_error(u_true, u_pred))