# # 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:
#
# - [`optimizing_julia.jl`](optimizing_julia.jl)
# - [`Project.toml`](Project.toml)
# - [`Manifest.toml`](Manifest.toml)
#
# This website renders the content of
# [`optimizing_julia.jl`](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](https://github.com/fredrikekre/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/
#
# if you care about performance.
# ## 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
read_dataset(io, "x"), read_dataset(io, "u0")
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()
# Task: Optimize the code!
#
# You can find one improved version in the file
# [`optimizing_julia_solution.jl`](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