Constructive Code Review

[ocots]

CTDirect.jl — Constructive Code Review

Scope: Based on the source files provided (DOCP_data.jl, DOCP_functions.jl, DOCP_variables.jl, collocation.jl, direct_shooting.jl, ode/common.jl, ode/euler.jl, ode/irk.jl, ode/irk_stagewise.jl, ode/midpoint.jl, ode/trapeze.jl, ode/variable.jl).

---

1. Architecture & Design

1.1 Scheme dispatch via a large if-else chain

Issue. The DOCP constructor contains a long if-else chain on a Symbol to select the discretization scheme:

if scheme == :trapeze
    ...
elseif scheme == :midpoint
    ...
elseif scheme == :euler || scheme == :euler_explicit || ...

This is the classic anti-pattern that Julia's multiple dispatch is specifically designed to avoid. Adding a new scheme requires modifying the constructor, and the mapping from symbol to type is implicit.

Proposal. Define a small trait/helper that maps a symbol to a concrete type, or better, let each scheme register itself. A clean approach is a two-argument constructor that directly accepts the scheme type:

# Each scheme struct is self-describing
DOCP(ocp, grid_size, control_steps, Trapeze, time_grid)   # type-based dispatch
DOCP(ocp, grid_size, control_steps, :trapeze, time_grid)  # keep for user API, but route through a lookup table

A lookup Dict{Symbol, Function} populated by each scheme file (or a scheme_from_symbol function defined in ode/common.jl) would centralize the mapping and make it open for extension without touching DOCP_data.jl.

---

1.2 DOCP is declared mutable but only bounds needs mutation

Issue. The DOCP struct is declared mutable struct, but inspection shows that only bounds.var_l / bounds.var_u and bounds.con_l / bounds.con_u are ever mutated (via __variables_bounds! and __constraints_bounds!), and even those are mutated inside DOCPbounds, which is already a heap-allocated struct. Making DOCP mutable forces the Julia compiler to heap-allocate it and inhibits some optimisations.

Proposal. Keep DOCP immutable and, if mutation of the bounds really is required post-construction, make DOCPbounds a mutable sub-struct (it effectively already is, since it holds Vectors). The DOCP struct itself can then be struct rather than mutable struct.

---

1.3 Asymmetric override of __variables_bounds! and __initial_guess

Issue. The general __variables_bounds! (in DOCP_variables.jl) loops over time steps and calls set_control_at_time_step!. Gauss_Legendre_2_Stagewise / Gauss_Legendre_3_Stagewise override this with a separate __variables_bounds! in irk_stagewise.jl because their control layout is different. The same asymmetry holds for __initial_guess. This creates a hidden coupling: a new scheme that has a non-standard variable layout silently falls back to the wrong general method unless the developer remembers to override it.

Proposal. Make __variables_bounds! and __initial_guess dispatch explicitly on the discretization type:

__variables_bounds!(docp::DOCP{<:Scheme})         # generic fallback
__variables_bounds!(docp::DOCP{<:GenericIRKStagewise})  # override

The dispatch is already implicit (it works today because the override exists), but making the intent explicit—and adding a comment pointing to the contract—prevents future regressions. Consider also whether __initial_guess should be a method on the discretizer rather than on DOCP, since the initial guess layout is entirely determined by the scheme.

---

1.4 Closure mutation for exa_getter

Issue. In collocation.jl:

exa_getter = nothing  # will be set in build_exa_model

function build_exa_model(...)
    # ...
    nlp, exa_getter = build_exa(...)   # mutates outer-scope variable
end

function build_exa_solution(...)
    if isnothing(exa_getter)
        error("build_exa_solution: exa_getter is nothing")
    end
    # uses exa_getter
end

build_exa_solution silently depends on build_exa_model having been called first. This is an implicit temporal coupling that cannot be expressed in the type system and produces a runtime error if the call order is violated. The pattern is also impossible to test in isolation.

Proposal. Return exa_getter from build_exa_model and pair it with the solution explicitly:

function build_exa_model(...) -> (ExaModels.ExaModel, GetterType)
    # return nlp AND getter
