diff --git a/mira/examples/decapodes/decapodes_examples.py b/mira/examples/decapodes/decapodes_examples.py index 1ca63ab74..89de8bf04 100644 --- a/mira/examples/decapodes/decapodes_examples.py +++ b/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 diff --git a/mira/metamodel/decapodes.py b/mira/metamodel/decapodes.py index cefd31d35..4b1dea325 100644 --- a/mira/metamodel/decapodes.py +++ b/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) diff --git a/tests/test_decapodes.py b/tests/test_decapodes.py index 05ef00448..d6e5b2b38 100644 --- a/tests/test_decapodes.py +++ b/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,8 +67,7 @@ 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(): @@ -70,47 +75,53 @@ def test_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