Single column parameter estimation with Reactant and Enzyme

This example shows how to estimate turbulence closure parameters and surface fluxes by differentiating through an Oceananigans simulation. We set up a "twin experiment": first we run a "nature run" with known parameters, then we pretend we don't know those parameters and recover them by minimizing the mismatch between the model and the nature run. We compute gradients with Enzyme automatic differentiation, compiled by Reactant.

This example demonstrates:

• How to run a single column model with TKEDissipationVerticalDiffusivity (k-ϵ) on ReactantState.

• How to assign closure parameters and boundary fluxes to a model inside a differentiable function.

• How to compute gradients with respect to closure parameters and boundary fluxes.

• How to recover the parameters of a nature run with a gradient-based optimizer.

Install dependencies

First let's make sure we have all required packages installed.

using Pkg
pkg"add Oceananigans, Enzyme, Reactant, CUDA, Optim, CairoMakie"

Reactant needs CUDA.jl to be loaded to compile Oceananigans' kernels, even when it compiles for the CPU.

Reactant can run with one of two runtimes, PJRT (the default) or IFRT. Oceananigans' Reactant tests use IFRT. To use IFRT, add the following to the Project.toml or LocalPreferences.toml of your project, and restart Julia:

[preferences.Reactant]
xla_runtime = "IFRT"

using Oceananigans
using Oceananigans.Units
using Enzyme
using Reactant
using CUDA
using CairoMakie
using Printf

A single column is small, so we run Reactant on the CPU even when a GPU is available,

Reactant.set_default_backend("cpu")

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1790311910.991865 2084408 pjrt_client.cc:579] PjRt-IFRT device count: total=1, addressable=1
I0000 00:00:1790311910.991907 2084408 pjrt_client.cc:583] Addressable PjRt-IFRT device: CpuDevice(id=0)

A single column model

We simulate a column of ocean mixed by a surface wind stress τˣ and cooled by a surface buoyancy flux Jᵇ. The k-ϵ model computes an eddy viscosity and an eddy diffusivity from the turbulent kinetic energy e and its dissipation rate ϵ,

where the "stability functions" 𝕊u and 𝕊c depend on the local shear and stratification. We estimate two constants that set the magnitude of the stability functions, Cu₀ for the eddy viscosity and Cc₀ for the eddy diffusivity, together with the surface fluxes τˣ and Jᵇ.

To assign new parameters inside a compiled function, the closure's numbers must be Reactant numbers, which we obtain with Reactant.to_rarray(closure; track_numbers=Number). Similarly, we represent the surface fluxes with Fields, whose values we can change.

using Oceananigans.Architectures: ReactantState

grid = RectilinearGrid(ReactantState(), size=32, z=(-128, 0), topology=(Flat, Flat, Bounded))
closure = Reactant.to_rarray(TKEDissipationVerticalDiffusivity(); track_numbers=Number)

τˣ = Field{Face, Center, Nothing}(grid)
Jᵇ = Field{Center, Center, Nothing}(grid)
u_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(τˣ))
b_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(Jᵇ))

model = HydrostaticFreeSurfaceModel(grid; closure,
                                    coriolis = FPlane(f=1e-4),
                                    tracers = :b,
                                    buoyancy = BuoyancyTracer(),
                                    boundary_conditions = (u=u_bcs, b=b_bcs))

HydrostaticFreeSurfaceModel{ReactantState, RectilinearGrid}(time = 0 seconds, iteration = ConcreteIFRTNumber(0))
├── grid: 1×1×32 RectilinearGrid{Float64, Flat, Flat, Bounded} on ReactantState with 0×0×3 halo
├── timestepper: QuasiAdamsBashforth2TimeStepper
├── tracers: (b, e, ϵ)
├── closure: TKEDissipationVerticalDiffusivity{VerticallyImplicitTimeDiscretization}
├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
├── advection scheme: 
│   ├── momentum: VectorInvariant
│   ├── b: Centered(order=2)
│   ├── e: Centered(order=2)
│   └── ϵ: Centered(order=2)
├── vertical_coordinate: ZCoordinate
└── coriolis: FPlane{Oceananigans.Advection.EnstrophyConserving{Float64}, Float64}

Assigning parameters

We work with parameters normalized by typical values, so that every parameter is order one,

scales = (Cu₀ = 0.1, Cc₀ = 0.1, τˣ = 1e-4, Jᵇ = 1e-8)

(Cu₀ = 0.1, Cc₀ = 0.1, τˣ = 0.0001, Jᵇ = 1.0e-8)

The function below assigns normalized_parameters to the model. It builds a new closure with the VariableStabilityFunctions constructor, which also computes another constant, 𝕊u₀, from Cu₀. Because we call the constructor inside the function we differentiate, the gradient accounts for the dependence of 𝕊u₀ on Cu₀.

using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: VariableStabilityFunctions
const vitd = VerticallyImplicitTimeDiscretization()

function assign_parameters!(model, normalized_parameters)
    Cu₀ = scales.Cu₀ * normalized_parameters.Cu₀
    Cc₀ = scales.Cc₀ * normalized_parameters.Cc₀
    FT = typeof(Cu₀)

    stability_functions = VariableStabilityFunctions(FT; Cu₀, Cc₀)
    model.closure = TKEDissipationVerticalDiffusivity(vitd, FT; stability_functions)

    τˣ = model.velocities.u.boundary_conditions.top.condition
    Jᵇ = model.tracers.b.boundary_conditions.top.condition
    parent(τˣ) .= scales.τˣ * normalized_parameters.τˣ
    parent(Jᵇ) .= scales.Jᵇ * normalized_parameters.Jᵇ

    return nothing
end

assign_parameters! (generic function with 1 method)

We represent parameters with a NamedTuple of Reactant numbers,

reactant_parameters(values) = NamedTuple{keys(scales)}(Tuple(Reactant.ConcreteRNumber.(values)))

reactant_parameters (generic function with 1 method)

Running the column

Each run starts from rest with constant stratification and lasts 12 hours. Because we reuse the same model for every run, we first reset! the model, which zeros its fields (including e and ϵ), its tendencies, and the closure's state saved from previous time steps, so that every run starts from the same state. The k-ϵ equations need a fairly short time step: with Δt = 10minutes, for example, the solution becomes noisy.

using Oceananigans.Models: reset!

N² = 1e-5
bᵢ = set!(CenterField(grid), z -> N² * z)

Δt = 1minute
Nt = 720

function run_column!(model, normalized_parameters, bᵢ, Δt, Nt)
    assign_parameters!(model, normalized_parameters)

    reset!(model)
    set!(model, b=bᵢ)

    @trace track_numbers=false for n = 1:Nt
        time_step!(model, Δt)
    end

    return nothing
end

run_column! (generic function with 1 method)

The nature run

