Single column parameter estimation with Reactant and Enzyme

We estimate turbulence closure parameters and surface fluxes in a "twin experiment": generate observations with known parameters, then fit a model to those observations. Enzyme computes gradients through the simulation, compiled by Reactant.

Install dependencies

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.

To use the IFRT runtime, as Oceananigans' Reactant tests do, add the following to your project's LocalPreferences.toml and restart Julia:

[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")

A single column model

Surface wind stress τˣ and cooling buoyancy flux Jᵇ mix an ocean column. The k-ϵ closure computes viscosity and diffusivity from turbulent kinetic energy e and its dissipation rate ϵ,

The stability functions 𝕊u and 𝕊c depend on shear and stratification. We estimate their magnitude parameters Cu₀ and Cc₀, along with the surface fluxes τˣ and Jᵇ.

To change parameters inside a compiled function, we store closure parameters as Reactant numbers and surface fluxes as mutable Fields.

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)

We rebuild the closure inside the differentiated function because VariableStabilityFunctions derives another constant, 𝕊u₀, from Cu₀. This includes that dependence in the gradient.

using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: VariableStabilityFunctions

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(VerticallyImplicitTimeDiscretization(), 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. We reset the fields, tendencies, closure state, and clock before each run. Reactant compiles run! and its time-stepping loop together. We use a one-minute step; ten-minute steps give noisy solutions with this closure.

using Oceananigans.Models: reset!

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

simulation = Simulation(model; Δt=1minute, stop_time=12hours, verbose=false)

function run_column!(simulation, normalized_parameters, bᵢ)
    model = simulation.model
    assign_parameters!(model, normalized_parameters)

    reset!(model)
    reset!(model.clock)
    set!(model, b=bᵢ)
    run!(simulation)

    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⁻³. The vector θ★ contains these values divided by scales, in the same order.

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

compiled_run_column! = @compile raise=true raise_first=true sync=true run_column!(
    simulation, reactant_parameters(θ★), bᵢ)

compiled_run_column!(simulation, reactant_parameters(θ★), bᵢ)

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

observations = deepcopy((u=model.velocities.u, v=model.velocities.v, b=model.tracers.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, simulation, bᵢ, observations)
    run_column!(simulation, normalized_parameters, bᵢ)

    model = simulation.model
    u, v = model.velocities
    b = model.tracers.b
    u★, v★, b★ = observations

    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)

Enzyme.gradient manages derivative buffers for the parameters and the mutated simulation. ReverseWithPrimal returns the cost too. We keep the parameter gradient and treat the initial condition and observations as constants.

function cost_and_gradient(parameters, simulation, bᵢ, observations)
    mode = Enzyme.set_strong_zero(Enzyme.ReverseWithPrimal)
    result = Enzyme.gradient(mode, cost, parameters, simulation, Enzyme.Const(bᵢ), Enzyme.Const(observations))
    return result.val, result.derivs[1]
end

cost_and_gradient (generic function with 1 method)

We start from a deliberately inaccurate guess,

θ₀ = [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.

compiled_cost = @compile raise=true raise_first=true sync=true cost(
    reactant_parameters(θ₀), simulation, bᵢ, observations)

compiled_cost_and_gradient = @compile raise=true raise_first=true sync=true cost_and_gradient(
    reactant_parameters(θ₀), simulation, bᵢ, observations)

Reactant compiled function cost_and_gradient (with tag ##cost_and_gradient_reactant#665)

Minimizing the cost

We use Optim.jl's BFGS optimizer. A combined cost-and-gradient function reuses the cost from differentiation; a cost-only request uses the forward calculation.

using Optim
using Optim.NLSolversBase: only_fg!

function cost_and_gradient!(F, G, θ)
    parameters = reactant_parameters(θ)
    if isnothing(G)
        return Float64(compiled_cost(parameters, simulation, bᵢ, observations))
    end
    value, gradient = compiled_cost_and_gradient(parameters, simulation, bᵢ, observations)
    G .= Float64.(collect(gradient))
    return isnothing(F) ? nothing : Float64(value)
end

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

 * Status: failure (reached maximum number of iterations)

 * Candidate solution
    Final objective value:     9.338283e-04

 * Found with
    Algorithm:     BFGS

 * Convergence measures
    |x - x'|               = 2.27e-02 ≰ 0.0e+00
    |x - x'|/|x'|          = 1.10e-02 ≰ 0.0e+00
    |f(x) - f(x')|         = 2.67e-04 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 2.85e-01 ≰ 0.0e+00
    |g(x)|                 = 3.92e-03 ≰ 1.0e-08

 * Work counters
    Seconds run:   1  (vs limit Inf)
    Iterations:    10
    f(x) calls:    14
    ∇f(x) calls:   14
    ∇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))]

11-element Vector{@NamedTuple{θ::Vector{Float64}, 𝒥::Float64}}:
 (θ = [0.6, 1.8, -0.5, 0.5], 𝒥 = 0.10753814307469763)
 (θ = [0.5433533552166119, 1.804315031110822, -1.487415123429518, 0.7615825783944596], 𝒥 = 0.06422337779398422)
 (θ = [0.5807307635138251, 1.7946120232248015, -1.0502019781298815, 0.6686777746643612], 𝒥 = 0.012127171515344275)
 (θ = [0.5791329077954477, 1.7934357592509598, -1.087197554246519, 0.6836050954929247], 𝒥 = 0.01167630075449744)
 (θ = [0.58148216918827, 1.7916096406637585, -1.0769223485331085, 0.6964648768697342], 𝒥 = 0.011356635106900262)
 (θ = [0.5927691014369907, 1.7809207191448528, -1.0481205495368886, 0.7766410393144456], 𝒥 = 0.010323540038151469)
 (θ = [0.6158079839900025, 1.7590676797246247, -1.0158061195879493, 0.9465620543420097], 𝒥 = 0.008641889044206392)
 (θ = [0.6848654607745668, 1.6948399937328658, -0.954329522392289, 1.4502339793887709], 𝒥 = 0.004604158605613909)
 (θ = [0.7597337430053009, 1.6280425833834091, -0.9168685140596078, 1.9729110046121865], 𝒥 = 0.0019196133218774582)
 (θ = [0.7783089741563496, 1.6127025411971836, -0.9267831515817962, 2.0847872185619196], 𝒥 = 0.0012003890913986994)
 (θ = [0.778624460324299, 1.6135586324027291, -0.9485584837394159, 2.0621223335468346], 𝒥 = 0.0009338283325708095)

Visualizing the parameter estimation

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

using Oceananigans.Architectures: on_architecture

z = on_architecture(CPU(), znodes(grid, Center()))

function final_state(θ)
    compiled_run_column!(simulation, reactant_parameters(θ), bᵢ)
    u = vec(on_architecture(CPU(), interior(model.velocities.u)))
    v = vec(on_architecture(CPU(), interior(model.velocities.v)))
    b = vec(on_architecture(CPU(), 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.828330698974153e-57, -8.639892656839315e-55, -1.5597002607069982e-52, -2.720736899784879e-50, -4.5992716822599976e-48, -7.529316368795246e-46, -1.1930668019606532e-43, -1.8288429372146635e-41, -2.71035234811193e-39, -3.88075490683321e-37, -5.364430732784074e-35, -7.153552863878164e-33, -9.196829899229518e-31, -1.1395703960883703e-28, -1.3612501552923554e-26, -1.5689910200761436e-24, -1.7472020382305014e-22, -1.881003353146055e-20, -1.9542191483084875e-18, -1.946524304591569e-16, -1.834983643023791e-14, -1.6051787560775865e-12, -1.2696312038628783e-10, -8.79502412084637e-9, -5.138469947083334e-7, -2.3305579043132505e-5, -0.014166397135379438, -0.031836875419803394, -0.030283632571008558, -0.02548866892539888, -0.016643798884779563, 0.003236163414431197], v = [7.813001302115706e-57, 1.373749621221405e-54, 2.432306146542253e-52, 4.15192585231236e-50, 6.849699771387683e-48, 1.090940028246282e-45, 1.6757776460826593e-43, 2.480117426053591e-41, 3.532599350889142e-39, 4.836902254699809e-37, 6.357958493011279e-35, 8.011170876833784e-33, 9.65958752672177e-31, 1.1123782586749847e-28, 1.2206810611173101e-26, 1.2731734704660473e-24, 1.2582901986998247e-22, 1.1737426684906428e-20, 1.0274625086670098e-18, 8.360634704320304e-17, 6.218296586765465e-15, 4.0966412158660164e-13, 2.2421334962493392e-11, 8.5719762646541e-10, 3.961027929735586e-9, -2.59932265486795e-6, -0.004419693050865934, -0.02061101631191612, -0.028340770088155068, -0.03400150486336239, -0.039552449977375906, -0.045739819925910516], 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, 1.0842021724855044e-19, 1.0842021724855044e-19, -9.75781955236954e-18, -1.2414114874959026e-16, 2.748848241043711e-14, -4.605965673969148e-13, 4.2922897115944825e-11, -7.424125522789878e-10, 4.5807760591708244e-8, -6.596958509778265e-7, 4.861940315181333e-5, 5.4846815668836533e-5, 1.937857719759361e-5, -1.9611485075264582e-5, -5.935022276083305e-5, -9.956767885180152e-5])
 (u = [-3.4282491633732205e-53, -7.872377141629999e-51, -1.8279133333981737e-48, -4.12536213716125e-46, -9.078212541591162e-44, -1.9470988828168567e-41, -4.068378709862096e-39, -8.27305303584165e-37, -1.6341790159043372e-34, -3.12583793307128e-32, -5.7634281635646526e-30, -1.0181550452652964e-27, -1.7106659445889973e-25, -2.7110011799161357e-23, -4.0182103911764564e-21, -5.53245056988518e-19, -7.071815525043126e-17, -8.523845394084119e-15, -1.0113391591760219e-12, -1.2414869070335164e-10, -1.540932760652173e-8, -1.6899199045810374e-6, -0.0015332234699545607, -0.04211783984406421, -0.05775342735019523, -0.05941543978727937, -0.057048819635057806, -0.05208100966713598, -0.044228975500766474, -0.03260877340713555, -0.014876713582795192, 0.01894454678735471], v = [2.8228432589272694e-53, 6.19224714274385e-51, 1.3714689172301228e-48, 2.9447704365056314e-46, 6.142981188966825e-44, 1.2432978503909685e-41, 2.4385254760795118e-39, 4.628877006645802e-37, 8.490194067905211e-35, 1.5013021302633507e-32, 2.5508961665165994e-30, 4.145112950331111e-28, 6.399772040366236e-26, 9.306983707708263e-24, 1.260453388091515e-21, 1.5652132555304703e-19, 1.7394191729857446e-17, 1.6463460924934094e-15, 1.1403571034142633e-13, 1.1167286519847775e-12, -1.3868196758244698e-9, -2.9042880351830294e-7, -0.0003901240059414713, -0.014334691593711627, -0.03293192981642873, -0.04420595090786853, -0.052142401900135346, -0.05945671573224279, -0.06672088496504774, -0.0739734509967224, -0.08115592795234809, -0.08834516928555708], b = [2.0347386354307474e-6, 1.2000364073110059e-8, 7.25774373174122e-9, 2.2752594079633592e-11, 1.6905135277744954e-11, 3.122415520584454e-14, 2.9540822912843545e-14, -2.3852447794681098e-17, -2.927345865710862e-17, -1.3877787807814457e-17, 2.667137344314341e-17, 6.364266752489911e-17, 7.806255641895632e-17, 7.806255641895632e-17, 7.361732751176575e-17, 4.7379634937616544e-17, -3.1051550219984847e-16, 5.105985079190134e-14, -1.985969982215563e-13, 4.732710686024266e-11, 2.074436603559642e-9, -6.025128793319323e-8, 1.1869918301719731e-5, 0.00012118048506306399, 0.00010206572849704338, 7.046229189074771e-5, 3.3252010981191704e-5, -5.431482713237569e-6, -4.472311287020294e-5, -8.434612567389533e-5, -0.0001242092811967295, -0.0001643672575361013])
 (u = [-1.1087524046675897e-54, -2.1062644801213776e-52, -4.101457289489421e-50, -7.851003410014196e-48, -1.4805790965871933e-45, -2.7450450127035465e-43, -4.991510414776549e-41, -8.879050723497113e-39, -1.5415022090578566e-36, -2.607328845919465e-34, -4.2919176576455663e-32, -6.871676230588414e-30, -1.069398478161319e-27, -1.6147152720730407e-25, -2.3555941823518075e-23, -3.2947089527291803e-21, -4.367217568354442e-19, -5.4024213448222666e-17, -6.121008883202487e-15, -6.21302737267227e-13, -5.503374289389769e-11, -4.119644627053547e-9, -2.5112261481916127e-7, -1.1706320656009547e-5, -0.022132564467768467, -0.04764809909124042, -0.0488984902648092, -0.04592711241466846, -0.0402584753656421, -0.031278917973320254, -0.017116634849811956, 0.011290937911387154], v = [1.6673774936827613e-54, 2.955699463128883e-52, 5.337166744967296e-50, 9.423644783928527e-48, 1.633277417175492e-45, 2.7776927565129422e-43, 4.630916044502733e-41, 7.55443003869887e-39, 1.2026499506104038e-36, 1.8624148378846972e-34, 2.795595337561005e-32, 4.052847570661097e-30, 5.654073949320311e-28, 7.56159450643552e-26, 9.649242204100925e-24, 1.1672209049748729e-21, 1.3251491576351368e-19, 1.3905617953534307e-17, 1.3177507166294575e-15, 1.0874928206757483e-13, 7.332561108373849e-12, 3.4654752315891133e-10, 4.082857289749848e-9, -1.1237253709817555e-6, -0.006787320066900276, -0.02689336556494583, -0.038112103973541624, -0.045082850731213195, -0.05161899550719958, -0.05813871200647633, -0.0646301621479878, -0.07140760323982787], b = [2.190063229868611e-6, 1.3992694618660331e-8, 8.387877782753059e-9, 2.8847011151328372e-11, 2.0944262600283703e-11, 4.335312490943988e-14, 3.923445769660194e-14, 3.187554387107383e-17, -3.903127820947816e-18, -4.694595406862234e-17, -5.4643789493269423e-17, -6.320898665590491e-17, -7.155734338404329e-17, -8.630249292984615e-17, -7.806255641895632e-17, -8.131516293641283e-18, 7.37257477290143e-18, -1.839891086707901e-16, 2.8975303059675106e-14, -4.4081903925741184e-13, 4.712974818353241e-11, -6.937384866871274e-10, 5.3449591592329496e-8, -6.568611002120517e-7, 8.052181651221834e-5, 8.562953589860757e-5, 5.6948693304646275e-5, 1.94584434004727e-5, -1.949479023281026e-5, -5.900787281783278e-5, -9.88464903223903e-5, -0.00013903497055367718])
 (u = [-1.4545330836939293e-54, -2.833009449119871e-52, -5.651191917376977e-50, -1.106375960694668e-47, -2.1292544231608716e-45, -4.019087069492883e-43, -7.424777460623262e-41, -1.3399356344228485e-38, -2.3588676300738413e-36, -4.046679867963547e-34, -6.759150293645651e-32, -1.0978129664282036e-29, -1.729547234030482e-27, -2.6311791778297566e-25, -3.837871440435958e-23, -5.314223418815604e-21, -6.898271030850285e-19, -8.270139886158048e-17, -9.001374585761147e-15, -8.722435608179531e-13, -7.362047436128556e-11, -5.302596276954622e-9, -3.297211867381328e-7, -2.0844639689248985e-5, -0.028521319989109915, -0.048738449902780624, -0.04969496566441427, -0.04656973848967755, -0.040693371915399956, -0.03143832892149176, -0.016907830140449595, 0.012079565149346093], v = [2.042204615894487e-54, 3.6952771001682755e-52, 6.8218350617930365e-50, 1.2327825186414496e-47, 2.1873272132311882e-45, 3.805738928460892e-43, 6.481726694196645e-41, 1.0780573178779244e-38, 1.746219103738272e-36, 2.746580465127761e-34, 4.182552482362202e-32, 6.148205581063648e-30, 8.69546407645393e-28, 1.178422520157707e-25, 1.5215948932054688e-23, 1.85645025570418e-21, 2.1147059635090156e-19, 2.2112863234322984e-17, 2.0724160383185546e-15, 1.679674504974956e-13, 1.10560393492268e-11, 5.012034703645976e-10, 1.3735415854349997e-9, -3.296165107801945e-6, -0.009056980372564243, -0.02835308959035654, -0.03938648106080834, -0.046502569918107736, -0.053132069008894695, -0.05971564568812078, -0.06625069286383044, -0.07304730282368906], b = [2.1827319069845502e-6, 1.3900594803039137e-8, 8.329939971954067e-9, 2.8568932809724945e-11, 2.072679614745221e-11, 4.2804952291031206e-14, 3.8713823813374404e-14, 2.7321894746634712e-17, 1.3552527156068805e-17, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -2.8189256484623115e-18, 1.5178830414797062e-17, -3.1799649718999845e-16, 4.407325199240475e-14, -6.383137291302021e-13, 6.389862629904014e-11, -8.672330789692012e-10, 6.421177909288374e-8, -7.181300127481012e-7, 9.51084322467104e-5, 8.460239801488882e-5, 5.444606358862529e-5, 1.6919725888660464e-5, -2.2025474081071932e-5, -6.15283574662013e-5, -0.00010135940495601503, -0.00014154302320537333])
 (u = [-1.4551314108692696e-54, -2.810223537513546e-52, -5.561127236950559e-50, -1.0807296238024793e-47, -2.066116935367805e-45, -3.877056636922193e-43, -7.125317706693103e-41, -1.2799014388376078e-38, -2.243392203325716e-36, -3.8324090432288293e-34, -6.374844520357264e-32, -1.0313027799708871e-29, -1.6190794991073815e-27, -2.4566284262577735e-25, -3.5784313272913477e-23, -4.956232327621621e-21, -6.446310727609562e-19, -7.756317144109252e-17, -8.484259441451785e-15, -8.269097887745582e-13, -7.016109774922905e-11, -5.050684611664837e-9, -3.039864438416499e-7, -1.609539015113887e-5, -0.02697769416883626, -0.048387430333869136, -0.04942085933420358, -0.04633832873118958, -0.040525457598325175, -0.031360657250293474, -0.016958841000926494, 0.011847605331496756], v = [2.0858123504669542e-54, 3.744797843990104e-52, 6.857580961246174e-50, 1.2290491270236953e-47, 2.162947503872346e-45, 3.733983751998306e-43, 6.313542832010329e-41, 1.0432284919024035e-38, 1.6800013960309994e-36, 2.6288352540517425e-34, 3.984735767658905e-32, 5.832511602283937e-30, 8.216151715878599e-28, 1.1093487809784187e-25, 1.4277148666946205e-23, 1.7373767909708416e-21, 1.9758255242029127e-19, 2.06517894680709e-17, 1.937284404233692e-15, 1.5739768725797722e-13, 1.0413313345661755e-11, 4.822052340695038e-10, 4.730088493250454e-9, -2.037957640478882e-6, -0.008504772469504564, -0.02800166524410141, -0.03906764093653383, -0.0461231500270958, -0.05270235218113845, -0.059242533442824805, -0.06573996504493795, -0.07251561819364559], b = [2.1919501772150324e-6, 1.402992967766345e-8, 8.396112008771106e-9, 2.8997270079891835e-11, 2.0966656579315956e-11, 4.3766856458460346e-14, 3.929430565652314e-14, 7.025630077706069e-17, 1.0050554138940626e-16, 4.7271214720367993e-17, 2.949029909160572e-17, 6.396792817664476e-18, 4.5428071027142636e-17, 4.0766001685454967e-17, 6.114900252818245e-17, 6.797947621484113e-17, 6.700369425960417e-17, -2.817841446289826e-16, 4.0816308666258294e-14, -6.051117723213095e-13, 6.027821053715063e-11, -8.397659537694689e-10, 6.192536767022456e-8, -7.152223936519788e-7, 9.164546771205041e-5, 8.45097963386806e-5, 5.4800231163736106e-5, 1.729201952428179e-5, -2.165579837449261e-5, -6.116485319821069e-5, -0.00010100401558771386, -0.00014120140346789522])
 (u = [-1.2328233837993419e-54, -2.4084649620289937e-52, -4.817836658108432e-50, -9.448764102637555e-48, -1.8190306528603512e-45, -3.428563589142489e-43, -6.31122994670172e-41, -1.1320568489001326e-38, -1.9752010832557864e-36, -3.3485691286070383e-34, -5.513071424852204e-32, -8.812903188892062e-30, -1.3669640658669132e-27, -2.0530320525538887e-25, -2.97166784045019e-23, -4.112010490379786e-21, -5.375851710218434e-19, -6.539995096120707e-17, -7.269272658249158e-15, -7.225246589463354e-13, -6.260636006866082e-11, -4.58605494292554e-9, -2.7449266574691143e-7, -1.2714191889055032e-5, -0.023388096613237184, -0.0473337517729441, -0.04849505444935308, -0.04552112156873662, -0.03989323321228923, -0.031011617297575293, -0.017036566505323383, 0.011190707064434438], v = [1.6453129245736095e-54, 2.989706443547995e-52, 5.54889426797518e-50, 1.0083445637702314e-47, 1.7985710343591925e-45, 3.1434859137571485e-43, 5.371428386471108e-41, 8.94838422835078e-39, 1.4487728393105842e-36, 2.272055133624164e-34, 3.4401982109160487e-32, 5.0134722557379625e-30, 7.010538730547041e-28, 9.374600571932707e-26, 1.1935046239269507e-23, 1.4372020328386048e-21, 1.6204393382353865e-19, 1.6844108981417674e-17, 1.5770291255269797e-15, 1.2826160184311258e-13, 8.503173873852743e-12, 3.937981790580542e-10, 4.319345940403177e-9, -1.2735409929602004e-6, -0.007305775471164086, -0.027101971839169403, -0.0381714697760242, -0.04504265936231621, -0.05145431599214494, -0.05783321602877544, -0.0641781328758401, -0.07086463076722016], b = [2.235169464756201e-6, 1.4657494839609778e-8, 8.703411028335162e-9, 3.1137631535671506e-11, 2.2082777446919e-11, 4.8552958528680357e-14, 4.2066176930699584e-14, 1.444157293750692e-16, 1.3097162243624894e-16, 7.26415455565288e-17, 2.4936649967166602e-18, 1.2034644114589099e-17, -8.565197162635485e-18, 1.1492543028346347e-17, 4.2717565595928875e-17, 8.456776945386935e-18, 2.677979366039196e-17, -1.8832591736073212e-16, 3.2757759178825285e-14, -5.038349295063971e-13, 5.116772516756615e-11, -7.537753318113787e-10, 5.583697664695089e-8, -6.847319731699921e-7, 8.228547910649322e-5, 8.398133311487404e-5, 5.508737016052416e-5, 1.7557364264015105e-5, -2.1430361323263156e-5, -6.098024392743323e-5, -0.00010086628058117565, -0.00014114087866625818])
 (u = [-1.5125765567682205e-54, -2.949605229096679e-52, -5.881610748459437e-50, -1.1471337575808776e-47, -2.1906082255576045e-45, -4.084938279191054e-43, -7.419275394203874e-41, -1.3093064489808667e-38, -2.2405145468740896e-36, -3.7126000059901095e-34, -5.953506886196869e-32, -9.240635342743349e-30, -1.3889825751907294e-27, -2.0219229233087563e-25, -2.8450689853706705e-23, -3.849434919601651e-21, -4.959709436534403e-19, -5.997928427173696e-17, -6.681484846912644e-15, -6.701339762612787e-13, -5.888803763232664e-11, -4.3857433056560985e-9, -2.6629862473951787e-7, -1.2366055740139495e-5, -0.020747901072942426, -0.04596484815531468, -0.04715739880558741, -0.0442911972950299, -0.03888940290519258, -0.0303878802166701, -0.017023980723662246, 0.010419217794954609], v = [1.8912961562334545e-54, 3.445790360877863e-52, 6.413626558278268e-50, 1.1674076913254893e-47, 2.081920998488024e-45, 3.629884721105613e-43, 6.171799166267559e-41, 1.0202860232955769e-38, 1.6345047405727476e-36, 2.5287150003751887e-34, 3.7649662672499217e-32, 5.376793860419249e-30, 7.341840141147688e-28, 9.554591777970469e-26, 1.1806479942692276e-23, 1.3779974105192772e-21, 1.5065274502750597e-19, 1.5218631346288058e-17, 1.389225656470933e-15, 1.1042612847799048e-13, 7.131965727886936e-12, 3.123248691723129e-10, 1.0220713229414688e-9, -1.3106024831109863e-6, -0.00654749201495973, -0.026286624090100835, -0.03722698368094291, -0.043839369352300796, -0.0499857472115716, -0.05608634068340704, -0.06215226287588159, -0.06866798555941493], b = [2.3220170757174904e-6, 1.5985724079094904e-8, 9.325447229095035e-9, 3.588573123554373e-11, 2.43785651388162e-11, 5.972522823527449e-14, 4.785863545742064e-14, 1.489693784995083e-16, 1.1373280789372942e-16, 1.336821278674627e-16, 1.3780209612290761e-16, 5.377642775528102e-17, 7.708677446371937e-17, 6.396792817664476e-17, 8.749511531958021e-17, 5.930585883495709e-17, 9.833713704443525e-17, -1.2587587222556706e-16, 2.7269202521051916e-14, -4.438055825617404e-13, 4.4480527930639696e-11, -7.03741992360684e-10, 5.071293825159454e-8, -6.649813001609133e-7, 7.35631999514951e-5, 8.245950206067136e-5, 5.395930855043313e-5, 1.6328285332446523e-5, -2.2745117288153785e-5, -6.237629437918199e-5, -0.00010235700890615894, -0.00014279303776932448])
 (u = [-3.562684391314721e-54, -6.881707146476817e-52, -1.3541128340390019e-49, -2.5922684179086746e-47, -4.834122525860176e-45, -8.763505717978834e-43, -1.5413569753485518e-40, -2.6250499348343354e-38, -4.321250717720359e-36, -6.866399159737597e-34, -1.0524194634088562e-31, -1.5560790368843396e-29, -2.2213412173305493e-27, -3.064628618880791e-25, -4.0862614926434884e-23, -5.251569870684012e-21, -6.458990736831318e-19, -7.5070501506978e-17, -8.096872215928095e-15, -7.917956422537303e-13, -6.824281481940386e-11, -5.007937930551067e-9, -3.0047532294028553e-7, -1.394947599469487e-5, -0.018778004994887652, -0.042897953667871906, -0.043935738717256805, -0.04125542355368903, -0.03633893822413588, -0.02870184334881752, -0.01677851883467971, 0.008809721513193286], v = [3.781836192571377e-54, 6.958885936732664e-52, 1.3036223208046452e-49, 2.37428731498229e-47, 4.209981386745383e-45, 7.25290906879501e-43, 1.211437910552237e-40, 1.9569667845381015e-38, 3.0491367357289246e-36, 4.568812175530446e-34, 6.563272201052249e-32, 9.010612606198593e-30, 1.1784583445292746e-27, 1.4633751366535173e-25, 1.7188370398659905e-23, 1.9000733531072638e-21, 1.9616039274943397e-19, 1.8671329492820013e-17, 1.6027857490979645e-15, 1.1927534821961065e-13, 7.0945424525678126e-12, 2.621342174948195e-10, -4.6408921627016565e-9, -1.7479413483585436e-6, -0.006284925956387863, -0.025302186140422477, -0.03567393900569944, -0.04161149108479537, -0.04703243770344607, -0.052363214157098756, -0.05765051943433698, -0.06364358536810366], b = [2.5694092078579287e-6, 2.031562991423007e-8, 1.1104688333804513e-8, 5.3219223299094676e-11, 3.114064149774276e-11, 1.0539724475122636e-13, 6.556387377454342e-14, 2.0903417885520525e-16, 1.6967763999398144e-16, 7.665309359472516e-17, 5.063224145507306e-17, -3.0357660829594124e-17, -1.2468324983583301e-17, 1.1275702593849246e-17, 5.605325231750058e-17, 7.296680620827445e-17, 4.3042826247674526e-17, -1.330316065639714e-16, 2.4116559443898566e-14, -4.700766685628022e-13, 3.975782322720306e-11, -7.386595417355896e-10, 4.575324025391956e-8, -6.879289659029292e-7, 6.285017449994003e-5, 7.674168058892237e-5, 4.830396476272929e-5, 1.0373113817348738e-5, -2.8939920618101055e-5, -6.880638127536433e-5, -0.00010907491803599668, -0.000150031022497084])
 (u = [-1.060124668627774e-53, -1.9871196084790933e-51, -3.7972303538375474e-49, -7.054860241222256e-47, -1.2755738407585852e-44, -2.2405323723413996e-42, -3.8171357973883e-40, -6.298313235543349e-38, -1.0052924530585802e-35, -1.5510150617803818e-33, -2.312638226365751e-31, -3.333452337439867e-29, -4.64685849550876e-27, -6.263300083930082e-25, -8.145716764093173e-23, -1.016910938107121e-20, -1.2073540711746125e-18, -1.3447602666333123e-16, -1.380237882514118e-14, -1.2772564194209021e-12, -1.038191201604177e-10, -7.182261488351072e-9, -4.074661395515537e-7, -1.81042897688379e-5, -0.020880743864788304, -0.040464634165896074, -0.04116104161994273, -0.03859516910702923, -0.034072447652773286, -0.027157025369143408, -0.016440988392790012, 0.007531157970582211], v = [1.0769443877459218e-53, 1.931606142822548e-51, 3.528701893079646e-49, 6.260653809130494e-47, 1.0798019466333686e-44, 1.8069881627512766e-42, 2.928193020598108e-40, 4.584784055915497e-38, 6.91940804931743e-36, 1.0040199197574665e-33, 1.3969838511251625e-31, 1.8588319681674068e-29, 2.3586355680591903e-27, 2.8449714218923575e-25, 3.2489814944660546e-23, 3.492616882843466e-21, 3.502225329625398e-19, 3.228156029163034e-17, 2.670022647909607e-15, 1.9003533721478604e-13, 1.0675443830566077e-11, 3.5626434599335446e-10, -9.195993930864416e-9, -2.522476147482308e-6, -0.007563474435530196, -0.025389158842415875, -0.035040094118190895, -0.040318699531118066, -0.045022286169330954, -0.04960065705758931, -0.054135685883818285, -0.05955484038236643], b = [2.8110310978201984e-6, 2.5485435069380266e-8, 1.2776622981174393e-8, 7.691252626104705e-11, 3.7547979048896396e-11, 1.7579579285331715e-13, 8.284952585091432e-14, 3.5518463170625125e-16, 1.2620113287731272e-16, -5.2150124496552763e-17, -6.591949208711867e-17, -6.093216209368535e-17, -7.806255641895632e-17, -7.806255641895632e-17, -2.45029690981724e-17, -1.214306433183765e-17, 5.659535340374333e-17, -3.423910460709223e-16, 3.069517296588886e-14, -7.072593379009451e-13, 4.623479935494598e-11, -9.50832448312535e-10, 4.865736140423956e-8, -7.892705384401292e-7, 6.278623985304924e-5, 6.961326102022819e-5, 4.077259615186044e-5, 2.5650667533769417e-6, -3.697473952226908e-5, -7.707800317714377e-5, -0.00011765474275555119, -0.00015921195624561004])
 (u = [-1.6686205176245078e-53, -3.11662093520416e-51, -5.932982461183175e-49, -1.0974264853476645e-46, -1.9740331990481058e-44, -3.447381814858062e-42, -5.836976349472106e-40, -9.570890881627321e-38, -1.5185859101432243e-35, -2.330752415056794e-33, -3.4605543152671463e-31, -4.971386252758892e-29, -6.908841083872254e-27, -9.274458447667424e-25, -1.1980919682817813e-22, -1.479107801678475e-20, -1.726832882670565e-18, -1.8799284965929665e-16, -1.875705705968186e-14, -1.6807265880542658e-12, -1.321621119053734e-10, -8.904852183515445e-9, -5.08883617924926e-7, -2.6329119506959204e-5, -0.023975426378661804, -0.04052973223113278, -0.041002051071246465, -0.03838482245282721, -0.03385660456482012, -0.026978526876909364, -0.016361371530538973, 0.00757149886190957], v = [1.66061096938527e-53, 2.971712139057197e-51, 5.416083260640652e-49, 9.582334017671423e-47, 1.6469246757840004e-44, 2.744276105428692e-42, 4.424569062999225e-40, 6.887613607544343e-38, 1.0328581175380323e-35, 1.4886168672714831e-33, 2.0572351684996954e-31, 2.719660742624375e-29, 3.430589411600445e-27, 4.1163740942026597e-25, 4.678686471652104e-23, 5.0053543955660164e-21, 4.990324464806408e-19, 4.564912759856853e-17, 3.737703773698972e-15, 2.6267755998354765e-13, 1.4526949604532075e-11, 4.664581015652964e-10, -1.6914630247068953e-8, -4.697805928074212e-6, -0.008928593850999955, -0.0260478228885608, -0.03552454571265517, -0.040734017888770055, -0.0453324377002087, -0.049786502953583535, -0.054188355693961604, -0.05950431506049898], b = [2.8668395537147662e-6, 2.6829435715398572e-8, 1.3145253052520775e-8, 8.346344991007337e-11, 3.893815890380903e-11, 1.9624688159247672e-13, 8.662450097507435e-14, 4.505944228849756e-16, 2.8341044788771086e-16, 1.4842727741326556e-16, 1.463672932855431e-16, 1.5395670849294163e-16, 1.5612511283791264e-16, 1.5612511283791264e-16, 9.660241356845845e-17, 7.546047120499111e-17, 7.470152968425126e-17, -4.578585774406285e-16, 4.08906849352908e-14, -9.27142368954692e-13, 5.6529666310391885e-11, -1.1258156470874553e-9, 5.471232497931229e-8, -8.503118770895228e-7, 6.873377922205165e-5, 6.754839768282765e-5, 3.821951016707897e-5, -2.0434861799554192e-8, -3.9590007021159145e-5, -7.973257399662677e-5, -0.00012036913410912665, -0.00016205682389039262])
 (u = [-2.1249202641728983e-53, -3.968641265108508e-51, -7.552267384776523e-49, -1.3961375163261153e-46, -2.5097345965314137e-44, -4.380909623575297e-42, -7.418064837250505e-40, -1.217524837341126e-37, -1.936179581017804e-35, -2.9829616638007873e-33, -4.452307392112877e-31, -6.436077100019644e-29, -8.99899729238418e-27, -1.213316018334397e-24, -1.5687579073822714e-22, -1.9287901133200292e-20, -2.2297010450099347e-18, -2.389891797694313e-16, -2.3368757406319024e-14, -2.0479206742362607e-12, -1.5822141431449216e-10, -1.0742153028185253e-8, -6.770535077527665e-7, -4.8198938936529514e-5, -0.026817091308095297, -0.0413308733619403, -0.041623514779421805, -0.03890202069863699, -0.034246374259625866, -0.027204720437811875, -0.016372919767871856, 0.007985175189129926], v = [2.1020610581234513e-53, 3.763965196268764e-51, 6.86213040972326e-49, 1.2140403855605192e-46, 2.085826767145008e-44, 3.4733858432414374e-42, 5.595608604118077e-40, 8.703860513197173e-38, 1.304620413423763e-35, 1.880628235740461e-33, 2.6019849863959855e-31, 3.4481084798905255e-29, 4.365857830885381e-27, 5.26449351230194e-25, 6.016800125466243e-23, 6.47043226944603e-21, 6.475308535701666e-19, 5.931202115063972e-17, 4.847853307469168e-15, 3.3870690687938424e-13, 1.8394130749726648e-11, 5.091222326121079e-10, -4.475533190483764e-8, -1.1200360709438414e-5, -0.010079638873645227, -0.026755963827003102, -0.036238517582821005, -0.041572935063778474, -0.04626266336151833, -0.05078845083400123, -0.05524556022264364, -0.06061613115168572], b = [2.868719821914152e-6, 2.6861373021713764e-8, 1.3164660763202371e-8, 8.36084530601855e-11, 3.9026911910489126e-11, 1.9670268018578962e-13, 8.686974750649057e-14, 3.642919299551295e-16, 1.4159680372660688e-16, 6.505213034913027e-19, 0.0, 0.0, 0.0, 0.0, 0.0, 9.64939933512099e-18, 3.946495907847236e-17, -6.862999751833243e-16, 5.1964400664405996e-14, -1.115037641212513e-12, 6.725966013526943e-11, -1.2562196184543253e-9, 6.076724895709906e-8, -8.578053492486019e-7, 7.50313022474396e-5, 6.746388501321861e-5, 3.752731472347678e-5, -6.852047063831085e-7, -4.022769573928809e-5, -8.034815258358941e-5, -0.00012096310478670732, -0.00016261819684241247])

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

