lululxvi / deepxde

A library for scientific machine learning and physics-informed learning
https://deepxde.readthedocs.io
GNU Lesser General Public License v2.1
2.55k stars 725 forks source link

The ODE system of the SIR model does not converge #816

Open mavilar opened 2 years ago

mavilar commented 2 years ago

I am trying to solve the system of ordinary differential equations of a SIR model like the following one:

$$ \begin{cases} \frac{dS}{dt} = - \tau I(t) S(t) \ \frac{dI}{dt} = \tau I(t) S(t) - \frac{1}{k} I(t) \ \frac{dR}{dt} = \frac{1}{k} I(t) \ \end{cases} $$

As I could not figure out how to solve the train convergence problem I was having, I tried to implement all the recommendations given in this question of the FAQ:

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

I was trying to get the solution of that ODE system for $t$ between $0$ and $60$, so the first thing was to transform it so that the solution fell between $0$ and $1$ instead, as this is one of the most repeated recommendations I could found.

After also following all the other recommendations in that section, my code is like follows:

tau = 0.8
k = 4

ts, te = 0, 1

s0 = 0.995
i0 = 0.005
r0 = 0

t_span = te - ts

def ode_system(t, sir):
    s = sir[:, 0:1]
    i = sir[:, 1:2]
    r = sir[:, 2:3]

    # Compute the derivatives of the input vector with automatic differentiation
    dS_t = dde.grad.jacobian(sir, t, i=0)
    dI_t = dde.grad.jacobian(sir, t, i=1)
    dR_t = dde.grad.jacobian(sir, t, i=2)

    return [
        dS_t + t_span * tau * i * s,
        dI_t - t_span * tau * i * s + t_span * i / k,
        dR_t - t_span * i / k,
    ]

geom = dde.geometry.TimeDomain(0, 1)

def boundary(x, on_initial):
    return on_initial and np.isclose(x[0], 0)

ic = [
    dde.icbc.IC(geom, lambda x: s0, boundary, component=0),
    dde.icbc.IC(geom, lambda x: i0, boundary, component=1),
    dde.icbc.IC(geom, lambda x: r0, boundary, component=2),
]

n = 16384
data = dde.data.PDE(
    geometry=geom,
    pde=ode_system,
    bcs=ic,
    num_domain=n,
    num_boundary=len(ic),
    num_test=n
)

layer_size = [1] + [64] * 7 + [len(ic)]
activation = "tanh"
initializer = "Glorot normal"
net = dde.nn.FNN(layer_size, activation, initializer)

model = dde.Model(data, net)
model.compile("adam", lr=0.0001)
epochs = 1000000
losshistory, train_state = model.train(
    epochs=epochs,
)
model.compile("L-BFGS-B")
losshistory, train_state = model.train()

Any idea of what may be happening?

lululxvi commented 2 years ago

The code seems OK. Not sure what the problem is.

Bensayah commented 2 years ago

Hi Mavilar and Lulu I think you will encounter the problem of trivial solution which will be in this system S(t)=1000, I(t)=0,R(t)=0. The question is how to avoide the trivial solution. There are some suggestions in FAQ e,g,. #356