Implement unit normalization

mira/dkg/model.py

@@ -21,11 +21,12 @@
 from fastapi.responses import FileResponse
 from pydantic import BaseModel, Field
 
-from mira.examples.sir import sir_bilayer, sir
+from mira.examples.sir import sir_bilayer, sir, sir_parameterized_init
 from mira.metamodel import (
     NaturalConversion, Template, ControlledConversion,
     stratify, Concept, ModelComparisonGraphdata, TemplateModelDelta,
     TemplateModel, Parameter, simplify_rate_laws, aggregate_parameters,
+    counts_to_dimensionless
 )
 from mira.modeling import Model
 from mira.modeling.askenet.petrinet import AskeNetPetriNetModel, ModelSpecification
@@ -87,6 +88,9 @@
 #: PetriNetModel json example
 petrinet_json = PetriNetModel(Model(sir)).to_pydantic()
 askenet_petrinet_json = AskeNetPetriNetModel(Model(sir)).to_pydantic()
+askenet_petrinet_json_units_values = AskeNetPetriNetModel(
+    Model(sir_parameterized_init)
+).to_pydantic()
 
 
 @model_blueprint.post(
@@ -221,6 +225,58 @@ def model_stratification(
     return template_model
 
 
+@model_blueprint.post(
+    "/counts_to_dimensionless_mira",
+    response_model=TemplateModel,
+    tags=["modeling"]
+)
+def dimension_transform(
+        query: Dict[str, Any] = Body(
+            ...,
+            example={
+                "model": sir_parameterized_init,
+                "counts_unit": "person",
+                "norm_factor": 1e5,
+            },
+        )
+):
+    """Convert all entity concentrations to dimensionless units"""
+    # convert to template model
+    tm_json = query.pop("model")
+    tm = TemplateModel.from_json(tm_json)
+    # The concepts should have their units' expressions as sympy.Expr,
+    # currently they are strings
+    tm_dimless = counts_to_dimensionless(tm=tm, **query)
+    return tm_dimless
+
+
+@model_blueprint.post(
+    "/counts_to_dimensionless_amr",
+    response_model=ModelSpecification,
+    tags=["modeling"]
+)
+def dimension_transform(
+        query: Dict[str, Any] = Body(
+            ...,
+            example={
+                "model": askenet_petrinet_json_units_values,
+                "counts_units": "persons",
+                "norm_factor": 1e5,
+            },
+        )
+):
+    """Convert all entity concentrations to dimensionless units"""
+    # convert to template model
+    amr_json = query.pop("model")
+    tm = template_model_from_askenet_json(amr_json)
+
+    # Create a dimensionless model
+    dimless_model = counts_to_dimensionless(tm=tm, **query)
+
+    # Transform back to askenet model
+    return AskeNetPetriNetModel(Model(dimless_model)).to_pydantic()
+
+
 @model_blueprint.get(
     "/biomodels/{model_id}",
     response_model=TemplateModel,

mira/examples/sir.py

@@ -6,7 +6,7 @@
 
 from mira.metamodel import ControlledConversion, NaturalConversion, \
     GroupedControlledConversion, TemplateModel, Initial, Parameter, \
-    safe_parse_expr
+    safe_parse_expr, Unit
 from .concepts import susceptible, infected, recovered, infected_symptomatic, \
     infected_asymptomatic
 
@@ -16,6 +16,7 @@
     "sir_2_city",
     "sir_bilayer",
     "sir_parameterized",
+    "sir_parameterized_init",
     "svir",
 ]
 
@@ -136,3 +137,22 @@
         infection_asymptomatic,
     ],
 )
+
+# SIR Parameterized Model with initial values and units, used as example in
+# docs and tests
+sir_parameterized_init = _d(sir_parameterized)
+sir_init_val_norm = 1e5
+for template in sir_parameterized_init.templates:
+    for concept in template.get_concepts():
+        concept.units = Unit(expression=sympy.Symbol('person'))
+sir_parameterized_init.initials['susceptible_population'].value = \
+    sir_init_val_norm - 1
+sir_parameterized_init.initials['infected_population'].value = 1
+sir_parameterized_init.initials['immune_population'].value = 0
+
+sir_parameterized_init.parameters['beta'].units = \
+    Unit(expression=1 / (sympy.Symbol('person') * sympy.Symbol('day')))
+old_beta = sir_parameterized_init.parameters['beta'].value
+
+for initial in sir_parameterized_init.initials.values():
+    initial.concept.units = Unit(expression=sympy.Symbol('person'))

mira/metamodel/__init__.py

@@ -6,4 +6,5 @@
 from .schema import *
 from .search import *
 from .ops import *
+from .units import *
 from .utils import *

mira/metamodel/ops.py

@@ -9,11 +9,16 @@
 
 from .template_model import TemplateModel, Initial, Parameter
 from .templates import *
+from .units import Unit, dimensionless_units
+from .utils import SympyExprStr
+
 
 __all__ = [
     "stratify",
     "simplify_rate_laws",
-    "aggregate_parameters"
+    "aggregate_parameters",
+    "get_term_roles",
+    "counts_to_dimensionless"
 ]
 
 
