Tag Archives: Julia

Accelerate Your Julia Code with Effective Profiling Methods

By: Great Lakes Consulting

Re-posted from: https://blog.glcs.io/profiling_allocations

This post was written by Steven Whitaker.

In this Julia for Devs post,we will discussusing Julia’s Profile standard libraryfor performance and allocation profiling.We will illustrate these toolswith example codeand then show how to improve the code.

This post will showcasethe powerful techniques we usedto significantly accelerateour clients simulations,resulting in a remarkable 25% reduction in run timeof code that was already highly optimized for performance.The impact of these techniqueson our client’s workis a testament to their effectivenessand should inspire youin your own projects.

Many new Julia developersface poor code performance,but these techniques can net10x or even 100x faster run times!

While we won’t be sharing client-specific code,we will provide similar examplesthat are practical and applicableto a wide range of simulations.This approachshould give you the confidenceto apply these techniquesin your work.

Here are some of the key ideas we will focus on:

  • Pinpoint areas of improvement with Profile.@profile.
  • Track down allocations with Profile.Allocs.@profile.
  • Improve run time and reduce allocations with @generated functions.

Let’s dive in!

Profiling in Julia

Profiling is a crucial toolfor locating performance bottlenecks.This knowledge is invaluableas it guides development effortsto the parts of the codethat will have the most significant impact on run time.Understanding this processwill keep you informedand in controlof your code’s performance.

Here is some example codewe will useto illustrate profiling in Julia:

using StaticArrays: SVectorfunction kernel_original(x::SVector{N}) where N    y = zeros(N + 1)    y[1] = x[1]    for i = 2:N        y[i] = 0.5 * y[i-1] + x[i]    end    y[end] = sum(@view(y[1:end-1]))    return SVector{N+1}(y)endfunction workflow_original()    total = 0.0    for i = 1:100        x = SVector{10, Float64}(rand(10))        y = kernel_original(x)        total += y[end]    end    return totalend

The main idea with this exampleis that there is a workflow,workflow_original,that calls out to a core computationfor many different input values.

To profile this code,we will use the Profile standard library.Since Profile implements a statistical profiler,we need to ensure the code we profileruns long enoughto reduce the impact of noise in the measurements(see the documentation for more info).So,let’s see how long workflow_original takes(using BenchmarkTools.jl):

julia> using BenchmarkToolsjulia> @btime workflow_original();  6.410 s (300 allocations: 25.00 KiB)

6 microseconds is very fastcompared to the default profiling sample delayof 1 millisecond.Therefore,we must run the workflow many timesto get enough samples.We have found that 1000–5000 samplesare usually plentyfor identifying performance bottlenecks,though slower code will likely need more samplesand/or a larger sample delay.

So,to get approximately 1000 samples,we will need to run the workflow\( 1000 \cdot \frac{1 \mathrm{ms}}{0.006 \mathrm{ms}} = 166,667 \) times.We’ll round up to 200,000 for good measure.

Let’s see how to profile the code:

julia> using Profilejulia> Profile.clear()julia> workflow_original();julia> Profile.@profile for _ in 1:200_000 workflow_original() end

A couple of notes:

  • The Profile.clear() stepis not strictly necessary.However, a habit of always clearing the profile dataensures one doesn’t inadvertently spoil their profileswith old data from previous profiling results.
  • We called workflow_original once before profilingto avoid profiling JIT compilation.

Now we want to display the profile data.One way is via Profile.print,which prints a textual representationof the profile to the console,but this isn’t the most efficient methodfor inspecting the data.Typically,profile data is visualizedusing a flamegraph.

The Profiling docs list several packagesfor visualizing profiles.We will use ProfileCanvas.jl,which creates an HTML filewe can view in a web browserand interactively inspect the profiling results:

julia> using ProfileCanvasjulia> ProfileCanvas.view()

This code creates and displays a flamegraph:

Profile of original workflow

One thing that stands out in this flamegraphis the three yellow rectangles.These indicate occurrences of garbage collection (GC),which implies the code allocated memory.Since these yellow blocksrepresent a decent portion of the run time(as indicated by the width of the rectangles),let’s investigate.

Allocation Profiling

We will use Julia’s allocation profiling toolsto further inspectthe allocations we know the workflow has.The process is similar to performance profiling,with the following differences:

  • We will use Profile.Allocs.@profileand pass the option sample_rate = 1.0to record all allocations.In a larger workflow with many more allocations,a smaller sample_rate is advised(the default value is 0.1).
  • We will run the workflow just one time.
  • We will visualize the results with ProfileCanvas.view_allocs.

Here’s how to profile allocations:

julia> using Profile, ProfileCanvasjulia> Profile.Allocs.clear()julia> workflow_original();julia> Profile.Allocs.@profile sample_rate = 1.0 workflow_original();julia> ProfileCanvas.view_allocs()

Allocation profile of original workflow

In this flamegraph,the widths of each rectanglecorrespond to how many allocations were made.In this example,each yellow blockis the same width,meaning they all contributedthe same number of allocations.

