I have a function that takes 4 inputs (S, E, I, R) and returns one of three values in (0, 1, 2). However, S+E+I+R=N, where N is a constant. I tried the following:
def dummy(S,E,I,R):
if S+E<500:
return 0
elif I+R >600:
return 1
else:
return 2
N = 1e5
States = []
Cateogory = []
S = np.arange(0.0, N, 1000.)
for s in S:
E = np.arange(0.0, N-s, 1000)
if not E.shape[0]>0:
E=np.arange(0.0, 1., 1000.)
for e in E:
I = np.arange(0.0, N-s-e, 1000)
if not I.shape[0]>0:
I=np.arange(0.0, 1., 1000.)
for i in I:
states.append([s,e,i,N-s-e-i])
Cateogory.append(dummy(s,e,i,N-s-e-i))
But its taking too much time to generate the data. Is there an optimal way to do this? After this I am planning to follow this post to make a 3D plot.