ax = Axis(top[1, 1]; xlabel="Iteration", ylabel="Parameter / nature run value", title="Parameters")
hlines!(ax, 1; color=:gray, linestyle=:dash)
for (i, name) in enumerate(keys(scales))
    ratio = [h.θ[i] / θ★[i] for h in history]
    points = @lift Point2f.(iterations[1:$n], ratio[1:$n])
    scatterlines!(ax, points; label=string(name))
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

---

Julia version and environment information

This example was executed with the following version of Julia:

using InteractiveUtils: versioninfo
versioninfo()

Julia Version 1.13.2
Commit 45b6d15e34b (2026-10-14 20:33 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.2+0.x64.linux.gnu/local/share/julia:/var/lib/buildkite-agent/.julia/juliaup/julia-1.13.2+0.x64.linux.gnu/share/julia
  JULIA_VERSION_ENZYME = 1.13.1
  JULIA_PKG_SERVER_REGISTRY_PREFERENCE = eager
  LD_LIBRARY_PATH = 
  JULIA_MAX_NUM_PRECOMPILE_FILES = 24
  JULIA_VERSION = 1.13.2
  JULIA_CUDA_USE_COMPAT = false
  JULIA_PROJECT = /var/lib/buildkite-agent/Oceananigans.jl-35983/docs/
  JULIA_DEBUG = Literate

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

import Pkg
Pkg.status()

Status `~/Oceananigans.jl-35983/docs/Project.toml`
  [79e6a3ab] Adapt v4.8.0
⌃ [052768ef] CUDA v6.4.1
  [13f3f980] CairoMakie v0.15.17
  [8e50ac2c] ConservativeRegridding v0.2.20
  [e30172f5] Documenter v1.20.0
  [daee34ce] DocumenterCitations v1.5.0
  [4710194d] DocumenterVitepress v0.3.9
  [7da242da] Enzyme v0.13.218
  [033835bb] JLD2 v0.6.7
⌅ [63c18a36] KernelAbstractions v0.9.45
  [98b081ad] Literate v2.21.0
  [da04e1cc] MPI v0.20.29
  [85f8d34a] NCDatasets v0.14.15
  [9e8cae18] Oceananigans v0.117.1 `..`
  [429524aa] Optim v2.3.2
  [f27b6e38] Polynomials v4.1.6
  [3c362404] Reactant v0.2.307
  [6038ab10] Rotations v1.7.1
  [d496a93d] SeawaterPolynomials v0.3.11
  [09ab397b] StructArrays v0.7.4
  [bdfc003b] TimesDates v0.3.3
  [0a941bbe] Zarr v0.11.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.