Now we need to determinewhether we can do anything about these allocations.

Moving up from the bottom of the flamegraphlets us trace the call stackto see where these allocations came from.For example,we can see that GenericMemorywas called from Array,which itself was called from Array,and so on.Eventually,we get to workflow_original.

If we hover our mouse cursor over that block,it will display the file and line number:

Mouse hover displaying file and line info

(In this case,I defined workflow_originalin the first REPL promptof a fresh Julia session,so that’s why REPL[1] shows up for the file name.)

Looking at these profile resultsshows us that the offending lines are:

  • x = SVector{10, Float64}(rand(10)) (in workflow_original)
  • y = zeros(N + 1) (in kernel_original)

This makes sense;in each of these lines,we explicitly create an array(via rand and zeros),which allocates memory.

But it turns outwe can eliminate these allocations!

Eliminating the first allocationrequires knowledgeof the StaticArrays.jl package.In particular,the @SVector macrocan create an SVectordirectly from standard array expressions(like rand)without allocating memory.So,instead of using the SVector constructor,we can write:

x = @SVector rand(10)

The second allocation,however,is a bit trickier to remove.

Reducing Allocations with @generated Functions

What makes it difficultto remove the allocationin kernel_originalis the elements of the vector ydepend on previous elementsof the vector.That means we have to compute one elementto compute the next,which means we have to store that element somehow.If we knew exactly how long y would be,we could store the computationsin local (scalar) variables.However, the function needs to work with any input size;we don’t want to create a separate methodfor each possible input size.

At least, not by hand.

It turns outJulia can automatethe creation of these specialized methods.The trick is to use the @generated macro.

A method annotated with @generateduses type informationto produce specialized implementationsof the methoddepending on the input types.And since type information is available at compile time,this specialization occurs at compile time,leading to run time improvements.

Using a @generated functionto replace kernel_originalwill allow us to, essentially,move the run time allocationto compile time.

Here’s what the new function looks like:

@generated function kernel_generated(x::SVector{N}) where N    assignments = [:(y1 = x[1])]    for i = 2:N        yprev = Symbol(:y, i - 1)        yi = Symbol(:y, i)        push!(assignments, :($yi = 0.5 * $yprev + x[$i]))    end    yend = Symbol(:y, N + 1)    sum_expr = reduce((a, b) -> :($a + $b), (Symbol(:y, i) for i = 1:N))    push!(assignments, :($yend = $sum_expr))    vars = ntuple(i -> Symbol(:y, i), N + 1)    return quote        $(assignments...)        return SVector{$(N + 1), Float64}($(vars...))    endend

The first thing you might notice,especially if you are unfamiliar with metaprogramming,is that this function looks quite differentfrom kernel_original.So, let’s unpack this a bit:

  • A @generated functionneeds to return an expression.The compiler will then compilethe code resulting from the expression.Finally, at run time,the compiled code will be used,not the code usedto generate the compiled expression.In other words,this function must returnthe specialized code itself,not the result of a run time computation.
  • The returned expressionis created with a quote block.This quote blockinterpolates (using $) expressionsbuilt up earlier in the function.The computationsto build up these expressionsoccur only at compile time.

If we want to inspectwhat the generated function looks like,we need to refactor the code just a bit.Essentially,we’ll create a regular Julia functionthat returns an expression,and the @generated functionwill just call that function:

function _gen_kernel_generated(::Type{<:SVector{N}}) where N    # Same code as `kernel_generated` above.end@generated kernel_generated(x::SVector) = _gen_kernel_generated(x)

Then we can call _gen_kernel_generated directlyto see what code actually runs:

julia> _gen_kernel_generated(typeof(@SVector rand(1)))quote    #= REPL[2]:18 =#    y1 = x[1]    y2 = y1    #= REPL[2]:19 =#    return SVector{2, Float64}(y1, y2)endjulia> _gen_kernel_generated(typeof(@SVector rand(2)))quote    #= REPL[2]:18 =#    y1 = x[1]    y2 = 0.5y1 + x[2]    y3 = y1 + y2    #= REPL[2]:19 =#    return SVector{3, Float64}(y1, y2, y3)endjulia> _gen_kernel_generated(typeof(@SVector rand(3)))quote    #= REPL[2]:18 =#    y1 = x[1]    y2 = 0.5y1 + x[2]    y3 = 0.5y2 + x[3]    y4 = (y1 + y2) + y3    #= REPL[2]:19 =#    return SVector{4, Float64}(y1, y2, y3, y4)end

Yep, the implementation looks right!

We can also seethat the generated codeavoids creating an arrayto store intermediate results,instead storing computationsin local variables.But we didn’t have to writeany of those methods ourselves!Using @generated allows usto maintain just one functionto generate all these specialized implementations.

Let’s benchmark the new implementation:

julia> @btime workflow_generated();  1.093 s (0 allocations: 0 bytes)

Nice, six times fasterand no allocations!

And here’s the profile:

Profile of optimized workflow

