"""
Utilities to compile functions, esp. from Sympy-constructed Matrix math.
"""
import numpy as np
from sympy import Matrix
from typing import Any, Callable
[docs]
def build_functions(arguments: str, expression: str) -> tuple[Callable, Callable]:
"""Build a function from its arguments and source code expression.
(This USED TO return a second function that had Numba set up to JIT-compile
the code on demand, but the compiled code time savings don't pay back the
compilation time in Python 3.9+.)
This still returns two functions so the JIT compiler could be restored
someday.
Args:
arguments: comma-separated lambda argument names
expression: expression to compile
Returns:
a function(arguments),
the same function(arguments) [was a JIT-enabled version]
"""
local_dict: dict[str, Any] = {}
expression = f"def f({arguments}):\n" + expression
exec(expression, globals(), local_dict)
f = local_dict["f"]
return f, f
[docs]
def _matrix_to_array(matrix: Matrix) -> str:
"""Convert a sympy Matrix expression to a function body."""
rows, cols = matrix.shape
_ = np # So the tools won't warn about unused np import.
function_str = [f"\tarr = np.zeros(({rows}, {cols}))\n"]
for i in range(rows):
for j in range(cols):
if matrix[i, j] != 0:
function_str.append(f"\tarr[{i}, {j}] = {matrix[i, j]}\n")
function_str.append("\treturn arr")
return "".join(function_str)
[docs]
def derivatives(matrix: Matrix) -> tuple[Callable, Callable]:
"""Build an optimized derivatives ODE function(y, t)."""
return build_functions("y, t", _matrix_to_array(matrix) + ".reshape(-1)")
[docs]
def derivatives_jacobian(jacobian_matrix: Matrix) -> tuple[Callable, Callable]:
"""Build an optimized derivatives ODE Jacobian function(y, t)."""
return build_functions("y, t", _matrix_to_array(jacobian_matrix))
[docs]
def rates(matrix: Matrix) -> tuple[Callable, Callable]:
"""Build an optimized rates function(t, y, kf, kr)."""
return build_functions("t, y, kf, kr", _matrix_to_array(matrix) + ".reshape(-1)")
[docs]
def rates_jacobian(jacobian_matrix: Matrix) -> tuple[Callable, Callable]:
"""Build an optimized rates Jacobian function(t, y, kf, kr)."""
return build_functions("t, y, kf, kr", _matrix_to_array(jacobian_matrix))