end
# The solution builder receives the getter as an explicit argument,
# or the DiscretizedModel carries both builders as a named pair.

If the DiscretizedModel API doesn't allow this yet, at minimum document the ordering constraint with an @assert and a clear error message.

---

1.5 VariableStepODE is half-finished and silently excluded

Issue. In CTDirect.jl the include is commented out:

#include("ode/variable.jl")

Meanwhile, the file itself contains:

println(sol[end])
error("type ?")

These are clearly debug remnants. The scheme is listed in the user-facing error message in the DOCP constructor (under the valid options string), which is misleading.

Proposal.

• Remove the scheme from the public-facing option list until it is ready.
• Either delete variable.jl or move it to a WIP/ or ext/ directory with a clear comment.
• Introduce a @warn "VariableStepODE is experimental..." if you want to expose it as an opt-in.

---

2. Code Quality

2.1 Symbol-dispatched getter in common.jl

Issue. getter(nlp_solution, docp; val::Symbol) is a ~80-line function that dispatches on a Symbol value at runtime:

if val == :costate || val == :mult_path_constraints || ...
elseif occursin("_l", String(val))
    ...
elseif val == :variable || val == :variable_l || ...
elseif val == :state || ...
elseif val == :control || ...

Problems:

• occursin("_l", String(val)) is fragile: any symbol that happens to contain _l (e.g., :nl_mult) would be misrouted. The test should at least be endswith.
• The function is hard to extend and impossible to specialise per scheme without adding more if branches.
• The ExaModels getter is a parallel function with a different API, leading to if isnothing(exa_getter) guards scattered through build_OCP_solution.

Proposal. Replace the symbol-dispatch approach with typed accessor functions:

get_costate(sol, docp)
get_state(sol, docp)
get_control(sol, docp)
get_path_mult(sol, docp)
# etc.

Each can dispatch on the solver backend (ADNLPModels vs ExaModels) via a trait or a wrapper type. This is idiomatic Julia and completely eliminates the string/symbol matching.

---

2.2 DOCPdims docstring example has wrong arity

Issue. The docstring for DOCPdims shows:

julia> DOCPdims(4, 2, 1, 3, 2, 1)
DOCPdims(4, 2, 1, 3, 2, 1)

But the struct has only 5 fields (NLP_x, NLP_u, NLP_v, path_cons, boundary_cons). The example passes 6 arguments and would fail at runtime.

Proposal. Fix the docstring. Consider enabling Documenter.jl doctest execution in CI to catch this class of error automatically.

---

2.3 Gauss_Legendre_1 is marked "test only" but is public

Issue.

"""
[test only] Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order ...
"""
struct Gauss_Legendre_1 <: GenericIRK

The struct is exported (implicitly, since CTDirect does not appear to have an explicit export list) and accessible to users, despite the "test only" disclaimer.

Proposal. Either:

• Prefix it with _ (_Gauss_Legendre_1) to signal private status, or
• Move it to a test file (test/ directory), or
• Remove it if Midpoint is already the canonical implementation.

---

2.4 stepPathConstraints! offset computation is fragile

Issue. In DOCP_functions.jl:

offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) 
(i <= docp.time.steps) && (offset += disc._state_stage_eqs_block)

The call-site must pass i in [1, docp.time.steps+1]. The condition i <= docp.time.steps silently handles the final-time edge case by not adding the dynamics block offset. This implicit behaviour is not explained and could easily be broken by a future refactor.

Proposal. Add an inline comment explaining the convention, or split into two clearly named helpers:

_path_constraint_offset_interior(disc, i)
_path_constraint_offset_final(disc, steps)

---

2.5 build_OCP_solution is too long

Issue. build_OCP_solution in DOCP_data.jl is ~80 lines mixing:

• Conversion of arrays
• Retrieval of primal variables
• Computation of the control time grid
• Retrieval of multipliers
• Final call to CTModels.build_solution

This makes it hard to test each concern in isolation and hard to read.

Proposal. Extract helpers:

_extract_primal(docp, solution, getter) -> (X, U, v, T_state, T_control)
_extract_dual(docp, solution, getter)   -> named tuple of multipliers

---

2.6 Magic constant 0.1 in initial guesses

