ormatex_py package¶

Submodules¶

ormatex_py.arnoldi_jax module¶

Arnoldi implementation in JAX.

ormatex_py.arnoldi_jax.arnoldi_mgs_lop(a_lo, a_scale, hs, qs, k, iom)¶

Parameters:

a_lo (LinOp) – linear operator
a_scale (float) – scaling factor
hs (Array) – hessenberg
qs (Array) – orthonormal basis of krylov subspace
k (int) – step of Arnoldi

Returns:

modified hessenberg qs: modified ONB bool. true if no happy breakdown

Return type:

hs

ormatex_py.arnoldi_jax.arnoldi_lop_jit(a_lo, a_scale, b, m, iom)¶

Parameters:

a_lo (LinOp) – linear operator
a_scale (float) – scaling factor
b (Array) – right hand side
m (int) – number of Krylov vectors to produce
iom (int) – number of (incomplete) orthogonalizations

Returns:

hessenberg qs: orthonormal basis of krylov subspace breakdown_k: iteration k where happy breakdown occurred, or m

Return type:

hs

ormatex_py.arnoldi_jax.arnoldi_lop(a_lo, a_scale, b, m, iom)¶

Parameters:

a_lo (LinOp) – linear operator
a_scale (float) – scaling factor
b (Array) – right hand side (it is assumed that b is not the zero vector)
m (int) – number of Krylov vectors to produce
iom (int) – number of (incomplete) orthogonalizations

Returns:

(reduced) hessenberg qs: (reduced) orthonormal basis of krylov subspace

Return type:

hs

ormatex_py.integrate_wrapper module¶

class ormatex_py.integrate_wrapper.IntegrateResult(t, y, callback_res, err_code)¶

Bases: object

Storage for ORMATEX integration results

t: list¶

y: list¶

callback_res: dict¶

err_code: int¶

property t_res¶

property y_res¶

property cb¶

ormatex_py.integrate_wrapper.integrate(ode_sys, y0, t0, dt, nsteps, method, **kwargs)¶: High level interface to all time integration methods in ORMATEX.

ormatex_py.integrate_wrapper.integrate_diffrax(ode_sys, y0, t0, dt, nsteps, method='implicit_euler', **kwargs)¶: Uses diffrax integrators to step adv diff system forward

ormatex_py.integrate_wrapper.integrate_ormatex(sys_int, y0, t0, dt, nsteps, method='exprb2', **kwargs)¶: Uses ormatex exponential integrators to step adv diff system forward

ormatex_py.matexp_krylov module¶

Krylov methods to calculate \(\mathrm{exp}(A \delta t)v\) and \(\phi_k(A \delta t)v\) with sparse A

ormatex_py.matexp_krylov.matexp_linop(a_lo, dt, v0, max_krylov_dim, iom=100)¶

Computes exp(A*dt)*v where A is a sparse linop

Return type:: Array

ormatex_py.matexp_krylov.phi_linop(a_lo, dt, v0, k, max_krylov_dim, iom=100)¶

Computes phi_k(A*dt)*v where A is a sparse linop

Return type:: Array

ormatex_py.matexp_krylov.kiops_fixedsteps(a_lo, dt, vb, max_krylov_dim, iom=100, n_steps=1)¶

Method based roughly on simplified KIOPS with fixes stepsize and not Krylov adaptivity. TODO: add adaptivity routines. This avoids the substepping procedure in phipm by computing the phi-vector products as:

\[ \begin{align}\begin{aligned}w(\tau) = \sum_{j=0}^p \tau^j \varphi_j(\tau A) b_j\\w(\tau) = \exp(\tau \tilde A)v\end{aligned}\end{align} \]

with \(v = [b_0, e_1]^T\) and

\[\tilde A = [[A, B],[0, K]]\]

where A = a_lo is NxN, B = vb[:0:-1] is Nxp, K = [[0, I_{p-1}],[0, 0]] is pxp, \(\tilde A\) is N+p x N+p

Return type:: Array

ormatex_py.matexp_leja module¶

