Added notebook for hierarchical sir model and AMR

docs/source/beta_mean_cycle_sir_model.json

@@ -0,0 +1 @@
+{"header": {"name": "Model", "schema": "https://raw.githubusercontent.com/DARPA-ASKEM/Model-Representations/petrinet_v0.6/petrinet/petrinet_schema.json", "schema_name": "petrinet", "description": "Model", "model_version": "0.1"}, "properties": {}, "model": {"states": [{"id": "S", "name": "S", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "I", "name": "I", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "R", "name": "R", "grounding": {"identifiers": {}, "modifiers": {}}}], "transitions": [{"id": "t1", "input": ["I", "S"], "output": ["I", "I"], "properties": {"name": "t1"}}, {"id": "t2", "input": ["I"], "output": ["R"], "properties": {"name": "t2"}}]}, "semantics": {"ode": {"rates": [{"target": "t1", "expression": "I*S*beta", "expression_mathml": "<apply><times/><ci>I</ci><ci>S</ci><ci>beta</ci></apply>"}, {"target": "t2", "expression": "I*gamma", "expression_mathml": "<apply><times/><ci>I</ci><ci>gamma</ci></apply>"}], "initials": [], "parameters": [{"id": "beta", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean*gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "gamma", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "beta_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "beta_mean", "beta": "10"}}}, {"id": "gamma_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "10", "beta": "10"}}}], "observables": [], "time": {"id": "t"}}}, "metadata": {"annotations": {}}}

docs/source/beta_mean_gamma_cycle_sir_model.json

@@ -0,0 +1 @@
+{"header": {"name": "Model", "schema": "https://raw.githubusercontent.com/DARPA-ASKEM/Model-Representations/petrinet_v0.6/petrinet/petrinet_schema.json", "schema_name": "petrinet", "description": "Model", "model_version": "0.1"}, "properties": {}, "model": {"states": [{"id": "S", "name": "S", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "I", "name": "I", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "R", "name": "R", "grounding": {"identifiers": {}, "modifiers": {}}}], "transitions": [{"id": "t1", "input": ["I", "S"], "output": ["I", "I"], "properties": {"name": "t1"}}, {"id": "t2", "input": ["I"], "output": ["R"], "properties": {"name": "t2"}}]}, "semantics": {"ode": {"rates": [{"target": "t1", "expression": "I*S*beta", "expression_mathml": "<apply><times/><ci>I</ci><ci>S</ci><ci>beta</ci></apply>"}, {"target": "t2", "expression": "I*gamma", "expression_mathml": "<apply><times/><ci>I</ci><ci>gamma</ci></apply>"}], "initials": [], "parameters": [{"id": "beta", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean*gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "gamma", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean", "scale": "0.0100000000000000"}}}, {"id": "beta_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "10*gamma", "beta": "10"}}}, {"id": "gamma_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "10", "beta": "10"}}}], "observables": [], "time": {"id": "t"}}}, "metadata": {"annotations": {}}}

docs/source/gamma_mean_beta_mean_cycle_sir_model.json

@@ -0,0 +1 @@
+{"header": {"name": "Model", "schema": "https://raw.githubusercontent.com/DARPA-ASKEM/Model-Representations/petrinet_v0.6/petrinet/petrinet_schema.json", "schema_name": "petrinet", "description": "Model", "model_version": "0.1"}, "properties": {}, "model": {"states": [{"id": "S", "name": "S", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "I", "name": "I", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "R", "name": "R", "grounding": {"identifiers": {}, "modifiers": {}}}], "transitions": [{"id": "t1", "input": ["I", "S"], "output": ["I", "I"], "properties": {"name": "t1"}}, {"id": "t2", "input": ["I"], "output": ["R"], "properties": {"name": "t2"}}]}, "semantics": {"ode": {"rates": [{"target": "t1", "expression": "I*S*beta", "expression_mathml": "<apply><times/><ci>I</ci><ci>S</ci><ci>beta</ci></apply>"}, {"target": "t2", "expression": "I*gamma", "expression_mathml": "<apply><times/><ci>I</ci><ci>gamma</ci></apply>"}], "initials": [], "parameters": [{"id": "beta", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean*gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "gamma", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "beta_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "gamma_mean", "beta": "10"}}}, {"id": "gamma_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "10*beta_mean", "beta": "10"}}}], "observables": [], "time": {"id": "t"}}}, "metadata": {"annotations": {}}}

docs/source/hierarchical_sir_model.json

@@ -0,0 +1 @@
+{"header": {"name": "Model", "schema": "https://raw.githubusercontent.com/DARPA-ASKEM/Model-Representations/petrinet_v0.6/petrinet/petrinet_schema.json", "schema_name": "petrinet", "description": "Model", "model_version": "0.1"}, "properties": {}, "model": {"states": [{"id": "S", "name": "S", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "I", "name": "I", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "R", "name": "R", "grounding": {"identifiers": {}, "modifiers": {}}}], "transitions": [{"id": "t1", "input": ["I", "S"], "output": ["I", "I"], "properties": {"name": "t1"}}, {"id": "t2", "input": ["I"], "output": ["R"], "properties": {"name": "t2"}}]}, "semantics": {"ode": {"rates": [{"target": "t1", "expression": "I*S*beta", "expression_mathml": "<apply><times/><ci>I</ci><ci>S</ci><ci>beta</ci></apply>"}, {"target": "t2", "expression": "I*gamma", "expression_mathml": "<apply><times/><ci>I</ci><ci>gamma</ci></apply>"}], "initials": [], "parameters": [{"id": "beta", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean*gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "gamma", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "beta_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "beta_mean", "beta": "10"}}}, {"id": "gamma_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "10", "beta": "10"}}}], "observables": [], "time": {"id": "t"}}}, "metadata": {"annotations": {}}}