(Click here to see the code for profiling.) juliausing StaticArrays: @SVector, SVector@generated function kernel_generated(x::SVector{N}) where N assignments = [:(y1 = x[1])] for i = 2:N yprev = Symbol(:y, i - 1) yi = Symbol(:y, i) push!(assignments, :($yi = 0.5 * $yprev + x[$i])) end yend = Symbol(:y, N + 1) sum_expr = reduce((a, b) -> :($a + $b), (Symbol(:y, i) for i = 1:N)) push!(assignments, :($yend = $sum_expr)) vars = ntuple(i -> Symbol(:y, i), N + 1) return quote $(assignments...) return SVector{$(N + 1), Float64}($(vars...)) endendfunction workflow_generated() total = 0.0 for i = 1:100 x = @SVector rand(10) y = kernel_generated(x) total += y[end] end return totalendusing BenchmarkTools@btime workflow_generated();using Profile, ProfileCanvasProfile.clear()Profile.@profile for _ in 1:200_000 workflow_generated() endProfileCanvas.view()Note,there’s no need for allocation profilingbecause there are no allocations.

There aren’t obvious places for improvement,so I’d say we did a pretty good joboptimizing the code.

Summary

In this post,we saw how to use the Profile standard libraryto profile the run timeand the allocationsof a piece of code.We also illustratedhow a @generated functioncan eliminate run time allocationsand speed up the code.These were key ideas we usedto help one of our clientsspeed up their simulations.

Do you need help pinpointing performance bottlenecksor tracking down allocationsin your code?Contact us, and we can help you out!

Additional Links

]]>

Implicit ODE Solvers Are Not Universally More Robust than Explicit ODE Solvers, Or Why No ODE Solver is Best

By: Christopher Rackauckas

Re-posted from: https://www.stochasticlifestyle.com/implicit-ode-solvers-are-not-universally-more-robust-than-explicit-ode-solvers-or-why-no-ode-solver-is-best/

A very common adage in ODE solvers is that if you run into trouble with an explicit method, usually some explicit Runge-Kutta method like RK4, then you should try an implicit method. Implicit methods, because they are doing more work, solving an implicit system via a Newton method having “better” stability, should be the thing you go to on the “hard” problems.

This is at least what I heard at first, and then I learned about edge cases. Specifically, you hear people say “but for hyperbolic PDEs you need to use explicit methods”. You might even intuit from this “PDEs can have special properties, so sometimes special things can happen with PDEs… but ODEs, that should use implicit methods if you need more robustness”. This turns out to not be true, and really understanding the ODEs will help us understand better why there are some PDE semidiscretizations that have this “special cutout”.

What I want to do in this blog post is more clearly define what “better stability” actually means, and show that it has certain consequences that can sometimes make explicit ODE solvers actually more robust on some problems. And not just some made-up problems, lots of real problems that show up in the real world.

A Quick Primer on Linear ODEs

First, let’s go through the logic of why implicit ODE solvers are considered to be more robust, which we want to define in some semi-rigorous way as “having a better chance to give an answer closer to the real answer”. In order to go from semi-rigorous into a rigorous definition, we can choose a test function, and what better test function to use than a linear ODE. So let’s define a linear ODE:

$$u’ = \lambda u$$

is the simplest ODE. We can even solve it analytically, $u(t) = \exp(\lambda t)u(0)$. For completeness, we can generalize this to a linear system of ODEs, where instead of having a scalar $u$ we can let $u$ be a vector, in which case the linear ODE has a matrix of parameters $A$, i.e.

$$u’ = Au$$

In this case, if $A$ is diagonalizable, $A = P^{-1}DP$, then we can replace $A$:

$$u’ = P^{-1}DP u$$

$$Pu’ = DPu$$

or if we let $w = Pu$, then

$$w’ = Dw$$

where $D$ is a diagonal matrix. This means that for every element of $w$ we have the equation:

$$w_i’ = \lambda_i w_i$$

where $w_i$ is the vector in the direction of the $i$th eigenvector of $A$, and $\lambda_i$ is the $i$th eigenvalue of $A$. Thus our simple linear ODE $u’ = \lambda u$ tells us about general linear systems along the eigenvectors. Importantly, since even for real $A$ we can have $\lambda$ be a complex number, i.e. real-valued matrices can have complex eigenvalues, it’s important to allow for $\lambda$ to be complex to understand all possible systems.

But why is this important for any other ODE? Well by the Hartman-Grobman theorem, for any sufficiently nice ODE:

$$u’ = f(u)$$

We can locally approximate the ODE by:

$$u’ = Au$$

where $A = f'(u)$, i.e. $A$ is the linear system defined by the Jacobian local to the point. This is effectively saying any “sufficiently nice” system (i.e. if $f$ isn’t some crazy absurd function and has properties like being differentiable), you can understand how things locally move by looking at the system approximated by a linear system, where the right linear approximation is given by the Jacobian. And we know that linear systems then boil down generally to just the scalar linear system, and so understanding the behavior of a solver on the scalar linear system tells us a lot about how it will do “for small enough h”.

