Merge pull request #358 from gyorilab/distribution_expr

Support expressions in parameter's of distributions

.github/workflows/tests.yml

@@ -22,7 +22,7 @@ jobs:
           pip install "tox<4.0.0"
       - name: Test with pytest
         run: |
-          export MIRA_REST_URL=http://34.230.33.149:8771
+          export MIRA_REST_URL=http://mira-epi-dkg-lb-c7b58edea41524e6.elb.us-east-1.amazonaws.com:8771
           tox -e py
 #      - name: Upload coverage report to codecov
 #        uses: codecov/codecov-action@v1

mira/metamodel/schema.json

@@ -946,15 +946,23 @@
       "properties": {
         "type": {
           "title": "Type",
-          "description": "The type of distribution, e.g. 'uniform', 'normal', etc.",
+          "description": "The type of distribution as provided by ProbOnto e.g. 'StandardUniform1', 'Beta1', etc.",
           "type": "string"
         },
         "parameters": {
           "title": "Parameters",
-          "description": "The parameters of the distribution.",
+          "description": "The parameters of the distribution keyed by parameter names controlled by ProbOnto and values that are either floating point values or symbolic expressions over other parameters.",
           "type": "object",
           "additionalProperties": {
-            "type": "number"
+            "anyOf": [
+              {
+                "type": "number"
+              },
+              {
+                "type": "string",
+                "example": "2*x"
+              }
+            ]
           }
         }
       },

mira/metamodel/template_model.py

@@ -104,10 +104,14 @@ class Distribution(BaseModel):
     """A distribution of values for a parameter."""
 
     type: str = Field(
-        description="The type of distribution, e.g. 'uniform', 'normal', etc."
+        description="The type of distribution as provided by ProbOnto "
+                    "e.g. 'StandardUniform1', 'Beta1', etc."
     )