Matrix function evaluation using leja point interpolation routines.

NOTE: When running on a CPU target, it is recommended to use the following env variables for best performance:

OMP_NUM_THREADS=4 XLA_FLAGS=–xla_cpu_use_thunk_runtime=false

Ref: https://github.com/jax-ml/jax/discussions/25711

ormatex_py.matexp_leja.gen_leja_conjugate(n=64, a=-1.0, b=1.0, c=1.0)¶

Generate the conjugate leja points for an ellipse contained in the square: [a,b] x [-c,c].

This is the ellipse with center shift = ((a+b)/2, 0.) and half axes ((b-a)/2, c).

Parameters:

n (int) – number of fast leja points to generate
a (float) – left boundary of box
b (float) – right boundary of box
c (float) – upper boundary of symmetric box

ormatex_py.matexp_leja.gen_leja_circle(n=64, conjugate=False)¶

Generate the fast leja points on a circle of radius 1 at zero.

Ref:: Baglama, J., D. Calvetti, and L. Reichel. “Fast leja points.” Electron. Trans. Numer. Anal 7.124-140 (1998): 119-120.

Parameters:: n (int) – number of fast leja points to generate

ormatex_py.matexp_leja.gen_leja_fast(a=-2.0, b=2.0, n=100)¶

Generate the fast leja points in [a, b].

Ref:: Baglama, J., D. Calvetti, and L. Reichel. “Fast leja points.” Electron. Trans. Numer. Anal 7.124-140 (1998): 119-120.

Parameters:

a (float) – left bounding point on the interval [a, b]
b (float) – right point
n (int) – number of fast leja points to generate

ormatex_py.matexp_leja.gen_complex_conj_leja_fast(beta, n=100)¶

Generates the leja point sequence on the imaginary axis

\(i[-\beta, \beta]\)

Parameters:

beta (float) – max imag eig magnitude
n (int) – number of fast leja points to generate

ormatex_py.matexp_leja.power_iter(a_lop, b0, iter, tol=0.05)¶

Performs power iteration to find dominant eigenvalue of system

Parameters:

a_lop (LinOp) – LinOp linear operator which must implement matvec.
b0 (Array) – initial eigen vector
iter (int) – max number of power iterations to perform.
tol (float) – tolerance for dominant (real) eigenvalue

ormatex_py.matexp_leja.complex_conj_imag_leja_expmv(a_lo, dt, u, shift, scale, imag_leja_x, coeffs, tol)¶

Interpolation approx \(\mathrm{exp}(\delta t A)u\) at the conjugate complex imaginary leja points. The complex conj. leja points are symmetric about the real line, and admit a 2 term recurrenace relationship to compute the polynomial terms as described in

Ref: Caliari, M., Kandolf, P., Ostermann, A. et al. Comparison of software for: computing the action of the matrix exponential. Bit Numer Math 54, 113–128 (2014). https://doi.org/10.1007/s10543-013-0446-0

Parameters:

a_lo (LinOp) – LinOp linear operator which must implement matvec.
dt (float) – (Sub)step size.
u (Array) – vector to which the matrix exponential is applied.
shift (float) – shift applied to leja points.
scale (float) – scale applied to leja points.
imag_leja_x (Array) – The unscaled complex leja sequence

ormatex_py.matexp_leja.decay_fun(e_new, e_old, gamma)¶

ormatex_py.matexp_leja.complex_conj_leja_expmv(a_lo, dt, u, shift, scale, leja_x, n_leja_real, coeffs, tol, abort_tol=10000000000.0)¶

Interpolation approx \(\mathrm{exp}(\delta t A)u\) at real and conjugate complex leja points.

leja_x[i] is real for 0 <= i < n_leja_real leja_x[i+1] = conj(leja_x[i]) for i >= n_leja_real

Note: this method does not check if leja_x has the assumed structure

Parameters:

