Search code examples
pythonpython-3.xscipydifferential-equationsodeint

Writing data from equation solver to file over loop only outputs last 2 elements


I am trying to solve this problem of writing to a file, the time runs from 0 -> 1000 with increments of 100, however when I run the script, I keep getting only the final 2 elements of the output. I think there is an issue with the loop but I'm not sure. I would like for all of the data to be present from t=0 to t=1000 and not just the final 2 elements. You can see in the file, only t=900 and t=1000 are recorded. Can't seem to see what to change.

import matplotlib.pyplot as plt
from scipy.integrate import odeint
import pandas as pd


plt.ion()
plt.rcParams['figure.figsize'] = 10, 8

P = 0      # birth rate
d = 0.0001  # natural death percent (per day)
B = 0.0095  # transmission percent  (per day)
G = 0.0001  # resurect percent (per day)
A = 0.0001  # destroy percent  (per day)

# solve the system dy/dt = f(y, t)
def f(y, t):
     Si = y[0]
     Zi = y[1]
     Ri = y[2]
     # the model equations (see Munz et al. 2009)
     f0 = P - B*Si*Zi - d*Si
     f1 = B*Si*Zi + G*Ri - A*Si*Zi
     f2 = d*Si + A*Si*Zi - G*Ri
     return [f0, f1, f2]

# initial conditions
S0 = 500.             # initial population
Z0 = 30                 # initial zombie population
R0 = 60                 # initial death population
y0 = [S0, Z0, R0]     # initial condition vector

#looping over some time instead of integrating in one go.
t_a = 0 
oput = 500
t_b = t_a + oput
delta_t = t_a + 100 
tend = 1000

dfs=[]
while t_a < tend: 

    t_c = t_a + delta_t 
    t=[t_a,t_c]
    y = odeint(f,y0,t,mxstep=10000) #Integrator
    t_a = t_c

    if(t_a > oput):
        t_b = t_b +oput

        S = y[:,0]
        R = y[:,1]
        Z = y[:,2]



dfs.append(pd.DataFrame({'t': t, 'Z': Z,'R': R}))
g=pd.concat(dfs,axis=0)
g.to_csv('example.csv',mode='w',index=False)

The other issue that arises is that I want to change my initial conditions to whatever the final element of the solver was with:

        S0 = S
        R0 = R
        Z0 = Z
        y0 = [S0,R0,Z0]

But I am met with this error:

ValueError: Initial condition y0 must be one-dimensional.

How can I work around this error if I want to update the initial conditions with the pop values in the loop?


Solution

  • You have only the last values because you do not output anything during the loop execution but just at the end of it.

    To solve your problem, I'd go like this, leaving to odeint the task of looping over time steps,

    In [21]: from scipy.integrate import odeint
    In [22]: P, D, B, G, A = 0, 0.0001, 0.0095, 0.0001, 0.0001
    In [23]: def fy(y, t):
        ...:      Si, Zi, Ri = y
        ...:      # the model equations (see Munz et al. 2009)
        ...:      return P-B*Si*Zi-D*Si, B*Si*Zi+G*Ri-A*Si*Zi, D*Si+A*Si*Zi-G*Ri
    In [24]: y0 = [500.0, 30.0, 60.0]
    In [25]: res = odeint(fy, y0, [i*100.0 for i in range(11)], mxstep=1000)
    In [26]: print('\n'.join(','.join('%+15.9f'%x for x in row) for row in res))
     +500.000000000,  +30.000000000,  +60.000000000
       +0.000000000, +525.356080701,  +64.643919299
       -0.000000000, +525.999298342,  +64.000701658
       +0.000000000, +526.636115893,  +63.363884107
       +0.000000000, +527.266597017,  +62.733402983
       +0.000000000, +527.890804714,  +62.109195286
       +0.000000000, +528.508801154,  +61.491198846
       +0.000000000, +529.120648555,  +60.879351445
       -0.000000000, +529.726408164,  +60.273591837
       -0.000000000, +530.326140479,  +59.673859521
       -0.000000000, +530.919905407,  +59.080094593
    

    If you want to save a CSV format to a file, print supports redirection to an open file, like print(..., file=open('res.txt'))

    Re "I am met with this error", you never update (at least in the code you've shown) the values of the initial conditions, so that it's difficult to guess what went wrong, but I suppose that you used the like of y0 = y, while it should be y0 = y[1] — if you print y you'll understand why.

    Finally, from what I can understand, maybe your time step is a bit longer than what is desirable, as the live population at the end of the first time step is 0 (more precisely, 2.72457580e-13) and it is possible that you want to follow the process rather than look at its final result.


    Addendum

    As a follow-up to your comment, you can loop over the results, examine them and take appropriate actions, e.g.

    res = odeint(...)
    
    for S, R, Z in res:
        output_style = 0
        if S<40 and R> 30 : output_style = 1 
        elif ............ : output_style = 2
        ....
        output(S, R, Z, output_style)