-    parameters: Dict[str, float] = Field(
-        description="The parameters of the distribution."
+    parameters: Dict[str, Union[float, SympyExprStr]] = Field(
+        description="The parameters of the distribution keyed by parameter names "
+                    "controlled by ProbOnto and values that are either floating "
+                    "point values or symbolic expressions over other "
+                    "parameters."
     )
 
 

mira/modeling/amr/petrinet.py

@@ -12,7 +12,7 @@
 
 from pydantic import BaseModel, Field
 
-from mira.metamodel import expression_to_mathml, TemplateModel
+from mira.metamodel import expression_to_mathml, TemplateModel, SympyExprStr
 from mira.sources.amr import sanity_check_amr
 
 from .. import Model
@@ -188,9 +188,16 @@ def __init__(self, model: Model):
             elif param.distribution.type is None:
                 logger.warning("can not add distribution without type: %s", param.distribution)
             else:
+                serialized_distr_parameters = {}
+                for param_key, param_value in param.distribution.parameters.items():
+                    if isinstance(param_value, SympyExprStr):
+                        serialized_distr_parameters[param_key] = \
+                            str(param_value.args[0])
+                    else:
+                        serialized_distr_parameters[param_key] = param_value
                 param_dict['distribution'] = {
                     'type': param.distribution.type,
-                    'parameters': param.distribution.parameters,
+                    'parameters': serialized_distr_parameters,
                 }
             if param.concept and param.concept.units:
                 param_dict['units'] = {

mira/sources/amr/petrinet.py

@@ -109,15 +109,16 @@ def template_model_from_amr_json(model_json) -> TemplateModel:
     # }
     ode_semantics = model_json.get("semantics", {}).get("ode", {})
     symbols = {state_id: sympy.Symbol(state_id) for state_id in concepts}
-    mira_parameters = {}
+
+    # We first make symbols for all the parameters
     for parameter in ode_semantics.get('parameters', []):
-        mira_parameters[parameter['id']] = parameter_to_mira(parameter)
         symbols[parameter['id']] = sympy.Symbol(parameter['id'])
 
-    param_values = {
-        p['id']: p['value'] for p in ode_semantics.get('parameters', [])
-        if p.get('value') is not None
-    }
+    # We then process the parameters into MIRA Parameter objects
+    mira_parameters = {}
+    for parameter in ode_semantics.get('parameters', []):
+        mira_parameters[parameter['id']] = \
+            parameter_to_mira(parameter, param_symbols=symbols)
 
     # Next we process initial conditions
     initials = {}

mira/sources/amr/regnet.py

@@ -10,7 +10,7 @@
 import requests
 
 from mira.metamodel import *
-from mira.sources.util import get_sympy
+from mira.sources.util import get_sympy, parameter_to_mira
 
 
 def model_from_url(url: str) -> TemplateModel:
@@ -83,11 +83,14 @@ def template_model_from_amr_json(model_json) -> TemplateModel:
     # Next, we capture all symbols in the model, including states and
     # parameters. We also extract parameters at this point.
     symbols = {state_id: sympy.Symbol(state_id) for state_id in concepts}
-    mira_parameters = {}
     for parameter in model.get('parameters', []):
-        mira_parameters[parameter['id']] = parameter_to_mira(parameter)
         symbols[parameter['id']] = sympy.Symbol(parameter['id'])
 
+    mira_parameters = {}
+    for parameter in model.get('parameters', []):
+        mira_parameters[parameter['id']] = \
+            parameter_to_mira(parameter, param_symbols=symbols)
+
     # Next we process any intrinsic positive/negative growth
     # at the vertex level into templates
     templates = []
@@ -213,12 +216,3 @@ def vertex_to_template(vertex, concept):
         template = NaturalDegradation(subject=concept)
     template.set_mass_action_rate_law(rate_constant)
     return template
-
-
-def parameter_to_mira(parameter):
-    """Return a MIRA parameter from a parameter"""
-    distr = Distribution(**parameter['distribution']) \
-        if parameter.get('distribution') else None
-    return Parameter(name=parameter['id'],
-                     value=parameter.get('value'),
-                     distribution=distr)

mira/sources/amr/stockflow.py

@@ -42,13 +42,17 @@ def template_model_from_amr_json(model_json) -> TemplateModel:
         all_stocks.add(stock['id'])
         symbols[stock['id']] = sympy.Symbol(stock['id'])
 
-    # Process parameters
+    # Process parameters, first to get all symbols, then
+    # processing the parameters to get the MIRA parameters
     ode_semantics = model_json.get("semantics", {}).get("ode", {})
-    mira_parameters = {}
     for parameter in ode_semantics.get('parameters', []):
-        mira_parameters[parameter['id']] = parameter_to_mira(parameter)
         symbols[parameter['id']] = sympy.Symbol(parameter['id'])
 
+    mira_parameters = {}
+    for parameter in ode_semantics.get('parameters', []):
+        mira_parameters[parameter['id']] = \
+            parameter_to_mira(parameter, param_symbols=symbols)
+
     # Process auxiliaries
     aux_expressions = {}
     for auxiliary in auxiliaries:

mira/sources/util.py

@@ -161,7 +161,7 @@ def transition_to_templates(
             )
 
 
-def parameter_to_mira(parameter) -> Parameter:
+def parameter_to_mira(parameter, param_symbols=None) -> Parameter:
     """
     Return a MIRA parameter from a mapping of MIRA Parameter attributes to
     values.
@@ -170,17 +170,33 @@ def parameter_to_mira(parameter) -> Parameter:
     ----------
     parameter : Dict[str,Any]
         A mapping containing MIRA Parameter attributes to values.
+    param_symbols : Optional[Dict]
+        An optional dict of all parameter symbols that are needed if expressions
+        are used in parameter distributions so that these can be
+        recognized as symbols.
 
     Returns
     -------
     :
         The corresponding MIRA Parameter.
     """
-    distr = (
-        Distribution(**parameter["distribution"])
-        if parameter.get("distribution")
-        else None
-    )
+    distr_json = parameter.get("distribution")
+    if distr_json:
+        distr_type = distr_json.get("type") if distr_json else None
+        # We need to check for symbolic expressions in parameters
+        processed_distr_parameters = {}
+        for param_key, param_value in distr_json.get("parameters", {}).items():
+            if isinstance(param_value, float):
+                processed_distr_parameters[param_key] = param_value
+            else:
+                processed_distr_parameters[param_key] = \
+                    safe_parse_expr(param_value)
+        distr = Distribution(
+            type=distr_type,
+            parameters=processed_distr_parameters,
+        )
+    else:
+        distr = None
     data = {
         "name": parameter["id"],
         "display_name": parameter.get("name"),

tests/test_distributions.py

@@ -0,0 +1,65 @@
+import sympy
+
+from mira.metamodel import *
+from mira.modeling import Model
+from mira.modeling.amr.petrinet import template_model_to_petrinet_json
+from mira.sources.amr import model_from_json
+
+
+def test_distribution_expressions():
+    beta_mean = Parameter(name='beta_mean',
+                          distribution=Distribution(type="Beta1",
+                                                    parameters={'alpha': sympy.Integer(1),
+                                                                'beta': sympy.Integer(10)}))
+    gamma_mean = Parameter(name='gamma_mean',
+                           distribution=Distribution(type="Beta1",
+                                                     parameters={'alpha': sympy.Integer(10),
+                                                                 'beta': sympy.Integer(10)}))
+    beta = Parameter(name='beta',
+                     distribution=Distribution(type="InverseGamma1",
+                                               parameters={'shape': sympy.Symbol('beta_mean'),
+                                                           'scale': sympy.Float(0.01)}))
+    gamma = Parameter(name='gamma',
+                      distribution=Distribution(type="InverseGamma1",
+                                                parameters={'shape': sympy.Symbol('gamma_mean'),
+                                                            'scale': sympy.Float(0.01)}))
+
+    # Make an SIR model with beta and gamma in rate laws
+    sir_model = TemplateModel(
+        templates=[
+            ControlledConversion(
+                subject=Concept(name='S'),
+                outcome=Concept(name='I'),
+                controller=Concept(name='I'),
+                rate_law=sympy.Symbol('S') * sympy.Symbol('I') * sympy.Symbol('beta')
+            ),
+            NaturalConversion(
+                subject=Concept(name='I'),
+                outcome=Concept(name='R'),
+                rate_law=sympy.Symbol('I') * sympy.Symbol('gamma')
+            ),
+        ],
+        parameters={
+            'beta': beta,
+            'gamma': gamma,
+            'beta_mean': beta_mean,
+            'gamma_mean': gamma_mean,
+            }
+    )
+
+    model = Model(sir_model)
+    pn_json = template_model_to_petrinet_json(sir_model)
+    params = pn_json['semantics']['ode']['parameters']
+    assert {p['id'] for p in params} == \
+        {'beta_mean', 'gamma_mean', 'beta', 'gamma'}
+    beta = [p for p in params if p['id'] == 'beta'][0]
+    assert beta['distribution']['type'] == 'InverseGamma1'
+    assert beta['distribution']['parameters']['shape'] == 'beta_mean'
+
+    # Now read the model back and check if it is deserialized
+    tm = model_from_json(pn_json)
+    assert tm.parameters['beta'].distribution.type == 'InverseGamma1'
+    assert isinstance(tm.parameters['beta'].distribution.parameters['shape'],
+                      SympyExprStr)
+    assert tm.parameters['beta'].distribution.parameters['shape'].args[0] == \
+        sympy.Symbol('beta_mean')