a_lo (LinOp) – LinOp linear operator which must implement matvec.
dt (float) – (Sub)step size.
u (Array) – vector to which the matrix function is applied.
shift (float) – shift applied to leja points.
scale (float) – scale applied to leja points.
leja_x (Array) – unscaled complex leja sequence.
n_leja_real (int) – number of real leja points at the beginning of leja_x (one or two, usually).
coeffs (Array) – divied differences of the function applied to the matrix.
tol (float) – convergence tolerance.
abort_tol (float) – aborts polynomial construction if estimated error measure is above this value.

ormatex_py.matexp_leja.real_leja_expmv(a_lo, dt, u, shift, scale, leja_x, coeffs, tol, abort_tol=10000000000.0)¶

Computes leja polynomial interpolation to approx \(\mathrm{exp}(\delta t A)u\).

Ref: L. Bergamaschi. M. Caliari. A. Martinez and M. Vianello.: Comparing Leja and Krylov Approximations of Large Scale Matrix Exponentials. Intl. Conf on Computational Science. 2006.

Parameters:

a_lo (LinOp) – LinOp linear operator which must implement matvec.
dt (float) – (Sub)step size.
u (Array) – vector to which the matrix exponential is applied.
shift (float) – shift applied to leja points.
scale (float) – scale applied to leja points.
leja_x (Array) – The unscaled leja points
tol (float) – Tolerance of the polynomial interpolation such that \(|| p - exp(\delta t A)u || < tol\) where p is the leja polynomial approximation to \(\mathrm{exp}(\delta t A)u\).
abort_tol (float) – aborts polynomial construction if estimated error measure is above this value.

ormatex_py.matexp_leja.leja_shift_scale(a_tilde_lo, dim, max_power_iter=20, b0=None, scale_factor=1.0)¶

Computes scaling and shifting parameters for leja interpolation of the matrix exponential using power iteration to estimate the eigenvalue with maximum real magnitude. TODO: use arnoldi or other method to estimate the eigenvalue with the maximum imag magnitude. WARNING: This routine does not work for systems with only imaginary eigenvalues.

Parameters:

max_power_iter (int) – maximum number of power iterations.
b0 – vector. Initial guess for principle eigenvector of A.
scale_factor (float) – saftey factor for largest eig magnitude to ensure leja points encompass the spectrum.

ormatex_py.matexp_leja.real_leja_expmv_substep(a_tilde_lo, tau_dt, v, leja_x, n, shift, scale, tol=1e-10, substep=True, dd_method='taylor')¶

Computes \(exp(\Delta t A)v\) using the real leja point interpolation method and substeps by:

\[v_{i+1} = \mathrm{exp}(\tau_i \tilde A) v_{i}\]

where \(\sum(\tau_i) = 1\)

Substepping is automatically applied by detecting divergence, or failure to converge the leja polynomial approximation.

Ref:: M. Caliari, M. Vianello, L. Bergamaschi. Interpolating discrete advection–diffusion propagators at Leja sequences. Journal of Computational and Applied Mathematics. Volume 172, Issue 1.

Parameters:

a_tilde_lo (LinOp) – LinOp linear operator.
tau_dt (float) – inital substep size in (0, 1].
v (Array) – vector to which the matrix exponential is applied
leja_x (Array) – leja points generated by the gen_leja_fast() function.
substep (bool) – Flag to enable automatic substepping. True by default.
dd_method (str) – method used to compute divided differences. One of [‘taylor’, ‘pade’]

ormatex_py.matexp_leja.complex_conj_leja_expmv_substep(a_tilde_lo, tau_dt, v, leja_x, n_leja_real, n, shift, scale, tol=1e-10, substep=True, leja_n_zeros=2, dd_method='taylor')¶

Computes \(exp(\Delta t A)v\) using complex conjugate leja points and with substeps by:

\[v_{i+1} = \mathrm{exp}(\tau_i \tilde A) v_{i}\]

where \(\sum(\tau_i) = 1\)

Substepping is automatically applied by detecting divergence, or failure to converge the leja polynomial approximation.

Parameters:

a_tilde_lo (LinOp) – LinOp linear operator.
tau_dt (float) – inital substep size in (0, 1].
v (Array) – vector to which the matrix exponential is applied
leja_x (Array) – leja points generated by the gen_leja_conjugate() function.
n_leja_real (int) – number of real points at the start of the leja sequence. This value is returned by gen_leja_conjugate(). n_leja_real==2 for a circle or ellipse, but can be equal to len(leja_x) in the case of purely real leja points. If n_leja_real is equal to len(leja_x) then this method is equal to the real_leja_expmv_substep() method.
substep (bool) – Flag to enable automatic substepping. True by default.
dd_method (str) – method used to compute divided differences. One of [‘taylor’, ‘pade’]
leja_n_zeros (int) – number of repeated zeros prepended to the leja sequence

ormatex_py.matexp_leja.build_a_tilde(a_lo, dt, vb)¶

Builds extended linear system.

Parameters:

a_lo (LinOp) – System Linear operator
dt (float) – time step size
vb (list[Array]) – list of rhs vectors to which \(A\) will be applied to

Builds the matrix:

\[\tilde A = [[A, B],[0, K]]\]

where A = a_lo is NxN, B = vb[:0:-1] is Nxp, K = [[0, I_{p-1}],[0, 0]] is pxp, \(\tilde A\) is N+p x N+p

ormatex_py.matexp_leja.plot_leja_conjugate_ellipse_error(a=0.0, b=0.0, c=4.0, eigJ=[], dt=1.0, leja_n_zeros=0, n_leja=50, dd_method='taylor', leja_tol=1e-06, v=None, dirname='./', xscale='linear')¶

Helper method to plot and report the error of the conjugate leja polynomial approximation.

Parameters:

a – left bound of the ellipse enclosing the spectrum
b – right bound of the ellipse enclosing the spectrum
c – upper imaginary bound of the ellipse enclosing the spectrum
eigJ – eigenvalues of the system jacobian
leja_n_zeros – number of zeros to include in the interpolation sequence
n_leja – number of points in the leja sequence
dd_method – divided diff method used to compute the coefficients of the leja polynomial. one of [‘taylor’, ‘pade’]
v – rhs vector (optional, if None v=unit vec)

Returns:

error of the leja polynomial matexp(eigJ)*v approximation

ormatex_py.matexp_phi module¶

Implements the phi-functions

ormatex_py.matexp_phi.f_phi_k(z, k)¶

Computes phi_k(Z) for dense Z

Return type:: Array

ormatex_py.matexp_phi.f_phi_k_inv(z, k, eps)¶

Computes phi_k(Z) for dense Z, using a formula involving the inverse of Z. Returns an error estimate to warn about nearly singular Z.

Return type:: (Array, float)

ormatex_py.matexp_phi.f_phi_k_ext(z, k, return_all=False)¶

Computes phi_k(Z) for dense Z, using the stable but more expensive extension formula

Return type:: Array

ormatex_py.matexp_phi.f_phi_k_poly_all(z, k, poly_deg=4)¶

Computes phi_k(Z) for dense Z, using a Taylor polynomial

Return type:: list[Array]

ormatex_py.matexp_phi.f_phi_k_sq_all(z, k)¶

Computes phi_k(Z) for dense Z, using the scaling and squaring relations

Return type:: list[Array]

ormatex_py.matexp_phi.f_phi_k_sq(z, k, return_all=False)¶

Return type:: Array

ormatex_py.matexp_phi.f_phi_k_appl(z, b, k)¶

Computes phi_k(Z)B for dense Z and dense B, using an extension formula

Return type:: Array

ormatex_py.matexp_phi.f_phi_k_pfd(z, b, k, method)¶

Computes phi_k(Z)B for dense Z and dense B, using a rational approximation and partial fraction expansion

Return type:: Array

ormatex_py.matexp_phi_pfd_dict module¶

Dictionary of poles and coefficients for partial fraction expansion based rational approximation methods for computing

\[\phi_k(Z)*v\]

products.

ormatex_py.matexp_poly module¶

This submodule holds polynomial approximation tools for the matrix exponential and the action of the matrix exponential on a vector. A selection of divided difference routines are provided to compute the polynomial coefficients of the matrix exponential and related phi-functions.