Okay, there are lots of unanswered questions, such as what if $A$ is not diagonalizable? What if $f$ is not differentiable? What if the system is very nonlinear so the Jacobian changes very rapidly? But under assumptions that things are nice enough, we can say that if a solver does well on $u’ = \lambda u$ then it is probably some idea of good.

Why implicit ODE solvers are “better”, i.e. more robust

So now we have a metric by which we can analyze ODEs: if they have good behavior on $u’ = \lambda u$, then they are likely to be good in general. So what does it mean to have good behavior on $u’ = \lambda u$? One nice property would be to at least be asymptotically correct for the most basic statement, i.e. does it go to zero when it should? If you have $u’ = \lambda u$ and $\lambda$ is negative, then the analytical solution $u(t) = \exp(\lambda t)u(0)$ goes to zero as $t$ goes to infinity. So a good question to ask is, for a given numerical method, for what values of $h$ (the time step size) does the numerical method give a solution that goes to zero, and for which $h$ does it get an infinitely incorrect answer?

To understand this, we just take a numerical method and plug in the test equation. So the first thing to look at is Euler’s method. For Euler’s method, we step forward by $h$ by assuming the derivative is constant along the interval, or:

$$u_{n+1} = u_n + hf(u_n)$$

When does this method give a solution that is asymptotically consistent? With a little bit of algebra:

$$u_{n+1} = u_n + h\lambda u_n$$

$$u_{n+1} = (1 + h\lambda) u_n$$

Let $z = h\lambda$, which means

$$u_{n+1} = (1 + z) u_n$$

This is a discrete dynamical system which has the analytical solution:

$$u_n = u_0 (1+z)^{n}$$

Note that if $1 + z > 1$, then $(1+z)^n$ keeps growing as $n$ increases, so this goes to infinity, while if $1 + z < 1$ it goes to zero. But since $\lambda$ can actually be a complex number, the analysis is a little bit more complex (pun intended), but it effectively means that if $z$ is in the unit circle shifted to the left in the complex plane by 1, then $u_n \rightarrow 0$. This gives us the definition of the stability region, $G(z)$ is the region for which $u_n \rightarrow 0$, and this is the shifted unit circle in the complex plane for explicit Euler.

This shows a pretty bad property for this method. For any given $\lambda$ with negative real part, there is a maximum $h$, actually $h = 1/\lambda$, such that for any larger step size we don’t just get a bad answer, we can get an infinitely bad answer, i.e. the analytical solution goes to zero but the numerical solution goes to infinity!

So, is there a method that doesn’t have this bad property? In comes the implicit methods. If you run the same analysis with implicit Euler,

$$u_{n+1} = u_n + hf(u_{n+1})$$

$$u_{n+1} = u_n + h\lambda u_{n+1}$$

$$(1-z) u_{n+1} = u_n$$

$$u_{n+1} = \frac{1}{1-z} u_n$$

Then we have almost an “inverse” answer, i.e. $G(z)$ is everything except the unit circle in the complex plane shifted to the right. This means that for any $\lambda$ with negative real part, for any $h$ the implicit Euler method has $u_n \rightarrow 0$, therefore it’s never infinitely wrong.

Therefore it’s just better, QED.

This then generalizes to more advanced methods. For example, the stability region of RK4

an explicit method has a maximum $h$, while the stability region of BDF2

an implicit method does not. You can even prove it’s impossible for any explicit method to have this “good” property, so “implicit methods are better”. QED times 2, done deal.

Wait a second, what about that other “wrongness”?

Any attentive student should immediately throw their hand up. “Teacher, given the $G(z)$ you said, you also have that for any $\lambda$ where $\text{Re}(\lambda)>1$, you also have that $u_n \rightarrow 0$, but in reality the analytical solution has $u(t) \rightarrow \infty$, so implicit Euler is infinitely wrong! And explicit Euler has the correct asymptotic behavior since it goes to infinity!”

That is completely correct! But it can be easy to brush this off with “practical concerns”. If you have a real model which has positive real eigenvalues like that, then it’s just going to explode to infinity. Those kinds of models aren’t really realistic? Energy goes to infinity, angular momentum goes to infinity, the chemical concentration goes to infinity: whatever you’re modeling just goes crazy! If you’re in this scenario, then your model is probably wrong. Or if the model isn’t wrong, the numerical methods aren’t very good anyways. If you analyze the error propagation properties, you’ll see the error of the numerical method also increases exponentially! So this is a case you shouldn’t be modeling anyways.

Seeing this robustness in practice

Therefore if you need a more accurate result, use an implicit method. And you don’t need to go to very difficult models to see this manifest in practice. Take the linear ODE:

$$T’ = 5(300-T)$$

with $T(0) = 320$. This is a simple model of cooling an object with a constant temperature influx. It’s easy to analytically solve, you just have an exponential fall in the temperature towards $T = 300$ the steady state. But when we solve it with an explicit method at default tolerances, that’s not what we see:

using OrdinaryDiffEq
function cooling(du,u,p,t)
    du[1] = 5.0*(300-u[1])
end
u0 = [310.0]
tspan = (0.0,10.0)
prob = ODEProblem(cooling, u0, tspan)
sol = solve(prob, Tsit5())
 
using Plots
plot(sol, title="RK Method, Cooling Problem")
savefig("rk_cooling.png")

We see that the explicit method gives oscillations in the solution! Meanwhile, if we take a “robust” implicit method like the BDF method from the classic C++ library SUNDIALS, we can solve this:

using Sundials
sol = solve(prob, CVODE_BDF())
plot(sol, title="BDF Method, Cooling Problem")
savefig("bdf_cooling.png")

Sure it’s not perfectly accurate, but at least it doesn’t give extremely wrong behavior. We can decrease tolerances to make this all go away,

But the main point is that the explicit method is just generally “less robust”, you have to be more careful, it can give things that are just qualitatively wrong.

This means that “good tools”, tools that have a reputation for robustness, they should default to just using implicit solvers because that’s going to be better. And you see that in tools like Modelica. For example, the Modelica University’s playground and other tools in the space like OpenModelica and Dymola, default to implicit solvers like DASSL. And you can see they do great on this problem by default!

Modelica tools gives a good answer out of the box.

So QED, that’s the “right thing to do”: if you want to be robust, stick to implicit methods.

But why oscillations?

Hold up a bit… why does the explicit method give oscillations? While we know that’s wrong, it would be good to understand why it gives the qualitatively wrong behavior that it does. It turns out that this falls right out of the definition of the method. If you go back to the definition of explicit Euler on the test problem, i.e.

$$u_{n+1} = u_n + hf(u_n)$$

then substitute in:

$$u_{n+1} = (1 + h\lambda) u_{n}$$

If we think about our stability criteria $G(z)$ another way, its boundaries are exactly the value by which the next $u_{n+1}$ would have a negative real part. So the analytical solution is supposed to go to zero, but the “bad” behavior is when we choose a step size $h$ such that if we extrapolate out with a straight line for $h$ long in time, then we will “jump” over this zero, something that doesn’t happen in the analytical solution. But now let’s think about what happens in that case. If you jump over zero, then $u_n < 0$ (think real right now), so therefore the derivative of the next update points in the other direction, i.e. we're still going towards zero, but now from the negative side we go up to zero. But since $\|1 + h\lambda\| > 1$, we have that $\|u_{n+1}\| > \|u_n\|$, i.e. the norm of the solution keeps growing. So you jump from positive to negative, then negative to positive, then positive to negative, where the jumps are growing each time. This is the phantom oscillations of the explicit ODE solver!

So what’s happening is the default tolerances of the explicit ODE solver were large enough that the chosen $h$s were in the range of the phantom oscillation behavior, and so you just need to cap $h$ below that value, which is dependent on the real part of the eigenvalue of $h$ (you can do the same analysis with complex numbers, but that just gives rotations in the complex plane along with the real part oscillation).

But if explicit methods give oscillations, what’s going on with implicit ODE solvers with large $h$? Let’s look at the update equation again:

$$u_{n+1} = \frac{1}{1-z} u_n$$

now instead of multiplying each time by $(1-z)$, we divide by it. This means that when $\lambda < 0$ (or $\text{Re}(\lambda) < 0$ to be more exact), then for any $h$ we have that $\|u_{n+1}\| < \|u_{n}\|$. Therefore, we might jump over the zero with a big enough $h$, but we are guaranteed that our "jump size" is always shrinking. Thus for any $h$, we will get to zero because we're always shrinking in absolute value. This means that implicit methods are working because they have a natural dampening effect. So:

Explicit methods introduce spurious oscillations, but implicit methods naturally damp oscillations

This explains in more detail why we saw what we saw: the explicit method when the error tolerance is sufficiently high will introduce oscillations that don’t exist, while the implicit method will not have this behavior. This is a more refined version of the “energy doesn’t go to infinity!”, now it’s “energy doesn’t come from nowhere in real systems”, and because of this implicit solvers give a better qualitative answer. This is why they are more robust, which is why robust software for real engineers just always default to them.

Wait a second… do we always want that?

You should now be the student in the front row raising your hand, “implicit methods are always dampening… is that actually a good idea? Are you sure that’s always correct?” And the answer is… well it’s not. And that then gives us exactly the failure case for which implicit methods are less robust. If you have a system that is supposed to actually oscillate, then this “hey let’s always dampen everything to make solving more robust” actually leads to very wrong answers!

To highlight this, let’s just take a simple oscillator. You can think of this as a harmonic oscillator, or you can think about it as a simple model of a planet going around a star. However you want to envision it, you can write it out as a system of ODEs:

$$u_1′ = 500u_2$$
$$u_2′ = -500u_1$$

This is the linear ODE $u’ = Au$ where $A = [0\ 500; -500\ 0]$, which has complex eigenvalues with zero real part. In other words, the analytical solution is $\sin(500t)$ and $\cos(500t)$, just a pure oscillation that just keeps going around and around in circles. If we solve this with an explicit ODE solver:

function f(du,u,p,t)
    du[1] = 500u[2]
    du[2] = -500u[1]
end
u0 = [1.0,1.0]
tspan = (0.0,1.0)
prob = ODEProblem(f, u0, tspan)
sol = solve(prob, Tsit5())
 
plot(sol, title="RK Method", idxs=(1,2))
savefig("rk_oscillate.png")

we can see that it generally gets the right answer. Over time you get some drift where the energy is slowly increasing due to numerical error in each step, but it’s going around in circles relatively well. However, our “robust implicit method”…

sol = solve(prob, CVODE_BDF())
plot(sol, title="BDF Method", idxs=(1,2))
savefig("bdf_oscillate.png")

is just not even close. And you can see that even our “robust Modelica tools” completely fall apart:

It says the answer goes to zero! Even when the analytical solution is just a circle! But we can understand why this is the case: the software developers made the implicit assumption that “dampening oscillations is always good, because generally that’s what happens in models, so let’s always do this by default so people get better answers”, and the result of this choice is that if someone puts in a model of the Earth going around the sun, then oops the Earth hits the sun pretty quickly.

Conclusion: ODE solvers make trade-offs, you need to make the right ones for your domain

This gives us the conclusion: there is no “better” or “more robust” ODE solver method, it’s domain-specific. This is why the Julia ODE solver package has hundreds of methods, because each domain can be asking for different properties that they want out of the method. Explicit methods are not generally faster, they are also something that tends to preserve (or generate) oscillations. Implicit methods are not generally more robust, they are methods which work by dampening transients, which is a good idea for some models but not for others. But then there’s a ton of nuance. For example, can you construct an explicit ODE solver so that on such oscillations you don’t get energy growth? You can! Anas5(w) is documented as “4th order Runge-Kutta method designed for periodic problems. Requires a periodicity estimate w which when accurate the method becomes 5th order (and is otherwise 4th order with less error for better estimates)”, i.e. if you give it a canonical frequency 500 it will be able to do extremely well on this problem (and being a bit off in that estimate still works, it just has energy growth that is small).

What about what was mentioned at the beginning of the article, “for hyperbolic PDEs you need to use explicit methods”? This isn’t a “special behavior” of PDEs, this is simply because for this domain, for example advective models of fluids, you want to conserve fluid as it moves. If you choose an implicit method, it “dampens” the solver, which means you get that as you integrate you get less and less fluid, breaking the conservation laws and giving qualitatively very incorrect solutions. If you use explicit methods, you don’t have this extraneous dampening, and this gives a better looking solution. But you can go even further and develop methods for which, if $h$ is sufficiently small, then you get little to no dampening. These are SSP methods, which we say are “for Hyperbolic PDEs (Conservation Laws)” but in reality what we mean is “when you don’t want things to dampen”.

But the point is, you can’t just say “if you want a better solution, use an implicit solver”. Maybe in some domains and for some problems that is true, but in other domains and problems that’s not true. And many numerical issues can stem from the implicit assumptions that follow from the choice being made for the integrator. Given all of this, it should be no surprise that much of the Modelica community has had many problems handling fluid models, the general flow of “everything is a DAE” → “always use an implicit solver” → “fluid models always dampen” → “we need to fix the dampening” could be fixed by making different assumptions at the solver level.

So, the next time someone tells you should just use ode15s or scipy.integrate.radau in order to make things robust without knowing anything about your problem, say “umm actually”.

Little Extra Details

The article is concluded. But here’s a few points I couldn’t fit into the narrative I want to mention:

Trapezoidal is cool

One piece I didn’t fit in here is that the Trapezoidal method is cool. The dampening property comes from L-stability, i.e. $G(z) \rightarrow 0$ as $\text{Re}(z) \rightarrow -\infty$. This is a stricter form of stability, since instead of just being stable for any finite $\lambda$, this also enforces that you are stable at the limit of bigger lambdas. “Most” implicit solvers that are used in practice, like Implicit Euler, have this property, and you can show the dampening is directly related to this property. But you can have an implicit method that isn’t L-stable. Some of these methods include Adams-Bashforth-Moulton methods, which are not even A-stable so they tend to have stability properties and act more like explicit methods. But the Trapezoidal method is A-stable without being L-stable, so it doesn’t tend to dampen while it tends to be also pretty stable. Though it’s not as stable, and the difference between “is stable for any linear ODE” versus “actually stable for nonlinear ODEs” (i.e. B-stability) is pronounced on real-world stiff problems. What this means in human terms is that the Trapezoidal method tends to not be stable enough for hard stiff problems, but it also doesn’t artificially dampen, so it can be a good default in cases where you know you have “some stiffness” but also want to keep some oscillations. One particular case of this is in some electrical circuit models with natural oscillators.

Lower order methods have purposes too