Issue. All __initial_guess functions start with:

NLP_X = 0.1 * ones(docp.dim_NLP_variables)

The value 0.1 is not explained. Is it intended to be nonzero to avoid degenerate Jacobians? Is it arbitrary?

Proposal. Extract a named constant or a configurable default:

const __DEFAULT_INITIAL_GUESS_VALUE = 0.1

and document the rationale in the function docstring.

---

3. Performance Concerns

3.1 Allocations in integral for GenericIRK and GenericIRKStagewise

Issue. In ode/irk.jl, integral allocates two work vectors per call:

work_xij = similar(xu, dims.NLP_x)
work_sumbk = similar(xu, 1)

and similarly in irk_stagewise.jl. These are called inside the NLP solver's objective evaluation, which runs at every iteration. For N = 250 steps with AD, this can add up.

Proposal. Pre-allocate these in setWorkArray and pass them through work, which already exists for this purpose. The Euler and Midpoint schemes already follow this pattern.

---

3.2 Allocation in stepStateConstraints! for GenericIRK

Issue.

@views @. work_xij[1:dims.NLP_x] = xi
for l in 1:disc.stage
    kil = get_stagevars_at_time_step(xu, docp, i, l)
    @views @. work_xij[1:dims.NLP_x] =
        work_xij[1:dims.NLP_x] + hi * disc.butcher_a[j, l] * kil[1:dims.NLP_x]
end

The comment in the code itself says # +++ still some allocations here. The kil[1:dims.NLP_x] slice creates a copy. Since get_stagevars_at_time_step already returns a @view, the slicing kil[1:dims.NLP_x] is redundant and should be removed.

Proposal.

kil = get_stagevars_at_time_step(xu, docp, i, l)  # already a view of length NLP_x
@. work_xij += hi * disc.butcher_a[j, l] * kil

---

3.3 integral for Trapeze evaluates f at every point independently

Issue. The trapeze integral calls f(ti, xi, ui, v) for each time step in a fresh loop, without reusing the dynamics already computed in setWorkArray. In contrast, stepStateConstraints! does use the work array. The inconsistency means the Lagrange cost evaluates the dynamics redundantly (once in the work array setup, once in integral).

Proposal. Extend the work array layout to also cache the Lagrange integrand, or have integral read from the same work array used for the dynamics. This is already partially done for Euler—the same approach can be applied to Trapeze.

---

4. Robustness & Testing

4.1 time_grid validation only checks strict monotonicity

Issue. In DOCPtime:

if !issorted(time_grid; lt=<=)
    throw(ArgumentError("given time grid is not strictly increasing. Aborting..."))
    return nothing  # dead code: throw already exits
end

The return nothing after throw is dead code. More importantly, there is no check that the first and last elements are finite, or that the grid has at least 2 points.

Proposal.

• Remove the dead return nothing.
• Add explicit checks:

length(time_grid) >= 2 || throw(ArgumentError("time grid must have at least 2 points"))
all(isfinite, time_grid) || throw(ArgumentError("time grid must contain finite values"))

---

4.2 add_nonzero_block! with symmetric flag can produce duplicate entries

Issue. add_nonzero_block!(Is, Js, i, j; sym=true) pushes both (i, j) and (j, i). When called for diagonal blocks (where i == j, or ranges overlap), this produces duplicate entries in the sparse pattern, which SparseArrays.sparse will silently sum. For a boolean pattern matrix this is harmless (1+1=2, which is truthy), but it is semantically wrong and may cause issues if the pattern is used for anything beyond boolean detection.

Proposal. Add a guard:

(sym && i != j) && (push!(Is, j); push!(Js, i))

