Skip to content

Question: Should we use let blocks in build_adnlp_model? #551

@ocots

Description

@ocots

Question: Should we use let blocks in build_adnlp_model?

Context

In [CTModels.jl/src/ocp/model.jl#L165-L188](https://github.com/control-toolbox/CTModels.jl/blob/ac5a9f37bd84aa7b82f5887c4ba6b51bf758fe2f/src/ocp/model.jl#L165-L188), the code uses let blocks when creating closures:

function make_path_cons_nl(constraints_number::Int, ...)
    let
        cn = constraints_number
        cd = constraints_dimensions
        cf = constraints_functions
        
        function path_cons_nl!(val, t, x, u, v)
            # Uses cn, cd, cf
        end
        return path_cons_nl!
    end
end

Why let blocks?

The let block ensures that captured variables are:

  1. Type-stable: Variables are captured with their concrete types, not as Box wrappers
  2. Performance-optimized: The compiler can better optimize closures with let-captured variables
  3. Predictable: Creates local immutable copies at capture time

Question

In [CTDirect.jl/src/collocation.jl#L204-L359](https://github.com/control-toolbox/CTDirect.jl/blob/da9c4e66303dc649556bea66a628b57f5d953c49/src/collocation.jl#L204-L359), build_adnlp_model creates closures that capture docp:

function build_adnlp_model(...)
    f = x -> CTDirect.DOCP_objective(x, docp)
    c! = (c, x) -> CTDirect.DOCP_constraints!(c, x, docp)
    # ...
end

Should we apply the same pattern here for consistency and performance?

function build_adnlp_model(...)
    let docp_local = docp
        f = x -> CTDirect.DOCP_objective(x, docp_local)
        c! = (c, x) -> CTDirect.DOCP_constraints!(c, x, docp_local)
        # ...
    end
end

This would ensure optimal performance in the NLP solver loop and maintain consistency with the rest of the codebase.

function (discretizer::Collocation)(ocp::AbstractOptimalControlProblem)
# common parts for builders
docp = get_docp(discretizer, ocp)
exa_getter = nothing # will be set in build_exa_model
# ==========================================================================================
# The needed builders for the construction of the final DiscretizedOptimalControlProblem
# ==========================================================================================
# +++ recheck kwargs passing / default with Olivier
function build_adnlp_model(
initial_guess::CTModels.AbstractOptimalControlInitialGuess;
adnlp_backend=__adnlp_backend(),
show_time=false,
kwargs...
)::ADNLPModels.ADNLPModel
# functions for objective and constraints
f = x -> CTDirect.DOCP_objective(x, docp)
c! = (c, x) -> CTDirect.DOCP_constraints!(c, x, docp)
# build initial guess
init = get_docp_initial_guess(:adnlp, docp, initial_guess)
# unused backends (option excluded_backend = [:jprod_backend, :jtprod_backend, :hprod_backend, :ghjvprod_backend] does not seem to work)
unused_backends = (
hprod_backend=ADNLPModels.EmptyADbackend,
jtprod_backend=ADNLPModels.EmptyADbackend,
jprod_backend=ADNLPModels.EmptyADbackend,
ghjvprod_backend=ADNLPModels.EmptyADbackend,
)
# set adnlp backends
if adnlp_backend == :manual
# build sparsity patterns for Jacobian and Hessian
J_backend = ADNLPModels.SparseADJacobian(
docp.dim_NLP_variables, f,
docp.dim_NLP_constraints, c!,
CTDirect.DOCP_Jacobian_pattern(docp),
)
H_backend = ADNLPModels.SparseReverseADHessian(
docp.dim_NLP_variables, f,
docp.dim_NLP_constraints, c!,
CTDirect.DOCP_Hessian_pattern(docp),
)
backend_options = (
gradient_backend=ADNLPModels.ReverseDiffADGradient,
jacobian_backend=J_backend,
hessian_backend=H_backend,
)
else
# use backend preset
backend_options = (backend=adnlp_backend,)
end
# build NLP
nlp = ADNLPModel!(
f,
init,
docp.bounds.var_l,
docp.bounds.var_u,
c!,
docp.bounds.con_l,
docp.bounds.con_u;
minimize=(!docp.flags.max),
backend_options...,
unused_backends...,
show_time=show_time,
)
return nlp
end
# Solution builder for ADNLPModels
function build_adnlp_solution(nlp_solution::SolverCore.AbstractExecutionStats)
# retrieve data from NLP solver
minimize = !docp.flags.max
objective, iterations, constraints_violation, message, status, successful = CTModels.extract_solver_infos(nlp_solution, minimize)
# retrieve time grid
T = get_time_grid(nlp_solution.solution, docp)
# build OCP solution from NLP solution
sol = CTDirect.build_OCP_solution(docp, nlp_solution, T,
objective, iterations, constraints_violation, message, status, successful)
return sol
end
# NLP builder for ExaModels
# +++ recheck kwargs passing / default with Olivier
function build_exa_model(
::Type{BaseType},
initial_guess::CTModels.AbstractOptimalControlInitialGuess;
exa_backend=CTDirect.__exa_backend(),
kwargs...
)::ExaModels.ExaModel where {BaseType<:AbstractFloat}
# recover discretization scheme and options
# since exa part does not reuse the docp struct
scheme = get_scheme(discretizer)
grid_size, time_grid = grid_options(discretizer)
# build initial guess
init = get_docp_initial_guess(:exa, docp, initial_guess)
# build Exa model and getters
# +++ later try to call Exa constructor here if possible, reusing existing functions...
build_exa = CTModels.get_build_examodel(ocp)
nlp, exa_getter = build_exa(;
grid_size=grid_size,
backend=exa_backend,
scheme=scheme,
init=init,
)
return nlp
end
# Solution builder for ExaModels
function build_exa_solution(nlp_solution::SolverCore.AbstractExecutionStats)
# NB exa_getter is set during build_exa_model call !
if isnothing(exa_getter)
error("build_exa_solution: exa_getter is nothing")
end
# retrieve data from NLP solver
minimize = !docp.flags.max
objective, iterations, constraints_violation, message, status, successful = CTModels.extract_solver_infos(nlp_solution, minimize)
# retrieve time grid
T = get_time_grid_exa(nlp_solution, docp, exa_getter)
# build OCP solution from NLP solution
sol = CTDirect.build_OCP_solution(docp, nlp_solution, T,
objective, iterations, constraints_violation, message, status, successful;
exa_getter=exa_getter)
return sol
end
#NB. it would be better to return builders as model/solution pairs since they are linked
return CTModels.DiscretizedOptimalControlProblem(
ocp,
CTModels.ADNLPModelBuilder(build_adnlp_model),
CTModels.ExaModelBuilder(build_exa_model),
CTModels.ADNLPSolutionBuilder(build_adnlp_solution),
CTModels.ExaSolutionBuilder(build_exa_solution),
)
end

Metadata

Metadata

Labels

internal devModification to the code is requested and should not affect user experiment

Type

No type
No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions