integrating solver/data into PI-DeepOnet

[erkara]

def pde_system(x, y, f):
    # 1D Convection-Diffusion
    c = y[:, 0:1]

    #constant velocity for now
    u = torch.tensor(0.1)
    du_x = 0.
    # derivatives
    dc_dt = dde.grad.jacobian(y, x, j=1)      # ∂c/∂t
    dc_x = dde.grad.jacobian(y, x, i=0, j=0)  # ∂c/∂x
    dc_xx = dde.grad.hessian(y, x, i=0, j=0)  # ∂²c/∂x²
    pde = dc_dt - D * dc_xx + u * dc_x + c * du_x - f
    return pde

geom = dde.geometry.Interval(0, 1)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
bc = dde.icbc.DirichletBC(geomtime, lambda *: 0, lambda *, on_boundary: on_boundary)
ic = dde.icbc.IC(geomtime, lambda *: 0, lambda *, on_initial: on_initial)
pde_system = dde.data.TimePDE(geomtime,pde_system,[bc, ic],num_domain=200,num_boundary=40,num_initial=20,num_test=500,)
num_eval_points = 50
evaluation_points = geom.uniform_points(num_eval_points, boundary=True)
func_space = dde.data.GRF(length_scale=0.2)
data=dde.data.PDEOperatorCartesianProd(pde_system,func_space,evaluation_points,num_function=1000,num_test=100,batch_size=50)
net = dde.nn.DeepONetCartesianProd([num_eval_points, 128, 128], [2, 128, 128],"tanh","Glorot normal")
#train
model = dde.Model(data, net)
model.compile("adam", lr=1e-3, decay=("inverse time", 1000, 0.5))
losshistory, train_state = model.train(iterations=20000)
dde.utils.plot_loss_history(losshistory)

Hi @lululxvi I have been working on the problem above where the task is to learn the mapping from f --> c. Numerical and PI-DeepOnet solutions are in good agreement.
In the actual application though, u is part of another PDE which uses the same source f(x) and that PDE is easy to solve. I am working on incorporating the right velocity from the sampled source function during training. I can think of two ways:

(1) If we assume we get the velocity like "u = solve_u(x,source_fun)", we can use this solver during training per sampled function to obtain the right velocity within pde_system. Apart from being slow, I also don't know how to make such solver compatible with DeepXDE,i,e, differentiability of inputs etc. (2) We can use the same solver to create a dataset of the form (f,u) before training. We can then load it to train PI-DeepOnet. However, this case requires no sampling of source function, thus no function space etc. I tried to create a custom PDEOperatorCartesianProd class to accommodate this case but I am going down to rabbit hole as this new class has lots of compatibility issues within DeepXDE.

Is there any way we can solve this problem with the existing methods in DeepXDE. If not, I really appreciate if you can point out to a direction I can handle this problem. Thank you!

[lululxvi]

1. You may need to modify the data class to call the solver to get u in each iteration.
2. For the current version of PI-DeepONet, you may have to modify the code to accept the dataset directly instead of function space. Or you can just use DeepONet.