ormatex_py.matexp_poly.newton_poly_div_diff(x, y)¶

Divided difference formula to compute coeffs of the newton polynomial given pairs of (x_i, y_i)

Parameters:

x (Array) – data support points, where the data is sampled.
y (Array) – value of the function at x points.

ormatex_py.matexp_poly.leja_coeffs_exp_dd(leja_x, shift, scale, h=1.0)¶

Dividied difference computation of the leja polynomial coefficients.

Parameters:

leja_x (Array) – The leja points
shift (float) – The shift applied to the leja points
scale (float) – The scale applied to the leja points
h (float) – step fraction

ormatex_py.matexp_poly.leja_coeffs_exp(leja_x, shift, scale, h=1.0)¶

Alternate method to compute the leja polynomial coefficients. Reduced roundoff error compared to divided difference formula. Ref:

M. Calari. Accurate evaluation of divided differences for polynomial interpolation of exponential propagators. Computing. 80. 2007.

Parameters:

leja_x (Array) – The leja points
shift (float) – The shift applied to the leja points
scale (float) – The scale applied to the leja points
h (float) – step fraction

ormatex_py.matexp_poly.expm_taylor(A, shift, scale, p=20)¶

Taylor expansion of the shifted matrix exp function for dense matricies A.

Where A is a matrix. We approximate the second term with the TS:

Parameters:

A (Array) – square matrix
shift (float) – argument shift
scale (float) – argument scale
p (int) – taylor series polynomial order

ormatex_py.matexp_poly.dd_exp_taylor(leja_x, shift, scale, h=1.0, p=20, mu_scale=1.0)¶

Divided differences of the shifted and scaled exponenial function using a taylor series approximation and scaling and squaring. Similar to leja_coeffs_exp but with custom TS impl.

Parameters:

leja_x – The leja points
shift – The shift applied to the leja points
scale – The scale applied to the leja points
h – time step fraction
p – taylor series polynomial order

ormatex_py.matexp_poly.leja_seq_pad_zeros(x, p=2)¶

Pads a leja sequence with leading (near) zeros. When used with the dd_exp_taylor method, The divided differences of the repeated zeros are the coefficients of a newton polynomial with terms that are approx. the taylor series.

Parameters:

x – leja sequence
l – number of (near) zeros to pad

ormatex_py.matexp_poly.leja_seq_reordering(x)¶

Reorders the points in sequence x into a leja sequence.

Parameters:: x – an unordered sequence of points in the complex plane

ormatex_py.matexp_poly.leja_coeffs_ts_ss(leja_x, shift, scale, h_in=1.0, l=0)¶: Divided differences of Xi, phi_l(Xi) by taylor series and scaling and squaring.

ormatex_py.matexp_poly.leja_coeffs_ts_phi(leja_x, shift, scale, h=1.0, k=0, full=False)¶

Divided difference of phi_k(h*(a+b*xi)) evaluated at xi in [-2, 2] by taylor series approximation. Uses the TS(II) method from Refs:

A. McCurdy, K. Ng and B. Parlett. Accurate Computation of Divided differences of the exponential function. Mathmatics of Computation. v 43. 1984.

M. Calari. Accurate evaluation of divided differences for polynomial interpolation of exponential propagators. Computing. 80. 2007.

Parameters:

leja_x (Array) – The leja points
shift (float) – The shift applied to the leja points
scale (float) – The scale applied to the leja points
h (float) – step fraction
k (int) – the phi function order

ormatex_py.matexp_poly.leja_coeffs_ts_phi_jax(leja_x, shift, scale, h=1.0, k=0)¶

JAX implementation of TS(II). See reference numpy implementation in :method:`leja_coeffs_ts_phi`

Parameters:

leja_x (Array) – The leja points
shift (float) – The shift applied to the leja points
scale (float) – The scale applied to the leja points
h (float) – time step fraction
k (int) – the phi function order

