Originally posted by **xuliang5115** April 9, 2024
Dear Pro.Lu:
I encountered some problems while training the model and would like your advice.
## Research objectives & Question
**Research objectives**
Train DeepONet to learning the mapping $[\xi, \eta]\rightarrow[r, z]$ which could transform grids between physical domain coordinates(Cartesian coordinates,$[\xi, \eta]$) and computational domain coordinates(Curvilinear body-fitted coordinates,$[r, z]$). The topography function $z(\xi)$ is the output of DeepNet on the boundary, and the uniform points in compational domain are the input of trunk-net.
Once the model is trained, we want the model to be able to compute physical problems in computational domain according to the chain rule display the results in curvilinear body-fitted coordinates.
**Question**
> We assume that the values generated by the function space are the function values of the boundary. However, the model predict results after training are not ideal, which only occurs when the PDE has a **second derivative**.
what i've tried:
1. add fourier feature layer
2. use hard constraint and add loss weights
3. change step of decay, more iterations
These don't seem to work, I don't understand why DeepONet cannot approximate this boundary condition ` bc_topography_top`, do you or community have any suggestions for this?
I try to do some test for this issue with another example script "stokes_aligned_zcs_pideeponet.py"(after small changes and shown in next block in the discussion area).
Many thanks.
## Results&Figures
### Loss History
### Model predict
### Topography function
### Code
```python
"""Backend supported: tensorflow, pytorch"""
# dde.backend.set_default_backend('pytorch')
try:
import deepxde as dde
dde.config.disable_xla_jit()
import numpy as np
import matplotlib.pyplot as plt
import os
import datetime
import torch
dim_x = 5 # number of feature
sin = torch.sin
cos = torch.cos
concat = torch.cat
# creat folder to save result
if not os.path.exists("./save"):
os.makedirs("./save")
if not os.path.exists("./figures"):
os.makedirs("./figures")
finally:
pass
# Generate grids in orthogonal body-fitted coordinates(OBFC)
Training = Train # switch 'Training to True' means Training the model.
# ZCS saved GPU memory a lot, but need {DeepXDE 1.11.0, python 3.9}{pytorch 2.2.* with cuda121}
Filename = datetime.datetime.now().strftime('%Y%b%d')
# Hyperparameters
iterations = 100000
lr, step, decayRate = 0.001, 200, 0.01
decay = ("inverse time", step, decayRate)
# PDE equation
def pde(uv, rz, _):
# input: \xi:u, \eta:v； output: x(\xi, \eta), y(\xi, \eta)
r, z = rz[..., 0:1], rz[..., 1:2]
grad_r = dde.zcs.LazyGrad(uv, r)
grad_z = dde.zcs.LazyGrad(uv, z)
# first order
dr_u = grad_r.compute((1, 0))
dr_v = grad_r.compute((0, 1))
dz_v = grad_z.compute((0, 1))
dz_u = grad_z.compute((1, 0))
# second order
dr_uu = grad_r.compute((2, 0))
dr_vv = grad_r.compute((0, 2))
dz_uu = grad_z.compute((2, 0))
dz_vv = grad_z.compute((0, 2))
# dr_uv = grad_r.compute((1, 1))
# dz_uv = grad_z.compute((1, 1))
# Laplace equation
g12 = dr_u * dr_v + dz_u * dz_v
g11 = dr_u ** 2 + dz_u ** 2
g22 = dr_v ** 2 + dz_v ** 2
f = torch.sqrt(g22) / torch.sqrt(g11)
CauchyRiemann = dr_v + f * dz_u + dz_v - f * dr_u
# Laplace_1 =
# Laplace_2 =
Laplace = dr_uu + dr_vv + dz_uu + dz_vv
return g12, CauchyRiemann, Laplace
# down-left and up-right
xi_max, eta_max = 1., 1.
xi_min, eta_min = 0, 0
geom2d = dde.geometry.Rectangle([0, 0], [xi_max, eta_max])
# Boundary Condition
# other boundary conditions will be enforced by output transform(hard-constraint)
def bc_topography_top(_, anx_var):
# using (perturbation / 10 + 1.) to avoid zero
# z(\xi) / z(\xi_0) -> anx_var / 10 + 1.
return anx_var / 10 + eta_max
bc_topography_z = dde.icbc.DirichletBC(
geom=geom2d,
func=bc_topography_top,
on_boundary=lambda x, on_boundary: np.isclose(x[1], eta_max),
component=1)
# PDE object
pde = dde.data.PDE(geom2d,
pde,
bcs=bc_topography_z,
num_domain=5000,
num_boundary=8000,
num_test=529)
# Function space, GRF or PowerSeries or Chebyshev, kernel={RBF or AE or ExpSineSquared}
func_space = dde.data.GRF(length_scale=0.2)
# Data
n_pts_bdry = 101 # using the size of bdry sample points, but this is unnecessary
eval_pts = np.linspace(0, 1, num=n_pts_bdry)[:, None]
data = dde.zcs.PDEOperatorCartesianProd(
pde, func_space, eval_pts, num_function=1000,
function_variables=[0], num_test=100, batch_size=50
)
# Net
net = dde.nn.DeepONetCartesianProd(
[n_pts_bdry, 128, 128, 128],
[dim_x, 128, 128, 128],
"tanh",
"Glorot normal",
num_outputs=2,
multi_output_strategy='independent'
)
# Output transform for other boundary conditions
def out_transform(inputs, outputs):
_xi, _eta = inputs[1][:, 0], inputs[1][:, 1]
# O-BFC r on left, right. if eta = 0, r = \xi; xi = 0, r=0; xi = xi_max, r=xi_max,
_r = outputs[:, :, 0] * (_xi * (xi_max - _xi)) * _eta[None, :] + _xi[None, :]
# O-BFC z on down. if xi = 0, z = \eta; eta = 0, z = 0.
_z = outputs[:, :, 1] * _eta * _xi[None, :] + _eta[None, :]
return dde.backend.stack((_r, _z), axis=2)
def periodic(x):
_xi, _eta = x[:, 0:1], x[:, 1:2]
_xi = _xi * 2 * np.pi
return concat([cos(_xi), sin(_xi), cos(2 * _xi), sin(2 * _xi), _eta], 1)
net.apply_feature_transform(periodic)
net.apply_output_transform(out_transform)
# Train Model
model = dde.zcs.Model(data, net)
if Training:
model.compile("adam", lr=lr, decay=decay, loss='MSE', loss_weights=[1, 1, 1e-2, 4])
# training the model
loss_history, train_state = model.train(iterations=iterations)
# save model if needed
model.save('./save/OBFC_{name}'.format(name=Filename), verbose=1)
dde.utils.plot_loss_history(loss_history, './figures/OBFC_{name}_loss.png'.format(name=Filename))
plt.show()
else:
# restore and predict
lr1, step, decayRate = 0.001, 100, 0.01
decay1 = ("inverse time", step, decayRate)
model.compile("adam", lr=lr1, decay=decay1, loss='MSE') # compile before restore, this is important.
model.restore("./save/OBFC_{name}-{iter}.pt".format(name=Filename, iter=iterations))
# Evaluation
lenR, lenZ = 40, 80
xi = np.linspace(0, 1, lenR)
eta = np.linspace(0, 1, lenZ)
features = func_space.random(1)
v_branch = func_space.eval_batch(features, eval_pts)
# v_branch[:] = 0.
xiv, etv = np.meshgrid(xi, eta)
x_trunk = np.vstack((np.ravel(xiv), np.ravel(etv))).T
model_out = model.predict((v_branch, x_trunk))
r_pred, z_pred = model_out[:, :, 0:1], model_out[:, :, 1:2]
r_new = r_pred.reshape(lenZ, lenR).T
z_new = z_pred.reshape(lenZ, lenR).T
# Figure 1
# row1.1: true result
fig, ax = plt.subplots(1, 2, sharey=True)
h0 = ax[0].plot(eval_pts, (v_branch / 10 + 1.).T, 'b', label='bd in Cart')
ax[0].vlines(0, ymin=0, ymax=v_branch[:, 0] / 10 + 1., label='left')
ax[0].vlines(1, ymin=0, ymax=v_branch[:, -1] / 10 + 1., label='right')
ax[0].hlines(0, xmin=0, xmax=1, label='top')
ax[0].set_xlabel('Range,after scale')
ax[0].set_ylabel('Depth,after scale')
ax[0].set_title('Cartesian Coordinate')
# row1.2: pred
for i in range(lenR):
ax[1].plot(r_new[i, :], z_new[i, :], 'k', linewidth=1.5)
for i in range(lenZ):
ax[1].plot(r_new[:, i], z_new[:, i], 'k', linewidth=1.5)
ax[1].set_xlabel('Range,after scale')
ax[1].set_title('O-BFC')
temp = 0.08
ax[1].set_xlim([0 - 0.02, 1 + 0.02])
ax[1].set_ylim([0 - 0.02, (v_branch / 10 + 1).max() + 0.02])
plt.tight_layout()
plt.show()
# Figure 2
fig1, ax2 = plt.subplots(1, 1)
ax2.plot(eval_pts, (v_branch / 10 + 1.).T, 'b', label='v_branch')
ax2.plot(r_new[:, -1], z_new[:, -1], 'r-', label='r_predict')
ax2.set_xlabel('Range')
plt.tight_layout()
plt.legend()
plt.show()
```
Discussed in https://github.com/lululxvi/deepxde/discussions/1704