Release Notes

Symjit is a lightweight just-in-time (JIT) compiler that directly translates sympy expressions into machine code without using a separate library such as LLVM. Its main utility is to generate fast numerical functions to feed into various numerical solvers provided by the NumPy/SciPy ecosystem, including numerical integration routines, ordinary differential equation (ODE) solvers, and image filtering functions. It is also able to generate LowLevelCallable functions for better coupling to certain fast Scipy numerical functions.

Symjit has two different code-generating backends. The default is a Rust library with minimum external dependencies. The second backend is written in plain Python, relying solely on the Python standard library and NumPy. Both backends generate AMD64 (also known as x86-64) and ARM64 (also known as aarch64) machine code on Linux, Windows, and Darwin (MacOS) platforms. Further architectures (e.g., RISC V) are planned.

As of version 2.0, the Rust backend generates AVX-compatible code by default for x86-64/AMD64 processors but can downgrade to SSE2 instructions if the processor does not support AVX or if explicitly requested by passing ty='amd-sse' to compile functions (see below). In version 1, the default was SSE instructions. Note that SSE2 instructions were introduced in 2000, meaning that virtually all current 64-bit x86-64 processors support them. Intel introduced the AVX instruction set in 2011; therefore, most processors support it. The Python backend uses only AVX instructions.

On ARM64 processors, both the Rust and Python backends generate code for the aarch64 instruction set. ARM32 and IA32 are not supported.

FuncBuilder is a companion package that provides a more general code generator akin to llvmlite. It is currently in the early stages of development.

Installing symjit from the conda-forge channel can be achieved by adding conda-forge to your channels with:

conda config --add channels conda-forge
conda config --set channel_priority strict

Once the conda-forge channel has been enabled, symjit can be installed with conda:

conda install symjit

or with mamba:

mamba install symjit

It is possible to list all of the versions of symjit available on your platform with conda:

conda search symjit --channel conda-forge

or with mamba:

mamba search symjit --channel conda-forge

Alternatively, mamba repoquery may provide more information:

# Search all versions available on your platform:
mamba repoquery search symjit --channel conda-forge

# List packages depending on `symjit`:
mamba repoquery whoneeds symjit --channel conda-forge

# List dependencies of `symjit`:
mamba repoquery depends symjit --channel conda-forge

You can also install symjit from pypi using pip:

python -m pip install symjit

However, the pip install may not include the correct binary Rust backend for different platforms and the conda-forge install is preferable. In addition, you can install symjit from the source by cloning symjit into symjit folder and then running

cd symjit
python -m pip install .

For the last option, you need a working Rust compiler and toolchains.

symjit is invoked by calling different compile_* functions. The most basic is compile_func, which behaves similarly to sympy lambdify function. While lambdify translates sympy expressions into regular Python functions, which in turn call numpy functions, compile_func returns a callable object Func, which is a thin wrapper over the jit code generated by the backends.

A simple example is

import numpy as np
from symjit import compile_func
from sympy import symbols

x, y = symbols('x y')
f = compile_func([x, y], [x+y, x*y])
assert(np.all(f(3, 5) == [8., 15.]))

compile_* functions support these operators and functions.

compile_func takes two mandatory arguments as compile_func(states, eqs). The first one, states, is a list or tuple of sympy symbols. The second argument, eqs, is a list, a tuple, or a single expression. We can think of states and eqs as corresponding to function signature and body.

If states has only one element, it can be passed directly. Similar to sympy lambdify, the output follows the form of the second argument to compile_func. Therefore, if f = compile_func([x, y], [x+y, x*y]), then f(2, 3) returns a list. On the other hand, if f = compile_func([x, y], (x+y, x*y)), the output will be (5, 6). The third form is a single scalar, such as if f = compile_func([x, y], sin(x+y)).

In addition, compile_func accepts a named argument params, which is a list of symbolic parameters. The output of compile_func, say f, is a callable object of type Func. The signature of f is f(x_1,...,x_n,p_1,...,p_m), where xs are the state variables and ps are the parameters. Therefore, n = len(states) and m = len(params). For example,

x, y, a = symbols('x y a')
f = compile_func([x, y], [(x+y)**a], params=[a])
assert(np.all(f(3., 5., 2.) == [64.]))  # 2. is the value of parameter a

By default, compile_func uses the Rust backend. However, we can force the use of the Python backend by passing backend='python' to compile_func. Moreover, if the binary library containing the Rust backend is unavailable or incompatible, symjit automatically switches to the Python backend.

compile_func helps generate functions to pass to numerical integration (quadrature) routines. The following example is adapted from scipy documentation:

import numpy as np
from scipy.integrate import nquad
from sympy import symbols, exp
from symjit import compile_func

N = 5
t, x = symbols("t x")
f = compile_func([t, x], exp(-t*x)/t**N)

sol = nquad(f, [[1, np.inf], [0, np.inf]])

np.testing.assert_approx_equal(sol[0], 1/N)

The output of the returned callable is a numpy array with dtype='double'. Note that you can call f by passing a list of numbers (say, f(1.0, 2.0)) or a list of numpy arrays (for example, f(np.asarray([1., 2.]), np.asarray([3., 4.])). However, broadcasting is not supported. All the parameters should be passed as scalars even if the state variables are arrays. For example,

import numpy as np
import matplotlib.pyplot as plt
from sympy import symbols
from symjit import compile_func

x, sigma = symbols('x sigma')
f = compile_func([x], [exp(-(x-100)**2/(2*sigma**2))], params=[sigma])

t = np.arange(0, 200)
y = f(t, 25.)[0]

plt.plot(t, y)

The following example uses the vectorization feature to calculate the Mandelbrot set.

# examples/mandelbrot.py
import numpy as np
import matplotlib.pyplot as plt
from sympy import symbols
from symjit import compile_func

x, y, a, b = symbols("x y a b")

A, B = np.meshgrid(np.arange(-2, 1, 0.002), np.arange(-1.5, 1.5, 0.002))
X = np.zeros_like(A)
Y = np.zeros_like(A)

f = compile_func([a, b, x, y], [x**2 - y**2 + a, 2*x*y + b])

for i in range(20):
    X, Y = f(A, B, X, Y)

Z = np.hypot(X, Y)

plt.imshow(Z < 2)

The output is:

The Rust backend supports different optimization methods, which can be controlled using compile_func arguments. See Optimization for details.

compile_ode returns a callable object (OdeFunc) suitable for passing to scipy.integrate.solve_ivp (the main Numpy/Scipy ODE solver). It takes three mandatory arguments as compile_ode(iv, states, odes). The first one (iv) is a single symbol that specifies the independent variable. The second argument, states, is a list of symbols defining the ODE state. The right-hand side of ODE equations is passed as the third argument, odes. It is a list of expressions that define the ODE by providing the derivative of each state variable with respect to the independent variable. In addition, similar to compile_func, compile_ode can accept an optional params. For example,

# examples/trig.py
import scipy.integrate
import matplotlib.pyplot as plt
import numpy as np
from sympy import symbols
from symjit import compile_ode

t, x, y = symbols('t x y')
f = compile_ode(t, (x, y), (y, -x))
t_eval=np.arange(0, 10, 0.01)
sol = scipy.integrate.solve_ivp(f, (0, 10), (0.0, 1.0), t_eval=t_eval)

plt.plot(t_eval, sol.y.T)

Here, the ODE definition is x' = y and y' = -x, which means y" = y. The solution is a*sin(t) + b*cos(t), where a and b are determined by the initial values. Given the initial values of 0 and 1 passed as the third argument of solve_ivp, the solutions are sin(t) and cos(t). We can confirm this by running the code. The output is

Note that OdeFunc conforms to the function form scipy.integrate.solve_ivp expects, i.e., it should be called as f(t, y, *args).

The following example is more complicated and showcases the Lorenz system, an important milestone in the historical development of chaos theory.

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from sympy import symbols

from symjit import compile_ode

t, x, y, z = symbols("t x y z")
sigma, rho, beta = symbols("sigma rho beta")

ode = (
    sigma * (y - x),
    x * (rho - z) - y,
    x * y - beta * z
    )

f = compile_ode(t, (x, y, z), ode, params=(sigma, rho, beta))

u0 = (1.0, 1.0, 1.0)
p = (10.0, 28.0, 8 / 3)
t_eval = np.arange(0, 100, 0.01)

sol = solve_ivp(f, (0, 100.0), u0, t_eval=t_eval, args=p)

plt.plot(sol.y[0, :], sol.y[2, :])

The result is the famous strange attractor:

The ODE examples discussed in the previous section are non-stiff and easy to solve using explicit methods. However, not all differential equations are so accommodating! Many important equations are stiff and usually require implicit methods. Many implicit ODE solvers use the system's Jacobian matrix to improve performance.

There are different techniques for calculating the Jacobian. In the last few years, automatic differentiation (AD) methods have gained popularity, working at the abstract syntax tree or lower level. However, if we define our model symbolically using a Computer Algebra System (CAS) such as sympy, we can calculate the Jacobian by differentiating the source symbolic expressions.

compile_jac is the symjit function to calculate the Jacobian of an ODE system. It has the same call signature as compile_ode, i.e., it is called compile_jac(iv, states, odes) with an optional argument params. The return value (of type JacFunc) is a callable similar to OdeFunc, which returns an n-by-n matrix J, where n is the number of states. The element at the ith row and jth column of J is the derivative of odes[i] w.r.t state[j] (this is the definition of Jacobian).

For example, we can consider the Van der Pol oscillator. This system has a control parameter (mu). For small values of mu, the ODE system is not stiff and can easily be solved using explicit methods.

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp
from sympy import symbols
from math import sqrt
from symjit import compile_ode, compile_jac

t, x, y, mu = symbols('t x y mu')
ode = [y, mu * ((1 - x*x) * y - x)]

f = compile_ode(t, [x, y], ode, params=[mu])
u0 = [0.0, sqrt(3.0)]
t_eval = np.arange(0, 10.0, 0.01)

sol1 = solve_ivp(f, (0, 10.0), u0, method='RK45', t_eval=t_eval, args=[5.0])

plt.plot(t_eval, sol1.y[0,:])

The output is

On the other hand, as mu is increased (for example, to 1e6), the system becomes very stiff. An explicit ODE solver, such as RK45 (Runge-Kutta 4/5), cannot solve this problem. Instead, we need an implicit method, such as the backward differentiation formula (BDF). BDF needs a Jacobian. If one is not provided, it numerically calculates one using the finite-difference method. However, this technique is both inaccurate and computationally intensive. It would be much better to give the solver a closed-form Jacobian. As mentioned above, calculate_jac exactly does this.

jac = compile_jac(t, [x, y], ode, params=[mu])
sol2 = solve_ivp(f, (0, 10.0), u0, method='BDF', t_eval=t_eval, args=[1e6], jac=jac)

plot.plot(t_eval, sol2.y[0,:])

The output of the stiff system is

All compile_* functions accept an optional parameter ty, which defines the type of the code to generate. Currently, the possible values are:

• amd: generates 64-bit AMD64/x86-64 code. If the processor supports AVX, then this is equivalent to passing amd-avx; otherwise, it is equal to amd-sse.
• amd-avx: generates 64-bit AMD64/x86-64 AVX code.
• amd-sse: generates 64-bit AMD64/x86-64 SSE code. It requires a minimum SSE2.1 specification, which should be easily fulfilled by all except the most ancient processors.
• arm generates 64-bit ARM64/aarch64 code. This option is mainly tested on Apple Silicon.
• bytecode: this option uses a generic and simple bytecode evaluator as a fallback option in case of unsupported instruction sets. The utility is to test correctness (see option debug below), not speed.
• native (default): selects the correct instruction set based on the current processor.
• debug: is useful for debugging the generated code. It runs both native and bytecode versions, compares the results, and panics if they are different.

Note that ty='wasm' is no longer supported in version 2. Also, as discussed above, compile_* functions accept a backend argument with possible values of rust and python.

You can inspect the generate intermediate form (IR) and machine code. Refer to Inspection for details.