Skip to content

Commit 180b56a

Browse files
author
agent@localhost
committed
Trajectory datavec now cached during NLP obj/jac/hess assembly callbacks, leaving only constant allocation overhead
1 parent 9b27482 commit 180b56a

1 file changed

Lines changed: 33 additions & 13 deletions

File tree

src/solvers/evaluator.jl

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,12 @@ mutable struct Evaluator <: MOI.AbstractNLPEvaluator
9090
_jacobian_ncols::Int
9191
_hessian_ncols::Int
9292

93+
# Reusable trajectory wrapper — datavec/global_data are rebound per callback
94+
# to avoid reconstructing a NamedTrajectory on every MOI evaluation.
95+
# Shares all structural metadata (components, dims, bounds, etc.) with
96+
# `trajectory`; only the data pointers differ.
97+
_cached_traj::NamedTrajectory
98+
9399
function Evaluator(prob::DirectTrajOptProblem; eval_hessian = true, verbose = false)
94100
t_start = time()
95101

@@ -246,6 +252,17 @@ mutable struct Evaluator <: MOI.AbstractNLPEvaluator
246252
println(" evaluator ready (total: $(round(time() - t_start, digits=3))s)")
247253
end
248254

255+
# Build a one-time cached trajectory wrapper. This is the ONLY call to
256+
# the NamedTrajectory copy constructor during the entire solve. The
257+
# datavec and global_data are owned copies (real Vector{Float64}),
258+
# ensuring the backing-store contract is preserved. Subsequent
259+
# _update_trajectory_cache! calls copyto! into these buffers.
260+
_cached_traj = NamedTrajectory(
261+
prob.trajectory;
262+
datavec = copy(prob.trajectory.datavec),
263+
global_data = copy(prob.trajectory.global_data),
264+
)
265+
249266
return new(
250267
prob.trajectory,
251268
prob.objective,
@@ -266,6 +283,7 @@ mutable struct Evaluator <: MOI.AbstractNLPEvaluator
266283
hessian_linear_map,
267284
jacobian_ncols,
268285
hessian_ncols,
286+
_cached_traj,
269287
)
270288
end
271289
end
@@ -444,21 +462,23 @@ end
444462
"""
445463
_update_trajectory_cache!(evaluator, Z⃗)
446464
447-
Update the cached trajectory in-place with new data from Z⃗.
448-
Avoids repeated allocation of NamedTrajectory wrappers.
465+
Copy the solver's current iterate `Z⃗` into the pre-allocated cached
466+
trajectory's backing vectors. The `_cached_traj.datavec` and
467+
`_cached_traj.global_data` are owned `Vector{Float64}` buffers allocated
468+
once at Evaluator construction; this call overwrites their contents via
469+
`copyto!` without allocating or rebinding.
470+
471+
The returned trajectory shares all structural metadata (components, dims,
472+
bounds, names, etc.) with `evaluator.trajectory` (= `prob.trajectory`).
449473
"""
450-
@inline @views function _update_trajectory_cache!(evaluator::Evaluator, Z⃗::AbstractVector)
474+
@inline function _update_trajectory_cache!(evaluator::Evaluator, Z⃗::AbstractVector)
451475
n_traj = evaluator.trajectory.dim * evaluator.trajectory.N
452-
453-
# Create trajectory wrapper with views (minimal allocation)
454-
# This is equivalent to the old approach but reuses structure
455-
traj = NamedTrajectory(
456-
evaluator.trajectory;
457-
datavec = Z⃗[1:n_traj],
458-
global_data = Z⃗[(n_traj+1):end],
459-
)
460-
461-
return traj
476+
copyto!(evaluator._cached_traj.datavec, 1, Z⃗, 1, n_traj)
477+
n_global = evaluator.trajectory.global_dim
478+
if n_global > 0
479+
copyto!(evaluator._cached_traj.global_data, 1, Z⃗, n_traj + 1, n_global)
480+
end
481+
return evaluator._cached_traj
462482
end
463483

464484
"""

0 commit comments

Comments
 (0)