# Introduction

*Symjit* is a lightweight just-in-time (JIT) compiler that directly translates *sympy* expressions into machine code. Its main utility is to generate fast numerical functions to feed into different numerical solvers, including numerical integration routines and ordinary differential equation (ODE) solvers.   

The Symjit core is a Rust library with minimum external dependency. It does not use a separate compiler, such as LLVM or GCC. Currently, it can generate AMD64 (aka x86-64) and ARM64 (aka aarch64) machine codes on Linux and Windows platforms. Further architectures and operating systems (RISC V, aarch64 on Mac OS) are planned.

# Installing symjit

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

In addition, you can install *symjit* from the source by cloning [symjit](https://github.com/siravan/symjit) into `symjit` folder and then running

```
pip install .
```

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

# Tutorial

## `compile_func`: a fast substitute for `lambdify`

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

A simple example is 

```python
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_func` takes two mandatory arguments as `compile_func(states, eqs)`. The first one, `states`, is a list or tuple of symbols. The second argument, `eqs`, is a list or tuple of expressions. If either `states` or `eqs` has only one element, that element can be passed directly. 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. The signature of `f` is `f(x_1,...,x_n,p_1,...,p_m)`, where `x`s are the state variables and `p`s are the parameters. Therefore, `n = len(states)` and `m = len(params)`. For example,

```python
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
```

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

```python
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([1., 2.], [3., 4.])`. However, broadcasting is not supported. Moreover, all the parameters should be passed as scaler even if the state variables are arrays. For example,

```python
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](https://en.wikipedia.org/wiki/Mandelbrot_set). 

```python
# 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:

![Mandelbrot](./figures/mandelbrot.png)

As of version 1.4, the vectorized version uses SIMD instructions if possible (currently supporting AVX instructions in X86-64 processors). SIMD code should improve the performance up to 4x (using 256-bit registers that encode and operate on four doubles simultaneously). 

## `compile_ode`: to solve ODEs

`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 w.r.t the independent variable. In addition, similar to `compile_func`, `compile_ode` can accept an optional `args`. For example,

```python
# 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  

![sin/cos functions](./figures/trig.png)

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](https://en.wikipedia.org/wiki/Lorenz_system), an important milestone in the historical development of chaos theory. 

```python
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 strange attractor](./figures/lorenz.png)


## `compile_jac`: calculating Jacobian

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](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) 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](https://en.wikipedia.org/wiki/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.

```python
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

![non-stiff Van der Pol](./figures/van_der_pol_non_stiff.png)

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 finite-difference. 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. 

```python
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

![non-stiff Van der Pol](./figures/van_der_pol_stiff.png)

# Backends

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

* `amd`: generates 64-bit AMD64/x86-64 code. It expects a minimum SSE2.1 spec, which should be easily fulfilled by all except the most ancient processors!
* `arm` generates 64-bit ARM64/aarch64 code. To test this instruction set, we use 64-bit Raspbian on Raspberry Pi 4/5 computers.
* `bytecode`: this is a generic fast bytecode as a fallback option in case the instruction set is not supported.
* `native` (**default**): selects the correct instruction set based on the current processor.
* `wasm` (**optional**): generates and runs WebAssembly code using the wasmtime library. This option is not included in the binary distributions. To enable wasm, you must make `symjit` from the source and add `features=["wasm"]` to `pyproject.toml`. 

To inspect the generated code, we must first dump the binary into a file by calling the `dump` function of various `Func` callables. The resulting file is a flat binary code with no header or other extras. Then, use the disassembler of your choice to inspect the code. For example,

```python
from symjit import compile_func
from sympy import symbols

x, y = symbols('x y')
f = compile_func([x, y], [x+y, x*y])
f.dump('test.bin')
```

On a Linux system, we can invoke `objdump` as below:

```
objdump -b binary -m i386:x86-64 -M intel -D test.bin
```

The output (assuming a Linux x86-64 machine) is 

```
   0:	55                   	push   rbp
   1:	53                   	push   rbx
   2:	48 89 fd             	mov    rbp,rdi
   5:	48 89 d3             	mov    rbx,rdx
   8:	f2 0f 10 4d 30       	movsd  xmm1,QWORD PTR [rbp+0x30]
   d:	f2 0f 10 45 28       	movsd  xmm0,QWORD PTR [rbp+0x28]
  12:	f2 0f 58 c1          	addsd  xmm0,xmm1
  16:	f2 0f 11 45 38       	movsd  QWORD PTR [rbp+0x38],xmm0
  1b:	f2 0f 10 4d 30       	movsd  xmm1,QWORD PTR [rbp+0x30]
  20:	f2 0f 10 45 28       	movsd  xmm0,QWORD PTR [rbp+0x28]
  25:	f2 0f 59 c1          	mulsd  xmm0,xmm1
  29:	f2 0f 11 45 40       	movsd  QWORD PTR [rbp+0x40],xmm0
  2e:	5b                   	pop    rbx
  2f:	5d                   	pop    rbp
  30:	c3                   	ret    
```

