Skip to content

Instantly share code, notes, and snippets.

@thearn
Created April 22, 2020 14:03
Show Gist options
  • Save thearn/27beebabb048f3df16f7be3e00bddc50 to your computer and use it in GitHub Desktop.
Save thearn/27beebabb048f3df16f7be3e00bddc50 to your computer and use it in GitHub Desktop.
import numpy as np
import openmdao.api as om
class SIRvec(om.ExplicitComponent):
"""Basic epidemiological infection model
S (suceptible), I (infected), R (recovered).
"""
def initialize(self):
self.options.declare('num_nodes', types=int)
self.options.declare('num_runs', types=int)
def setup(self):
nn = self.options['num_nodes']
nr = self.options['num_runs']
# States
self.add_input('S',
val=np.zeros((nr, nn)))
self.add_input('I',
val=np.zeros((nr, nn)))
self.add_input('R',
val=np.zeros((nr, nn)))
# ROCs
self.add_output('Sdot', val=np.zeros((nr, nn)))
self.add_output('Idot', val=np.zeros((nr, nn)))
self.add_output('Rdot', val=np.zeros((nr, nn)))
# Params
self.add_input('beta',
val = np.zeros((nr, nn)))
self.add_input('gamma',
val = np.zeros((nr, nn)))
#self.add_output('max_I', val=0.0)
#arange = np.arange(self.options['num_nodes'], dtype=int)
arange = np.arange(nn + nr, dtype=int)
self.declare_partials('Sdot', ['beta', 'gamma', 'S', 'I', 'R'], rows=arange, cols=arange)
self.declare_partials('Idot', ['beta', 'gamma', 'S', 'I', 'R'], rows=arange, cols=arange)
self.declare_partials('Rdot', ['gamma', 'I'], rows=arange, cols=arange)
#self.declare_partials('max_I', 'I')
def compute(self, inputs, outputs):
S = inputs['S']
I = inputs['I']
R = inputs['R']
gamma = inputs['gamma']
beta = inputs['beta']
# time derivatives of the states of an SIR model
# substitution dynamic control 'beta' for constant 'beta'.
outputs['Sdot'] = - beta * S * I
outputs['Idot'] = beta * S * I - gamma * I
outputs['Rdot'] = gamma * I
def compute_partials(self, inputs, jacobian):
S = inputs['S']
I = inputs['I']
R = inputs['R']
gamma = inputs['gamma']
beta = inputs['beta']
# derivatives of the ODE equations of state
jacobian['Sdot', 'S'] = (-I * beta).flatten()
jacobian['Sdot', 'I'] = (-S * beta).flatten()
jacobian['Sdot', 'beta'] = ( - S * I).flatten()
jacobian['Idot', 'gamma'] = (-I).flatten()
jacobian['Idot', 'S'] = I * (beta).flatten()
jacobian['Idot', 'I'] = S * (beta - gamma).flatten()
jacobian['Idot', 'beta'] = (S * I).flatten()
jacobian['Rdot', 'gamma'] = (I).flatten()
jacobian['Rdot', 'I'] = (gamma).flatten()
if __name__ == '__main__':
np.random.seed(0)
p = om.Problem()
p.model = om.Group()
nr = 4
n = 35
p.model.add_subsystem('test', SIRvec(num_nodes=n, num_runs=nr), promotes=['*'])
p.setup()
np.random.seed(0)
p['S'] = np.random.uniform(1, 1000, (nr, n))
p['I'] = np.random.uniform(1, 1000, (nr, n))
p['R'] = np.random.uniform(1, 1000, (nr, n))
p['beta'] = np.random.uniform(0, 2, (nr, n))
p['gamma'] = np.random.uniform(0, 2, (nr, n))
#p['t'] = np.linspace(0, 100, n)
p.run_model()
x = p.check_partials(compact_print=True)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment