Skip to content

Commit 6e2d750

Browse files
ChrisRackauckasChrisRackauckas
committed
Add rootfinding blog post
1 parent ad8aad9 commit 6e2d750

File tree

1 file changed

+197
-0
lines changed

1 file changed

+197
-0
lines changed

‎news/2024/08/25/rootfinding_enzyme.md

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
@def rss_pubdate = Date(2024,8,25)
2+
@def rss = """SciML Update: Symbolic Solvers, Direct Enzyme on ODEs, and More"""
3+
@def published = " 25 August 2024 "
4+
@def title = "SciML Update: Symbolic Solvers, Direct Enzyme on ODEs, and More"
5+
@def authors = """<a href="https://github.com/ChrisRackauckas">Chris Rackauckas</a>"""
6+
7+
# SciML Update: Symbolic Solvers, Direct Enzyme on ODEs, and More
8+
9+
In this SciML update we have plenty of new features to check out including new symbolic
10+
solvers, direct Enzyme support on OrdinaryDiffEq.jl, and major loading improvements
11+
across the ecosystem. Time for the fun details!
12+
13+
## Symbolic Solvers: Solve Systems of Equations Symbolically with Symbolics.jl
14+
15+
A long-time requested feature for Symbolics.jl has been to add symbolic solvers. This is
16+
functionality to say give a system like `x^2 - y = 5, sin(y) + cos(x) = 2`, and let it
17+
spit out a symbolic exact solution. Well with this release it's finally here! Thanks to a
18+
Google Summer of Code project by Yassin ElBedwihy, the
19+
[core pull-request adding the symbolic solver](https://github.com/JuliaSymbolics/Symbolics.jl/pull/1192)
20+
has finally landed. With this PR, you can do things like:
21+
22+
```julia
23+
using Symbolics, Groebner
24+
@variables x y z;
25+
eqs = [x^2 + y + z - 1, x + y^2 + z - 1, x + y + z^2 - 1]
26+
Symbolics.symbolic_solve(eqs, [x,y,z])
27+
```
28+
29+
and get the symbolic solution to the set of polynomial equations. Want to know the answer
30+
to the above? Well install it and find it!
31+
32+
It comes complete with extensions that make use of the [Nemo](https://github.com/Nemocas/Nemo.jl)
33+
Computer Algebra System (CAS), which is an abstract algebra system developed by other smart
34+
Julia developers that specializes in "abstract algebra" like computational group theory,
35+
computational ring theory, and more. It mixes some of these techniques in with rule-based
36+
techniques to give a solver that is "best of both worlds" and is easily extendable with
37+
other submodules. For Groebner basis calculations, it uses the
38+
[Groebner.jl](https://github.com/sumiya11/Groebner.jl) package which is
39+
[demonstrated to be one of the most efficient implementations](https://arxiv.org/abs/2304.06935)
40+
and thus serves as a solid building block.
41+
42+
In more detail, The `symbolic_solve` function uses 4 hidden solvers in order to solve the
43+
user's input. Its base, `solve_univar`, uses analytic solutions up to polynomials of
44+
degree 4 and factoring as its method for solving univariate polynomials. The function's `solve_multipoly` uses GCD on the input polynomials then throws passes the result
45+
to `solve_univar`. The function's `solve_multivar` uses Groebner basis and a separating
46+
form in order to create linear equations in the input variables and a single high degree
47+
equation in the separating variable. Each equation resulting from the basis is then passed
48+
to `solve_univar`. We can see that essentially, `solve_univar` is the building block of
49+
`symbolic_solve`. If the input is not a valid polynomial and can not be solved by the
50+
algorithm above, `symbolic_solve` passes it to `ia_solve`, which attempts solving by
51+
attraction and isolation. This only works when the input is a single expression
52+
and the user wants the answer in terms of a single variable. Say `log(x) - a == 0`
53+
gives us `[e^a]`. This attraction isolation is then extendable via rules, so down the line
54+
we can add extensions to handle cases like LambertW.jl and SpecialFunctions.jl detection.
55+
56+
Its current feature completeness can be summarized as:
57+
58+
- [x] Linear and polynomial equations
59+
- [x] Systems of linear and polynomial equations
60+
- [x] Some transcendental functions
61+
- [x] Systems of linear equations with parameters (via `symbolic_linear_solve`)
62+
- [ ] Systems of polynomial equations with parameters
63+
- [ ] Inequalities
64+
- [ ] Differential Equations (ODEs)
65+
- [ ] Integrals
66+
67+
With plans to continue developing and handle the next cases soon after.
68+
69+
## Symbolics v6 / SymbolicUtils v3 / TermInterface v2: Core Interface Improvements
70+
71+
A major was released on the Symbolics stack this month, signifying a breaking change. The
72+
major breaking change here is the re-adoption of
73+
[TermInterface.jl](https://github.com/JuliaSymbolics/TermInterface.jl),
74+
which is a core interface for symbolic terms. By having all symbolic libraries extend term
75+
interface, be it Metatheory.jl, Symbolics.jl, SymPy.jl, and more, the core interface gives
76+
a common specification for building and translating terms, making all of them interopable.
77+
This means that by re-adopting TermInterface, we now have bidirectional translation to and
78+
from SymPy as a well-maintained part of the interfaces, and moving between rule-based
79+
approaches and E-graphs based approaches is a standard part of the interfaces.
80+
81+
At the same time that we tacked TermInterface v2, we also tackled some of the long-standing
82+
problems in the Symbolics ecosystem. In particular, for a symbolic term like `ex = f(x,y)`,
83+
while `arguments(ex) == [x,y]`, if you build a symbolic expression of `ex = x + y + z` the
84+
internal data structures can be optimized by assuming no ordering in such a commutative
85+
operation. However before we guarenteed argument order on `arguments(ex)`, so then
86+
`arguments(ex) = [x,y,z]` was enforced to always be sorted lexicographically. However, it
87+
turns out that after building out the symbolic stack that this choice is one of the most
88+
costly in the entire ecosystem! Thus we have changed `arguments(ex)` to not guarantee a
89+
sorting order, allowing symbolic manipulations which specialize on commutativity to skip
90+
spending 99% of their time calculating lexicographic sorts. In order to allow for printing
91+
in a stable manner, we created the new interface functions such as `sorted_arguments(ex)`
92+
which guarantee a sorting and thus take the performance hit, and this is used so that
93+
things like Latex outputs and displays are more stable.
94+
95+
Another major change was a breaking change to the `maketerm` syntax to remove the `symtype`
96+
argument. While symbolic terms like those in Symbolics.jl can still be typed, i.e.
97+
`@variables x::Complex` which changes their behavior in things like simplification rules,
98+
this information is now captured in the metadata instead of the term itself. This unifies
99+
more of the implementation between Symbolics.jl and Metatheory.jl to better allow usage
100+
of E-graphs on symbolic terms.
101+
102+
SymbolicUtils.jl v3 and Symboilcs v6 thus take these interface changes as their breaking
103+
bits. Most code should actually not be broken by these changes, we only had to update the
104+
code of approximately 10% of the upstream libraries to allow for this change, and many of
105+
those only because we the symbolics developers are using some deeper features. So it's
106+
generally a small break but we get some major performance improvements and nice new features
107+
that unify the ecosystem.
108+
109+
We had a few major Symbolics.jl updates, with this year having Symbolics v4, v5, and now v6,
110+
a few with SymbolicUtils, and new a few with TermInterface. We are happy to report that we
111+
believe TermInterface.jl is now finally stable, so these updates should have reached their
112+
conclusion. There is a common core change to SymbolicUtils.jl which should be a major
113+
performance improvement by changing the structure of the BasicSymbolic, however we believe this
114+
is likely to be non-breaking. Thus the 2023-2024 stream of majors on Symbolics seems to have
115+
come to its end and Symbolics is now in more of a feature building phase.
116+
117+
## Direct ODE Support with Enzyme
118+
119+
There are two kinds of adjoints, one is the continuous adjoint approaches which define a new
120+
ODE problem to solve, and another which uses automatic differentiation directly through the solver.
121+
There are pros and cons of each, as described in
122+
[our recent review on differentiation of ODEs](https://arxiv.org/abs/2406.09699) in exquisite
123+
detail.
124+
125+
However, improving support for both is always on the menu. With discrete adjoints, AD support
126+
directly through the solver has been supported by ForwardDiff, ReverseDiff, and Tracker since
127+
almost the dawn of DifferentialEquations.jl. Limited support for Zygote through specific
128+
solvers, such as those in SimpleDiffEq.jl, has also always existed as well. But Enzyme is a
129+
much more powerful system. We have used it rather extensively in the continuous adjoint
130+
infrastructure for almost half a decade now, but it always lacked the capability to directly
131+
differentiate the complexity of the ODE solvers...
132+
133+
Until now. As part of the JuliaCon 2024 Hackathon, the remaining issues were worked through
134+
and now the explicit methods in OrdinaryDiffEq.jl can be directly differentiated with
135+
Enzyme. This can be seen by using the AD passthrough setting to avoid the SciMLSensitivity
136+
adjoint catching as follows:
137+
138+
```julia
139+
using Enzyme, OrdinaryDiffEq, StaticArrays
140+
141+
function lorenz!(du, u, p, t)
142+
du[1] = 10.0(u[2] - u[1])
143+
du[2] = u[1] * (28.0 - u[3]) - u[2]
144+
du[3] = u[1] * u[2] - (8 / 3) * u[3]
145+
end
146+
147+
const _saveat = SA[0.0,0.25,0.5,0.75,1.0,1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0]
148+
149+
function f(y::Array{Float64}, u0::Array{Float64})
150+
tspan = (0.0, 3.0)
151+
prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
152+
sol = DiffEqBase.solve(prob, Tsit5(), saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough())
153+
y .= sol[1,:]
154+
return nothing
155+
end;
156+
u0 = [1.0; 0.0; 0.0]
157+
d_u0 = zeros(3)
158+
y = zeros(13)
159+
dy = zeros(13)
160+
161+
Enzyme.autodiff(Reverse, f, Duplicated(y, dy), Duplicated(u0, d_u0));
162+
```
163+
164+
This is a major leap forward in Enzyme support. SciMLSensitivity will soon update to include
165+
a EnzymeAdjoint option which makes use of this direct differentiation mode with automation
166+
of the process (such as unwrapping of function pointers for `SciMLBase.FullSpecialize`, so
167+
more standard definitions work). There are still a few details to work through in order to
168+
get compatability with all features, in particular
169+
[support for ranges is the leading issue](https://github.com/EnzymeAD/Enzyme.jl/issues/274),
170+
but these are actively being worked on.
171+
172+
Support for implicit methods is currently lacking in this form because Enzyme is incompatible
173+
with PreallocationTools.jl structures due to being unable to being able to prove the
174+
lack of aliasing. Thus Enzyme adjoint sensitivity support for implicit methods will directly
175+
come with the OrdinaryDiffEq v7 changes to the `autodiff` API which is planning to change
176+
from the lagacy ForwardDiff-based Jacobian interface to using DifferentiationInterface.jl
177+
for the Jacobian specification, which would then allow for Enzyme-based Jacobians and
178+
will work due to Enzyme-over-Enzyme support. This is expected over the next month as mentioned
179+
at JuliaCon.
180+
181+
## SciMLSensitivity Adjoint Support for General SciMLStructures Types
182+
183+
For many years SciMLSensitivity.jl only supported AbstractArray parameter types for its
184+
continuous adjoints because it required being able to solve differential equations
185+
based on the object type. This was relaxed a bit with the introduction of `GaussAdjoint`
186+
as the new standard adjoint method in 2023, but there were still limitations. Now thanks
187+
to a [massive effort by Dhariya](https://github.com/SciML/SciMLSensitivity.jl/pull/1057),
188+
this limitation has been lifted. Now any type which defines the
189+
[SciMLStructures interface](https://docs.sciml.ai/SciMLStructures/stable/) for
190+
tunable canonicalization can automatically be supported.
191+
192+
One major result of this is that the `MTKParameters` object from ModelingToolkit v9 is
193+
supported, which means that general adjoint differentiation is now compatible with MTK
194+
models. This includes compatability with the
195+
[SymbolicIndexingInterface](https://docs.sciml.ai/SymbolicIndexingInterface/stable/),
196+
meaning that lazy observed quantities can be differented with respect to, allowing all
197+
of the symbolic simplifications of MTK to be used within the context of AD.

0 commit comments

Comments
 (0)