The nature run uses the default closure parameters Cu₀ = 0.1067 and Cc₀ = 0.1120, a wind stress τˣ = -10⁻⁴ m² s⁻², and a cooling buoyancy flux Jᵇ = 2 × 10⁻⁸ m² s⁻³. We collect the normalized parameters in the vector θ = (Cu₀, Cc₀, τˣ, Jᵇ) ./ scales, and convert them to Reactant numbers with reactant_parameters.

θ★ = [1.067, 1.120, -1, 2]

compiled_run_column! = @compile raise=true raise_first=true sync=true run_column!(
    model, reactant_parameters(θ★), bᵢ, Δt, Nt)

compiled_run_column!(model, reactant_parameters(θ★), bᵢ, Δt, Nt)

We save the final state of the nature run as our "observations",

u★ = XFaceField(grid)
v★ = YFaceField(grid)
b★ = CenterField(grid)
set!(u★, model.velocities.u)
set!(v★, model.velocities.v)
set!(b★, model.tracers.b)

obs = (; u★, v★, b★)

NamedTuple with 3 Fields on 1×1×32 RectilinearGrid{Float64, Flat, Flat, Bounded} on ReactantState with 0×0×3 halo:
├── u★: 1×1×32 Field{Face, Center, Center} on RectilinearGrid on ReactantState
├── v★: 1×1×32 Field{Center, Face, Center} on RectilinearGrid on ReactantState
└── b★: 1×1×32 Field{Center, Center, Center} on RectilinearGrid on ReactantState

Cost function and its gradient

The cost function is the normalized mean square difference between the final state of the model and the observations,

using Statistics: mean

function cost(normalized_parameters, model, bᵢ, obs, Δt, Nt)
    run_column!(model, normalized_parameters, bᵢ, Δt, Nt)

    u, v = model.velocities
    b = model.tracers.b
    u★, v★, b★ = obs

    U² = 1e-2
    B² = (N² * 10)^2

    𝒥u = mean((u - u★)^2) / U²
    𝒥v = mean((v - v★)^2) / U²
    𝒥b = mean((b - b★)^2) / B²

    return 𝒥u + 𝒥v + 𝒥b
end

cost (generic function with 1 method)

To compute the gradient, we differentiate cost in reverse mode. Enzyme accumulates the gradient of the cost with respect to the parameters into the "shadow" parameters cost_gradient. The model is also mutated by cost, so we give it a shadow as well.

function cost_and_gradient!(cost_gradient, parameters, model, shadow, bᵢ, obs, Δt, Nt)
    mode = Enzyme.set_strong_zero(Enzyme.ReverseWithPrimal)
    _, 𝒥 = Enzyme.autodiff(mode, cost, Enzyme.Active,
                           Enzyme.Duplicated(parameters, cost_gradient),
                           Enzyme.Duplicated(model, shadow),
                           Enzyme.Const(bᵢ),
                           Enzyme.Const(obs),
                           Enzyme.Const(Δt),
                           Enzyme.Const(Nt))
    return 𝒥
end

cost_and_gradient! (generic function with 1 method)

We start from an initial guess that's quite different from the nature run,

θ₀ = [0.6, 1.8, -0.5, 0.5]

4-element Vector{Float64}:
  0.6
  1.8
 -0.5
  0.5

and compile the cost and its gradient. Compiling the gradient takes a few minutes.

shadow = Enzyme.make_zero(model)
cost_gradient = reactant_parameters(zeros(4))

compiled_cost = @compile raise=true raise_first=true sync=true cost(
    reactant_parameters(θ₀), model, bᵢ, obs, Δt, Nt)

compiled_cost_and_gradient! = @compile raise=true raise_first=true sync=true cost_and_gradient!(
    cost_gradient, reactant_parameters(θ₀), model, shadow, bᵢ, obs, Δt, Nt)

𝒥(θ) = Float64(compiled_cost(reactant_parameters(θ), model, bᵢ, obs, Δt, Nt))

function cost_and_gradient(θ)
    shadow = Enzyme.make_zero(model)
    cost_gradient = reactant_parameters(zeros(4))
    𝒥θ = compiled_cost_and_gradient!(cost_gradient, reactant_parameters(θ), model, shadow, bᵢ, obs, Δt, Nt)
    return Float64(𝒥θ), Float64.(collect(cost_gradient))
end

cost_and_gradient (generic function with 1 method)

Minimizing the cost

We minimize the cost with the BFGS algorithm implemented by Optim.jl. BFGS is a quasi-Newton flavor of gradient descent: rather than stepping along -∇𝒥, it steps along -H ∇𝒥, where H is an estimate of the inverse Hessian of 𝒥 that BFGS builds from the change in the gradient between iterations. Optim needs the cost 𝒥, and a function that computes the gradient in place,

using Optim

function ∇𝒥!(G, θ)
    _, ∇𝒥θ = cost_and_gradient(θ)
    G .= ∇𝒥θ
    return G
end

