How to make the Optim.jl zero allocated?

Xu_Shan · 2025-02-27T09:07:54.555Z

Hi guys,

Right now I am using Optim.jl to optimise my problem, however, @btime shows the whole script has around hundreds allocations…which is quite a lot. So is there any trick to reduce the allocations to zero? and also increases the speed of the script? because it would be in a for-loop which would be executed for many many times…thanks!

Here is my script:

using Optim
using BenchmarkTools

# Define the conductance function
function conductance(conductance_max, shape_para, 
                    water_potential, potential_50_loss_conductance, o_one)
    return conductance_max / (o_one + exp(shape_para * (water_potential - potential_50_loss_conductance)))
end

# Define the loss function
function pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, p_up_to_stem,
                            kRootMax, pShapeRoot, pWaterRootHalfLoss,
                            areaLeaf, ΔhRoot, capRoot, o_one)

    k_root_update = conductance(kRootMax, pShapeRoot, 
                    pWaterRoot_new, pWaterRootHalfLoss, o_one)
    
    sap_root_update = (p_soil - pWaterRoot_new - ΔhRoot) * k_root_update * areaLeaf
    ΔrootW = capRoot * (pWaterRoot_new - pWaterRoot) * areaLeaf
    
    return (sap_root_update - ΔrootW - p_up_to_stem)^2  # Squared error for minimization
end

# Wrapper function to optimize
function optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem,
                            kRootMax, pShapeRoot, pWaterRootHalfLoss,
                            areaLeaf, ΔhRoot, capRoot, o_one)
    
    # Define the function to minimize
    objective(pWaterRoot_new) = pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, 
                                                    p_up_to_stem, kRootMax, pShapeRoot, 
                                                    pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
                                                    capRoot, o_one)
    
    # Set initial guess and bounds
    pWaterRoot_new_init = [pWaterRoot]  # Initial guess
    lower_bound = -5.0  # Define reasonable lower bound
    upper_bound = 0.0   # Define reasonable upper bound

    opt_options = Optim.Options(g_tol=1e-8)

    # Perform the optimization
    result = optimize(pWaterRoot_new -> pWaterRoot_optimize(pWaterRoot_new[1], pWaterRoot, p_soil, 
                                                            p_up_to_stem, kRootMax, pShapeRoot, 
                                                            pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
                                                            capRoot, o_one), 
                    lower_bound, upper_bound, pWaterRoot_new_init, Fminbox(LBFGS()), opt_options) # , Fminbox(LBFGS())
    
    return Optim.minimizer(result), Optim.minimum(result)
end

# Example usage with sample values
pWaterRoot=-1.0
    p_soil=-0.5 
    p_up_to_stem=0.1
    kRootMax=10.0
    pShapeRoot=5.0
    pWaterRootHalfLoss=-1.2
    areaLeaf=1.0
    ΔhRoot=0.2
    capRoot=0.05
    o_one=1.0
pWaterRoot_new_opt, min_value = optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem,
                                                    kRootMax, pShapeRoot, pWaterRootHalfLoss,
                                                    areaLeaf, ΔhRoot, capRoot, o_one
                                                    )

println("Optimal pWaterRoot_new: ", pWaterRoot_new_opt)
println("Minimum loss value: ", min_value)

Here is the allocation results:

  39.363 μs (625 allocations: 28.69 KiB)
([-1.657185450744575e-9], 0.027992088033827963)

JADekker · 2025-03-31T18:27:28.907Z

Xu_Shan:

using Optim
using BenchmarkTools

# Define the conductance function
function conductance(conductance_max, shape_para, 
                    water_potential, potential_50_loss_conductance, o_one)
    return conductance_max / (o_one + exp(shape_para * (water_potential - potential_50_loss_conductance)))
end

# Define the loss function
function pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, p_up_to_stem,
                            kRootMax, pShapeRoot, pWaterRootHalfLoss,
                            areaLeaf, ΔhRoot, capRoot, o_one)

    k_root_update = conductance(kRootMax, pShapeRoot, 
                    pWaterRoot_new, pWaterRootHalfLoss, o_one)
    
    sap_root_update = (p_soil - pWaterRoot_new - ΔhRoot) * k_root_update * areaLeaf
    ΔrootW = capRoot * (pWaterRoot_new - pWaterRoot) * areaLeaf
    
    return (sap_root_update - ΔrootW - p_up_to_stem)^2  # Squared error for minimization