ormatex_py.matexp_poly.leja_coeffs_ts_phi_substep(leja_x, shift, scale, k=0)¶: Substepping procedure for the divided differences using the ts_phi method.

ormatex_py.matexp_poly.leja_coeffs_dd_phi(leja_x, shift, scale, h=1.0, l=0)¶

Fast divided difference calculation for phi_l(Z) evaluated at z=(a+b*xi) using the dd_phi routine from Zivovich (2019). Avoids expm call from leja_coeffs_exp method.

Ref:: Fast and accurate computation of divided differences for analytic functions, with an application to the exponential function. F. Zivcovich. Dolomites Research Notes on Approximation. v 12. 2019.

Parameters:

leja_x (Array) – The leja points
shift (float) – The shift applied to the leja points
scale (float) – The scale applied to the leja points
h (float) – time step fraction
l (int) – the phi function order

ormatex_py.ode_exp module¶

Exponential time integration methods

class ormatex_py.ode_exp.ExpRBIntegrator(sys, t0, y0, method='epi2', **kwargs)¶

Bases: IntegrateSys

reset_ic(t0, y0)¶

step(dt, frhs_kwargs={})¶

Take a time step

Parameters:: dt (float) – time step size
Return type:: StepResult

accept_step(s)¶

default implementation, maybe overridden

Parameters:: s (StepResult) – Step result

class ormatex_py.ode_exp.ExpLejaIntegrator(sys, t0, y0, method='epi2_leja', **kwargs)¶

Bases: IntegrateSys

step(dt, frhs_kwargs={})¶

Take a time step

Parameters:: dt (float) – time step size
Return type:: StepResult

accept_step(s)¶

default implementation, maybe overridden

Parameters:: s (StepResult) – Step result

class ormatex_py.ode_exp.ExpSplitIntegrator(sys, t0, y0, method='exp3', **kwargs)¶

Bases: IntegrateSys

reset_ic(t0, y0)¶

step(dt, frhs_kwargs={})¶

Take a time step

Parameters:: dt (float) – time step size
Return type:: StepResult

accept_step(s)¶

default implementation, maybe overridden

Parameters:: s (StepResult) – Step result

ormatex_py.ode_explicit module¶

class ormatex_py.ode_explicit.RKIntegrator(sys, t0, y0, method='rk4', **kwargs)¶

Bases: IntegrateSys

step(dt)¶

Take a time step

Parameters:: dt (float) – time step size
Return type:: StepResult

accept_step(s)¶

default implementation, maybe overridden

Parameters:: s (StepResult) – Step result

ormatex_py.ode_sys module¶

Defines an interface for coupled systems of ODEs. This interface is compatible with all time integration methods in this ormatex_py package. JAX allows user defined types as PyTrees: See: https://jax.readthedocs.io/en/latest/pytrees.html NOTE: we use equinox to make jax compatibility simpler by removing boilerplate to register objects as pytrees. from discussion here: https://github.com/jax-ml/jax/discussions/10598

class ormatex_py.ode_sys.LinOp¶

Bases: Module

matvec(v)¶

Return type:: Array

matvec_npcompat(v)¶

Return type:: ndarray

dense()¶

Return type:: Array

class ormatex_py.ode_sys.EyeLinOp(n)¶

Bases: LinOp

Identity LinOp

n: int¶

class ormatex_py.ode_sys.DiagLinOp(d)¶

Bases: LinOp

Diagonal linear operator with diagonal d

d: Array¶

class ormatex_py.ode_sys.MatrixLinOp(a)¶

Bases: LinOp

Helper class to wrap a jax.Array as a LinOp similar to scipy.sparse.linalg.aslinearoperator()

a: Array | JAXSparse¶

class ormatex_py.ode_sys.AugMatrixLinOp(a_lo, dt, B, K)¶

Bases: LinOp

Helper class to create a larger linear operator out of a block matrix comprised of components:

a_lo_aug =
    [[dt*A,  B],
     [0,     K]]

where A is the original NxN linop dt is a scalar factor applied to A. b is a Nxp dense matrix. K is a pxp sparse matrix.

a_lo: LinOp¶

dt: float¶

B: Array¶

K: Array¶

class ormatex_py.ode_sys.SysJacLinOp(t, u, frhs, frhs_kwargs={}, **kwargs)¶

Bases: LinOp

Extend base LinOp class with pure virtual method for frhs time derivative and frhs caching.

class ormatex_py.ode_sys.CustomJacLinOp(t, u, frhs, f_du, f_dt=None, frhs_kwargs={}, **kwargs)¶: Bases: SysJacLinOp

class ormatex_py.ode_sys.AdJacLinOp(t, u, frhs, frhs_kwargs={}, **kwargs)¶: Bases: SysJacLinOp

class ormatex_py.ode_sys.FdJacLinOp(t, u, frhs, frhs_kwargs={}, scale=1.0, gamma=0.0, **kwargs)¶

Bases: SysJacLinOp

property shape: (<class 'int'>, <class 'int'>)¶: LinOp shape

class ormatex_py.ode_sys.OdeSys(*args, **kwargs)¶

Bases: Module

frhs(t, u, **kwargs)¶

Define the RHS of the system. dU/dt = F(t, U)

Ex: Could look like: F(t, U) = M^-1 [AU + BU + S] where M=const, A, B are matricies and S=S(u,t) is a vec.

Parameters:

t (float) – current time
u (Array) – current system state.

Return type:

Array

fjac(t, u, **kwargs)¶

Define the Jacobian of the system as a linear operator

Ex: Could look like: dF(t, U)/dU|_(t,u) = M^-1 (A + B + dS/dU|_(t,u)) where M, A, B are matricies and S is a vec. and M is const.

The default implementation uses jax.linearize.

Parameters:

t (float) – current time
u (Array) – current system state.

Return type:

LinOp

fm(t, u, **kwargs)¶

Method that returns the mass matrix LinOp

Return type:: LinOp

class ormatex_py.ode_sys.OdeSplitSys(*args, **kwargs)¶

Bases: OdeSys

Define a split ODE system.

dU/dt = F(t, U) = L(t, U) @ U + R(t, U),

where L(t, U) is a (potentially time and state dependent) LinOp, different from the Jacobian.

fl(t, u, **kwargs)¶

Define the LinOp L(t,u) of the system. :type t: float :param t: current time :type u: Array :param u: current system state.

Return type:: LinOp

class ormatex_py.ode_sys.OdeSysNp(inner)¶

Bases: OdeSys

inner: OdeSys¶

class ormatex_py.ode_sys.StepResult(t, dt, y, err)¶

Bases: object

t: float¶

dt: float¶

y: Array¶

err: float¶

class ormatex_py.ode_sys.IntegrateSys(sys, t0, y0, order, method, *args, **kwargs)¶

Bases: object

Defines interface to time integrators

sys: OdeSys¶

t: float¶

order: int¶

method: str¶

t_hist: deque[float]¶

y_hist: deque[Array]¶

reset_ic(t0, y0)¶

abstractmethod step(dt)¶

Take a time step

Parameters:: dt (float) – time step size
Return type:: StepResult

accept_step(s)¶

default implementation, maybe overridden

Parameters:: s (StepResult) – Step result

property time: float¶: current time

property state: Array¶: current state

ormatex_py.ode_utils module¶

Helper functions for ODE integrators.

ormatex_py.ode_utils.stack_u(u, n)¶

ormatex_py.ode_utils.flatten_u(u)¶

ormatex_py package¶

Submodules¶

ormatex_py.arnoldi_jax module¶

ormatex_py.integrate_wrapper module¶

ormatex_py.matexp_krylov module¶

ormatex_py.matexp_leja module¶

ormatex_py.matexp_phi module¶

ormatex_py.matexp_phi_pfd_dict module¶

ormatex_py.matexp_poly module¶

ormatex_py.ode_exp module¶

ormatex_py.ode_explicit module¶

ormatex_py.ode_sys module¶

ormatex_py.ode_utils module¶

Module contents¶

ORMATEX

Navigation

Related Topics