Merge pull request #313 from gyorilab/observables_fix

Fix observable handling in stock and flow
gyorilab · Mar 25, 2024 · 469f077 · 469f077
2 parents 75da2b4 + ca415d4
commit 469f077
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 6 deletions.
diff --git a/mira/modeling/amr/stockflow.py b/mira/modeling/amr/stockflow.py
@@ -208,12 +208,12 @@ def to_json(self):
                 'flows': self.flows,
                 'stocks': self.stocks,
                 'auxiliaries': self.auxiliaries,
+                'observables': self.observables,
                 'links': self.links
             },
             'semantics': {'ode': {
                 'parameters': self.parameters,
                 'initials': self.initials,
-                'observables': self.observables,
                 'time': self.time if self.time else {'id': 't'}
             }},
             'metadata': self.metadata,

diff --git a/mira/sources/amr/stockflow.py b/mira/sources/amr/stockflow.py
@@ -78,7 +78,7 @@ def template_model_from_amr_json(model_json) -> TemplateModel:
         observable = Observable(name=observable['id'],
                                 expression=observable_expr,
                                 display_name=observable.get('name'))
-        observables[observable.name] = observable
+        tm_observables[observable.name] = observable
 
     time = ode_semantics.get("time")
     if time:
@@ -139,7 +139,7 @@ def template_model_from_amr_json(model_json) -> TemplateModel:
     return TemplateModel(templates=templates,
                          parameters=mira_parameters,
                          initials=initials,
-                         observable=tm_observables,
+                         observables=tm_observables,
                          annotations=anns,
                          time=model_time)
 

diff --git a/tests/test_modeling/test_regnet.py b/tests/test_modeling/test_regnet.py
@@ -4,7 +4,7 @@
 
 from mira.modeling.amr.regnet import AMRRegNetModel, \
     template_model_to_regnet_json
-
+from mira.sources.amr.regnet import template_model_from_amr_json
 
 def test_regnet_end_to_end():
     url = 'https://raw.githubusercontent.com/DARPA-ASKEM/' \
@@ -89,3 +89,13 @@ def test_custom_rates():
     assert r['initial'] == 3.0
     assert r['rate_constant'] == 'V'
     assert r['sign'] is False
+
+
+def test_regnet_observable_roundtrip():
+    import sympy
+    t = NaturalDegradation(subject=Concept(name='x'))
+    observable = Observable(name='x', expression=sympy.Symbol('x'))
+    tm = TemplateModel(templates=[t], observables={'xo': observable})
+    rj = template_model_to_regnet_json(tm)
+    tm2 = template_model_from_amr_json(rj)
+    assert tm2.observables
diff --git a/tests/test_modeling/test_stockflow.py b/tests/test_modeling/test_stockflow.py
@@ -2,8 +2,7 @@
 from copy import deepcopy as _d
 from mira.sources.amr.stockflow import *
 from mira.modeling.amr.stockflow import *
-from mira.metamodel.io import expression_to_mathml
-from mira.metamodel.utils import safe_parse_expr
+from mira.metamodel import *
 
 stockflow_example = 'https://raw.githubusercontent.com/DARPA-ASKEM/' \
                     'Model-Representations/7f5e377225675259baa6486c64102f559edfd79f/stockflow/examples/sir.json'
@@ -95,3 +94,13 @@ def list_of_dict_to_dict(l):
         assert initial_id in initials_dict
         assert initial['expression'] == initials_dict[initial_id]['expression']
         assert initial['expression_mathml'] == initials_dict[initial_id]['expression_mathml']
+
+
+def test_stockflow_observable_roundtrip():
+    import sympy
+    t = NaturalDegradation(subject=Concept(name='x'))
+    observable = Observable(name='x', expression=sympy.Symbol('x'))
+    tm = TemplateModel(templates=[t], observables={'xo': observable})
+    rj = template_model_to_stockflow_json(tm)
+    tm2 = template_model_from_amr_json(rj)
+    assert tm2.observables