Xu Shan
Xu Shan

Reputation: 315

How make the Optim.jl zero allocated?

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)

Upvotes: 0

Views: 27

Answers (0)

Related Questions