Merge pull request #266 from gyorilab/add-budyko-sellers-decapode

Handle digesting budyko-sellers decapode.

mira/examples/decapodes/decapodes_examples.py

@@ -11,6 +11,11 @@
     "/sa_climate_modeling/examples/climate/ice_dynamics.json"
 )
 
+BUDYKO_SELLERS_EXAMPLE_JSON_URL = (
+    "https://raw.githubusercontent.com/AlgebraicJulia/Decapodes.jl"
+    "/gh-pages/dev/budyko_sellers.json"
+)
+
 HERE = Path(__file__).parent
 EXAMPLES = HERE / "decapodes_vs_decaexpr_composite"
 DECAPODE_OSCILLATOR = EXAMPLES / "d1_oscillator_decapode.json"
@@ -42,6 +47,29 @@ def get_ice_decapode_json():
     return requests.get(ICE_DYNAMICS_EXAMPLE_JSON_URL).json()
 
 
+def get_budyko_sellers_example() -> Decapode:
+    """Returns the budyko sellers decapode example
+
+    Returns
+    -------
+    :
+        The budyko-sellers example as a Decapode object
+    """
+    res_json = requests.get(BUDYKO_SELLERS_EXAMPLE_JSON_URL).json()
+    return process_decapode(res_json)
+
+
+def get_budyko_seller_decapode_json():
+    """Return the budyko sellers decapode example as json
+
+    Returns
+    -------
+    : JSON
+        The budyko-sellers example as a json object
+    """
+    return requests.get(BUDYKO_SELLERS_EXAMPLE_JSON_URL).json()
+
+
 def get_oscillator_decapode() -> Decapode:
     """Return the oscillator decapode example
 

mira/metamodel/decapodes.py

@@ -29,15 +29,15 @@ def expand_variable(variable, var_produced_map):
     elif isinstance(var_prod, Op2):
         arg1 = expand_variable(var_prod.proj1, var_produced_map)
         arg2 = expand_variable(var_prod.proj2, var_produced_map)
-        if var_prod.function_str == "/":
+        if var_prod.function_str == "/" or var_prod.function_str == "./":
             return arg1 / arg2
-        elif var_prod.function_str == "*":
+        elif var_prod.function_str == "*" or var_prod.function_str == ".*":
             return arg1 * arg2
-        elif var_prod.function_str == "+":
+        elif var_prod.function_str == "+" or var_prod.function_str == ".+":
             return arg1 + arg2
-        elif var_prod.function_str == "-":
+        elif var_prod.function_str == "-" or var_prod.function_str == ".-":
             return arg1 - arg2
-        elif var_prod.function_str == "^":
+        elif var_prod.function_str == "^" or var_prod.function_str == ".^":
             return arg1**arg2
         else:
             return sympy.Function(var_prod.function_str)(arg1, arg2)

tests/test_decapodes.py

@@ -4,7 +4,9 @@
     get_friction_decapode,
     get_friction_decapode_json,
     get_ice_dynamics_example,
-    get_ice_decapode_json
+    get_ice_decapode_json,
+    get_budyko_sellers_example,
+    get_budyko_seller_decapode_json,
 )
 
 from mira.metamodel.decapodes import Variable, RootVariable
@@ -18,28 +20,30 @@ def test_oscillator_decapode():
     json_decapode = get_oscillator_decapode_json()
 
     assert len(decapode.variables) == 6
-    for var_json, var_obj in zip(json_decapode['Var'],
-                                 decapode.variables.values()):
-        assert var_json['name'] == var_obj.name
-        assert var_json['type'] == var_obj.type
+    for var_json, var_obj in zip(
+        json_decapode["Var"], decapode.variables.values()
+    ):
+        assert var_json["name"] == var_obj.name
+        assert var_json["type"] == var_obj.type
 
     assert len(decapode.tangent_variables) == 2
-    for tvar_json, tvar_obj in zip(json_decapode['TVar'],
-                                   decapode.tangent_variables.values()):
-        assert tvar_obj.incl_var == decapode.variables[tvar_json['incl']]
+    for tvar_json, tvar_obj in zip(
+        json_decapode["TVar"], decapode.tangent_variables.values()
+    ):
+        assert tvar_obj.incl_var == decapode.variables[tvar_json["incl"]]
 
     assert len(decapode.op1s) == 2
-    for op1_json, op1_obj in zip(json_decapode['Op1'], decapode.op1s.values()):
-        assert decapode.variables[op1_json['src']] == op1_obj.src
-        assert decapode.variables[op1_json['tgt']] == op1_obj.tgt
-        assert op1_json['op1'] == op1_obj.function_str
+    for op1_json, op1_obj in zip(json_decapode["Op1"], decapode.op1s.values()):
+        assert decapode.variables[op1_json["src"]] == op1_obj.src
+        assert decapode.variables[op1_json["tgt"]] == op1_obj.tgt
+        assert op1_json["op1"] == op1_obj.function_str
 
     assert len(decapode.op2s) == 2
-    for op2_json, op2_obj in zip(json_decapode['Op2'], decapode.op2s.values()):
-        assert decapode.variables[op2_json['proj1']] == op2_obj.proj1
-        assert decapode.variables[op2_json['proj2']] == op2_obj.proj2
-        assert decapode.variables[op2_json['res']] == op2_obj.res
-        assert op2_json['op2'] == op2_obj.function_str
+    for op2_json, op2_obj in zip(json_decapode["Op2"], decapode.op2s.values()):
+        assert decapode.variables[op2_json["proj1"]] == op2_obj.proj1
+        assert decapode.variables[op2_json["proj2"]] == op2_obj.proj2
+        assert decapode.variables[op2_json["res"]] == op2_obj.res
+        assert op2_json["op2"] == op2_obj.function_str
 
     assert len(decapode.summations) == 0
 
@@ -48,9 +52,11 @@ def test_oscillator_decapode():
             assert isinstance(var, RootVariable)
             assert len(var.expression) == 2
             assert var.expression[0] == DERIVATIVE_FUNCTION(
-                DERIVATIVE_FUNCTION(sympy.Symbol('X')))
-            assert var.expression[1] == sympy.Symbol('-1') * sympy.sympify(
-                'X*k')
+                DERIVATIVE_FUNCTION(sympy.Symbol("X"))
+            )
+            assert var.expression[1] == sympy.Symbol("-1") * sympy.sympify(
+                "X*k"
+            )
         else:
             assert isinstance(var, Variable)
             if isinstance(var.expression, sympy.Symbol):
@@ -61,56 +67,61 @@ def test_oscillator_decapode():
             elif var.id == 2:
                 DERIVATIVE_FUNCTION(var.name)
             elif var.id == 4:
-                assert var.expression == sympy.Symbol('-1') * sympy.Symbol(
-                    'k')
+                assert var.expression == sympy.Symbol("-1") * sympy.Symbol("k")
 
 
 def test_friction_decapode():
     decapode = get_friction_decapode()
     json_decapode = get_friction_decapode_json()
 
     assert len(decapode.variables) == 9
-    for var_json, var_obj in zip(json_decapode['Var'],
-                                 decapode.variables.values()):
-        assert var_json['name'] == var_obj.name
-        assert var_json['type'] == var_obj.type
+    for var_json, var_obj in zip(
+        json_decapode["Var"], decapode.variables.values()
+    ):
+        assert var_json["name"] == var_obj.name
+        assert var_json["type"] == var_obj.type
 
     assert len(decapode.tangent_variables) == 1
-    for tvar_json, tvar_obj in zip(json_decapode['TVar'],
-                                   decapode.tangent_variables.values()):
-        assert tvar_obj.incl_var == decapode.variables[tvar_json['incl']]
+    for tvar_json, tvar_obj in zip(
+        json_decapode["TVar"], decapode.tangent_variables.values()
+    ):
+        assert tvar_obj.incl_var == decapode.variables[tvar_json["incl"]]
         assert isinstance(tvar_obj.incl_var, RootVariable)
 
     assert len(decapode.op1s) == 1
-    for op1_json, op1_obj in zip(json_decapode['Op1'], decapode.op1s.values()):
-        assert decapode.variables[op1_json['src']] == op1_obj.src
-        assert decapode.variables[op1_json['tgt']] == op1_obj.tgt
-        assert op1_json['op1'] == op1_obj.function_str
+    for op1_json, op1_obj in zip(json_decapode["Op1"], decapode.op1s.values()):
+        assert decapode.variables[op1_json["src"]] == op1_obj.src
+        assert decapode.variables[op1_json["tgt"]] == op1_obj.tgt
+        assert op1_json["op1"] == op1_obj.function_str
 
     assert len(decapode.op2s) == 3
-    for op2_json, op2_obj in zip(json_decapode['Op2'], decapode.op2s.values()):
-        assert decapode.variables[op2_json['proj1']] == op2_obj.proj1
-        assert decapode.variables[op2_json['proj2']] == op2_obj.proj2
-        assert decapode.variables[op2_json['res']] == op2_obj.res
-        assert op2_json['op2'] == op2_obj.function_str
+    for op2_json, op2_obj in zip(json_decapode["Op2"], decapode.op2s.values()):
+        assert decapode.variables[op2_json["proj1"]] == op2_obj.proj1
+        assert decapode.variables[op2_json["proj2"]] == op2_obj.proj2
+        assert decapode.variables[op2_json["res"]] == op2_obj.res
+        assert op2_json["op2"] == op2_obj.function_str
 
     assert len(decapode.summations) == 1
-    assert decapode.summations[1].sum == decapode.variables[json_decapode["Σ"][
-        0]['sum']]
+    assert (
+        decapode.summations[1].sum
+        == decapode.variables[json_decapode["Σ"][0]["sum"]]
+    )
 
     assert len(decapode.summations[1].summands) == 2
-    for summand_json, summand_var in zip(json_decapode['Summand'],
-                                         decapode.summations[1].summands):
-        assert decapode.variables[summand_json['summand']] == summand_var
+    for summand_json, summand_var in zip(
+        json_decapode["Summand"], decapode.summations[1].summands
+    ):
+        assert decapode.variables[summand_json["summand"]] == summand_var
 
     for var in decapode.variables.values():
         if var.id == 6:
             assert isinstance(var, RootVariable)
             assert len(var.expression) == 2
-            assert var.expression[0] == DERIVATIVE_FUNCTION(sympy.Symbol('Q'))
+            assert var.expression[0] == DERIVATIVE_FUNCTION(sympy.Symbol("Q"))
             assert var.expression[1] == (
-                sympy.sympify('V*κ') + sympy.Symbol('λ') *
-                (sympy.Symbol('Q') - sympy.Symbol('Q₀')))
+                sympy.sympify("V*κ")
+                + sympy.Symbol("λ") * (sympy.Symbol("Q") - sympy.Symbol("Q₀"))
+            )
         else:
             assert isinstance(var, Variable)
             if isinstance(var.expression, sympy.Symbol):
@@ -119,65 +130,116 @@ def test_friction_decapode():
             # Check for expressions of non-base level variables that aren't
             # root variables
             elif var.id == 7:
-                assert var.expression == sympy.sympify('V*κ')
+                assert var.expression == sympy.sympify("V*κ")
             elif var.id == 8:
-                assert var.expression == sympy.Symbol('λ') * (sympy.Symbol(
-                    'Q') - sympy.Symbol('Q₀'))
+                assert var.expression == sympy.Symbol("λ") * (
+                    sympy.Symbol("Q") - sympy.Symbol("Q₀")
+                )
             elif var.id == 9:
-                assert var.expression == (sympy.Symbol(
-                    'Q') - sympy.Symbol('Q₀'))
+                assert var.expression == (
+                    sympy.Symbol("Q") - sympy.Symbol("Q₀")
+                )
 
 
 def test_ice_decapode():
-    json_decapode = get_ice_decapode_json()
-
     decapode = get_ice_dynamics_example()
+    json_decapode = get_ice_decapode_json()
 
     assert len(decapode.variables) == 30
-    for var_json, var_obj in zip(json_decapode['Var'],
-                                 decapode.variables.values()):
-        assert var_json['name'] == var_obj.name
-        assert var_json['type'] == var_obj.type
+    for var_json, var_obj in zip(
+        json_decapode["Var"], decapode.variables.values()
+    ):
+        assert var_json["name"] == var_obj.name
+        assert var_json["type"] == var_obj.type
 
     assert len(decapode.tangent_variables) == 1
-    for tvar_json, tvar_obj in zip(json_decapode['TVar'],
-                                   decapode.tangent_variables.values()):
-        assert tvar_obj.incl_var == decapode.variables[tvar_json['incl']]
+    for tvar_json, tvar_obj in zip(
+        json_decapode["TVar"], decapode.tangent_variables.values()
+    ):
+        assert tvar_obj.incl_var == decapode.variables[tvar_json["incl"]]
         assert isinstance(tvar_obj.incl_var, RootVariable)
 
     assert len(decapode.op1s) == 10
-    for op1_json, op1_obj in zip(json_decapode['Op1'], decapode.op1s.values()):
-        assert decapode.variables[op1_json['src']] == op1_obj.src
-        assert decapode.variables[op1_json['tgt']] == op1_obj.tgt
-        assert op1_json['op1'] == op1_obj.function_str
+    for op1_json, op1_obj in zip(json_decapode["Op1"], decapode.op1s.values()):
+        assert decapode.variables[op1_json["src"]] == op1_obj.src
+        assert decapode.variables[op1_json["tgt"]] == op1_obj.tgt
+        assert op1_json["op1"] == op1_obj.function_str
 
     assert len(decapode.op2s) == 11
-    for op2_json, op2_obj in zip(json_decapode['Op2'], decapode.op2s.values()):
-        assert decapode.variables[op2_json['proj1']] == op2_obj.proj1
-        assert decapode.variables[op2_json['proj2']] == op2_obj.proj2
-        assert decapode.variables[op2_json['res']] == op2_obj.res
-        assert op2_json['op2'] == op2_obj.function_str
+    for op2_json, op2_obj in zip(json_decapode["Op2"], decapode.op2s.values()):
+        assert decapode.variables[op2_json["proj1"]] == op2_obj.proj1
+        assert decapode.variables[op2_json["proj2"]] == op2_obj.proj2
+        assert decapode.variables[op2_json["res"]] == op2_obj.res
+        assert op2_json["op2"] == op2_obj.function_str
 
     assert len(decapode.summations) == 2
-    assert decapode.summations[1].sum == decapode.variables[json_decapode["Σ"][
-        0]['sum']]
+    assert (
+        decapode.summations[1].sum
+        == decapode.variables[json_decapode["Σ"][0]["sum"]]
+    )
 
     assert len(decapode.summations[1].summands) == 2
-    for summand_json, summand_var in zip(json_decapode['Summand'],
-                                         decapode.summations[1].summands):
-        assert decapode.variables[summand_json['summand']] == summand_var
+    for summand_json, summand_var in zip(
+        json_decapode["Summand"], decapode.summations[1].summands
+    ):
+        assert decapode.variables[summand_json["summand"]] == summand_var
 
     assert len(decapode.summations[2].summands) == 2
-    for summand_json, summand_var in zip(json_decapode['Summand'],
-                                         decapode.summations[1].summands):
-        assert decapode.variables[summand_json['summand']] == summand_var
+    for summand_json, summand_var in zip(
+        json_decapode["Summand"], decapode.summations[1].summands
+    ):
+        assert decapode.variables[summand_json["summand"]] == summand_var
 
     for var in decapode.variables.values():
         if var.id == 4:
             assert isinstance(var, RootVariable)
             assert len(var.expression) == 2
-            assert var.expression[0] == DERIVATIVE_FUNCTION(sympy.Symbol('h'))
+            assert var.expression[0] == DERIVATIVE_FUNCTION(sympy.Symbol("h"))
         else:
             assert isinstance(var, Variable)
             if isinstance(var.expression, sympy.Symbol):
                 assert var.expression == sympy.Symbol(var.name)
+
+
+def test_budyko_sellers_decapode():
+    decapode = get_budyko_sellers_example()
+    json_decapode = get_budyko_seller_decapode_json()
+
+    assert len(decapode.variables) == 25
+    for var_json, var_obj in zip(
+        json_decapode["Var"], decapode.variables.values()
+    ):
+        assert var_json["name"] == var_obj.name
+        assert var_json["type"] == var_obj.type
+
+    assert len(decapode.tangent_variables) == 1
+    for tvar_json, tvar_obj in zip(
+        json_decapode["TVar"], decapode.tangent_variables.values()
+    ):
+        assert tvar_obj.incl_var == decapode.variables[tvar_json["incl"]]
+        assert isinstance(tvar_obj.incl_var, RootVariable)
+
+    assert len(decapode.op1s) == 5
+    for op1_json, op1_obj in zip(json_decapode["Op1"], decapode.op1s.values()):
+        assert decapode.variables[op1_json["src"]] == op1_obj.src
+        assert decapode.variables[op1_json["tgt"]] == op1_obj.tgt
+        assert op1_json["op1"] == op1_obj.function_str
+
+    assert len(decapode.op2s) == 10
+    for op2_json, op2_obj in zip(json_decapode["Op2"], decapode.op2s.values()):
+        assert decapode.variables[op2_json["proj1"]] == op2_obj.proj1
+        assert decapode.variables[op2_json["proj2"]] == op2_obj.proj2
+        assert decapode.variables[op2_json["res"]] == op2_obj.res
+        assert op2_json["op2"] == op2_obj.function_str
+
+    assert len(decapode.summations) == 1
+    assert (
+        decapode.summations[1].sum
+        == decapode.variables[json_decapode["Σ"][0]["sum"]]
+    )
+
+    assert len(decapode.summations[1].summands) == 2
+    for summand_json, summand_var in zip(
+        json_decapode["Summand"], decapode.summations[1].summands
+    ):
+        assert decapode.variables[summand_json["summand"]] == summand_var