docs/source/multilevel_sir_model.json

@@ -0,0 +1 @@
+{"header": {"name": "Model", "schema": "https://raw.githubusercontent.com/DARPA-ASKEM/Model-Representations/petrinet_v0.6/petrinet/petrinet_schema.json", "schema_name": "petrinet", "description": "Model", "model_version": "0.1"}, "properties": {}, "model": {"states": [{"id": "S", "name": "S", "grounding": {"identifiers": {"ido": "0000514"}, "modifiers": {}}, "units": {"expression": "person", "expression_mathml": "<ci>person</ci>"}}, {"id": "I", "name": "I", "grounding": {"identifiers": {"ido": "0000511"}, "modifiers": {}}, "units": {"expression": "person", "expression_mathml": "<ci>person</ci>"}}, {"id": "R", "name": "R", "grounding": {"identifiers": {"ido": "0000592"}, "modifiers": {}}, "units": {"expression": "person", "expression_mathml": "<ci>person</ci>"}}], "transitions": [{"id": "t1", "input": ["I", "S"], "output": ["I", "I"], "properties": {"name": "t1"}}, {"id": "t2", "input": ["I"], "output": ["R"], "properties": {"name": "t2"}}]}, "semantics": {"ode": {"rates": [{"target": "t1", "expression": "I*S*beta", "expression_mathml": "<apply><times/><ci>I</ci><ci>S</ci><ci>beta</ci></apply>"}, {"target": "t2", "expression": "I*gamma", "expression_mathml": "<apply><times/><ci>I</ci><ci>gamma</ci></apply>"}], "initials": [{"target": "S", "expression": "1.0", "expression_mathml": "<cn>1.0</cn>"}, {"target": "I", "expression": "0.0", "expression_mathml": "<cn>0.0</cn>"}, {"target": "R", "expression": "0.0", "expression_mathml": "<cn>0.0</cn>"}], "parameters": [{"id": "beta", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean*gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "gamma", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "beta_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "1.0*gamma_mean", "beta": "10.0000000000000"}}}, {"id": "gamma_mean", "distribution": {"type": "Beta1", "parameters": {"alpha": "10.0000000000000", "beta": "10.0000000000000"}}}], "observables": [], "time": {"id": "t"}}}, "metadata": {"annotations": {}}}

docs/source/multilevel_sir_nodist_model.json

@@ -0,0 +1 @@
+{"header": {"name": "Model", "schema": "https://raw.githubusercontent.com/DARPA-ASKEM/Model-Representations/petrinet_v0.6/petrinet/petrinet_schema.json", "schema_name": "petrinet", "description": "Model", "model_version": "0.1"}, "properties": {}, "model": {"states": [{"id": "S", "name": "S", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "I", "name": "I", "grounding": {"identifiers": {}, "modifiers": {}}}, {"id": "R", "name": "R", "grounding": {"identifiers": {}, "modifiers": {}}}], "transitions": [{"id": "t1", "input": ["I", "S"], "output": ["I", "I"], "properties": {"name": "t1"}}, {"id": "t2", "input": ["I"], "output": ["R"], "properties": {"name": "t2"}}]}, "semantics": {"ode": {"rates": [{"target": "t1", "expression": "I*S*beta", "expression_mathml": "<apply><times/><ci>I</ci><ci>S</ci><ci>beta</ci></apply>"}, {"target": "t2", "expression": "I*gamma", "expression_mathml": "<apply><times/><ci>I</ci><ci>gamma</ci></apply>"}], "initials": [{"target": "S", "expression": "1.0", "expression_mathml": "<cn>1.0</cn>"}, {"target": "I", "expression": "0.0", "expression_mathml": "<cn>0.0</cn>"}, {"target": "R", "expression": "0.0", "expression_mathml": "<cn>0.0</cn>"}], "parameters": [{"id": "beta", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "beta_mean*gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "gamma", "distribution": {"type": "InverseGamma1", "parameters": {"shape": "gamma_mean", "scale": "0.0100000000000000"}}}, {"id": "beta_mean", "value": 1.0}, {"id": "gamma_mean", "value": 2.0}], "observables": [], "time": {"id": "t"}}}, "metadata": {"annotations": {}}}

pyciemss/compiled_dynamics.py

@@ -21,15 +21,16 @@ class CompiledDynamics(pyro.nn.PyroModule):
     def __init__(self, src, **kwargs):
         super().__init__()
         self.src = src
-
+        params = _compile_param_values(self.src)
         try:
             params = _compile_param_values(self.src)
         except Exception as e:
             raise ValueError(
                 "The model parameters could not be compiled. Please check the model definition."
             ) from e
 
-        for k, v in params.items():
+        for k in _sort_dependencies(self.src):
+            v = params[get_name(k)]
             if hasattr(self, get_name(k)):
                 continue
 
@@ -189,6 +190,11 @@ def _compile_deriv(src) -> Callable[..., Tuple[torch.Tensor]]:
     raise NotImplementedError
 
 
+@functools.singledispatch
+def _sort_dependencies(src) -> list:
+    raise NotImplementedError
+
+
 @functools.singledispatch
 def _compile_initial_state(src) -> Callable[..., Tuple[torch.Tensor]]:
     raise NotImplementedError

pyciemss/mira_integration/compiled_dynamics.py

@@ -3,13 +3,15 @@
 import numbers
 from typing import Callable, Dict, Optional, Tuple, TypeVar, Union
 
+import networkx as nx
 import mira
 import mira.metamodel
 import mira.modeling
 import mira.sources
 import mira.sources.amr
 import numpy
 import pyro
+import pyro.nn
 import sympy
 import sympytorch
 import torch
@@ -20,16 +22,50 @@
     _compile_initial_state,
     _compile_observables,
     _compile_param_values,
+    _sort_dependencies,
     eval_deriv,
     eval_initial_state,
     eval_observables,
     get_name,
 )
+
 from pyciemss.mira_integration.distributions import mira_distribution_to_pyro
 
 S = TypeVar("S")
 T = TypeVar("T")
 
+@_sort_dependencies.register(mira.modeling.Model)
+def sort_mira_dependencies(src: mira.modeling.Model) -> list:
+    """
+    Sort the model parameters of a MIRA TemplateModel by their distribution parameter dependencies.
+
+    Parameters
+    ----------
+    src : mira.modeling.Model
+        The MIRA Model to sort.
+
+    Returns
+    -------
+    list
+        A list of parameter names in the order in which they must be evaluated.
+    """
+    dependencies = nx.DiGraph()
+    for param_name, param_info in src.parameters.items():
+        param_name = get_name(param_info)
+        #if param_info.placeholder:
+        #    continue
+        param_dist = getattr(param_info, "distribution", None)
+
+        if param_dist is None:
+            dependencies.add_node(param_name)
+        else:
+            for k, v in param_dist.parameters.items():
+                # Check to see if the distribution parameters are sympy expressions
+                # and add their free symbols to the dependency graph
+                if isinstance(v, mira.metamodel.utils.SympyExprStr):
+                    for free_symbol in v.free_symbols:
+                        dependencies.add_edge(str(free_symbol), str(param_name))
+    return list(nx.topological_sort(dependencies))
 
 @_compile_deriv.register(mira.modeling.Model)
 def _compile_deriv_mira(src: mira.modeling.Model) -> Callable[..., Tuple[torch.Tensor]]:
@@ -84,29 +120,39 @@ def _compile_observables_mira(
 def _compile_param_values_mira(
     src: mira.modeling.Model,
 ) -> Dict[str, Union[torch.Tensor, pyro.nn.PyroParam, pyro.nn.PyroSample]]:
-    values = {}
-    for param_info in src.parameters.values():
-        param_name = get_name(param_info)
-
+    param_values = {}
+    sorted_dependencies = _sort_dependencies(src)
+    for param_name in sorted_dependencies:
+        param_info = src.parameters[param_name]
         if param_info.placeholder:
             continue
-
+            
         param_dist = getattr(param_info, "distribution", None)
         if param_dist is None:
-            param_value = param_info.value
+            param_value = float(param_info.value)
         else:
-            param_value = mira_distribution_to_pyro(param_dist)
+            idx = sorted_dependencies.index(param_name)
+            upstream_dependencies = sorted_dependencies[:idx]
+
+            def param_value(model: pyro.nn.PyroModule) -> torch.Tensor:
+                return mira_distribution_to_pyro(
+                    param_dist, {
+                        k: getattr(model, f"persistent_{k}")
+                        for k in upstream_dependencies
+                    })
 
         if isinstance(param_value, torch.nn.Parameter):
-            values[param_name] = pyro.nn.PyroParam(param_value)
-        elif isinstance(param_value, pyro.distributions.Distribution):
-            values[param_name] = pyro.nn.PyroSample(param_value)
+            param_values[param_name] = pyro.nn.PyroParam(param_value)
+        elif isinstance(param_value, (pyro.distributions.distribution.Distribution, Callable)):
+            param_values[param_name] = pyro.nn.PyroSample(param_value)
         elif isinstance(param_value, (numbers.Number, numpy.ndarray, torch.Tensor)):
-            values[param_name] = torch.as_tensor(param_value, dtype=torch.float32)
+            param_values[param_name] = torch.as_tensor(param_value, dtype=torch.float32)
+        elif isinstance(param_value, Callable):
+            param_values[param_name] = pyro.nn.PyroSample(param_value)
         else:
             raise TypeError(f"Unknown parameter type: {type(param_value)}")
 
-    return values
+    return param_values
 
 
 @eval_deriv.register(mira.modeling.Model)
@@ -198,6 +244,9 @@ def _get_name_mira_transition(trans: mira.modeling.Transition) -> str:
 def _get_name_mira_modelparameter(param: mira.modeling.ModelParameter) -> str:
     return str(param.key)
 
+@get_name.register
+def _get_name_mira_metamodel_parameter(param: mira.metamodel.Parameter) -> str:
+    return str(param.name)
 
 @get_name.register
 def _get_name_mira_model_observable(obs: mira.modeling.ModelObservable) -> str: