0

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.

kosa
  • 262
  • 4
  • 17

1 Answers1

0

You can speed up things using two ways:

  1. Multiprocessing toolbox
  2. using permutations from itertool (as shown below) from itertools import permutations

from itertools import permutations

def random_uniform_state():
    popu=1e5
    S = np.random.uniform(low=0.0, high=popu)
    E = np.random.uniform(low=0.0, high=popu-S)
    I = np.random.uniform(low=0.0, high=popu-(S+E))
    R = popu-(S+E+I)
    state = [S,E,I,R]
    perms = permutations([S,E,I,R])
    states = []
    for state in perms:
        states.append(state)
    return states
kosa
  • 262
  • 4
  • 17