diff --git a/pybamm/expression_tree/operations/evaluate_julia.py b/pybamm/expression_tree/operations/evaluate_julia.py index 615f3a20b2..d3b2607e91 100644 --- a/pybamm/expression_tree/operations/evaluate_julia.py +++ b/pybamm/expression_tree/operations/evaluate_julia.py @@ -270,7 +270,7 @@ def to_julia(symbol, debug=False): return constant_values, "\n".join(variable_lines) -def get_julia_function(symbol): +def get_julia_function(symbol, funcname="f"): """ Converts a pybamm expression tree into pure julia code that will calculate the result of calling `evaluate(t, y)` on the given expression tree. @@ -290,18 +290,26 @@ def get_julia_function(symbol): constants, var_str = to_julia(symbol, debug=False) # extract constants in generated function - const_str = "" + const_str = "const cs=(\n" for symbol_id, const_value in constants.items(): const_name = id_to_julia_variable(symbol_id, True) - const_str = const_str + "{} = {}\n".format(const_name, const_value) + const_str += " {} = {},\n".format(const_name, const_value) + const_str += ")\n" # indent code var_str = " " + var_str var_str = var_str.replace("\n", "\n ") + # add "c." to constant names + var_str = var_str.replace("const", "c.const") # add function def and sparse arrays to first line - imports = "begin\nusing SparseArrays\n" - julia_str = imports + const_str + "\nfunction f_pybamm(t, y, p)\n" + var_str + imports = "begin\nusing SparseArrays, LinearAlgebra\n" + julia_str = ( + imports + + const_str + + f"\nfunction {funcname}_with_consts(dy, y, p, t, c)\n" + + var_str + ) # calculate the final variable that will output the result result_var = id_to_julia_variable(symbol.id, symbol.is_constant()) @@ -309,11 +317,15 @@ def get_julia_function(symbol): result_value = symbol.evaluate() # add return line - # two "end"s: one to close the function, one to close the "begin" if symbol.is_constant() and isinstance(result_value, numbers.Number): - julia_str = julia_str + "\n return " + str(result_value) + "\nend\nend" + julia_str = julia_str + "\n return " + str(result_value) + "\nend\n" else: - julia_str = julia_str + "\n return " + result_var + "\nend\nend" + julia_str = julia_str + "\n return " + result_var + "\nend\n" + + julia_str += f"{funcname}(dy, y, p, t) = {funcname}_with_consts(dy, y, p, t, cs)\n" + + # close the "begin" + julia_str += "end" return julia_str diff --git a/tests/unit/test_expression_tree/test_operations/quick_julia_test.py b/tests/unit/test_expression_tree/test_operations/quick_julia_test.py new file mode 100644 index 0000000000..2629f773af --- /dev/null +++ b/tests/unit/test_expression_tree/test_operations/quick_julia_test.py @@ -0,0 +1,160 @@ +# +# Test for the evaluate-to-Julia functions +# +import pybamm + +from tests import ( + get_mesh_for_testing, + get_1p1d_mesh_for_testing, + get_discretisation_for_testing, + get_1p1d_discretisation_for_testing, +) +import unittest +import numpy as np +import scipy.sparse +from collections import OrderedDict + +from julia import Main + + +a = pybamm.StateVector(slice(0, 1)) +b = pybamm.StateVector(slice(1, 2)) + +y_tests = [np.array([[2], [3]]), np.array([[1], [3]])] +t_tests = [1, 2] + +# test a * b +# expr = a * b +# evaluator_str = pybamm.get_julia_function(expr) +# print(evaluator_str) + +# test something with a matrix multiplication +A = pybamm.Matrix([[1, 2], [3, 4]]) +expr = A @ pybamm.StateVector(slice(0, 2)) +evaluator_str = pybamm.get_julia_function(expr) +print(evaluator_str) + +# # test something with a heaviside +# a = pybamm.Vector([1, 2]) +# expr = a <= pybamm.StateVector(slice(0, 2)) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# # note 1D arrays are flattened in Julia +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).flatten()) + +# expr = a > pybamm.StateVector(slice(0, 2)) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# # note 1D arrays are flattened in Julia +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).flatten()) + +# # # test something with a minimum or maximum +# # a = pybamm.Vector([1, 2]) +# # expr = pybamm.minimum(a, pybamm.StateVector(slice(0, 2))) +# # evaluator_str = pybamm.get_julia_function(expr) +# # evaluator = Main.eval(evaluator_str) +# # for t, y in zip(t_tests, y_tests): +# # result = evaluator(t,y,None) +# # np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) + +# # expr = pybamm.maximum(a, pybamm.StateVector(slice(0, 2))) +# # evaluator_str = pybamm.get_julia_function(expr) +# # evaluator = Main.eval(evaluator_str) +# # for t, y in zip(t_tests, y_tests): +# # result = evaluator(t,y,None) +# # np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) + +# # test something with an index +# expr = pybamm.Index(A @ pybamm.StateVector(slice(0, 2)), 0) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# self.assertEqual(result, expr.evaluate(t=t, y=y)) + +# # test something with a sparse matrix multiplication +# A = pybamm.Matrix([[1, 2], [3, 4]]) +# B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) +# C = pybamm.Matrix(scipy.sparse.coo_matrix(np.array([[1, 0], [0, 4]]))) +# expr = A @ B @ C @ pybamm.StateVector(slice(0, 2)) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# # note 1D arrays are flattened in Julia +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).flatten()) + +# expr = B @ pybamm.StateVector(slice(0, 2)) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# # note 1D arrays are flattened in Julia +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).flatten()) + +# # test numpy concatenation +# a = pybamm.StateVector(slice(0, 1)) +# b = pybamm.StateVector(slice(1, 2)) +# c = pybamm.StateVector(slice(2, 3)) + +# y_tests = [np.array([[2], [3], [4]]), np.array([[1], [3], [2]])] +# t_tests = [1, 2] + +# expr = pybamm.NumpyConcatenation(a, b) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# # note 1D arrays are flattened in Julia +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).flatten()) + +# expr = pybamm.NumpyConcatenation(a, c) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None) +# # note 1D arrays are flattened in Julia +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).flatten()) + +# # test sparse stack +# A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) +# B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[2, 0], [5, 0]]))) +# a = pybamm.StateVector(slice(0, 1)) +# expr = pybamm.SparseStack(A, a * B) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t, y, None).toarray() +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).toarray()) + +# # test Inner +# expr = pybamm.Inner(a, b) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t,y,None) +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) + +# v = pybamm.StateVector(slice(0, 2)) +# A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) +# expr = pybamm.Inner(A, v) +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t,y,None).toarray() +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).toarray()) + +# y_tests = [np.array([[2], [3], [4], [5]]), np.array([[1], [3], [2], [1]])] +# t_tests = [1, 2] +# a = pybamm.StateVector(slice(0, 1), slice(3, 4)) +# b = pybamm.StateVector(slice(1, 3)) +# expr = a * b +# evaluator_str = pybamm.get_julia_function(expr) +# evaluator = Main.eval(evaluator_str) +# for t, y in zip(t_tests, y_tests): +# result = evaluator(t,y,None) +# np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) diff --git a/tests/unit/test_expression_tree/test_operations/test.py b/tests/unit/test_expression_tree/test_operations/test.py new file mode 100644 index 0000000000..58554252ae --- /dev/null +++ b/tests/unit/test_expression_tree/test_operations/test.py @@ -0,0 +1,25 @@ +from julia import Main +import numpy as np + +f_b = Main.eval( + """ +begin +function f_b(dy,y,a,b) + dy[1] = a*y[1] + dy[2] = b*y[2] + dy +end +end +""" +) +dy = [0, 0] +y = [1, 3] +print(dy) # returns [0 0] +print(f_b(dy, y, 5, 3)) # returns [5 9] +print(dy) # returns [0 0] (expected [5 9]) + +Main.dy = [0, 0] +Main.y = [1, 3] +print(Main.dy) # returns [0 0] +print(Main.f_b(Main.dy, Main.y, 5, 3)) # returns [5 9] +print(Main.dy) # returns [0 0] (expected [5 9])