end

# Wrapper function to optimize
function optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem,
                            kRootMax, pShapeRoot, pWaterRootHalfLoss,
                            areaLeaf, ΔhRoot, capRoot, o_one)
    
    # Define the function to minimize
    objective(pWaterRoot_new) = pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, 
                                                    p_up_to_stem, kRootMax, pShapeRoot, 
                                                    pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
                                                    capRoot, o_one)
    
    # Set initial guess and bounds
    pWaterRoot_new_init = [pWaterRoot]  # Initial guess
    lower_bound = -5.0  # Define reasonable lower bound
    upper_bound = 0.0   # Define reasonable upper bound

    opt_options = Optim.Options(g_tol=1e-8)

    # Perform the optimization
    result = optimize(pWaterRoot_new -> pWaterRoot_optimize(pWaterRoot_new[1], pWaterRoot, p_soil, 
                                                            p_up_to_stem, kRootMax, pShapeRoot, 
                                                            pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
                                                            capRoot, o_one), 
                    lower_bound, upper_bound, pWaterRoot_new_init, Fminbox(LBFGS()), opt_options) # , Fminbox(LBFGS())
    
    return Optim.minimizer(result), Optim.minimum(result)
end

# Example usage with sample values
pWaterRoot=-1.0
    p_soil=-0.5 
    p_up_to_stem=0.1
    kRootMax=10.0
    pShapeRoot=5.0
    pWaterRootHalfLoss=-1.2
    areaLeaf=1.0
    ΔhRoot=0.2
    capRoot=0.05
    o_one=1.0
pWaterRoot_new_opt, min_value = optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem,
                                                    kRootMax, pShapeRoot, pWaterRootHalfLoss,
                                                    areaLeaf, ΔhRoot, capRoot, o_one
                                                    )

println("Optimal pWaterRoot_new: ", pWaterRoot_new_opt)
println("Minimum loss value: ", min_value)

I’m not sure how to minimize allocations in LBFGS here. However, in your example you may wish to use the Brent method for a univariate function on an interval: this gives a good speed-up for me:

using Optim
using BenchmarkTools

# Define the conductance function
function conductance(conductance_max, shape_para, 
                    water_potential, potential_50_loss_conductance, o_one)
    return conductance_max / (o_one + exp(shape_para * (water_potential - potential_50_loss_conductance)))
end

# Define the loss function
function pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, p_up_to_stem,
                            kRootMax, pShapeRoot, pWaterRootHalfLoss,
                            areaLeaf, ΔhRoot, capRoot, o_one)

    k_root_update = conductance(kRootMax, pShapeRoot, 
                    pWaterRoot_new, pWaterRootHalfLoss, o_one)
    
    sap_root_update = (p_soil - pWaterRoot_new - ΔhRoot) * k_root_update * areaLeaf
    ΔrootW = capRoot * (pWaterRoot_new - pWaterRoot) * areaLeaf
    
    return (sap_root_update - ΔrootW - p_up_to_stem)^2  # Squared error for minimization
end