options = Optim.Options(iterations=10, store_trace=true, extended_trace=true)
bfgs_result = optimize(𝒥, ∇𝒥!, θ₀, BFGS(), options)

 * Status: failure (reached maximum number of iterations)

 * Candidate solution
    Final objective value:     3.351246e-04

 * Found with
    Algorithm:     BFGS

 * Convergence measures
    |x - x'|               = 1.07e-02 ≰ 0.0e+00
    |x - x'|/|x'|          = 5.13e-03 ≰ 0.0e+00
    |f(x) - f(x')|         = 2.49e-05 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 7.44e-02 ≰ 0.0e+00
    |g(x)|                 = 1.63e-03 ≰ 1.0e-08

 * Work counters
    Seconds run:   15  (vs limit Inf)
    Iterations:    10
    f(x) calls:    15
    ∇f(x) calls:   15
    ∇f(x)ᵀv calls: 0

The trace records the parameters and cost at each iteration,

history = [(θ=θ, 𝒥=𝒥θ) for (θ, 𝒥θ) in zip(Optim.x_trace(bfgs_result), Optim.f_trace(bfgs_result))]

for (n, h) in enumerate(history)
    @info @sprintf("iteration %2d: 𝒥 = %.2e, θ / θ★ = %s", n - 1, h.𝒥, string(round.(h.θ ./ θ★, digits=3)))
end

[ Info: iteration  0: 𝒥 = 1.08e-01, θ / θ★ = [0.562, 1.607, 0.5, 0.25]
[ Info: iteration  1: 𝒥 = 6.24e-02, θ / θ★ = [0.509, 1.611, 1.476, 0.381]
[ Info: iteration  2: 𝒥 = 1.27e-02, θ / θ★ = [0.544, 1.602, 1.039, 0.334]
[ Info: iteration  3: 𝒥 = 1.23e-02, θ / θ★ = [0.54, 1.6, 1.147, 0.353]
[ Info: iteration  4: 𝒥 = 6.81e-03, θ / θ★ = [0.632, 1.495, 1.086, 0.536]
[ Info: iteration  5: 𝒥 = 2.90e-03, θ / θ★ = [0.809, 1.294, 0.926, 0.884]
[ Info: iteration  6: 𝒥 = 6.29e-04, θ / θ★ = [0.838, 1.258, 0.974, 0.953]
[ Info: iteration  7: 𝒥 = 4.57e-04, θ / θ★ = [0.871, 1.22, 0.965, 1.02]
[ Info: iteration  8: 𝒥 = 4.07e-04, θ / θ★ = [0.881, 1.207, 0.968, 1.039]
[ Info: iteration  9: 𝒥 = 3.60e-04, θ / θ★ = [0.887, 1.198, 0.975, 1.05]
[ Info: iteration 10: 𝒥 = 3.35e-04, θ / θ★ = [0.886, 1.199, 0.98, 1.044]

Visualizing the parameter estimation

We compute the final state of the model at each iteration,

z = Array(znodes(grid, Center()))

function final_state(θ)
    compiled_run_column!(model, reactant_parameters(θ), bᵢ, Δt, Nt)
    u = Array(interior(model.velocities.u))[:]
    v = Array(interior(model.velocities.v))[:]
    b = Array(interior(model.tracers.b))[:] .- N² .* z
    return (; u, v, b)
end

nature_state = final_state(θ★)
iteration_states = [final_state(h.θ) for h in history]

11-element Vector{@NamedTuple{u::Vector{Float64}, v::Vector{Float64}, b::Vector{Float64}}}:
 (u = [-4.9016413739267804e-57, -8.769963353184815e-55, -1.5831404817445888e-52, -2.761850356044449e-50, -4.669663534195564e-48, -7.646833200145853e-46, -1.2121688026912987e-43, -1.8590135769667686e-41, -2.7565431097950974e-39, -3.9491176104862746e-37, -5.461975516985426e-35, -7.287416614902618e-33, -9.373192066080938e-31, -1.1618545526804226e-28, -1.3882501612364635e-26, -1.6003689474480441e-24, -1.7821664235206388e-22, -1.9182607117505854e-20, -1.9919222290312614e-18, -1.9823031367074232e-16, -1.86622519703178e-14, -1.6296530003445893e-12, -1.2862848941986687e-10, -8.88939516252812e-9, -5.18066124929221e-7, -2.343935193319183e-5, -0.01421961219006483, -0.031859991404355636, -0.030307927040972904, -0.02551300381163179, -0.01666741816362195, 0.003212278755645892], v = [7.982764747851425e-57, 1.4030044074674024e-54, 2.483229925610387e-52, 4.2376971036585036e-50, 6.989912114742729e-48, 1.1131723908400092e-45, 1.7099488506883795e-43, 2.5309801957028754e-41, 3.6058199764887464e-39, 4.938668240297314e-37, 6.494214874757762e-35, 8.186466720055352e-33, 9.875658156245394e-31, 1.1378185434539593e-28, 1.2492088607133418e-26, 1.303557226142498e-24, 1.2889474012895332e-22, 1.2029685316231719e-20, 1.0536881983547745e-18, 8.580772539361372e-17, 6.3892968804339095e-15, 4.217450530350118e-13, 2.3177758496051782e-11, 8.976503728924055e-10, 5.716032190445799e-9, -2.5444660228089183e-6, -0.0044014795383430436, -0.020544965636934077, -0.02826181902946271, -0.033919338397839637, -0.03946874266011554, -0.045654625576282956], b = [2.2749471783388257e-6, 1.5121756733729383e-8, 9.05373786662042e-9, 3.2442662488765817e-11, 2.3470614557777325e-11, 5.0806798004843223e-14, 4.566117449422702e-14, 6.353424730765056e-17, 6.505213034913027e-19, -6.364266752489911e-17, -1.3010426069826053e-17, -2.0925101928970236e-17, -2.5478751053409354e-17, -1.6479873021779667e-17, 0.0, -4.391018798566293e-17, 0.0, 0.0, -9.64939933512099e-18, -1.2912847874302358e-16, 2.7896847158703775e-14, -4.694688648249068e-13, 4.3348559930402736e-11, -7.500178683412378e-10, 4.606183277974709e-8, -6.627158511894301e-7, 4.875270686956475e-5, 5.4821845291630564e-5, 1.934249318107564e-5, -1.9649172961818575e-5, -5.93884782116295e-5, -9.960621172232776e-5])
 (u = [-3.197462829291986e-53, -7.304916542019159e-51, -1.6863690309196203e-48, -3.781084710680777e-46, -8.259915048921153e-44, -1.7575229407924974e-41, -3.6418479109185985e-39, -7.345520468357516e-37, -1.4403580574511722e-34, -2.739199944685055e-32, -5.032403646645819e-30, -8.881323814200575e-28, -1.4947457433680857e-25, -2.3784432770305623e-23, -3.5445938672674645e-21, -4.903670262112162e-19, -6.266014456232991e-17, -7.453888233942201e-15, -8.567263476700669e-13, -1.0162121403272574e-10, -1.259226438046018e-8, -1.4261420562620975e-6, -0.0012475181454689205, -0.04068860769269932, -0.05742257519364053, -0.05926016040047185, -0.0569383693107012, -0.052009963981526026, -0.04420217148366566, -0.03263571644331002, -0.014975000794896179, 0.01872185830816641], v = [2.684534746162925e-53, 5.863994082401422e-51, 1.2928868452692665e-48, 2.7622696174642548e-46, 5.730463315612586e-44, 1.1526479170400763e-41, 2.2452521472878666e-39, 4.2303531043977005e-37, 7.699073375493003e-35, 1.3509868140422752e-32, 2.2794520905141273e-30, 3.6830189984122895e-28, 5.665105422387894e-26, 8.228092833139315e-24, 1.116168015364611e-21, 1.3933629817054592e-19, 1.5655565992753862e-17, 1.517652070827641e-15, 1.1253363283823935e-13, 2.7887268795085235e-12, -1.0240784547146617e-9, -2.4388777648158506e-7, -0.00032222332898324636, -0.013518416860273719, -0.03236597789471324, -0.04378077263172095, -0.051683564041453245, -0.058956763372893, -0.06618891278081024, -0.07341826345073106, -0.08058424285352551, -0.08776100720965534], b = [2.033112394231101e-6, 1.1980292317654434e-8, 7.246371666559323e-9, 2.2693631045525575e-11, 1.686597931628564e-11, 3.112007179728593e-14, 2.944172683427837e-14, 4.228388472693467e-17, -2.4936649967166602e-17, -8.098990228466718e-17, -4.434386885465713e-17, 0.0, -9.64939933512099e-18, 0.0, 6.505213034913027e-19, 3.0140820395097023e-17, -3.6765295668983455e-16, 4.944655795924291e-14, -3.778545401914024e-13, 5.399857915412004e-11, 1.1391231286147704e-9, -3.6590701322949774e-8, 9.420681079598103e-6, 0.0001202894432656275, 0.00010195686074347102, 7.086489999663675e-5, 3.372965646378676e-5, -4.933201573686604e-6, -4.421796289406709e-5, -8.383870539317586e-5, -0.00012370143148779692, -0.00016386010410496314])
 (u = [-8.271970453970968e-55, -1.6035100706599247e-52, -3.1854618935042687e-50, -6.214982123571104e-48, -1.1927086337916209e-45, -2.2454055389235977e-43, -4.1349659964953476e-41, -7.426662018156968e-39, -1.2975517323572595e-36, -2.2009455661700463e-34, -3.6205084249216446e-32, -5.774389975966369e-30, -8.931586061083086e-28, -1.3395278964380466e-25, -1.9439006610717836e-23, -2.7145706031643807e-21, -3.611403155996945e-19, -4.5100676721391276e-17, -5.187168920020039e-15, -5.369061862428099e-13, -4.865587297646881e-11, -3.73251438650908e-9, -2.3285586303868205e-7, -1.0982237804795422e-5, -0.02031463117478161, -0.047366178398171276, -0.04868017473182789, -0.0457544857693887, -0.04014877595235844, -0.03125251200509727, -0.01720080805185312, 0.011023524693412124], v = [1.1781022916107365e-54, 2.1249670667803178e-52, 3.9106897729529025e-50, 7.046143256025163e-48, 1.2469551914226478e-45, 2.1650930931249695e-43, 3.681782221071544e-41, 6.11589503450362e-39, 9.891206763876373e-37, 1.551663810398814e-34, 2.3518033693641687e-32, 3.4307087966820164e-30, 4.798998552049039e-28, 6.414108734499997e-26, 8.158234083553205e-24, 9.821388020209469e-22, 1.1096931212825364e-19, 1.1607675982205237e-17, 1.0995496403146415e-15, 9.097564206731418e-14, 6.157862118450582e-12, 2.901896487128913e-10, 2.82101925608861e-9, -1.0417790624001826e-6, -0.0061307351256337335, -0.026322147261739235, -0.03756593512574135, -0.044501280273404706, -0.05101892756347002, -0.05752688022512077, -0.0640109372842295, -0.07078484037758782], b = [2.1907432423793217e-6, 1.4001248792739479e-8, 8.393264291266592e-9, 2.8872879347802138e-11, 2.0964527206249195e-11, 4.3416008635444037e-14, 3.928866780522622e-14, -1.0408340855860843e-17, 4.401860820291148e-17, 0.0, 0.0, 0.0, 0.0, 2.7430314963883262e-17, 0.0, 0.0, 2.4936649967166602e-18, -1.1904539853890839e-16, 2.4532459397264006e-14, -3.7102254862170225e-13, 4.1809186050300456e-11, -6.280965945916137e-10, 4.9730558794027976e-8, -6.267653534716157e-7, 7.57194770913299e-5, 8.616385408012647e-5, 5.770763712065617e-5, 2.019478168250612e-5, -1.8772880486727923e-5, -5.8295218416437683e-5, -9.814017391247746e-5, -0.00013833388872236974])
 (u = [-2.4247909829777624e-54, -4.874840736962524e-52, -1.0028057046753584e-49, -2.0232905730225332e-47, -4.011963320437807e-45, -7.805886280584578e-43, -1.4880725599029817e-40, -2.775010271937199e-38, -5.052172403821301e-36, -8.95556612921445e-34, -1.5398996500039947e-31, -2.555718701782353e-29, -4.0682474168762194e-27, -6.164143754881476e-25, -8.813298209729589e-23, -1.1779425190266776e-20, -1.4581426136962015e-18, -1.6604153826463509e-16, -1.74360680098202e-14, -1.734462293678728e-12, -1.7388957234054325e-10, -1.8242880382417252e-8, -1.8007681065760757e-6, -0.0011100601914395029, -0.038430651985070646, -0.05062361456833265, -0.05079439265827457, -0.04743233936254137, -0.04126467217860264, -0.031631554078985255, -0.016588027121187895, 0.013191177275707563], v = [3.1123507786400216e-54, 5.794439420800749e-52, 1.1019536972038682e-49, 2.0530124407885288e-47, 3.756416452563016e-45, 6.7396965146639716e-43, 1.1836794152886814e-40, 2.0307754296757602e-38, 3.395489931845123e-36, 5.517786408181866e-34, 8.685570072060979e-32, 1.318745018063015e-29, 1.9207437210832804e-27, 2.664736308701117e-25, 3.489967269251741e-23, 4.266583706895366e-21, 4.799291903669645e-19, 4.868779575995092e-17, 4.30506348247211e-15, 3.0513704151003924e-13, 1.1831941422080505e-11, -1.0399049355757136e-9, -2.9915746766021815e-7, -0.0003020198407263961, -0.01348521662929522, -0.030741986246191462, -0.040876523499165894, -0.0482973222575921, -0.055206797593701576, -0.06199703627150707, -0.06868238926099675, -0.07556751015680957], b = [2.171106422803543e-6, 1.3753068173454758e-8, 8.239542363501495e-9, 2.8119466824461092e-11, 2.0390301860157e-11, 4.1934337946525346e-14, 3.784928100103446e-14, 8.239936510889834e-18, -4.174178364069192e-17, -9.93129189996722e-17, -7.806255641895632e-17, -7.806255641895632e-17, -7.806255641895632e-17, -7.806255641895632e-17, -6.971419969081794e-17, -3.71881345162528e-17, 1.3118846287074604e-17, -6.076953176781252e-16, 7.37217361809761e-14, -7.928353069550087e-13, 7.459473528237046e-11, 4.731367959638459e-10, -2.6015855866346003e-8, 6.917921367138836e-6, 0.00010781912174508493, 8.521724325859267e-5, 5.108060854492396e-5, 1.3073551967875027e-5, -2.6005259909528205e-5, -6.555776773525897e-5, -0.0001054105185637016, -0.00014560479368338456])
 (u = [-1.111859837560944e-53, -2.1336331648471292e-51, -4.181892559904703e-49, -7.997980760811635e-47, -1.4947637023012334e-44, -2.7259334417162617e-42, -4.845797280326319e-40, -8.389894533645167e-38, -1.4135863115344117e-35, -2.315019458320499e-33, -3.678060427441285e-31, -5.651425780652738e-29, -8.358606736889107e-27, -1.1823487426591399e-24, -1.5864966420647483e-22, -1.9998312560877145e-20, -2.342492838745977e-18, -2.521136925956474e-16, -2.4700800459480853e-14, -2.204386090748163e-12, -1.8539418488728178e-10, -1.626923967768826e-8, -1.5645214134589397e-6, -0.0004151746995729549, -0.03433516219777938, -0.04780722191136722, -0.04806550682449063, -0.04481053982999262, -0.0390857376195313, -0.030359624815486562, -0.016961680178483988, 0.011230183651395508], v = [1.259904154294353e-53, 2.2669849048015895e-51, 4.170029107317421e-49, 7.489527862194549e-47, 1.3147441045675802e-44, 2.2506443594862375e-42, 3.748993006288982e-40, 6.063656555184577e-38, 9.502945557571347e-36, 1.4400182836759789e-33, 2.1050181289252685e-31, 2.959951543156567e-29, 3.988391157443328e-27, 5.122662745667023e-25, 6.225900646550863e-23, 7.089595588992836e-21, 7.464870984321375e-19, 7.138239673620225e-17, 6.033004022992809e-15, 4.272711704803108e-13, 2.122480346169747e-11, -1.91324378878382e-10, -2.3620370235116175e-7, -0.0001205271515951342, -0.011909205202858117, -0.029264935673130543, -0.03979424066478579, -0.04661801046047513, -0.05272727181831293, -0.05861890952143576, -0.06435911779728129, -0.07073331339111578], b = [2.5130353904403557e-6, 1.9483040601198445e-8, 1.0586214055753954e-8, 5.0102344365324236e-11, 2.895129275204411e-11, 9.746522165732241e-14, 5.941991690350257e-14, 1.5547459153442134e-16, 7.37257477290143e-17, -6.179952383167375e-18, 0.0, 9.432558900623889e-18, 3.5453411040275995e-17, 1.0842021724855044e-17, 0.0, 4.119968255444917e-18, 4.9873299934333204e-17, -7.840950111415168e-16, 7.63470233214325e-14, -1.2193515511529918e-12, 9.050708836902399e-11, -1.1482215907243631e-9, 5.1938780478420245e-8, 8.093099675233269e-7, 9.728494062092848e-5, 7.938226702490216e-5, 4.7997577711293335e-5, 1.0346201103140053e-5, -2.8699227942061172e-5, -6.832801481916098e-5, -0.00010834707605249823, -0.00014891401194244602])
 (u = [-3.885034956584103e-53, -7.017119969470931e-51, -1.2920388523842569e-48, -2.313227070310186e-46, -4.0277774847595356e-44, -6.809694674069492e-42, -1.116388253985832e-39, -1.7726069136536594e-37, -2.723550197129081e-35, -4.0473833500058885e-33, -5.816918569304154e-31, -8.086179659370504e-29, -1.0870624884412596e-26, -1.4115005260607548e-24, -1.7642665950159613e-22, -2.1092400080927092e-20, -2.3879585458122826e-18, -2.525368502638964e-16, -2.4520939768194826e-14, -2.1410249159925647e-12, -1.6396109308586657e-10, -1.068798290887064e-8, -5.739631552754351e-7, -2.4926023121638395e-5, -0.021172284958535806, -0.040658476401281146, -0.041438066119574964, -0.038715735984170115, -0.034113683077588726, -0.027289230956029375, -0.016966930668712195, 0.006701434190993619], v = [3.932622759040065e-53, 6.815308966918635e-51, 1.2027964238328184e-48, 2.0615619586471226e-46, 3.4319457278348906e-44, 5.538899346755158e-42, 8.650375553584811e-40, 1.304628973063795e-37, 1.8959905657657934e-35, 2.6491421680088506e-33, 3.5505869721393416e-31, 4.5541262930633436e-29, 5.576110079695402e-27, 6.498073157909345e-25, 7.178125633111695e-23, 7.471179592875719e-21, 7.258430291003131e-19, 6.486686156375665e-17, 5.212219122186927e-15, 3.628098920919954e-13, 2.039319830311657e-11, 7.632790936940354e-10, -9.899969950173872e-10, -3.0328544736893376e-6, -0.0071415807083441704, -0.02478286023956643, -0.03570155135890721, -0.04115449063307272, -0.04582134903549309, -0.050222244646661224, -0.05447499750117919, -0.05970450995304669], b = [2.997364951847197e-6, 3.1926787627057154e-8, 1.3124671139393834e-8, 1.0888304320666453e-10, 3.6997745145317396e-11, 2.721143622930189e-13, 7.834466582423705e-14, 5.956606735635361e-16, 2.4557179206796675e-16, 1.5233040523421337e-16, 1.5482407023093003e-16, 1.4029576111962427e-16, 7.112366251504909e-17, 3.458604930228759e-17, 8.359198749863239e-17, 8.131516293641283e-17, 8.543513119185775e-17, -3.754592123317302e-16, 2.6204299147236654e-14, -8.154340917776448e-13, 4.0180865525010073e-11, -1.0519713907483977e-9, 4.319173987362519e-8, -8.631183252820625e-7, 5.7769794055010054e-5, 6.944913384175733e-5, 4.4648704762999665e-5, 7.010330379191888e-6, -3.232916039544719e-5, -7.23620512702243e-5, -0.00011294904796554455, -0.00015466971286774673])
 (u = [-1.0399563219422958e-52, -1.878917737627551e-50, -3.4589175849434866e-48, -6.191060893910417e-46, -1.0778873824342594e-43, -1.8238422648068045e-41, -2.99762173282627e-39, -4.783690808524904e-37, -7.409007575367737e-35, -1.1129220035840383e-32, -1.6191985721141823e-30, -2.276301816003312e-28, -3.080228680329687e-26, -3.9896072535237654e-24, -4.9094879581671305e-22, -5.686996489727266e-20, -6.133939025572413e-18, -6.085427544335285e-16, -5.4828891459650067e-14, -4.441274666581482e-12, -3.2457164388472e-10, -2.24570136312573e-8, -1.678280657455927e-6, -0.0001548941066290062, -0.03067435663706566, -0.042011012689089036, -0.041949560957398595, -0.038984013890482776, -0.034211035593774815, -0.027229939292378966, -0.01673772510678821, 0.0072452971025700665], v = [1.0100390365092026e-52, 1.7525970295546738e-50, 3.095650812585154e-48, 5.308178468716343e-46, 8.834867511919914e-44, 1.4249384599052813e-41, 2.2237464510129205e-39, 3.3528988421784483e-37, 4.876973511736907e-35, 6.832612080384128e-33, 9.2031103277565e-31, 1.188984249723591e-28, 1.4686006859319333e-26, 1.7262770203194925e-24, 1.9184489038910972e-22, 1.9974343210524073e-20, 1.924369424603893e-18, 1.6865204573267539e-16, 1.3114159743187783e-14, 8.662307050481727e-13, 4.347946694299726e-11, 7.962376663182945e-10, -1.8545947388664266e-7, -3.938181147178822e-5, -0.011190447902764392, -0.026959128574737524, -0.03748185534947611, -0.04298749972734365, -0.047613537097395846, -0.05193303194425507, -0.056085946073839835, -0.061183284022481015], b = [3.046702897705573e-6, 3.38422747182647e-8, 1.3154100060716253e-8, 1.1791605040480102e-10, 3.6516888905074874e-11, 2.973717192628067e-13, 7.615327639320935e-14, 5.139118297581291e-16, 7.405100838075995e-17, 6.071532165918825e-18, -2.3852447794681098e-18, -4.5428071027142636e-17, -5.561957144850638e-17, -3.187554387107383e-17, 2.3310346708438345e-17, 4.336808689942018e-19, -1.1058862159352145e-17, -9.78926141537162e-16, 5.177737578965225e-14, -1.4097133833532172e-12, 6.538248246848802e-11, -1.4498379429525934e-9, 5.75407863456388e-8, -8.052037263504494e-7, 7.540479879052981e-5, 6.646025012969713e-5, 3.9762315391402166e-5, 2.1464269870808815e-6, -3.719049622050871e-5, -7.725522665976975e-5, -0.00011791973572143379, -0.0001598418998482085])
 (u = [-1.5068127846713254e-52, -2.694835165767626e-50, -4.913810304304289e-48, -8.720043680893749e-46, -1.5065728074082e-43, -2.5323690632658973e-41, -4.1393933561782767e-39, -6.576943027152152e-37, -1.0151042783492711e-34, -1.5202368913961812e-32, -2.2048673764252408e-30, -3.087235845070274e-28, -4.154213276215059e-26, -5.339089316768571e-24, -6.503631797387058e-22, -7.440452323791097e-20, -7.913507225168318e-18, -7.742076344804015e-16, -6.904109645984971e-14, -5.60156529374491e-12, -4.2213248582989493e-10, -3.160143244475771e-8, -2.529012346307954e-6, -0.0003177037674630412, -0.03145005902447202, -0.04147253561053737, -0.041175409609579806, -0.03820966956510668, -0.033537826432249215, -0.026753452929515146, -0.0165670072352808, 0.006698215359001076], v = [1.4674449753710646e-52, 2.519402478835386e-50, 4.403955100892838e-48, 7.476951091636888e-46, 1.2326459678599043e-43, 1.9703722738767464e-41, 3.049895139790586e-39, 4.56529842010043e-37, 6.599207208731258e-35, 9.197338041240534e-33, 1.233447455658298e-30, 1.587484948578446e-28, 1.9534760207684014e-26, 2.286348729427691e-24, 2.5268787999084025e-22, 2.6117056615607112e-20, 2.491968603122619e-18, 2.1563706821205317e-16, 1.6471642036954005e-14, 1.0543083314087436e-12, 4.818016086277138e-11, 2.9399559372463165e-11, -3.4920964769543515e-7, -8.959287714764645e-5, -0.011703354435981472, -0.02696609155845569, -0.03736928249855488, -0.04270950766702868, -0.0471502173885233, -0.05127532613676507, -0.055240015593609054, -0.06010274588071191], b = [3.0950410119824894e-6, 3.590752390932907e-8, 1.310679039042903e-8, 1.2721130464125874e-10, 3.5641363124727654e-11, 3.216433196173707e-13, 7.282022207455441e-14, 7.507015842289633e-16, 1.5027042110649091e-16, 5.355958732078392e-17, 6.266688556966216e-17, 7.19910242530375e-17, 9.031404096804252e-17, 1.1405806854547507e-16, 1.3975366003338152e-16, 1.4072944198861848e-16, 1.4864411784776266e-16, -9.392443420241925e-16, 5.00682394849461e-14, -1.431553877176417e-12, 6.314828168837647e-11, -1.4067692437921657e-9, 4.720073196808365e-8, -3.6358761302862554e-7, 7.441656207878041e-5, 6.458428265370043e-5, 3.799094746635901e-5, 3.1916295811380813e-7, -3.908559046669703e-5, -7.923670428186155e-5, -0.00012002788821256743, -0.00016222066449025327])
 (u = [-1.7179399080149983e-52, -3.0813716968957176e-50, -5.633513399027339e-48, -1.0021992482882972e-45, -1.7353572058678215e-43, -2.922647501856014e-41, -4.7853192831998287e-39, -7.613158254207318e-37, -1.17598624975109e-34, -1.7613897479851043e-32, -2.5526000701499938e-30, -3.5672449513281277e-28, -4.784618017644341e-26, -6.1209689822484755e-24, -7.411765124460525e-22, -8.41934736634669e-20, -8.884686674048925e-18, -8.625817792392369e-16, -7.651501397236794e-14, -6.222102315275765e-12, -4.778389874908771e-10, -3.687048935296814e-8, -2.8326019243823955e-6, -0.0004950947075373166, -0.03216157592269388, -0.041518865226654364, -0.041064169903588944, -0.03806801555366623, -0.033400401596690735, -0.02664540443067022, -0.016512729259005827, 0.006606372543131606], v = [1.643682996702548e-52, 2.831265212671428e-50, 4.964790326895552e-48, 8.455148466835362e-46, 1.3979144753442887e-43, 2.2404304870258813e-41, 3.476138498048626e-39, 5.21424252375594e-37, 7.550763560656666e-35, 1.0538480597796234e-32, 1.4146419943467831e-30, 1.8212637787452395e-28, 2.2400427645991742e-26, 2.6178347719073856e-24, 2.8855853376332404e-22, 2.970799867472182e-20, 2.8196504130531785e-18, 2.4229590965126575e-16, 1.8322930565615898e-14, 1.1502618485214018e-12, 4.920052005772474e-11, -5.435716635839974e-10, -4.098289389070262e-7, -0.0001454730446315154, -0.012090079273876831, -0.027178629856167584, -0.037473570440978425, -0.042785943094897264, -0.047188803310838384, -0.05127012077973221, -0.05518996039819058, -0.05999225154503015], b = [3.1072121769531134e-6, 3.650984785373915e-8, 1.3056990859176831e-8, 1.297586519569821e-10, 3.521878036601145e-11, 3.276335366203531e-13, 7.138755732383206e-14, 6.7914424084492e-16, 1.336821278674627e-16, -1.9949319973733282e-17, -4.911435841359335e-17, -4.7488055154865094e-17, -5.724587470723463e-17, -7.28583859910259e-17, -5.616167253474913e-17, -3.241764495731658e-17, 3.393552799879629e-17, -1.036063596027148e-15, 5.1016157444350174e-14, -1.4643908915340509e-12, 6.127174890504725e-11, -1.3249225541559287e-9, 1.728733475694997e-8, 3.631013595205724e-7, 7.463756634299687e-5, 6.413296521035679e-5, 3.7269383391538526e-5, -4.247706102490959e-7, -3.984770681289172e-5, -8.00227686370354e-5, -0.00012085108748288651, -0.00016312554129741587])
 (u = [-2.0968429723376048e-52, -3.7555826926736724e-50, -6.852824465161359e-48, -1.2162695581064036e-45, -2.1002734388881475e-43, -3.5263189527582066e-41, -5.754091589138241e-39, -9.120482608583306e-37, -1.4031228860462705e-34, -2.0922397801437863e-32, -3.017020766084017e-30, -4.192741780001554e-28, -5.58832700231972e-26, -7.099609206845292e-24, -8.533253315572163e-22, -9.622463475308384e-20, -1.009358399703501e-17, -9.781145019278049e-16, -8.747078357953271e-14, -7.315331545172385e-12, -5.922567595007808e-10, -4.775150039795989e-8, -3.6210500955295115e-6, -0.0008048324406538335, -0.033271534131606155, -0.041817890438451834, -0.041157332473479545, -0.03809831409195092, -0.033394669793557305, -0.026614801717873338, -0.01646179658218563, 0.006671282509999266], v = [1.9970031529292313e-52, 3.438686449020943e-50, 6.025028794095489e-48, 1.0248034827593099e-45, 1.69141139980524e-43, 2.7048511847909753e-41, 4.185608013807381e-39, 6.259284242650113e-37, 9.033029271678306e-35, 1.2559573774296906e-32, 1.678958460101082e-30, 2.1517321375004726e-28, 2.633235546320843e-26, 3.060210627771735e-24, 3.3522018481489e-22, 3.426883717562479e-20, 3.225743437641418e-18, 2.7425053984056877e-16, 2.0384630858199245e-14, 1.229467583923724e-12, 4.4821471172030435e-11, -1.8573103455354868e-9, -5.698047923216553e-7, -0.00023455744787922295, -0.012682917171866655, -0.027595445311958806, -0.037712687427845185, -0.04303402650977138, -0.04743806707268687, -0.05151341010991593, -0.0554227886597674, -0.060205020380797135], b = [3.112396286415785e-6, 3.6843400258185915e-8, 1.299819880093249e-8, 1.3106268501039686e-10, 3.4836440806693125e-11, 3.3020895046087517e-13, 7.014441111286018e-14, 6.353424730765056e-16, 3.2851325826310784e-17, -8.001412032943023e-17, -8.955509944730267e-17, -7.806255641895632e-17, -7.806255641895632e-17, -7.611099250848241e-17, -7.806255641895632e-17, -7.990570011218168e-17, -1.3769367590565906e-17, -1.1721309686740788e-15, 5.253393206561263e-14, -1.4586163223431758e-12, 5.231287072363783e-11, -9.92435411194953e-10, -3.125266697984099e-8, 1.6176944465922261e-6, 7.553437354452058e-5, 6.405319277255847e-5, 3.6603752311091206e-5, -1.116556720827459e-6, -4.054801481961863e-5, -8.073370705484081e-5, -0.00012158048553716435, -0.0001638983115925196])
 (u = [-2.0959116778279995e-52, -3.752733091887228e-50, -6.852174210192125e-48, -1.2182914697354874e-45, -2.1099261641288477e-43, -3.5570825878494187e-41, -5.8345350165329604e-39, -9.304560718655105e-37, -1.4410300287144422e-34, -2.1634535010087692e-32, -3.1398723457446292e-30, -4.38785691159531e-28, -5.873680620884493e-26, -7.483541381971602e-24, -9.008275532673882e-22, -1.0164598919568489e-19, -1.0671497963811946e-17, -1.0374043599847725e-15, -9.359621751278302e-14, -7.966446791984271e-12, -6.592972097152703e-10, -5.3661854351440945e-8, -4.055095843826338e-6, -0.0009507937690573675, -0.03369604408033994, -0.0420382318832992, -0.041322169229287845, -0.038233342503307184, -0.03349666303904367, -0.026675298127638643, -0.01646785694128416, 0.006779436469105604], v = [2.008372860685943e-52, 3.4497014137353405e-50, 6.03373628725076e-48, 1.0253530582212963e-45, 1.6924220693395657e-43, 2.7095237727465515e-41, 4.20234448319586e-39, 6.305831651123303e-37, 9.141377715866687e-35, 1.2779713025780004e-32, 1.718871877267842e-30, 2.2170400077535054e-28, 2.7301311621424293e-26, 3.1906129167256094e-24, 3.5108637280498613e-22, 3.6001287144962596e-20, 3.3931334215929332e-18, 2.8813884700865167e-16, 2.129714233796717e-14, 1.2627208054623812e-12, 4.2658995345587096e-11, -2.4464323258672535e-9, -6.473615103468636e-7, -0.0002718059630058246, -0.012888743130497645, -0.02778605207884922, -0.037854368398311126, -0.04320865240952337, -0.047642739002261014, -0.05174470204779401, -0.05567676822215159, -0.06048413826782202], b = [3.108385042020767e-6, 3.6721789127731885e-8, 1.2976466902316769e-8, 1.3047295145326987e-10, 3.476728909108939e-11, 3.284575302714421e-13, 6.998850284045677e-14, 5.748439918518145e-16, 6.14742631799281e-17, -9.616873269946424e-17, -1.0137290312739466e-16, -9.573505183047004e-17, -9.020562075079397e-17, -7.936359902593892e-17, -7.903833837419327e-17, -7.068998164605489e-17, -2.8297676701871666e-17, -1.1808045860539629e-15, 5.284585703063671e-14, -1.4292078805155928e-12, 4.43403058952678e-11, -7.48488011750436e-10, -5.6630108443735117e-8, 2.24814446183187e-6, 7.603446740507013e-5, 6.424809912111893e-5, 3.655934998238907e-5, -1.1570037261105686e-6, -4.057917874401742e-5, -8.075535874635586e-5, -0.0001215909672702564, -0.00016388850860333566])

and plot the parameters, the cost, and the profiles of velocity and the buoyancy anomaly b - N² z. We plot the state of the optimization at iteration n, starting with the final iteration.

fig = Figure(size=(1200, 800))
top = fig[2, 1] = GridLayout()
bottom = fig[3, 1] = GridLayout()

n = Observable(length(history))

title = @lift @sprintf("Iteration %d, 𝒥 = %.1e", $n - 1, history[$n].𝒥)
Label(fig[1, 1], title, fontsize=20, tellwidth=false)

iterations = 0:length(history)-1
labels = ["Cu₀", "Cc₀", "τˣ", "Jᵇ"]

ax = Axis(top[1, 1]; xlabel="Iteration", ylabel="Parameter / nature run value", title="Parameters")
hlines!(ax, 1; color=:gray, linestyle=:dash)
for i in 1:4
    ratio = [h.θ[i] / θ★[i] for h in history]
    points = @lift Point2f.(iterations[1:$n], ratio[1:$n])
    scatterlines!(ax, points; label=labels[i])
end
xlims!(ax, -0.5, length(history) - 0.5)
ylims!(ax, 0, 2)
axislegend(ax, position=:rt)

costs = [h.𝒥 for h in history]
ax = Axis(top[1, 2]; xlabel="Iteration", ylabel="𝒥", yscale=log10, title="Cost")
cost_points = @lift Point2f.(iterations[1:$n], costs[1:$n])
scatterlines!(ax, cost_points)
xlims!(ax, -0.5, length(history) - 0.5)
ylims!(ax, minimum(costs) / 2, 2 * maximum(costs))

axu = Axis(bottom[1, 1]; xlabel="u (m s⁻¹)", ylabel="z (m)", title="u")
axv = Axis(bottom[1, 2]; xlabel="v (m s⁻¹)", title="v")
axb = Axis(bottom[1, 3]; xlabel="b - N² z (m s⁻²)", title="Buoyancy anomaly", xticks=LinearTicks(3))

for (ax, name) in zip((axu, axv, axb), (:u, :v, :b))
    lines!(ax, getproperty(nature_state, name), z; linewidth=4, label="nature run")
    lines!(ax, getproperty(first(iteration_states), name), z; linewidth=2, linestyle=:dot, color=(:gray, 0.6), label="initial guess")
    profile = @lift getproperty(iteration_states[$n], name)
    lines!(ax, profile, z; linewidth=3, linestyle=:dash, label="estimate")

    all_profiles = [getproperty(state, name) for state in (nature_state, iteration_states...)]
    xmin = minimum(minimum, all_profiles)
    xmax = maximum(maximum, all_profiles)
    δx = (xmax - xmin) / 20
    xlims!(ax, xmin - δx, xmax + δx)
    ylims!(ax, -80, 0)
end

axislegend(axu, position=:lb)

Finally, we animate the progress of the optimization,

CairoMakie.record(fig, "single_column_parameter_estimation.mp4", 1:length(history), framerate=2) do i
    n[] = i
end

Comparing BFGS with gradient descent

Optim makes it easy to try other optimizers with the same compiled cost and gradient. We repeat the optimization with plain gradient descent, starting from the same initial guess and taking the same number of iterations,

gradient_descent_result = optimize(𝒥, ∇𝒥!, θ₀, GradientDescent(), options)

 * Status: failure (reached maximum number of iterations)

 * Candidate solution
    Final objective value:     6.099023e-03

 * Found with
    Algorithm:     Gradient Descent

 * Convergence measures
    |x - x'|               = 1.26e-02 ≰ 0.0e+00
    |x - x'|/|x'|          = 7.42e-03 ≰ 0.0e+00
    |f(x) - f(x')|         = 1.79e-04 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 2.94e-02 ≰ 0.0e+00
    |g(x)|                 = 1.07e-02 ≰ 1.0e-08

 * Work counters
    Seconds run:   17  (vs limit Inf)
    Iterations:    10
    f(x) calls:    17
    ∇f(x) calls:   17
    ∇f(x)ᵀv calls: 0

The optimizers' line searches evaluate the cost and gradient a different number of times per iteration, so we also count evaluations, and measure how long each evaluation takes,

forward_time = @elapsed 𝒥(θ₀)
gradient_time = @elapsed cost_and_gradient(θ₀)

@info @sprintf("Evaluating the cost takes %.2f s and evaluating its gradient takes %.2f s", forward_time, gradient_time)

for (name, result) in (("BFGS", bfgs_result), ("Gradient descent", gradient_descent_result))
    θ = Optim.minimizer(result)
    @info @sprintf("%s: 𝒥 = %.2e after %d iterations with %d cost and %d gradient evaluations; θ / θ★ = %s",
                   name, Optim.minimum(result), Optim.iterations(result),
                   Optim.f_calls(result), Optim.g_calls(result), string(round.(θ ./ θ★, digits=3)))
end

fig = Figure(size=(600, 400))
ax = Axis(fig[1, 1]; xlabel="Iteration", ylabel="𝒥", yscale=log10, title="Cost")

for (label, result) in (("BFGS", bfgs_result), ("gradient descent", gradient_descent_result))
    cost_trace = Optim.f_trace(result)
    scatterlines!(ax, 0:length(cost_trace)-1, cost_trace; label)
end

axislegend(ax)

---

Julia version and environment information

This example was executed with the following version of Julia:

using InteractiveUtils: versioninfo
versioninfo()

Julia Version 1.13.0
Commit d1c37793dd2 (2026-09-09 19:00 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 128 × AMD EPYC 9374F 32-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-20.1.8 (ORCJIT, znver4)
  GC: Built with stock GC
Threads: 1 default, 1 interactive, 1 GC (on 128 virtual cores)
Environment:
  JULIA_LOAD_PATH = @:@v#.#:@stdlib
  JULIA_DEPOT_PATH = /var/lib/buildkite-agent/.julia:/var/lib/buildkite-agent/.julia/juliaup/julia-1.13.0+0.x64.linux.gnu/local/share/julia:/var/lib/buildkite-agent/.julia/juliaup/julia-1.13.0+0.x64.linux.gnu/share/julia
  JULIA_VERSION_ENZYME = 1.11.9
  JULIA_PKG_SERVER_REGISTRY_PREFERENCE = eager
  LD_LIBRARY_PATH = 
  JULIA_MAX_NUM_PRECOMPILE_FILES = 24
  JULIA_VERSION = 1.13.0
  JULIA_CUDA_USE_COMPAT = false
  JULIA_PROJECT = /var/lib/buildkite-agent/Oceananigans.jl-33941/docs/
  JULIA_DEBUG = Literate

These were the top-level packages installed in the environment:

import Pkg
Pkg.status()

Status `~/Oceananigans.jl-33941/docs/Project.toml`
  [79e6a3ab] Adapt v4.7.1
⌃ [052768ef] CUDA v5.11.3
  [13f3f980] CairoMakie v0.15.15
⌅ [e30172f5] Documenter v1.17.0
  [daee34ce] DocumenterCitations v1.5.0
  [4710194d] DocumenterVitepress v0.3.6
  [7da242da] Enzyme v0.13.205
  [033835bb] JLD2 v0.6.7
  [63c18a36] KernelAbstractions v0.9.42
  [98b081ad] Literate v2.21.0
  [da04e1cc] MPI v0.20.27
  [85f8d34a] NCDatasets v0.14.15
  [9e8cae18] Oceananigans v0.113.2 `..`
  [429524aa] Optim v2.3.2
  [f27b6e38] Polynomials v4.1.3
  [3c362404] Reactant v0.2.288
  [6038ab10] Rotations v1.7.1
  [d496a93d] SeawaterPolynomials v0.3.10
  [09ab397b] StructArrays v0.7.3
  [bdfc003b] TimesDates v0.3.3
  [0a941bbe] Zarr v0.10.2
  [b77e0a4c] InteractiveUtils v1.11.0
  [37e2e46d] LinearAlgebra v1.13.0
  [44cfe95a] Pkg v1.13.0
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated`

---

This page was generated using Literate.jl.