“All ODE solvers have a purpose”, I give some talks that give the justification for many high order methods, so in general “higher order is good if you solve with stricter tolerances and need more precision”. But lower order methods can be better because the higher order methods require that more derivatives of $f$ are defined, and if that’s not the case (like derivative discontinuities), then lower order methods will be more efficient. So even implicit Euler has cases where it’s better than higher order BDF methods, and it has to do with “how nice” $f$ is.

BDF methods like DASSL are actually α-stable

I said that generally implicit methods that you use are A-stable. That’s also a small lie to make the narrative simpler. The BDF methods which Sundials, DASSL, LSODE, FBDF, etc. use are actually α-stable, which means that they are actually missing some angle α of the complex plane for stability. The stability regions look like this:

So these BDF methods are actually pretty bad for other reasons on very oscillatory problems! Meanwhile, things like Rosenbrock methods can also solve DAEs while actually being L-stable, which can make them more stable in many situations where there’s oscillations towards a steady state. So there’s a trade-off there… again every method has a purpose. But this is another “ode15s is more stable than ode23s”… “well actually…”

The post Implicit ODE Solvers Are Not Universally More Robust than Explicit ODE Solvers, Or Why No ODE Solver is Best appeared first on Stochastic Lifestyle.

Enhance Model Accuracy Using Universal Differential Equations

By: Great Lakes Consulting

Re-posted from: https://blog.glcs.io/ude_symbolic_regression

This post was written by Steven Whitaker.

In this Julia for Devs post,we will explore the fascinating world of using universal differential equations (UDEs)for model discovery.A UDE is a differential equationwhere part (or all) of itis defined by a universal approximator(typically a neural network).These UDEs play a pivotal rolein uncovering missing aspectsof established models.

UDEs enable seamless integrationof known mechanistic modelswith data-driven modeling approaches.Rather than omit or guess at unknown aspects of a model,a neural network can be embedded into the model,which can then be trained using observed data.

Once trained,the embedded neural networkcan be approximated by a mathematical expression.This allows us to replace the neural networkwith an interpretable formula,improving transparencywhile filling in the missing aspectsof the original modelwith new structure learned from observed data.

This post will illustratekey ideasand design decisions madewhile helping one of our clientsuncover missing dynamics in their model,a complex system representingthe human body under treatmentof a rare disease.

By replacing just one aspect of our client’s modelwith a neural network,we achieved a significant milestone:a 30% reduction in the error of the model predictions.This success story is a testament to the powerof UDEs in enhancing model performance.

We will not share client-specific code,but will provide similar examplesto illustrate our approach.

Here are some of the key ideas we will focus on:

  • Ensure nonnegativity of state variables (as needed).
  • Appropriately initialize neural network weights.
  • Normalize neural network inputs.
  • Make dynamics modular.
  • Fit symbolic expression.

Note that we will discussfully connected neural networks,so some comments belowmay not apply if other architectures are used.

And with that,let’s dive in!

Ensure Nonnegativity of States

When working with a modelrepresenting a real phenomenon,chances are that at least some of the model’s state variablesare naturally nonnegative.For example,the energy of a systemor the concentration of a chemicalcannot be negative quantities.

However,a differential equation solversees only math,not the physical phenomenonthe math represents.So,we have to tell the solverto respect the natural constraintsthat exist.

One way to enforce nonnegativityis to pass the isoutofdomain optionto solve,where isoutofdomain is a functionthat returns truewhen at least one of the provided state variablesviolates the nonnegativity constraint.For example,if all state variables should be nonnegative,the function can be defined as:

isoutofdomain(u, p, t) = any(<(0), u)

As another example,if only the first and third states should be nonnegative:

isoutofdomain(u, p, t) = u[1] < 0 || u[3] < 0

For more informationand other toolsfor tackling nonnegativity,check out the SciML docs.

The ideaof telling the solverabout the nonnegativity constraintsmay seem obvious.Still, it has important implicationsfor the initialization of neural networksembedded into the model.Care should be takento ensure the initial network weights(before training)do not cause states to become negative.See the next section below for more details.

Appropriately Initialize Neural Network Weights

To learn missing dynamics,the neural network in a UDEneeds to be trained.

Training the neural network can occuronly if the UDE can be solvedusing the initial network weights.Otherwise,the gradient of the loss,which is used by optimization algorithmsto train the neural network,cannot be computed.

As a result,it is necessaryto choose a good initializationfor the neural network weights.Here are some approaches for initializationand how well they work:

  • Random initialization:Typically,neural network weightsare initialized randomly.However,this results in random UDE dynamics,so whether a state variablestays positive(or satisfies any other constraints)is left to chance.And if constraints are violated,but we try to enforce themvia, e.g., isoutofdomain,the solver will be unable to solve the UDE,preventing computation of the loss gradient.So,random initialization might work,but it might not.
  • Zero initialization:Don’t do this.Network weights won’t ever update during trainingbecause the gradients with respect to those weightswill always be zero.
  • Constant initialization:Don’t do this, either.Initializing all the weightsto the same, non-zero valuewill result in some amount of learning,but it severely limitsthe expressivity of the network.
  • Pre-training:If the neural network in the UDEis used to replace an expressionthat is believed to be incorrect(but otherwise gives a usable model),one way to initialize the network weightsfor UDE trainingis to initialize the weights randomlybut then train the network directly(i.e., not the UDE as a whole)to fit the expression the network is to replace.This works because no constraints are involvedduring pre-training(so randomly initializing the weights is okay),but when training the UDE,the network weights have been setto avoid running into issues with constraints.
  • Smoothed Collocation:Smoothed collocation is another approachfor pre-training the neural networkwithout involving constraints.See the SciML docs for more info.

