Having trouble with nan when training the Schrodinger equation after Madelung transformation

[ChenRui077]

Dear Prof. Lu @lululxvi, I'm trying to solve a Schrodinger equation after Madelung transformation, but while training the model I got several NAN values. The equation is as follows: with discontinuous Initial Value Condition (for example rho = 1 while x<0, rho = 0.25 while x>=0, so is t, and I use a tanh function to fit it). The range of x is [-30,30] and the range of t is [0,10].

Kindly suggest how to fix it.

Code is:

x_lower = -30
x_upper = 30
t_lower = 0
t_upper = 10

x = np.linspace(x_lower, x_upper, 512)
t = np.linspace(t_lower, t_upper, 201)
X, T = np.meshgrid(x, t)

X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))

space_domain = dde.geometry.Interval(x_lower, x_upper)
time_domain = dde.geometry.TimeDomain(t_lower, t_upper)
geomtime = dde.geometry.GeometryXTime(space_domain, time_domain)

# The "physics-informed" part of the loss
def pde(x, y):

    u = y[:, 0:1]
    v = y[:, 1:2]
    w = y[:, 2:3]

    # In 'jacobian', i is the output component and j is the input component
    u_t = dde.grad.jacobian(y, x, i=0, j=1)
    v_t = dde.grad.jacobian(y, x, i=1, j=1)

    u_x = dde.grad.jacobian(y, x, i=0, j=0)
    v_x = dde.grad.jacobian(y, x, i=1, j=0)
    w_x = dde.grad.jacobian(y, x, i=2, j=0)

    # In 'hessian', i and j are both input components. (The Hessian could be in principle something like d^2y/dxdt, d^2y/d^2x etc)
    # The output component is selected by "component"
    u_xx = dde.grad.hessian(y, x, component=0, i=0, j=0)
    v_xx = dde.grad.hessian(y, x, component=1, i=0, j=0)

    #u_xxx = dde.grad.jacobian(u_xx, x, i=0, j=0)    
    #func_x = (u_xxx * u - u_xx * u_x)/u**2 - (u_x * u_xx * u**2 - u_x**3 * u)/u**4
    #(episilon**2/4) * func_x

    episilon = 0.02
    f_u = u_t + u_x * v + u * v_x
    f_v = v_t + v * v_x + u_x - (episilon**2/4) * w_x
    f_w = w - u_xx / u + (u_x**2) / (2*u**2)

    return [f_u, f_v, f_w]

# Boundary and Initial conditions

# no boundary conditions
# Initial conditions

rho_left = 1
rho_right = 0.25
v_left = 0
v_right = 0

def init_cond_u(x):
    return ((rho_left-rho_right)/2)*(-np.tanh(x[:, 0:1])+1)+rho_right

def init_cond_v(x):
    return ((v_left-v_right)/2)*(-np.tanh(x[:, 0:1])+1)+v_right

ic_u = dde.IC(geomtime, init_cond_u, lambda _, on_initial: on_initial, component=0)
ic_v = dde.IC(geomtime, init_cond_v, lambda _, on_initial: on_initial, component=1)

#build the model
data = dde.data.TimePDE(
    geomtime,
    pde,
    [ic_u, ic_v],
    num_domain=10000,
    # num_boundary=20,
    num_initial=200,
    train_distribution="pseudo",
)

# Network architecture
net = dde.maps.FNN([2] + [100] * 4 + [3], "tanh", "Glorot normal")

model = dde.Model(data, net)

# To employ a GPU accelerated system is highly encouraged.

model.compile("adam", lr=1e-3, loss="MSE")
model.train(epochs=10000, display_every=1000)

and the error messages are:

Compiling model...
Building feed-forward neural network...
'build' took 0.088686 s

/usr/local/lib/python3.7/dist-packages/deepxde/nn/tensorflow_compat_v1/fnn.py:110: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  kernel_constraint=self.kernel_constraint,
/usr/local/lib/python3.7/dist-packages/keras/legacy_tf_layers/core.py:261: UserWarning: `layer.apply` is deprecated and will be removed in a future version. Please use `layer.__call__` method instead.
  return layer.apply(inputs)
'compile' took 4.130018 s

Initializing variables...
Training model...

Step      Train loss                                            Test loss                                             Test metric
0         [1.73e-03, 2.51e-02, 6.42e+04, 1.18e+00, 7.32e-01]    [1.73e-03, 2.51e-02, 6.42e+04, 1.18e+00, 7.32e-01]    []  
1000      [nan, nan, nan, nan, nan]                             [nan, nan, nan, nan, nan]                             []  

Best model at step 0:
  train loss: 6.42e+04
  test loss: 6.42e+04
  test metric: []

'train' took 37.702558 s

(<deepxde.model.LossHistory at 0x7f39bfaa9110>,
 <deepxde.model.TrainState at 0x7f39bfaa93d0>)

[lululxvi]

Maybe the loss value 6.42e+04 is too large and makes the network training blow up.

[ariehlev]

Hi all, I constrained it so that the model would not initialise until all the losses in steps 0 and 1 are less than 1 (I also varied this number to see if it made a difference, but I never had it over 4). However, I still get the problem that at Step 1000, I get NaNs in the train and test losses. Any ideas? Thanks in advance!

[lululxvi]

See FAQ "Q: I failed to train the network or get the right solution, e.g., large training loss, unbalanced losses."