#1129 starting to optimize generated function

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,30 +290,42 @@ 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())
     if symbol.is_constant():
         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
 

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))

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])