thearn · April 22, 2020 14:03 · Apr 22, 2020
diff --git a/SIRvec.py b/SIRvec.py
@@ -0,0 +1,120 @@
+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)
+
+
+
+
+
+
+
No results found