For our client,randomly initialized weightsfailed to give a usable gradientfor training,so we decided on pre-trainingbecause they had a reasonable guessthat we could fit the initial network weights to.

Normalize Neural Network Inputs

One key to ensuring UDE training proceeds nicelyis to normalize the inputsto the neural networkto ensure all inputsare roughly of the same magnitude.

Normalizing is essentialbecause the gradient of the training loss functionis proportional to the network inputs.If one input tends to be 1000 times largerthan another,the gradient will also tend to be that much larger,dominating the learning algorithm.

Effect of input scale imbalance on training

To illustrate how to normalize,suppose we have the following neural network(created using Lux.jl):

nn = Chain(    Dense(2, 5, tanh),    Dense(5, 5, tanh),    Dense(5, 1),)

Then we can add a function to the Chainthat performs the normalization:

nn_normalized = Chain(    x -> x ./ normalization,    # Other layers as before)

In this example,normalization is a Vectorcontaining the valuesto scale each input variable.And of course,we’re not limited to scaling the inputs;we can implement whatever function we wantto process the inputs(such as z-score standardization).

The neural network in our client’s UDEdid not learn without normalization.After normalizing the inputs,we got the neural networkto learn meaningful dynamics.

Make Dynamics Modular

While not strictly necessary for UDEs,making the dynamics modularvastly simplifies comparingbetween different approaches.

For our client,we wanted to comparethree versions of the dynamics:the original dynamics,the UDE(where part of the original dynamicswas replaced with a neural network),and the updated dynamics(where the trained neural networkwas replaced with a symbolic expressionfit to the neural network).

Rather than duplicate the dynamics function three timeswith each slight variation,we allowed the variable part of the dynamicsto be passed in as an optionto the dynamics function.To illustrate:

function f!(du, u, p, t; g = original_g)    # ... some code ...    # `a` and `b` are values computed above    # or pulled from `u`, etc.    val = g(a, b, p)    # `val` is used to update `du` in some way.    # ... some code ...end

However,ODEProblems don’t allow passing optionsto the dynamics function.But that’s not a problem;we can create a closureto set the optionbefore handing the function to an ODEProblem.For example,if we want g to use a Lux.jl neural network:

nn = Chain(...)(p_nn, st) = Lux.setup(...)# Be sure to include `p_nn` in the `ODEProblem` parameters `p`.g_nn = (a, b, p) -> nn([a, b], p.p_nn, st)[1][1]f_nn! = (du, u, p, t) -> f!(du, u, p, t; g = g_nn)

Setting up our client’s code like thisreduced code duplicationwhile allowing easy comparisonsacross different versions of the dynamics.

Fit Symbolic Expression

One of the key stepsto learning interpretable dynamicsfrom a UDEis to fit a symbolic expressionto the trained neural network.Doing so allows us to replacethe black box neural networkwith a mathematical expressionwe can understand.

The first stepis to collect the data to usefor the fitting.If the neural network takes two inputs,we will need a Matrix of values:

X = [    a1 a2 ... aN;    b1 b2 ... bN;]

The (ai, bi) pairsshould cover the expected range of input values.One way to determine this rangeis to solve the UDEand save off the valuesthat end up going into the neural network.

To get the target values,evaluate the neural networkwith the X we just created:

y = nn(X, p_nn, st)[1] |> vec

(Here we assume the neural networkoutputs a single numberper input pair.)

Now we can figure outa mathematical expressionrelating X to y.One way to do thisis to use SymbolicRegression.jl.

After obtaining an Expression object eq(called tree in the README for SymbolicRegression.jl),we can use it in our dynamics functionas described in the previous section:

# `p` needs to be an input even though it's not used.g_eq = (a, b, p) -> eq([a; b;;])[1]f_eq! = (du, u, p, t) -> f!(du, u, p, t; g = g_eq)

Using SymbolicRegression.jlto fit the trained UDE neural network,we could provide our clientwith a precise mathematical expressionrepresenting the dynamics that were missingfrom their original differential equation.

Summary

In this post,we discussed vital aspectsof implementing UDEsto learn missing dynamics.We learned some tipsfor helping UDE training,including appropriately initializing network weightsand normalizing network inputs.We also saw the benefitsof code modularityand how that can easily allow usto insert a learned mathematical expressioninto the dynamics.For our client,we were able to use these ideasto train a UDEto achieve 30% less errorin the model predictions.

Do you want to leverage UDEsto improve your models’ predictive capabilities?Contact us, and we can help you out!

Additional Links

]]>