symjit 1.2__tar.gz

symjit-1.2/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "symjit"
+version = "0.3.0"
+edition = "2021"
+
+[dependencies]
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+memmap2 = { version = "0.9", optional = true }
+rand = "0.8"
+anyhow = "1"
+wasmtime = { version = "28.0", optional = true }
+cffi = "*"
+libc = "0.2"
+windows-sys = { version = "0.59", features = ["Win32_System_Memory"] }
+region = "3.0.2"
+wasmtime-jit-icache-coherence = "*"
+
+[features]
+wasm = ["dep:wasmtime"]
+selinux-fix = ['memmap2']
+default = []
+
+[lib]
+name = "_lib"                 
+path = "rust/lib.rs"
+crate-type = ["cdylib"]
+

symjit-1.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 siravan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

symjit-1.2/MANIFEST.in

@@ -0,0 +1,3 @@
+# MANIFEST.in
+include Cargo.toml
+recursive-include rust *.rs

symjit-1.2/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.2
+Name: symjit
+Version: 1.2
+License-File: LICENSE

symjit-1.2/README.md

@@ -0,0 +1,218 @@
+# 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.
+
+# Installation
+
+You can install *symjit* using pip command as 
+
+```
+pip install symjit
+```
+or from the source by cloning 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. For example,
+
+```python
+x, y, a = symbols('x y a')
+f = compile_func([x, y], [(x+y)**a], args=[a])
+assert(np.all(f([3, 5], args=[2]) == [64.]))
+```
+
+`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)
+```
+
+## `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    
+```
+

symjit-1.2/pyproject.toml

@@ -0,0 +1,23 @@
+# pyproject.toml
+[build-system]
+requires = ["setuptools", "setuptools-rust"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "symjit"
+version = "1.2"
+
+[tool.setuptools.packages]
+# Pure Python packages/modules
+find = { where = ["python"] }
+
+[[tool.setuptools-rust.ext-modules]]
+# Private Rust extension module to be nested into the Python package
+target = "symjit._lib"   # The last part of the name (e.g. "_lib") has to match lib.name in Cargo.toml,
+                             # but you can add a prefix to nest it inside of a Python package.
+path = "Cargo.toml"   
+binding = "NoBinding" 
+features = []
+debug = false
+
+

symjit-1.2/python/symjit/__init__.py

@@ -0,0 +1,172 @@
+import json
+import numpy as np
+
+from . import engine
+from . import structure
+
+lib = engine.Engine()   # interface to the rust codegen engine
+
+
+def from_raw_parts(ptr, count):
+    return np.ctypeslib.as_array(ptr, shape=(count,))
+
+
+class BaseFunc:
+    def __init__(self, model, ty="native"):
+        self.p = lib._compile(model.encode("utf-8"), ty.encode("utf8"))
+        status = lib._check_status(self.p)
+        if status != b"Success":
+            raise ValueError(status)
+        self.populate()
+        self.model = model  # for debugging
+
+    def __del__(self):
+        # lib._finalize(self.p)   
+        pass  
+        
+    def get_u0(self):
+        u0 = np.zeros(self.count_states, dtype="double")
+        lib._fill_u0(self.p, np.ctypeslib.as_ctypes(u0), self.count_states)
+        return u0
+
+    def get_p(self):
+        p = np.zeros(self.count_params, dtype="double")
+        lib._fill_p(self.p, np.ctypeslib.as_ctypes(p), self.count_params)
+        return p
+
+    def populate(self):
+        self.count_states = lib._count_states(self.p)
+        self.count_params = lib._count_params(self.p)
+        self.count_obs = lib._count_obs(self.p)
+        self.count_diffs = lib._count_diffs(self.p)
+
+        self._states = from_raw_parts(lib._ptr_states(self.p), self.count_states)
+        self._params = from_raw_parts(lib._ptr_params(self.p), self.count_params)
+        self._obs = from_raw_parts(lib._ptr_obs(self.p), self.count_obs)
+        self._diffs = from_raw_parts(lib._ptr_diffs(self.p), self.count_diffs)
+        
+    def dump(self, name):
+        lib._dump(self.p, name.encode("utf-8"))
+        
+
+class Func(BaseFunc):
+    def __init__(self, model, ty="native"):
+        super().__init__(model, ty=ty)
+
+    def __call__(self, *args):
+        u = np.array(args, dtype="double")
+        self._states[:] = u
+        status = lib._execute(self.p, 0.0)
+
+        if not status:
+            raise ValueError("cannot execute the model")
+
+        return self._obs.copy()
+
+
+class OdeFunc(BaseFunc):
+    def __init__(self, model, ty="native"):
+        super().__init__(model, ty=ty)
+
+    def __call__(self, t, y, *args):
+        y = np.array(y, dtype="double")
+        self._states[:] = y
+
+        if len(args) > 0:
+            p = np.array(args, dtype="double")
+            self._params[:] = p
+
+        status = lib._execute(self.p, t)
+
+        if not status:
+            raise ValueError("cannot execute the model")
+
+        return self._diffs.copy()
+
+
+class JacFunc(BaseFunc):
+    def __init__(self, model, ty="native"):
+        super().__init__(model, ty=ty)
+
+    def __call__(self, t, y, *args):
+        y = np.array(y, dtype="double")
+        self._states[:] = y
+
+        if len(args) > 0:
+            p = np.array(args, dtype="double")
+            self._params[:] = p
+
+        status = lib._execute(self.p, t)
+
+        if not status:
+            raise ValueError("cannot execute the model")
+
+        jac = self._obs.copy()        
+        return jac.reshape((self.count_states, self.count_states))
+
+
+def compile_func(states, eqs, params=None):
+    """Compile a list of symbolic expression into an executable form.
+    compile_func tries to mimic sympy lambdify, but instead of generating
+    a standard python funciton, it returns a callable (Func object) that
+    is a thin wrapper over compiled machine-code.
+    
+    Parameters
+    ==========
+    
+    states: a single symbol or a list/tuple of symbols
+    eqs: a single symbolic expression or a list/tuple of symbolic expressions
+    params (optional): a list/tuple of additional symbols as parameters to the model
+    
+    ==> returns a Func object
+    
+    >>> 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.]))
+    """
+    model = structure.model(states, eqs, params)
+    return Func(json.dumps(model))
+
+
+def compile_ode(iv, states, odes, params=None):
+    """Compile a symbolic ODE model into an executable form suitable for 
+    passung to scipy.integrate.solve_ivp.    
+    
+    Parameters
+    ==========
+    
+    iv: a single symbol, the independent variable.
+    states: a single symbol or a list/tuple of symbols
+    odes: a single symbolic expression or a list/tuple of symbolic expressions,
+        representing the derivative of the state with respect to iv
+    params (optional): a list/tuple of additional symbols as parameters to the model
+    
+    invariant => len(states) == len(odes)
+    
+    ==> returns an OdeFunc object
+    
+    >>> import scipy.integrate
+    >>> 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)
+
+    >>> np.testing.assert_allclose(sol.y[0,:], np.sin(t_eval), atol=0.005)
+    """
+    model = structure.model_ode(iv, states, odes, params)
+    return OdeFunc(json.dumps(model))
+    
+def compile_jac(iv, states, odes, params=None):
+    model = structure.model_jac(iv, states, odes, params)
+    return JacFunc(json.dumps(model))
+
+def compile_json(model):
+    return OdeFunc(model)

symjit-1.2/python/symjit/engine.py

@@ -0,0 +1,124 @@
+import os
+import sys
+import ctypes
+import platform
+
+def find_dll(substr):
+    files = os.listdir(os.path.dirname(__file__))
+    matches = list(filter(lambda s: s.find(substr) >= 0, files))
+    if len(matches) == 0:
+        return None
+    else:
+        return matches[0]
+
+
+dll_name = None
+
+if sys.platform == "linux" and platform.machine() == "x86_64":
+    dll_name = find_dll("x86_64-linux")
+if sys.platform == "linux" and platform.machine() == "aarch64":
+    dll_name = find_dll("aarch64-linux")
+elif sys.platform == "win32":
+    dll_name = find_dll("win_amd64")
+
+if dll_name is None:
+    raise ValueError("unsupported platform")
+
+print(dll_name)
+
+dll_path = os.path.join(os.path.dirname(__file__), dll_name)
+dll = ctypes.CDLL(dll_path)
+
+
+class Engine:
+    def __init__(self):
+        self._info = dll.info
+        self._info.argtypes = []
+        self._info.restype = ctypes.c_char_p
+
+        self._check_status = dll.check_status
+        self._check_status.argtypes = [ctypes.c_void_p]
+        self._check_status.restype = ctypes.c_char_p
+
+        self._count_states = dll.count_states
+        self._count_states.argtypes = [ctypes.c_void_p]
+        self._count_states.restype = ctypes.c_size_t
+
+        self._count_params = dll.count_params
+        self._count_params.argtypes = [ctypes.c_void_p]
+        self._count_params.restype = ctypes.c_size_t
+
+        self._count_obs = dll.count_obs
+        self._count_obs.argtypes = [ctypes.c_void_p]
+        self._count_obs.restype = ctypes.c_size_t
+
+        self._count_diffs = dll.count_diffs
+        self._count_diffs.argtypes = [ctypes.c_void_p]
+        self._count_diffs.restype = ctypes.c_size_t
+
+        self._run = dll.run
+        self._run.argtypes = [
+            ctypes.c_void_p,  # handle
+            ctypes.POINTER(ctypes.c_double),  # du
+            ctypes.POINTER(ctypes.c_double),  # u
+            ctypes.c_size_t,  # ns
+            ctypes.POINTER(ctypes.c_double),  # p
+            ctypes.c_size_t,  # np
+            ctypes.c_double,  # t
+        ]
+        self._run.restype = ctypes.c_bool
+
+        self._execute = dll.execute
+        self._execute.argtypes = [
+            ctypes.c_void_p,  # handle
+            ctypes.c_double,  # t
+        ]
+        self._execute.restype = ctypes.c_bool
+
+        self._fill_u0 = dll.fill_u0
+        self._fill_u0.argtypes = [
+            ctypes.c_void_p,  # handle
+            ctypes.POINTER(ctypes.c_double),  # u0
+            ctypes.c_size_t,  # ns
+        ]
+        self._fill_u0.restype = ctypes.c_bool
+
+        self._fill_p = dll.fill_p
+        self._fill_p.argtypes = [
+            ctypes.c_void_p,  # handle
+            ctypes.POINTER(ctypes.c_double),  # p
+            ctypes.c_size_t,  # np
+        ]
+        self._fill_p.restype = ctypes.c_bool
+
+        self._compile = dll.compile
+        self._compile.argtypes = [ctypes.c_char_p, ctypes.c_char_p]
+        self._compile.restype = ctypes.c_void_p
+
+        self._dump = dll.dump
+        self._dump.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
+        self._dump.restype = None
+
+        self._finalize = dll.finalize
+        self._finalize.argtypes = [ctypes.c_void_p]
+        self._finalize.restype = None
+
+        self._ptr_states = dll.ptr_states
+        self._ptr_states.argtypes = [ctypes.c_void_p]
+        self._ptr_states.restype = ctypes.POINTER(ctypes.c_double)
+
+        self._ptr_params = dll.ptr_params
+        self._ptr_params.argtypes = [ctypes.c_void_p]
+        self._ptr_params.restype = ctypes.POINTER(ctypes.c_double)
+
+        self._ptr_obs = dll.ptr_obs
+        self._ptr_obs.argtypes = [ctypes.c_void_p]
+        self._ptr_obs.restype = ctypes.POINTER(ctypes.c_double)
+
+        self._ptr_diffs = dll.ptr_diffs
+        self._ptr_diffs.argtypes = [ctypes.c_void_p]
+        self._ptr_diffs.restype = ctypes.POINTER(ctypes.c_double)
+
+    def info(self):
+        return self._info()
+