Oak Ridge MATrix EXponential tools.
ORMATEX contains methods to compute the matrix exponential:
ORMATEX is a mixed Rust and Python package that provides an extensible foundation to construct advanced exponential integrators.
The current set of implemented and planned time integration methods in each language:
Exponential integrators:
- EPI2 = EXPRB2 = EPIRK2
- EPI3
- EPIRK4
- EXPRB3
- EXPRB4 (exponential rosenbrock order 4)
Classic integrators:
- RK1,RK2,RK3,RK4
- Backward Euler
- BDF2
- Crank-Nicolson
- DIRK
- SDIRK
Exponential integrators:
Jacobian based:
- EPI2 = EXPRB2 = EPIRK2
- EPI3
- EPI4
- EPIRK4
- EXPRB3
- EXPRB4
Splitting linear operator based:
- EXP1
- EXP2
- EXP3
- EXP4
Classic integrators:
- all explicit and implicit integrators supported by diffrax (through interface with diffrax)
- jax
- numpy
- scipy
- pytest
- python3.8+
- equinox
- matplotlib
- scikit-fem
- diffrax
For a local development install, run:
pip install -e .
After running the above, the python unit tests can be executed. From the project base directory (the directory this readme is located in), run:
pytest
Examples are provided in the ormatex_py/progression directory.
Imports
import jax
from jax import numpy as jnp
from ormatex_py.ode_sys import OdeSys, MatrixLinOp
from ormatex_py import integrate_wrapper
Define the system
class LotkaVolterraAD(OdeSys):
alpha: float
beta: float
delta: float
gamma: float
def __init__(self, a=1.0, b=1.0, d=1.0, g=1.0, **kwargs):
super().__init__()
self.alpha = a
self.beta = b
self.delta = d
self.gamma = g
@jax.jit
def _frhs(self, t, x, **kwargs):
prey_t = self.alpha * x[0] - self.beta * x[0] * x[1]
pred_t = self.delta * x[0] * x[1] - self.gamma * x[1]
return jnp.array([prey_t, pred_t])
Initialize the system and integrate
method = 'epi3'
sys = LotkaVolterraAD()
y0 = jnp.array([0.1, 0.2])
t0 = 0.0
dt = 0.2
nsteps = 100
t_res, y_res = integrate_wrapper.integrate(sys, y0, t0, dt, nsteps, method, max_krylov_dim=2, iom=2)
Optionally, an explicit Jacobian can be supplied. If not supplied, as above, automatic differentiation will be used.
class LotkaVolterra(OdeSys):
alpha: float
beta: float
delta: float
gamma: float
def __init__(self, a=1.0, b=1.0, d=1.0, g=1.0, **kwargs):
super().__init__()
self.alpha = a
self.beta = b
self.delta = d
self.gamma = g
@jax.jit
def _frhs(self, t, x, **kwargs):
prey_t = self.alpha * x[0] - self.beta * x[0] * x[1]
pred_t = self.delta * x[0] * x[1] - self.gamma * x[1]
return jnp.array([prey_t, pred_t])
@jax.jit
def _fjac(self, t, x, **kwargs):
jac = jnp.array([
[self.alpha - self.beta * x[1], - self.beta*x[0]],
[self.delta*x[1], self.delta*x[0] - self.gamma]
])
return MatrixLinOp(jac)
Download rustup: https://www.rust-lang.org/tools/install
Then, get rust dev stuff:
rustup toolchain install stable
rustup handles installing the rust toolchain. For improved editing, install the language server for rust (lsp), rust-analyzer:
rustup component add rust-analyzer
To update rust toolchain
rustup update
After setting up rust and cargo, to create a debug build run:
cargo build
To run tests:
cargo test
For an optimized build run:
cargo build --release
To use the rust-based ORMATEX integrators from a python interface, build the python-rust bindings with:
pip install maturin
maturin develop --release
Ensure to use the --release flag for an optimized build. Forgetting this flag will build in debug mode and will result in significantly degraded performance.
After running the above maturin develop --release command from the directory containing this README.md file (the root ORMATEX project directory), the following example can be run:
cd examples
python ex_ormatex_rspy.py
Run the examples with
cargo run --example ex_sys_1 --release
cargo run --example ex_sys_2 --release
Expected resulting images from running the first example of the Lotka-Volterra system integrated with EPI3:
Expected result from the Bateman system in the second example integrated with EXPRB2:
William Gurecky (gureckywl@ornl.gov)
Konstantin Pieper (pieperk@ornl.gov)
Copyright© 2024-present, UT-Battelle, LLC
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
Incomplete orthogonalization procedure (IOM2), faster Arnoldi:
[1] Saad, Yousef. "Variations on Arnoldi's method for computing eigenelements of large unsymmetric matrices." Linear algebra and its applications 34 (1980): 269-295.
Exponential (EPI) applied to shallow water equations:
[2] Gaudreault, Stéphane, and Janusz A. Pudykiewicz.
An efficient exponential time integration method for the numerical solution of the
shallow water equations on the sphere.
Journal of Computational Physics 322 (2016): 827-848.
Exponential Rosenbrock integrators applied to atmosphere and ocean:
[3] Luan, Vu Thai, Janusz A. Pudykiewicz, and Daniel R. Reynolds. "Further development of efficient and accurate time integration schemes for meteorological models." Journal of Computational Physics 376 (2019): 817-837.