# Optimizing Julia code

This session provides an introduction to optimizing Julia code. The examples are developed with Julia v1.9.3. You can download all files from the summer school website:

This website renders the content of `optimizing_julia.jl`.

First, we install all required packages

``````import Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()``````

The markdown file is created from the source code using Literate.jl. You can create the markdown file via

``````using Literate
Literate.markdown(joinpath(@__DIR__, "optimizing_julia.jl");
flavor = Literate.CommonMarkFlavor())``````

## Basics

The Julia manual contains basic information about performance. In particular, you should be familiar with

• https://docs.julialang.org/en/v1/manual/performance-tips/
• https://docs.julialang.org/en/v1/manual/style-guide/

## Example: A linear advection simulation

This is a classical partial differential equation (PDE) simulation setup. If you are not familiar with it, just ignore what’s going on - but we need an example. This one is similar to several problematic cases I have seen in practice.

First, we generate initial data and store it in a file.

``````using HDF5
x = range(-1.0, 1.0, length = 1000)
dx = step(x)
h5open("initial_data.h5", "w") do io
u0 = sinpi.(x)
write_dataset(io, "x", collect(x))
write_dataset(io, "u0", u0)
end``````

Next, we write our simulation code as a function - as we have learned from the performance tips!

``````function simulate()
x, u0 = h5open("initial_data.h5", "r") do io
end

t = 0
u = u0
while t < t_end
u = u + dt * rhs(u)
t = t + dt
end

return t, x, u, u0
end

function rhs(u)
du = similar(u)
for i in eachindex(u)
im1 = i == firstindex(u) ? lastindex(u) : (i - 1)
du[i] = -(u[i] - u[im1]) / dx
end
return du
end``````

Now, we can define our parameters, run a simulation, and plot the results.

``````using Plots, LaTeXStrings

t_end = 2.5
dt = 0.9 * dx
t, x, u, u0 = simulate()
plot(x, u0; label = L"u_0", xguide = L"x")
plot!(x, u; label = L"u")``````

Maybe you can already spot problems in the code above. Anyway, before you begin optimizing some code, you should test it and make sure it’s correct. Let’s assume this is the case here. You should write tests making sure that the code continues to work as expected while we optimize it. We will not do this here right now.

## Profiling

First, we need to measure the performance of our code to see whether it’s already reasonable. For this, you can use `@time` - or better BenchmarkTools.jl.

``````using BenchmarkTools
@benchmark simulate()``````

This shows quite a lot of allocations - typically a bad sign if you are working with numerical simulations.

There are also profiling tools available in Julia. Some of them are already loaded in the Visual Studio Code extension for Julia. If you prefer working in the REPL, you can install ProfileView.jl (`@profview simulate()`) and PProf.jl (`pprof()` after creating a profile via `@profview`).

``@profview simulate()``

You can find one improved version in the file `optimizing_julia_solution.jl`.

## Introduction to generic Julia code and AD

One of the strengths of Julia is that you can use quite a few modern tools like AD. However, you need to learn Julia a bit to use everything efficiently. Here, we concentrate on ForwardDiff.jl.

``````using ForwardDiff
using Statistics

function foo0(x)
vector = zeros(typeof(x), 100)
foo0!(vector, x)
end

function foo0!(vector, scalar)
for i in eachindex(vector)
vector[i] = atan(i * scalar) / π
end

for _ in 1:1000
for i in eachindex(vector)
vector[i] = cos(vector[i])
end
end

return mean(vector)
end

let x = 2 * π
@show foo0(x)
@show ForwardDiff.derivative(foo0, x)
@benchmark foo0(\$x)
end``````

The code above is basically a fixed point iteration. By doing some analysis, one could figure out that it should converge to the solution `x` of `x == cos(x)`, the “Dottie number”. See https://en.wikipedia.org/wiki/Dottie_number

``````using SimpleNonlinearSolve

function dottie_number()
f(u, p) = cos(u) - u
bounds = (0.0, 2.0)
prob = IntervalNonlinearProblem(f, bounds)
sol = solve(prob, ITP())
return sol.u
end

dottie_number()``````

Next, we introduce a custom struct. Can you spot the problems?

``````struct MyData1
scalar
vector
end

function foo1(x)
data = MyData1(x, zeros(typeof(x), 100))
foo1!(data)
end

function foo1!(data)
(; scalar, vector) = data

for i in eachindex(vector)
vector[i] = atan(i * scalar) / π
end

for _ in 1:1000
for i in eachindex(vector)
vector[i] = cos(vector[i])
end
end

return mean(vector)
end

let x = 2 * π
@show foo1(x)
@show ForwardDiff.derivative(foo1, x)
@benchmark foo1(\$x)
end``````

We can fix type instabilities by introducing types explicitly. But we lose the possibility to use AD this way!

``````struct MyData2
scalar::Float64
vector::Vector{Float64}
end

function foo2(x)
data = MyData2(x, zeros(typeof(x), 100))
foo1!(data)
end

let x = 2 * π
@show foo2(x)
@benchmark foo2(\$x)
end

let x = 2 * π
ForwardDiff.derivative(foo2, x)
end``````

Thus, we need parametric types!

``````struct MyData3{T}
scalar::T
vector::Vector{T}
end

function foo3(x)
data = MyData3(x, zeros(typeof(x), 100))
foo1!(data)
end

let x = 2 * π
@show foo3(x)
@show ForwardDiff.derivative(foo3, x)
@benchmark foo3(\$x)
end``````