@@ -393,3 +398,69 @@ def get_term_roles(term, template, parameters):
         else:
             term_roles['other'].append(symbol.name)
     return dict(term_roles)
+
+
+def counts_to_dimensionless(tm: TemplateModel,
+                            counts_unit: str,
+                            norm_factor: float):
+    """Convert all entity concentrations to dimensionless units.
+
+    Parameters
+    ----------
+    tm :
+        A template model.
+    counts_unit :
+        The unit of the counts.
+    norm_factor :
+        The normalization factor to convert counts to concentration.
+
+    Returns
+    -------
+    :
+        A template model with all entity concentrations converted to
+        dimensionless units.
+    """
+    # Make a deepcopy up front so we don't change the original template model
+    tm = deepcopy(tm)
+    # Make a symbol of the counts unit for calculations
+    counts_unit_symbol = sympy.Symbol(counts_unit)
+
+    initials_normalized = set()
+    # First we normalize concepts and their initials
+    for template in tm.templates:
+        # Since concepts can be distributed across templates, we have to go
+        # template by template
+        for concept in template.get_concepts():
+            if concept.units:
+                # We figure out what the exponent of the counts unit is
+                # if it appears in the units of the concept
+                (coeff, exponent) = \
+                    concept.units.expression.args[0].as_coeff_exponent(counts_unit_symbol)
+                # If the exponent is other than zero then normalization is needed
+                if exponent:
+                    concept.units.expression = \
+                        SympyExprStr(concept.units.expression.args[0] /
+                                     (counts_unit_symbol ** exponent))
+                    # We not try to see if there is a corresponding initial condition
+                    # for the concept and if so, we normalize it as well
+                    if concept.name in tm.initials and concept.name not in initials_normalized:
+                        init = tm.initials[concept.name]
+                        if init.value is not None:
+                            init.value /= (norm_factor ** exponent)
+                            if init.concept.units:
+                                init.concept.units.expression = \
+                                    SympyExprStr(init.concept.units.expression.args[0] /
+                                                 (counts_unit_symbol ** exponent))
+                            initials_normalized.add(concept.name)
+    # Now we do the same for parameters
+    for p_name, p in tm.parameters.items():
+        if p.units:
+            (coeff, exponent) = \
+                p.units.expression.args[0].as_coeff_exponent(counts_unit_symbol)
+            if exponent:
+                p.units.expression = \
+                    SympyExprStr(p.units.expression.args[0] /
+                                 (counts_unit_symbol ** exponent))
+                p.value /= (norm_factor ** exponent)
+
+    return tm

mira/metamodel/template_model.py

@@ -3,14 +3,15 @@
 
 import datetime
 import sys
-from typing import List, Dict, Set, Optional, Mapping, Tuple
+from typing import List, Dict, Set, Optional, Mapping, Tuple, Any
 
 import networkx as nx
 import sympy
 from pydantic import BaseModel, Field
 
 from .templates import *
-from .utils import safe_parse_expr
+from .units import Unit
+from .utils import safe_parse_expr, SympyExprStr
 
 
 class Initial(BaseModel):
@@ -19,6 +20,23 @@ class Initial(BaseModel):
     concept: Concept
     value: float
 
+    class Config:
+        arbitrary_types_allowed = True
+        json_encoders = {
+            SympyExprStr: lambda e: str(e),
+        }
+        json_decoders = {
+            SympyExprStr: lambda e: sympy.parse_expr(e)
+        }
+
+    @classmethod
+    def from_json(cls, data: Dict[str, Any]) -> "Initial":
+        value = data.pop('value')
+        concept_json = data.pop('concept')
+        # Get Concept
+        concept = Concept.from_json(concept_json)
+        return cls(concept=concept, value=value)
+
 
 class Distribution(BaseModel):
     """A distribution of values for a parameter."""
@@ -363,24 +381,28 @@ def from_json(cls, data) -> "TemplateModel":
             for concept in template.get_concepts()
         }
 
-        initials = {
-            name: (
-                Initial(
+        initials = {}
+        for name, value in data.get('initials', {}).items():
+            if isinstance(value, float):
+                # If the data is just a float, upgrade it to
+                # a :class:`Initial` instance
+                initials[name] = Initial(
                     concept=concepts[name],
                     value=value,
                 )
-                # If the data is just a float, upgrade it to
-                # a :class:`Initial` instance
-                if isinstance(value, float)
+            else:
                 # If the data is not a float, assume it's JSON
-                # for a :class:`Initial` instance
-                else value
-            )
-            for name, value in data.get('initials', {}).items()
+                # for a :class:`Initial` instance and parse it to Initial
+                initials[name] = Initial.from_json(value)
+
+        # Handle parameters
+        parameters = {
+            par_key: Parameter.from_json(par_dict)
+            for par_key, par_dict in data.get('parameters', {}).items()
         }
 
         return cls(templates=templates,
-                   parameters=data.get('parameters', {}),
+                   parameters=parameters,
                    initials=initials,
                    annotations=data.get('annotations'))