(and the equivalent for ranges: only add the symmetric block if it doesn't overlap the original).

---

4.3 DirectShooting has empty build_exa_model / build_exa_solution

Issue.

function build_exa_model(::Type{BaseType}, ...) where {BaseType<:AbstractFloat}
    # try to call constructor here with constraints component wise ?
end

function build_exa_solution(nlp_solution::SolverCore.AbstractExecutionStats)
end

These silently return nothing, which will produce confusing errors downstream if a user requests ExaModels with DirectShooting.

Proposal. Add explicit error or @warn + return nothing with a message:

function build_exa_model(...)
    error("ExaModels backend is not yet supported for DirectShooting. Use Collocation instead.")
end

---

5. Documentation

5.1 DOCPtime has an undocumented control_steps field

The DOCPtime docstring lists fields steps, normalized_grid, fixed_grid but not control_steps, which was added later. Similarly the constructor example shows the 3-argument form DOCPtime(ocp, 10, nothing) but the actual signature is DOCPtime(ocp, grid_size, control_steps, time_grid).

5.2 The relationship between DirectShooting.control_steps and DOCPtime.control_steps is not documented

It is not immediately obvious to a reader what "multiple controls per time step" means in the context of direct shooting, how the controls are interpolated, or how they interact with the Lagrange cost evaluation in Midpoint. A top-level doc page (or at minimum a block comment in direct_shooting.jl) explaining the parametrisation would help contributors.

5.3 No module-level docstring

CTDirect.jl (the module file) has no @doc or """...""" module docstring. Adding one would make ?CTDirect useful in the REPL.

---

Summary Table

#	Location	Severity	Category	One-liner
1.1	DOCP_data.jl	Medium	Design	Replace if-else scheme dispatch with a lookup table or typed dispatch
1.2	DOCP_data.jl	Low	Design	Make DOCP immutable; DOCPbounds vectors are already mutable
1.3	DOCP_variables.jl / irk_stagewise.jl	Medium	Design	Make __variables_bounds! / __initial_guess dispatch explicit
1.4	collocation.jl	High	Design	Eliminate closure mutation for exa_getter
1.5	ode/variable.jl	High	Correctness	Remove debug println/error; hide unfinished scheme from public API
2.1	ode/common.jl	Medium	Quality	Replace symbol-dispatch getter with typed accessor functions
2.2	DOCP_data.jl	Low	Docs	Fix DOCPdims docstring (6 args shown, 5 in struct)
2.3	ode/irk.jl	Low	Quality	Hide or remove Gauss_Legendre_1 (marked "test only" but public)
2.4	DOCP_functions.jl	Low	Quality	Document / refactor fragile stepPathConstraints! offset logic
2.5	DOCP_data.jl	Low	Quality	Decompose build_OCP_solution into smaller helpers
2.6	Multiple	Low	Quality	Name the 0.1 default initial guess value
3.1	ode/irk.jl, ode/irk_stagewise.jl	Medium	Performance	Pre-allocate work_xij / work_sumbk in setWorkArray
3.2	ode/irk.jl	Low	Performance	Remove redundant kil[1:NLP_x] slicing (already a view)
3.3	ode/trapeze.jl	Low	Performance	Reuse dynamics work array in integral
4.1	DOCP_data.jl	Low	Robustness	Remove dead return nothing; add finite/length checks on time grid
4.2	ode/common.jl	Low	Correctness	Guard add_nonzero_block! sym path against duplicate diagonal entries
4.3	direct_shooting.jl	Medium	Robustness	Raise explicit error in empty build_exa_model / build_exa_solution
5.1	DOCP_data.jl	Low	Docs	Update DOCPtime docstring and constructor example
5.2	direct_shooting.jl	Low	Docs	Document control_steps semantics and control interpolation
5.3	CTDirect.jl	Low	Docs	Add module-level docstring

[ocots]

CTDirect.jl — Architecture Review (SOLID & Design Principles)

This is a follow-up to the first review comment which covered concrete code-level issues. This comment zooms out to architectural patterns, using SOLID principles and common Julia design idioms as the lens.

---

S — Single Responsibility Principle

DOCP_data.jl carries too many responsibilities

The file defines the core data structs (DOCPFlags, DOCPdims, DOCPtime, DOCPbounds, DOCP), but it also contains:

• the DOCP constructor with full scheme dispatch logic,
• the get_time_grid runtime accessor (called during NLP iterations),
• the build_OCP_solution solution reconstruction function (80 lines).

These are three distinct concerns: data layout, runtime variable access, and post-processing. A file named DOCP_data.jl should only define what a DOCP is, not what it does during solving or after.

Proposal — split by responsibility:

DOCP_data.jl       → struct definitions only (DOCPFlags, DOCPdims, etc.)
DOCP_accessors.jl  → get_time_grid, get_OCP_variable, etc.  (already DOCP_variables.jl for some)
DOCP_solution.jl   → build_OCP_solution, get_time_grid_exa

DOCP_variables.jl already exists and partially follows this split — completing it would make the structure fully coherent.

---

common.jl mixes variable accessors, constraint building, and sparsity pattern utilities

ode/common.jl contains: the getter function, the get_OCP_* / set_* family, setWorkArray, runningCost, and the add_nonzero_block! utilities. The last group (sparsity helpers) has no dependency on the ODE or DOCP logic; it is pure sparse-matrix bookkeeping and belongs in its own file (sparsity_utils.jl or similar), usable by all scheme files without pulling in ODE-specific code.

---

O — Open/Closed Principle

The DOCP constructor is closed to extension

As noted in §1.1 of the previous comment, the if-else chain on scheme::Symbol means that adding a new scheme requires modifying DOCP_data.jl. This violates OCP: the core data structure should not need to change when a new discretization method is introduced.

A concrete, Julia-idiomatic solution uses a scheme registry:

# In ode/common.jl — a Dict populated by each scheme file
const _SCHEME_CONSTRUCTORS = Dict{Symbol, Function}()

# Each scheme file registers itself:
# ode/trapeze.jl
_SCHEME_CONSTRUCTORS[:trapeze] = (dims, time) -> Trapeze(dims, time)

# ode/euler.jl
for alias in (:euler, :euler_explicit, :euler_forward)
    _SCHEME_CONSTRUCTORS[alias] = (dims, time) -> Euler(dims, time; explicit=true)
end

The DOCP constructor becomes:

builder = get(_SCHEME_CONSTRUCTORS, scheme, nothing)
isnothing(builder) && error("Unknown scheme: $scheme. Valid options: $(keys(_SCHEME_CONSTRUCTORS))")
discretization, dim_NLP_variables, dim_NLP_constraints = builder(dims, time)

Adding a new scheme now means adding one file and one registration call — no modification to the constructor.

---

L — Liskov Substitution Principle

DirectShooting <: AbstractDiscretizer silently breaks the contract

Collocation and DirectShooting both implement AbstractDiscretizer. The contract implied by the type hierarchy is that calling discretizer(ocp) returns a fully functional DiscretizedModel. But DirectShooting's ExaModels builders are empty:

function build_exa_model(::Type{BaseType}, ...) where {BaseType}
    # try to call constructor here with constraints component wise ?
end

Any code that calls build_exa_model on a DiscretizedModel built from DirectShooting will get nothing silently, violating the substitutability guarantee. The first comment (§4.3) suggested adding an explicit error; the LSP framing makes the reason clearer: a subtype must honour the parent contract, even if only to say "this operation is not supported" loudly.

Proposal. Define a fallback at the abstract level:

function build_exa_model(::AbstractDiscretizer, ::Type, ::Any; kwargs...)
    error("ExaModels backend not implemented for $(typeof(discretizer))")
end

This way the error fires at the right level with a clear message, and DirectShooting doesn't need its own stub.

---

I — Interface Segregation Principle

DiscretizedModel bundles four builders that users only partially need

DiscretizedModel is constructed with:

CTSolvers.DiscretizedModel(
    ocp,
    ADNLPModelBuilder(build_adnlp_model),
    ExaModelBuilder(build_exa_model),
    ADNLPSolutionBuilder(build_adnlp_solution),
    ExaSolutionBuilder(build_exa_solution),
)

A user who only ever uses ADNLPModels still carries the ExaModels builders (and their closures) everywhere. The interface is broader than what any single client needs.

This is a design decision that may be imposed by CTSolvers, so it might not be directly actionable in CTDirect. But if the CTSolvers API is flexible enough, consider a narrower interface — e.g., registering builders lazily or via keyword argument — so that a backend-specific build path only pays the cost of what it uses.

---

D — Dependency Inversion Principle

build_OCP_solution depends on a concrete solver type

function build_OCP_solution(docp::DOCP, nlp_solution::SolverCore.AbstractExecutionStats, ...)

At first glance this looks fine — AbstractExecutionStats is an abstract type. But inside the function, the code immediately accesses concrete fields:

solution    = Array(nlp_solution.solution)
multipliers = Array(nlp_solution.multipliers)
multipliers_L = Array(nlp_solution.multipliers_L)
multipliers_U = Array(nlp_solution.multipliers_U)

These field accesses are an implicit contract with SolverCore's concrete layout. If SolverCore ever changes field names, or if a third-party solver provides an AbstractExecutionStats without those fields, the function breaks silently at runtime.

Proposal. Introduce a thin extraction layer that depends on the abstract interface (or on explicit accessor functions from SolverCore), not on field names:

struct RawNLPResult
    primal::Vector
    mult_lower::Vector
    mult_upper::Vector
    mult_constraints::Vector
end

function extract_raw(sol::SolverCore.AbstractExecutionStats)
    return RawNLPResult(sol.solution, sol.multipliers_L, sol.multipliers_U, sol.multipliers)
end

build_OCP_solution then depends only on RawNLPResult — a type you control. The conversion from SolverCore types lives in one isolated place.

---

Multiple Dispatch vs. Symbol Dispatch

The getter(nlp_solution, docp; val::Symbol) function (already flagged in §2.1 of the first comment) is worth re-examining through the dispatch lens. Julia's whole identity is built around types carrying behaviour — a val::Symbol is just a runtime string in disguise, and it forfeits every benefit of the dispatch system: no compile-time specialisation, no method table lookup, no extensibility without editing the function.

The idiomatic replacement is not just "cleaner" — it unlocks scheme-specific overrides for free:

struct StateVal end;   struct ControlVal end;   struct CostateVal end
# ...

get_nlp_result(sol, docp, ::StateVal)   = # extract state
get_nlp_result(sol, docp, ::CostateVal) = # extract costate

# A scheme that stores costates differently overrides one method:
get_nlp_result(sol, docp::DOCP{<:GenericIRK}, ::CostateVal) = # IRK-specific extraction

No if-else, no occursin, fully open to extension.

---

Primitive Obsession

The scheme is passed as a Symbol all the way from the user API into the DOCP constructor and stored nowhere — it is consumed and discarded. The discretization type (e.g. Trapeze, Midpoint) is what actually carries the scheme's identity. Using Symbol as the primary representation is a classic case of primitive obsession: a domain concept (discretization scheme) represented by a primitive value instead of a domain type.

The scheme symbol is a valid user-facing convenience (:trapeze is more ergonomic than CTDirect.Trapeze()). The fix is to keep the symbol only at the API boundary and convert it immediately:

# Public API: symbol accepted for ergonomics
Collocation(; scheme=:midpoint, ...)

# Internally: convert once, carry the type
discretization_type = scheme_type(scheme)  # :midpoint → Midpoint

scheme_type is the registry function from the OCP section above — the symbol lives in exactly one place.

---

Summary of New Points

Principle	Location	Issue	Proposal
SRP	DOCP_data.jl	Mixes data layout, accessors, and post-processing	Split into DOCP_data, DOCP_accessors, DOCP_solution
SRP	ode/common.jl	Sparsity helpers mixed with ODE logic	Move add_nonzero_block! to sparsity_utils.jl
OCP	DOCP_data.jl constructor	Adding a scheme requires modifying the constructor	Scheme registry dict populated by each scheme file
LSP	direct_shooting.jl	Empty Exa builders silently break the discretizer contract	Abstract-level fallback error in AbstractDiscretizer
ISP	collocation.jl	DiscretizedModel bundles all backends unconditionally	Lazy or keyword-based builder registration (if CTSolvers API allows)
DIP	DOCP_data.jl	build_OCP_solution depends on SolverCore field layout	Extract RawNLPResult adapter struct
Dispatch	ode/common.jl	Symbol-based getter bypasses multiple dispatch	Typed singleton structs + dispatch
Primitives	collocation.jl, DOCP_data.jl	Symbol carries scheme identity beyond the API boundary	Convert symbol to type at API boundary via registry