This package contains bindings for the PETSc TS time integrators to allow them to be used with the SciML common interface. PETSc's linear and nonlinear solvers are already reachable from SciML through LinearSolve.jl and NonlinearSolve.jl; this package covers the third layer, TS. For more information on using the solvers from this package, see the DifferentialEquations.jl documentation.
using Pkg
Pkg.add("PETScDiffEq")Precompiling the package runs a few small solves through PETSc, so that the first solve
of a session compiles much less. To precompile without them:
using PETScDiffEq, Preferences
set_preferences!(PETScDiffEq, "precompile_workload" => false; force = true)Precompiling under mpiexec or srun skips these solves, and later sessions reuse that
build until the package or one of its dependencies changes, so load the package once
without the launcher before the first parallel run.
This library adds the common interface to PETSc's TS solvers.
See the DifferentialEquations.jl documentation for details on the interface.
Following the Lorenz example from
the ODE tutorial, we can
solve this using TSRK via the following:
using PETScDiffEq, SciMLBase
function lorenz(du, u, p, t)
du[1] = 10.0(u[2] - u[1])
du[2] = u[1] * (28.0 - u[3]) - u[2]
du[3] = u[1] * u[2] - (8 / 3) * u[3]
end
u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 100.0)
prob = SciMLBase.ODEProblem(lorenz, u0, tspan)
sol = SciMLBase.solve(prob, TSRK("5dp"); dt = 0.01, abstol = 1e-8, reltol = 1e-8)dt sets the first step. An adaptive solve can leave it out, and the first step is then
Hairer and Wanner's estimate as OrdinaryDiffEq uses it, taken with the abstol and reltol
keywords (SciML's defaults, abstol = 1e-6 and reltol = 1e-3, when not given). DAE and mass-matrix
problems start from a small step instead. A solve PETSc steps at a fixed size needs dt:
adaptive = false, -ts_adapt_type none, the fixed-step families (TSImplicit and TSDAE
other than bdf, TSIRK, TSMPRK), and subtypes registered without an embedded error
estimate. What counts is the type PETSc runs after petsc_options: TSImplicit("beuler", ["-ts_type", "bdf"]) runs without dt and TSImplicit("bdf", ["-ts_type", "beuler"]) needs
it. A TSGeneric naming a type no other constructor covers needs it too, since whether that
type adapts is not known here.
TSRK(subtype), explicit Runge-KuttaTSRosW(subtype), linearly implicit Rosenbrock-WTSImplicit(subtype; order), backward Euler, Crank-Nicolson, theta and BDFTSIRK(nstages), Gauss-Legendre implicit Runge-Kutta of order2 * nstagesTSARKIMEX(subtype), additive Runge-Kutta IMEX, for aSplitODEProblemTSDAE(subtype), the same implicit methods applied to aDAEProblemTSBasicSymplectic(subtype), symplectic splitting methods for aDynamicalODEProblemorSecondOrderODEProblemTSAlpha2(), generalized-alpha for aSecondOrderODEProblemTSGeneric(ts_type), a pass-through to any other PETScTSTypeby name
Each has a docstring covering its subtypes, whether it adapts and what it requires, so
?TSRosW at the REPL is the reference. One default worth knowing: PETSc's BDF is order 2,
so pass TSImplicit("bdf"; order = 5) when comparing against a higher-order method. Every solver takes petsc_options, a vector of
command-line style tokens passed to PETSc for that solve, which are parsed after the
options this package sets and so take precedence.
The options available in solve are documented
at the common solver options page.
This package supports dt, adaptive, dtmin, force_dtmin, dtmax, reltol and
abstol (either may be a vector of per-component tolerances), saveat, save_everystep,
save_start, save_end, save_on, save_idxs, dense, callback, tstops,
d_discontinuities, unstable_check, isoutofdomain, timeseries_errors,
dense_errors, verbose and initializealg. The warning a solve that ends early gives is logged at the
instability level of a DEVerbosity, so verbose = DEVerbosity(SciMLLogging.None())
silences it, as do SciMLLogging.None() and false. Keywords it cannot honour emit a
warning rather than being silently dropped.
Saving follows OrdinaryDiffEq: a saveat keeps only its own points, adding t0 or tf
only when it names them or save_start or save_end asks, and save_everystep = true
saves every step alongside it. dtmin is a floor the solve keeps: once PETSc proposes a
smaller step, the solve ends there with ReturnCode.DtLessThanMin. A step shortened to
land on a stop or the final time does not count, and a fixed-step solve has no floor, as in
OrdinaryDiffEq. With force_dtmin = true the solve goes on at dtmin instead, and the floor
wins over a smaller dtmax. d_discontinuities are stepped onto, and, as SciML defines them,
the step after each starts one ULP past it, so the right-hand side there sees the new regime
when written as if t > t_d. isoutofdomain(u, p, t) is asked after each step of an adaptive solve, and a
step that leaves the domain is taken again at a fifth of its size, as OrdinaryDiffEq takes it;
one that cannot be made small enough ends the solve with Unstable, or DtLessThanMin at
dtmin. An adaptive step whose error estimate is NaN or infinite, as when an explicit
method's right-hand side returns NaN or the state overflows, is taken again smaller the same
way, and both kinds of retry count in stats.nreject. An adaptive implicit step whose Newton
or linear solve fails, as when its right-hand side returns NaN, or whose Newton matrix has a
zero pivot, is taken again at PETSc's -ts_adapt_scale_solve_failed share of its size, a
quarter by default, as many times as maxiters allows, and counts in
stats.nnonlinconvfail rather than stats.nreject, as OrdinaryDiffEq counts a failed Newton
solve. A method that runs no Newton iteration, which is TSRosW or any type given
-snes_type ksponly, counts these failures in stats.nreject instead, as OrdinaryDiffEq's
Rosenbrock methods do. A zero pivot in a Newton method still counts in
stats.nnonlinconvfail, where OrdinaryDiffEq counts it in stats.nreject. A fixed-step
solve ends at its first failed Newton or linear solve with ConvergenceFailure, as
OrdinaryDiffEq's Newton-based methods do with adaptive = false.
A solve that stops short of the final time says why in its retcode: Unstable when the
state stops being finite, a step overflows, turns NaN or fails its Newton or linear solve at
every size tried, with a warning, an adaptive step is too small to move t, or unstable_check(dt, u, p, t)
returns true, which is asked before each step with the step about to be taken, as
OrdinaryDiffEq asks it, ConvergenceFailure when a fixed-step nonlinear
solve fails, DtLessThanMin as above, MaxIters after maxiters step attempts, accepted,
rejected or failed, as OrdinaryDiffEq counts them, and
Failure for a zero pivot in a fixed-step solve, with a warning, or another step PETSc cannot
take. Where
petsc_options asks PETSc to raise, with -ksp_error_if_not_converged,
-snes_error_if_not_converged or -ts_error_if_step_fails, it raises instead.
A state between step ends, for saveat, integrator(t) or a ContinuousCallback, comes
from PETSc's own interpolant for TSRK("5dp"), TSRosW("ra34pw2"), TSARKIMEX("4") and
"5", TSImplicit("bdf") and TSDAE("bdf"), and from the cubic Hermite interpolant dense
output uses for everything else, TSGeneric and a type petsc_options changes included.
With a mass matrix or a DAEProblem only PETSc's is available, and a type that has none
raises an ArgumentError when such a state is needed. integrator(t, Val{1}) is the slope
of the cubic Hermite interpolant through the step's ends, the interpolant's own derivative
where the package interpolates itself, and it raises with a mass matrix or a DAEProblem.
integrator(t; idxs), a vector of times and integrator(out, t) take the forms
OrdinaryDiffEq's integrator does.
ODEProblem, SplitODEProblem, DAEProblem, DynamicalODEProblem and
SecondOrderODEProblem are supported, in place or out of place, along with
ODEFunction's jac, jac_prototype and mass_matrix. Supply a jac_prototype for
anything sparse: without one the Jacobian is dense and forces a dense factorization.
Without a jac, the implicit algorithms build the Jacobian with ForwardDiff, as
OrdinaryDiffEq does, and colour a sparse jac_prototype, so a tridiagonal problem costs
one dual evaluation of f per Jacobian rather than one evaluation per state. The
prototype has to hold every entry the Jacobian can have: one it leaves out is left out of
the Jacobian, which then costs Newton iterations. Passing
autodiff = PETScDiffEq.AutoFiniteDiff() to the algorithm leaves the Jacobian to PETSc's
finite differences, coloured by a sparse prototype too. On a badly scaled stiff problem such as
Robertson's, those are far enough off that the solve reports success with an answer
wrong in its first digit, so keep them for a right-hand side ForwardDiff cannot run.
DiscreteCallback, ContinuousCallback, VectorContinuousCallback and CallbackSet
all work, as does the integrator interface through init, step!, solve!, reinit!,
terminate! and initialize_dae!. After a step the running solution's retcode is Success, as
OrdinaryDiffEq's is. check_error gives Success while the integrator can go on and the
retcode it stopped with after that. In a callback's finalize it gives the retcode passed to
terminate!, and Success for a solve that ended any other way, where OrdinaryDiffEq's
also gives MaxIters or Unstable. check_error! and postamble! work as SciMLBase
defines them, postamble! finishing the integrator where it is. Like terminate!, it saves
the point it stops at as OrdinaryDiffEq does: under a saveat that does not name that point,
only with save_end = true or when nothing is saved yet. Unlike OrdinaryDiffEq's, a finished
integrator cannot step again. auto_dt_reset! takes
the step init would take from the current state, and reinit! does the same with
reset_dt = true, or keeps the proposed step with reset_dt = false. A fixed-step
integrator goes on at its fixed size through both, as OrdinaryDiffEq's does, and only dt
shows the estimate. get_proposed_dt is signed, negative on a reversed span, and
set_proposed_dt! also takes another integrator whose proposed step it copies. set_abstol!
and set_reltol! hold for the steps after, until reinit! goes back to the tolerances init
was given, where OrdinaryDiffEq's reinit! keeps them. change_t_via_interpolation! with
Val{true} drops what was saved past the new time, and saves the new end when
save_everystep asks for every step. resize!, deleteat! and addat! raise an
ArgumentError, since PETSc sizes its vectors and solvers when the integrator is made.
A SecondOrderODEProblem, u'' = f(u', u, p, t), and a DynamicalODEProblem, whose state is
a velocity v and a position u with v' = f1(v, u, p, t) and u' = f2(v, u, p, t), have
two PETSc integrators of their own. Every other serial algorithm solves them as the
first-order system [v; u]' = [f1; f2], as OrdinaryDiffEq's general methods do.
TSBasicSymplectic steps the velocity and the position in turn with f1 and f2, so, as for
OrdinaryDiffEq's symplectic methods, f1 must not depend on v nor f2 on u. Its
subtypes "sieuler", "velverlet", "3" and "4" are of order 1 to 4 and step at a fixed
dt, and their energy error stays bounded over long times where a non-symplectic method's
grows:
using PETScDiffEq, SciMLBase
pendulum!(ddu, du, u, p, t) = (ddu .= -sin.(u); nothing)
prob = SecondOrderODEProblem(pendulum!, [0.0], [2.0], (0.0, 1000.0))
sol = solve(prob, TSBasicSymplectic("4"); dt = 0.1)
energy(s) = s.x[1][1]^2 / 2 - cos(s.x[2][1])
maximum(abs(energy(s) - energy(sol.u[1])) for s in sol.u) # 6.4e-6, as large at t = 1000 as at t = 500f1 is given the time of the positions it is evaluated at, so a force that depends on t
keeps each method's order; PETSc's own time for that update runs ahead of the positions.
"velverlet" then gives OrdinaryDiffEq's VelocityVerlet answer to round-off.
TSAlpha2 is PETSc's implicit generalized-alpha method of order 2, for stiff second-order
systems such as structural dynamics. radius, from 0 to 1, is its spectral radius at an
infinite step: below 1 it damps the frequencies the step cannot resolve, which radius 1, the
default, carries on undamped. It adapts on PETSc's error estimate with scalar abstol and
reltol, or steps at dt with adaptive = false. Its Jacobian, from a jac or from
autodiff as for the other implicit families, is that of the first-order system, 2n by
2n, and a sparse jac_prototype of that system keeps the n by n matrix PETSc factors
sparse, whether the Jacobian comes from a jac or is coloured, as for TSImplicit.
K, C = [1.0e6 0.0; 0.0 1.0], [200.0 0.0; 0.0 0.02]
spring!(ddu, du, u, p, t) = (ddu .= -K * u .- C * du; nothing)
prob = SecondOrderODEProblem(spring!, [0.0, 0.0], [1.0, 1.0], (0.0, 5.0))
sol = solve(prob, TSAlpha2(; radius = 0.5); dt = 0.1, adaptive = false)The states come back as OrdinaryDiffEq returns them: sol.u[i], sol(t), integrator.u and
get_du(integrator) are ArrayPartition(v, u), so sol.u[i].x[2] is the position, while
save_idxs indexes the flat [v; u] and saves plain vectors. Between steps sol(t) is the
cubic Hermite interpolant of the velocity and the position, which is also what OrdinaryDiffEq
gives for VelocityVerlet. stats.nf counts evaluations of f1, or of the whole system,
and stats.nf2 those of f2 alone.
These problems run on MPI.COMM_SELF only and take no mass matrix. Both algorithms take a
reversed tspan. PETSc only steps forward, so TSAlpha2 then integrates w(s) = u(-s),
whose velocity is -u', and gives back u': the states, jac and callbacks are those of the
problem as written. PETScAdjoint differentiates them through the first-order form
with TSRK, TSARKIMEX or TSImplicit's "beuler", "cn" or "theta"; PETSc has no
adjoint for TSBasicSymplectic or TSAlpha2.
A DAEProblem, or an ODEProblem whose mass matrix has zero rows and columns, has to start
where its algebraic equations hold. The initializealg keyword picks how, with the algorithms
OrdinaryDiffEq takes, which come from DiffEqBase:
CheckInit(), the default, evaluates the residual atu0, and atdu0for aDAEProblem, and throws aCheckInitFailureErrorwhen its RMS norm exceedsabstol, taken per component whenabstolis a vector.BrownFullBasicInit()keeps the differential variables and solves for the algebraic ones, and for aDAEProblemalso for the derivatives of the differential ones, which needsdifferential_vars. Its ownabstol,1e-10unless given, decides whether to solve.ShampineCollocationInit(initdt)takes one backward Euler step and starts from where it lands. The step isinitdtas given. Without one it runs towardtf, and for anODEProblemit is OrdinaryDiffEq's:dt / 5, at mostdtmax, whendtis given, and a thousandth of the span otherwise. For aDAEProblemit is a tenth ofdtmax, the span by default, att0 = 0, and the smaller of that and|t0| / 1000elsewhere. OrdinaryDiffEq ignoresinitdtfor aDAEProblem, and steps differently there whent0is negative or the span runs backward.NoInit()starts fromu0as given.
The two that solve use PETSc's SNES with the Jacobian the solve itself uses: the problem's
jac, the ForwardDiff one or PETSc's finite differences, sparse under a jac_prototype. So
they take no nlsolve, and when SNES fails the solve returns at t0 with
ReturnCode.InitialFailure. A start that already passes the check is left untouched, and
reinit! initializes again unless given reinit_dae = false. They do not run on a
communicator other than MPI.COMM_SELF, where CheckInit() still checks the whole state.
The default is CheckInit() even for a problem carrying ModelingToolkit's initialization
data, which OrdinaryDiffEq would solve with OverrideInit(); this package does not solve that
system and refuses OverrideInit() on such a problem.
initialize_dae!(integrator, initializealg) runs the same on the integrator's current state
and time, with the initializealg the solve was given unless another is passed, and writes
the result into PETSc. It takes du0 from the problem for a DAEProblem, the current
abstol and, for ShampineCollocationInit() on an ODEProblem, the current dt / 5, as
OrdinaryDiffEq's does. When SNES fails the integrator finishes where it is with
ReturnCode.InitialFailure, and on an ODEProblem without a singular mass matrix it does
nothing. After a callback's affect! runs without calling
derivative_discontinuity!(integrator, false), or its initialize calls
derivative_discontinuity!(integrator, true), the integrator is initialized again with the
callback's initializealg, or the solve's when the callback has none, as OrdinaryDiffEq does.
So with the default a callback has to leave the algebraic equations satisfied or the solve
throws CheckInitFailureError, and under BrownFullBasicInit() the algebraic variables are
solved for again. On a DAEProblem, whose derivative PETSc keeps to itself, CheckInit()
after a callback takes a state the callback left alone as consistent and checks a changed one
against the problem's du0, which is right at t0 only.
A Float64 state runs in PETSc's double-precision build and a ComplexF64 state in its
double-precision complex build. A Float32 or ComplexF32 state runs in the matching
single-precision build when the span is Float32 as well, following OrdinaryDiffEq's advice
to give a single-precision problem a Float32 span, and in the double-precision build when
the span is Float64, with the saved states given back in single precision. DiffEqBase
promotes the span to the type of dt, so a single-precision solve takes dt = 0.01f0
rather than dt = 0.01, and it makes a whole-number span Float64. Any other real state,
whole numbers included, is solved in Float64 and comes back in it. Where PETSc.jl has not
loaded the single-precision build, as with a library set with PETSc.set_library!, a
single-precision state runs in the double-precision one, and a problem whose build is not
loaded at all is refused with an ArgumentError that names it. On 32-bit x86 a
single-precision state always runs in the double-precision build: there PETSc_jll's
single-precision builds end BDF and ARKIMEX solves in failure at stops that its
double-precision build takes.
Times are in the type of the clock PETSc steps on: Float32 for a single-precision state
with a Float32 span and Float64 otherwise. That covers sol.t and the integrator's t
and dt, and the integrator's u is in the type PETSc steps. OrdinaryDiffEq gives the
times in the span's type whatever the state; here a Float64 state with a Float32 span is
stepped on PETSc's Float64 clock, and rounding its times to Float32 would give states
close together the same time.
Single precision has seven digits, which bounds the clock as well as the state. The first
step, and the step after a stop, are at least four ulps of t, since PETSc cannot take a
step whose stages have no room between its ends; a single-precision solve is stepped by
this package's integrator, which lands on each stop and on the final time itself, where
PETSc's own landing would refuse a step under about 1e-6 before a final time below 1.
Tolerances finer than single precision can resolve, about 1e-7, are accepted but buy
nothing past its rounding, and can drive PETSc's adaptive step below what the clock can
tell apart, which ends the solve with a failed retcode. PETSc's single build also takes the
norms its Newton and Krylov iterations stop on in single precision, and a vector whose
entries are all below about 1e-19 in size has a norm of zero there. An implicit solve of
such a state can then stop without moving it and still report success, so this package
warns when a solve starts on such a state, or fails on one; a state that decays there from
above is within any coarser tolerance. Rescale the problem, or give it a Float64 span.
With a complex state, times, dt, saveat, tstops and the tolerances stay real, and
PETSc's error norms take each component's modulus; a tolerance given as a complex number
with a zero imaginary part is taken as its real part. A ContinuousCallback's condition has
to return a real number, such as real(u[1]) - 0.5, since a root is a sign change. The
implicit methods' Newton iteration needs a holomorphic f, one that does not go through
conj, abs, real or imag of the state. ForwardDiff takes no complex numbers, so
without a jac the Jacobian is differentiated along the real parts of the state, which for
a holomorphic f is its complex Jacobian, and a sparse jac_prototype, or the pattern a
sparse backend is given, is coloured as for a real state. A check at the start compares the
derivatives along the real and the imaginary parts and refuses an f that is not
holomorphic; it is best effort, and can miss a term too small to show near the initial
state. AutoFiniteDiff() and a hand-written jac are not
checked. The explicit methods take any f.
With SciMLSensitivity loaded, adjoint_sensitivities(sol, alg; sensealg = PETScAdjoint(), ...)
runs PETSc's own discrete adjoint for TSRK, TSARKIMEX and TSImplicit's "beuler",
"cn" and "theta", for discrete costs, integral costs through PETSc's quadrature TS, or
both; TSARKIMEX takes a SplitODEProblem too, and discrete costs only. The keywords
that set the steps have to be repeated from solve. It runs in PETSc's double
real build, so a Float32 problem is differentiated in Float64 and its gradients come
back as Float32, and a complex one is refused. ?PETScAdjoint and the documentation cover
what it needs, what it refuses and how to check jac and paramjac.
Every algorithm takes a comm keyword. With a communicator
other than the default MPI.COMM_SELF the solve runs distributed over it: every rank of
comm calls solve with the same arguments, and u0 is the block of the state that rank
owns, the blocks following each other in rank order. Each rank's sol.u holds its own rows,
and sol.t is the same on every rank.
f(du, u, p, t) sees only its rank's rows, so it fetches what it needs from the other ranks
itself, and it has to be collective: the package calls it the same number of times, in the
same order, on every rank. A 1-D heat equation with eight rows on each rank:
using MPI, PETScDiffEq, SciMLBase
MPI.Init()
comm = MPI.COMM_WORLD
rank, nranks = MPI.Comm_rank(comm), MPI.Comm_size(comm)
dx = 1 / (8nranks + 1)
x = (8rank .+ (1:8)) .* dx
function heat!(du, u, p, t)
left = rank == 0 ? MPI.PROC_NULL : rank - 1
right = rank == nranks - 1 ? MPI.PROC_NULL : rank + 1
gl, gr = zeros(1), zeros(1)
MPI.Sendrecv!(u[1:1], gr, comm; dest = left, source = right)
MPI.Sendrecv!(u[end:end], gl, comm; dest = right, source = left)
for i in eachindex(u)
l = i == 1 ? gl[1] : u[i - 1]
r = i == length(u) ? gr[1] : u[i + 1]
du[i] = (l - 2u[i] + r) / dx^2
end
end
sol = solve(ODEProblem(heat!, sinpi.(x), (0.0, 0.1)), TSRK("5dp"; comm))Run it with the mpiexec MPI.jl provides, MPI.mpiexec(), as in mpiexec -n 4 julia heat.jl.
PETSc starts up collectively over MPI.COMM_WORLD, so under mpiexec a rank that solves on
its own, on MPI.COMM_SELF too, hangs unless an earlier solve or PETSc.initialize has
started PETSc on every rank.
saveat, tstops, d_discontinuities, a fixed dt and dense output work as in a serial
solve. Vector abstol and reltol, save_idxs and p are per rank.
unstable_check and isoutofdomain are asked on each rank's rows, and true on any rank
counts on all of them. A step that turns NaN or overflows on any rank's rows is taken again
smaller on every rank, and a fixed-step solve stops at the first state that is not finite on
some rank, as in a serial solve. When f, jac or one of those checks throws on some ranks,
those ranks go on with NaN until the ranks next agree, at the end of the step or when its
nonlinear solve fails, and then every rank throws rather than retrying the step, so an f
that throws has to do so after its own communication.
Callbacks and the integrator interface run distributed too, as long as every rank makes the
same calls with the same arguments in the same order: init, step!, solve!, reinit!,
terminate!, set_u!, initialize_dae!, add_tstop!, add_saveat!, savevalues!,
change_t_via_interpolation!, set_proposed_dt!, set_abstol!, set_reltol!,
postamble!, auto_dt_reset!, integrator(t) and get_du are all collective.
integrator.u holds the rank's own rows, and so
does the state given to set_u! or reinit!. set_proposed_dt! takes the smallest step any
rank proposes.
A callback's condition should give the same value on every rank, which usually means it
reduces over the ranks itself, for instance with MPI.Allreduce. The package reduces the
conditions as well, so ranks whose conditions disagree still stay together: a
DiscreteCallback fires when its condition is true on any rank, and a ContinuousCallback
or VectorContinuousCallback fires at the earliest event any rank finds, with that rank's
crossing. The affect then runs on every rank whatever its own condition gave, so it has to be
collective as well: an affect that calls terminate! has to call it on every rank. A condition,
affect, initialize or finalize that throws on some ranks makes every rank throw, as f
does, so an affect that throws has to do so after its own communication.
The implicit algorithms build their Jacobian as a distributed PETSc matrix whose pattern
comes from the problem's jac_prototype, which then holds this rank's rows only: it is
length(u0) by the length of the whole state, with global column indices. A jac fills
those rows, and is collective like f. For the heat equation above:
using SparseArrays
N = 8nranks
rows = 8rank .+ (1:8)
near(i) = max(1, i - 1):min(N, i + 1)
proto = sparse(
[k for (k, i) in enumerate(rows) for _ in near(i)], [j for i in rows for j in near(i)],
1.0, 8, N,
)
function heat_jac!(J, u, p, t)
for (k, i) in enumerate(rows), j in near(i)
J[k, j] = (i == j ? -2 : 1) / dx^2
end
end
f = ODEFunction(heat!; jac = heat_jac!, jac_prototype = proto)
sol = solve(ODEProblem(f, sinpi.(x), (0.0, 0.1)), TSImplicit("bdf"; comm))Without a jac, autodiff defaults to AutoFiniteDiff() on such a comm: PETSc colours the
prototype's pattern and differences f, calling it the same number of times on every rank.
ForwardDiff and the other autodiff backends are refused there, since the number of times
they call f differs between ranks. So are a jac or colouring without a sparse prototype,
a dense mass matrix, and TSIRK without a jac.
A mass matrix is a Diagonal of this rank's entries, or a sparse matrix holding this rank's
rows with global column indices, as the prototype does, such as a finite element mass
matrix. PETSc assembles a sparse one into a distributed matrix, and its pattern joins the
prototype's in the Jacobian a*M - J of the implicit solve. A rank whose rows of the mass
matrix are the identity can leave it as I, whatever the other ranks give.
TSIRK also needs each rank to hold PETSc's own share of the state, split evenly with the
first ranks taking one row more, since PETSc lays out its stage vector that way.
PETSc solves the linear systems of a distributed solve with GMRES and block Jacobi, one
ILU(0) block on each rank, to a relative tolerance of 1e-5, so such a solve agrees with a
serial one to that accuracy rather than to round-off. Options such as -ksp_rtol or
-sub_pc_type in petsc_options change that solver.
TSMPRK's slow and medium index the rank's own rows, and either may be empty on some
ranks as long as some rank names a slow row and, for "2a23" and "2a33", a medium one. An
implicit TSGeneric runs distributed for "beuler", "cn", "theta", "bdf", "rosw",
"arkimex", "irk", "alpha" and "dirk", as TSImplicit does. Other implicit types are
refused: "glle"'s step control follows the round-off of the distributed linear solve, so it
takes other steps than a serial solve and ends with another error, larger or smaller.
PETScAdjoint runs distributed too, for TSRK, TSARKIMEX on an ODEProblem and
TSImplicit's "beuler", "cn" and "theta". It needs the problem's jac, filling this
rank's rows of a sparse prototype as above, and when there are parameters a paramjac
filling this rank's rows, since automatic differentiation would call f a different number
of times on each rank; both are collective like f. An explicit method takes such a jac
in its own solve too, and ignores it there. dgdu_discrete gets this rank's rows of the
state and writes their gradient, and dgdp_discrete gives this rank's share of the cost's
direct derivative with respect to p, which the ranks add up. du0 comes back as this
rank's rows and dp as the whole gradient, the same on every rank. The cost times,
no_start, the length of p and whether dgdp_discrete is given have to agree across the
ranks. A jac, paramjac, cost function or f that throws on some ranks makes every rank
throw, as in a solve. The transposed linear solves of TSImplicit and TSARKIMEX use the
solver above; ["-ksp_type", "preonly", "-pc_type", "redundant"] in petsc_options solves
them directly.
A PETSc DM can do the halo exchange instead. Build a DMDA with PETSc.jl and pass it as dm,
which every algorithm that takes comm takes as well. The solve then runs on the DM's
communicator, which comm may name too but not contradict, and u0 is the block of the grid
this rank owns, in the DM's order. f(du, u, p, t) gets u ghosted: before every call to f,
including the package's own calls for the first step size, dense output, callbacks, get_du
and the Jacobian, the state is scattered into a local vector from DMGetLocalVector with
DMGlobalToLocalBegin and DMGlobalToLocalEnd, so u also holds the neighbouring ranks'
points within the stencil width. du is the owned block. PETScDiffEq.reshape_local_array(x, dm)
views either one by grid point in global numbering, as x[c, i] on a 1-D grid and x[c, i, j]
on a 2-D one, where c is the degree of freedom at the point. It is PETSc.jl's
reshape_local_array, which PETSc.jl 0.4 calls reshapelocalarray. With DM_BOUNDARY_GHOSTED
the ghost points past the edge of the grid read zero, so this heat equation is zero at both
ends:
using MPI, PETScDiffEq, SciMLBase
using PETScDiffEq: PETSc, LibPETSc
MPI.Init()
petsclib = PETSc.getlib(; PetscScalar = Float64)
PETSc.initialize(petsclib)
N = 64
dx = 1 / (N + 1)
da = PETSc.DMDA(petsclib, MPI.COMM_WORLD, (LibPETSc.DM_BOUNDARY_GHOSTED,), (N,), 1, 1)
function heat!(du, u, da, t)
U = PETScDiffEq.reshape_local_array(u, da)
D = PETScDiffEq.reshape_local_array(du, da)
for i in axes(D, 2)
D[1, i] = (U[1, i - 1] - 2U[1, i] + U[1, i + 1]) / dx^2
end
end
xs, _, _, xm = LibPETSc.DMDAGetCorners(petsclib, da)
prob = ODEProblem(heat!, sinpi.((xs .+ (1:xm)) .* dx), (0.0, 0.1), da)
sol = solve(prob, TSRK("5dp"; dm = da))
sol_bdf = solve(prob, TSImplicit("bdf"; dm = da))With a dm the implicit algorithms need no jac_prototype. Their Jacobian is the DM's own
matrix from DMCreateMatrix, whose pattern comes from the DM's stencil. Without a jac PETSc
fills it by colouring it and differencing f, so autodiff defaults to AutoFiniteDiff()
and the other backends are refused; the stencil has to cover every point f reads. Colouring
calls f once per colour at every Jacobian, and the colours grow with the stencil width and
the degrees of freedom at a point, so a jac(J, u, p, t) can fill the matrix instead. It gets
u ghosted, as f does, and J is the DM's matrix as a PETSc.jl Mat, zeroed before the
call and assembled after it, so jac writes df/du into it through PETSc's matrix API,
each rank its own rows: set_stencil_values!(J, rows, cols, vals) writes a block by grid
index through MatSetValuesStencil, in the numbering reshape_local_array uses, and
J[i, j] = v writes one entry by global index through MatSetValues, for those who have the
global numbering; LibPETSc has the rest. A column at a ghost point past the edge of a
DM_BOUNDARY_GHOSTED grid is dropped, having no global entry, and a set and an add cannot
follow each other without a PETSc.assemble!(J) between them. The package turns the matrix
into PETSc's shift * M - J itself. A DAEProblem's jac(J, du, u, p, gamma, t) gets u
ghosted and du owned, and writes PETSc's whole dG/du + gamma dG/du'. jac runs on every
rank at every Jacobian, so anything collective in it has to be called on all of them in the
same order, and it has to be in place: an out-of-place one is refused, as is a
jac_prototype, since the DM gives the pattern, and a jac on an explicit method, which
would never use it. autodiff is ignored with a jac. The heat equation above with its
Jacobian, whose columns past the ends of the grid are dropped:
function heat_jac!(J, u, da, t)
for i in (xs + 1):(xs + xm)
set_stencil_values!(J, (1, i), ((1, i - 1), (1, i), (1, i + 1)), (1, -2, 1) ./ dx^2)
end
end
fn = ODEFunction(heat!; jac = heat_jac!)
sol_jac = solve(ODEProblem(fn, sinpi.((xs .+ (1:xm)) .* dx), (0.0, 0.1), da), TSImplicit("bdf"; dm = da))The rest works as it does without a DM: TSRK, TSRosW, TSImplicit, TSDAE,
TSARKIMEX and TSGeneric(ts_type; explicit = true), saveat, dense output, callbacks and
the integrator interface, a Diagonal mass matrix, a SplitODEProblem, whose f2 gets u
ghosted as f does, and a DAEProblem, whose residual f(r, du, u, p, t) gets u ghosted
and du owned. Everything else the package calls, such as a callback, unstable_check or
isoutofdomain, sees the owned block. The TS works on a copy of the DM from DMClone, so the
DM itself stays free for further solves. A DMDA on MPI.COMM_SELF, or on a single rank, gives
a serial solve. Only a DMDA is taken so far.
A solve with a dm refuses TSIRK, TSMPRK, an implicit TSGeneric and PETScAdjoint
with an ArgumentError. A distributed solve, with a dm or without, is refused inside
Threads.@threads on more than one thread, as EnsembleThreads runs its trajectories:
nothing there keeps the ranks' solves in the same order, and ranks taking them in different
orders run different solves as one and can return wrong results without an error.
Distributed solves running at once from Threads.@spawn tasks are not refused, so the caller
has to keep them in the same order on every rank. An ensemble of distributed solves runs with
EnsembleSerial().
A DynamicalODEProblem or SecondOrderODEProblem does not run distributed yet, whatever the
algorithm, so TSBasicSymplectic and TSAlpha2 run on MPI.COMM_SELF only. PETSc TS is
built for large distributed problems, and reaching it from the SciML interface is what this
package is for; use OrdinaryDiffEq.jl for serial problems where it applies.
On 32-bit Julia, use Julia 1.10, or add PETSc_jll = "~3.22" to your own compat: PETSc_jll
3.25 has no 32-bit builds, and newer Julia versions would otherwise resolve it.
Solves from several threads, such as an EnsembleThreads ensemble, are safe but run one
at a time: PETSc's options and MPI are shared by the whole process. Distributed ones are not,
as the MPI section says.
Finish or terminate every integrator you start. One dropped part way is released by a finalizer, and if that finalizer runs at process exit, after MPI has shut down, PETSc's own object finalizers print an MPI warning and the process exits non-zero.
MIT. See LICENSE.