simple nonlinear 1st order ODE system: what am I doing wrong?

[rayannott]

I picked the original paper for review during my masters seminar "Deep Learning in Physics" at the Technical University of Munich. I wanted to try out some other equations and systems, but all of them failed to match the "exact" numerical solution (RK45).

This problem is similar to #1583

For example, here is my code for the SIR-model simulation.

This sets up the reference solution using scipy.integrate.solve_ivp.

from scipy.integrate import solve_ivp
import numpy as np

beta = 0.5
gamma = 0.1

time_domain = (0, 30)
initial_condition = (0.99, 0.01, 0.0) # 99% are susceptible and 1% are infected

def rhs(t, r):
    """ODE system:
    x' = -beta * x * y
    y' = beta * x * y - gamma * y
    z' = gamma * y
    """
    x, y, z = r
    return [-beta * x * y, beta * x * y - gamma * y, gamma * y]

def gen_truedata(t):
    sol = solve_ivp(rhs, time_domain, initial_condition, t_eval=t)
    x_true, y_true, z_true = sol.y
    x_true = x_true.reshape(-1, 1)
    y_true = y_true.reshape(-1, 1)
    z_true = z_true.reshape(-1, 1)
    return x_true, y_true, z_true

This is the model creation and training:

import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np
import torch

def ode_system(x, y):
    s = y[:, 0]
    i = y[:, 1]
    r = y[:, 2]
    ds_t = dde.grad.jacobian(y, x, i=0)
    di_t = dde.grad.jacobian(y, x, i=1)
    dr_t = dde.grad.jacobian(y, x, i=2)
    return [
        ds_t + beta * s * i,
        di_t - beta * s * i + gamma * i,
        dr_t - gamma * i,
    ]

def boundary(_, on_initial):
    return on_initial

geom = dde.geometry.TimeDomain(*time_domain)
ic1 = dde.icbc.IC(geom, lambda x: initial_condition[0], boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda x: initial_condition[1], boundary, component=1)
ic3 = dde.icbc.IC(geom, lambda x: initial_condition[2], boundary, component=2)
data = dde.data.PDE(geom, ode_system, [ic1, ic2, ic3], 40, 3, num_test=200)

layer_size = [1] + [100] * 5 + [3] 
activation = "tanh"
initializer = "Glorot normal"
net = dde.nn.FNN(layer_size, activation, initializer)


model = dde.Model(data, net)

model.compile("adam", lr=0.001)
losshistory, train_state = model.train(iterations=10000)
dde.saveplot(losshistory, train_state, issave=True, isplot=True)

The loss looks great, but the solution doesn't match the reference:

plt.xlabel("t")
plt.ylabel("population")

t = np.linspace(0, 30, 500)
s_true, i_true, r_true = gen_truedata(t)
plt.plot(t, s_true, color="black", label="s_true")
plt.plot(t, i_true, color="blue", label="i_true")
plt.plot(t, r_true, color="green", label="r_true")


t_pred = np.linspace(0, 30, 500).reshape(-1, 1)
sol_pred = model.predict(t_pred)
s_pred = sol_pred[:, 0]
i_pred = sol_pred[:, 1]
r_pred = sol_pred[:, 2]
plt.plot(t_pred, s_pred, color="black", linestyle="dashed", label="s_pred")
plt.plot(t_pred, i_pred, color="blue", linestyle="dashed", label="i_pred")
plt.plot(t_pred, r_pred, color="green", linestyle="dashed", label="r_pred")
plt.legend()
plt.show()

[xuliang5115]

I think you should post the loss_history, and try more iterations may be help?
Generally, the evaluation results are closer to the true value after the loss function drops by $10^4-10^5$ relative to the initial loss.（based on my experience）

[rayannott]

Thanks for the reply, @xuliang5115.
Here is the loss I get:

And, like you said, it drops by about 4-6 orders of magnitude.
I increased the number of iterations to 30K, and yet, the wrong result persists.

[xuliang5115]

I think you should check the PDE and BCs, the defined problem may have multi solution, the networks chooses one.
Also, you should try data scaling.

[aoguedao]

Hi @rayannott, I have worked with some SIR models before but mostly for inverse problems. I usually use 3 layers with 32-64 neurons and relu as activation. Also you could try to add some input transform, Dr. Lu has a good example in this repo https://github.com/lu-group/sbinn

I hope it helps!