rajeshrinet / pyross

PyRoss: inference, forecasts, and optimised control of epidemiological models in Python. github.com/rajeshrinet/pyross
https://pyross.readthedocs.io
MIT License
169 stars 57 forks source link

Simulate until steady state? #41

Closed sdwfrost closed 2 years ago

sdwfrost commented 2 years ago

If I just use solve_ivp in SciPy, I can run until steady state using callbacks. See this SO post and the following code for an SIR model. Can I do this in Pyross too?

import numpy as np
from scipy.integrate import solve_ivp

# Model definition 
def model(t, u, p):
    b,g = p
    S,I,R = u
    dS = -b*S*I
    dI = b*S*I-g*I
    dR = g*I 
    return [dS,dI,dR]

# Stop integration at steady state
def steady_state(t,u,p,f,tol):
    global flag
    du = f(t,u,p)
    condition = np.max(np.abs(du))<tol
    if flag == 1:
        if condition:
            test = [0]
        else:
            test = [1]
        flag = 0
    else:
        if condition:
            test = np.array([0])
        else:
            test = np.array([1])
    return test

# Define terminal condition and type-change flag
tol = 1e-6
limit = lambda t, u: steady_state(t,u,p,model,tol)
limit.terminal = True                                                                                                                                                                                   
global flag
flag = 1

# Extract inputs from JSON
tspan = [0,40]
u0 = [0.99,0.01,0.0]
p = [0.5,0.25]

# Run model
sol = solve_ivp(lambda t, u: model(t, u, p),
                tspan,
                u0,
                events = limit)
rajeshrinet commented 2 years ago

@sdwfrost thanks for sharing your code on how to simulate until a steady-state has been reached. I have now added an example on how to implement such a scenario using PyRoss for an arbitrary model at: https://github.com/rajeshrinet/pyross/blob/master/examples/others/simulate_until_steady-state.ipynb

sdwfrost commented 2 years ago

Thanks @rajeshrinet! Much appreciated.