# Wrapper function to optimize
function optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem,
                            kRootMax, pShapeRoot, pWaterRootHalfLoss,
                            areaLeaf, ΔhRoot, capRoot, o_one, method)
    
    # Define the function to minimize
    objective(pWaterRoot_new) = pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, 
                                                    p_up_to_stem, kRootMax, pShapeRoot, 
                                                    pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
                                                    capRoot, o_one)
    
    # Set initial guess and bounds
    pWaterRoot_new_init = [pWaterRoot]  # Initial guess
    lower_bound = -5.0  # Define reasonable lower bound
    upper_bound = 0.0   # Define reasonable upper bound

    opt_options = Optim.Options(g_tol=1e-8)

    # Perform the optimization
    if method == 1
        result = optimize(pWaterRoot_new -> pWaterRoot_optimize(pWaterRoot_new[1], pWaterRoot, p_soil, 
                                                                p_up_to_stem, kRootMax, pShapeRoot, 
                                                                pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
                                                                capRoot, o_one), 
                        lower_bound, upper_bound, pWaterRoot_new_init, Fminbox(LBFGS()), opt_options) # , Fminbox(LBFGS())
    elseif method == 2
        result = optimize(pWaterRoot_new -> pWaterRoot_optimize(pWaterRoot_new, pWaterRoot, p_soil, 
        p_up_to_stem, kRootMax, pShapeRoot, 
        pWaterRootHalfLoss, areaLeaf, ΔhRoot, 
        capRoot, o_one), 
        lower_bound, upper_bound)#, opt_options)
    end
    return Optim.minimizer(result), Optim.minimum(result)
end


function test()
    # Example usage with sample values
    pWaterRoot=-1.0
        p_soil=-0.5 
        p_up_to_stem=0.1
        kRootMax=10.0
        pShapeRoot=5.0
        pWaterRootHalfLoss=-1.2
        areaLeaf=1.0
        ΔhRoot=0.2
        capRoot=0.05
        o_one=1.0
    pWaterRoot_new_opt, min_value = optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem,
                                                        kRootMax, pShapeRoot, pWaterRootHalfLoss,
                                                        areaLeaf, ΔhRoot, capRoot, o_one, 1)

    println("Optimal pWaterRoot_new: ", pWaterRoot_new_opt)
    println("Minimum loss value: ", min_value)

    pWaterRoot_new_opt, min_value = optimize_pWaterRoot(pWaterRoot, p_soil, p_up_to_stem, kRootMax, pShapeRoot, pWaterRootHalfLoss, areaLeaf, ΔhRoot, capRoot, o_one, 2)

    println("Optimal pWaterRoot_new: ", pWaterRoot_new_opt)
    println("Minimum loss value: ", min_value)

    # Benchmarking
    display(@benchmark optimize_pWaterRoot($pWaterRoot, $p_soil, $p_up_to_stem, $kRootMax, $pShapeRoot, $pWaterRootHalfLoss, $areaLeaf, $ΔhRoot, $capRoot, $o_one, 1))
    display(@benchmark optimize_pWaterRoot($pWaterRoot, $p_soil, $p_up_to_stem, $kRootMax, $pShapeRoot, $pWaterRootHalfLoss, $areaLeaf, $ΔhRoot, $capRoot, $o_one, 2))
end

test()

gives

# Using Fminbox(LBFGS)
Optimal pWaterRoot_new: [-1.657185450744575e-9]
Minimum loss value: 0.027992088033827963
# Using Brent
Optimal pWaterRoot_new: -0.7946930792350235
Minimum loss value: 2.6441897161267073e-18

# Benchmark LBFGS
BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  28.584 μs …  34.151 ms  ┊ GC (min … max): 0.00% … 99.79%
 Time  (median):     31.375 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   35.411 μs ± 341.202 μs  ┊ GC (mean ± σ):  9.62% ±  1.00%

       ▃▅██▅▃▆▂                                                 
  ▂▄█████████████▇▆▅▄▃▃▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▃
  28.6 μs         Histogram: frequency by time           46 μs <

 Memory estimate: 28.78 KiB, allocs estimate: 625.

# Benchmark Brent
BenchmarkTools.Trial: 10000 samples with 159 evaluations per sample.
 Range (min … max):  664.572 ns …  1.449 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     672.956 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   686.921 ns ± 44.431 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▅██▆▄▂▃▂▁▁ ▂▂▁▁                                              ▂
  ████████████████▇▇▆▇▇█▇▇▇▇▇▇▇▇▆▇▆▆▇▇▆▇▇▇▅▆▆▅▆▆▅▅▆▆▆▅▅▅▅▅▆▄▆▅ █
  665 ns        Histogram: log(frequency) by time       889 ns <

 Memory estimate: 176 bytes, allocs estimate: 2.

which is roughly a factor 50 speed-up and factor 300 allocation reduction.