diff --git a/Project.toml b/Project.toml index 7ed1bb39..b7129b0a 100644 --- a/Project.toml +++ b/Project.toml @@ -4,7 +4,12 @@ version = "0.10.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" @@ -13,26 +18,23 @@ Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" -oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" # oneAPI version >= 2.6 is needed to use sort! on GPU arrays -# OptimalControl = "5f98b655-cc9a-415a-b60e-744165666948" +oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] ExaModelsIpopt = ["MathOptInterface", "NLPModelsIpopt"] ExaModelsJuMP = "JuMP" ExaModelsKernelAbstractions = "KernelAbstractions" -ExaModelsOneAPI = "oneAPI" ExaModelsMOI = "MathOptInterface" ExaModelsMadNLP = ["MadNLP", "MathOptInterface"] ExaModelsMetal = "Metal" +ExaModelsOneAPI = "oneAPI" ExaModelsOpenCL = "OpenCL" ExaModelsSpecialFunctions = "SpecialFunctions" -# ExaModelsOptimalControl = ["OptimalControl", "LinearAlgebra"] [compat] Adapt = "4" @@ -44,8 +46,12 @@ MathOptInterface = "1.19" Metal = "1.9" NLPModels = "0.21" NLPModelsIpopt = "0.11" +NLPModelsJuMP = "0.13.5" +NLPModelsTest = "0.10.9" OpenCL = "0.10" +Percival = "0.7.6" +PowerModels = "0.21.5" SolverCore = "0.3" SpecialFunctions = "2" julia = "1.9" -oneAPI = "2" \ No newline at end of file +oneAPI = "2" diff --git a/docs/make.jl b/docs/make.jl index 7de79027..15cfd6cf 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -15,6 +15,8 @@ if !(@isdefined _PAGES) "performance.md", "gpu.md", "parameters.md", + "two_stage.md", + "batch.md", "develop.md", "quad.md", @@ -38,6 +40,8 @@ if !(@isdefined _JL_FILENAMES) "gpu.jl", "performance.jl", "parameters.jl", + "two_stage.jl", + "batch.jl", ] end diff --git a/docs/src/batch.jl b/docs/src/batch.jl new file mode 100644 index 00000000..a2d3dc28 --- /dev/null +++ b/docs/src/batch.jl @@ -0,0 +1,131 @@ +# # [Batch Optimization](@id batch) +# ExaModels supports batch optimization through `BatchExaCore`. This feature +# enables efficient evaluation of multiple fully independent optimization instances +# that share identical structure but differ in parameter values. +# +# Unlike `TwoStageExaModel`, which couples instances through shared design variables, +# batch models treat each instance as completely independent. The key advantage +# is that all instances share one compiled expression pattern and are fused into a +# single model for efficient SIMD evaluation. + +# ## Problem Formulation +# A batch optimization problem solves `ns` independent instances simultaneously: +# ```math +# \begin{aligned} +# \min_{v_i} \quad & f(v_i; \theta_i), \quad i = 1, \ldots, S \\ +# \text{s.t.} \quad & h(v_i; \theta_i) = 0 \\ +# & v_i \in \mathcal{V} +# \end{aligned} +# ``` +# where each instance has the same structure but different parameters $\theta_i$. + +# ## Building a Batch Model +# Use `BatchExaCore(ns)` to create a core for `ns` instances. +# Variables, parameters, objectives, and constraints are defined once — +# the batch structure replicates them across all instances automatically. + +using ExaModels, NLPModelsIpopt +import NLPModels + +# Define the problem dimensions and instance parameters: +ns = 3 ## number of instances +nv = 1 ## variables per instance + +# Create a batch core: +c = BatchExaCore(ns) + +# Add variables and parameters: +@add_var(c, v, nv) +@add_par(c, θ, [2.0]) + +# Define objectives and constraints — these apply to every instance: +@add_obj(c, (v[j] - θ[1])^2 for j in 1:nv) +@add_con(c, g, v[j] for j in 1:nv; lcon = 0.0) + +# Build the model: +model = ExaModel(c) + +# ## Batch API (NLPModels) +# Batch models implement `AbstractNLPModel` with matrix-valued variables. +# All evaluation functions use matrices of size `(dim, ns)`: + +println("Variables per instance: ", NLPModels.get_nvar(model)) +println("Constraints per instance: ", NLPModels.get_ncon(model)) +println("Number of instances: ", ExaModels.get_nbatch(model)) + +# Evaluate objectives for all instances at once: +bx = reshape([1.0, 3.0, 5.0], nv, ns) +bf = zeros(ns) +NLPModels.obj!(model, bx, bf) +println("\nObjective values: ", bf) +## instance 1: (1-2)² = 1, instance 2: (3-4)² = 1, instance 3: (5-6)² = 1 + +# Evaluate gradients: +bg = zeros(nv, ns) +NLPModels.grad!(model, bx, bg) +println("Gradients: ", bg) + +# Evaluate constraints: +bc = zeros(NLPModels.get_ncon(model), ns) +NLPModels.cons!(model, bx, bc) +println("Constraints: ", bc) + +# ## Solving Batch Models +# There is currently no dedicated batch solver. To solve all instances, wrap +# the batch model in `FlatNLPModel`, which concatenates the `ns` instances +# into a single standard `AbstractNLPModel` that any NLPModels-compatible +# solver can consume: +flat = FlatNLPModel(model) +result = ipopt(flat; print_level = 0) +println("\nSolution status: ", result.status) + +# Extract per-instance solutions: +x_sol = result.solution +for i in 1:ns + v_sol = x_sol[ExaModels.var_indices(model, i)] + println("Instance $i: v* = ", round(v_sol[1], digits = 4)) +end + +# ## Per-Instance Parameters +# By default, `@add_par` replicates the same value across all instances. +# To give each instance its own parameter values, pass an `npar × ns` +# matrix to `set_value!` after building the model. Each column supplies the +# parameter vector for one instance. + +c3 = BatchExaCore(ns) +@add_var(c3, w, nv) +@add_par(c3, α, [0.0]) ## placeholder — will be overwritten per-instance +@add_obj(c3, (w[j] - α[1])^2 for j in 1:nv) +model3 = ExaModel(c3) + +# Set different targets for each instance (npar=1 row, ns=3 columns): +ExaModels.set_value!(model3, α, [2.0 4.0 6.0]) + +flat3 = FlatNLPModel(model3) +result3 = ipopt(flat3; print_level = 0) +for i in 1:ns + v_sol = result3.solution[ExaModels.var_indices(model3, i)] + println("Instance $i: w* = ", round(v_sol[1], digits = 4)) +end + +# For problems with multiple parameters per instance, supply a matrix with one +# row per parameter and one column per instance. For example, with `npar = 2` +# parameters and `ns = 3` instances: +# +# ```julia +# c4 = BatchExaCore(3) +# @add_var(c4, x, 2) +# @add_par(c4, p, [0.0, 0.0]) +# @add_obj(c4, sum((x[j] - p[j])^2 for j in 1:2)) +# model4 = ExaModel(c4) +# +# # 2×3 matrix: column i = parameter vector for instance i +# ExaModels.set_value!(model4, p, [1.0 3.0 5.0; +# 2.0 4.0 6.0]) +# ``` +# +# The parameter matrix `model.θ` always has shape `(npar, ns)` and can be +# inspected directly: + +println("\nParameter matrix (npar × ns): ", size(model3.θ)) +println("Values: ", model3.θ) diff --git a/docs/src/core.md b/docs/src/core.md index d41a1ed9..7b6ccbde 100644 --- a/docs/src/core.md +++ b/docs/src/core.md @@ -2,3 +2,4 @@ ```@autodocs Modules = [ExaModels] ``` + diff --git a/docs/src/parameters.jl b/docs/src/parameters.jl index 93477e4b..635c8154 100644 --- a/docs/src/parameters.jl +++ b/docs/src/parameters.jl @@ -35,11 +35,11 @@ result1 = ipopt(m_param) println("Original objective: $(result1.objective)") # Now change the penalty coefficient and solve again: -set_parameter!(c_param, θ, [200.0, 1.0]) # Double the penalty coefficient +set_value!(m_param, θ, [200.0, 1.0]) # Double the penalty coefficient result2 = ipopt(m_param) println("Modified penalty objective: $(result2.objective)") # Try a different offset parameter: -set_parameter!(c_param, θ, [200.0, 0.5]) # Change the offset in the objective +set_value!(m_param, θ, [200.0, 0.5]) # Change the offset in the objective result3 = ipopt(m_param) println("Modified offset objective: $(result3.objective)") diff --git a/docs/src/parameters.md b/docs/src/parameters.md index d09d7348..9ffb2bdb 100644 --- a/docs/src/parameters.md +++ b/docs/src/parameters.md @@ -31,7 +31,7 @@ An ExaCore Adding parameters is very similar to adding variables -- just pass a vector of values denoting the initial values. ````julia -@add_parameter(c_param, θ, [100.0, 1.0]) # [penalty_coeff, offset] +@add_par(c_param, θ, [100.0, 1.0]) # [penalty_coeff, offset] ```` ```` @@ -45,7 +45,7 @@ Define the variables as before: ````julia N = 10 -@add_variable(c_param, x_p, N; start = (mod(i, 2) == 1 ? -1.2 : 1.0 for i = 1:N)) +@add_var(c_param, x_p, N; start = (mod(i, 2) == 1 ? -1.2 : 1.0 for i = 1:N)) ```` ```` @@ -58,7 +58,7 @@ Variable Now we can use the parameters in our objective function just like variables: ````julia -@add_objective(c_param, θ[1] * (x_p[i-1]^2 - x_p[i])^2 + (x_p[i-1] - θ[2])^2 for i = 2:N) +@add_obj(c_param, θ[1] * (x_p[i-1]^2 - x_p[i])^2 + (x_p[i-1] - θ[2])^2 for i = 2:N) ```` ```` @@ -73,7 +73,7 @@ Objective Add the same constraints as before: ````julia -@add_constraint( +@add_con( c_param, 3x_p[i+1]^3 + 2 * x_p[i+2] - 5 + sin(x_p[i+1] - x_p[i+2])sin(x_p[i+1] + x_p[i+2]) + @@ -178,7 +178,7 @@ Original objective: 6.232458632437464 Now change the penalty coefficient and solve again: ````julia -set_parameter!(c_param, θ, [200.0, 1.0]) # Double the penalty coefficient +set_value!(m_param, θ, [200.0, 1.0]) # Double the penalty coefficient result2 = ipopt(m_param) println("Modified penalty objective: $(result2.objective)") ```` @@ -237,7 +237,7 @@ Modified penalty objective: 8.647439751691499 Try a different offset parameter: ````julia -set_parameter!(c_param, θ, [200.0, 0.5]) # Change the offset in the objective +set_value!(m_param, θ, [200.0, 0.5]) # Change the offset in the objective result3 = ipopt(m_param) println("Modified offset objective: $(result3.objective)") ```` diff --git a/docs/src/two_stage.jl b/docs/src/two_stage.jl index 16e6c79b..c820dcd8 100644 --- a/docs/src/two_stage.jl +++ b/docs/src/two_stage.jl @@ -13,23 +13,24 @@ nv = 2 ## recourse variables per scenario nd = 1 ## design variables weight = 1.0 / ns -# Two annotate the scenario for each variable and constraint, we can use the `scenario` we need to start with a special ExaCore that supports such scenario annotations, which can be created by calling `TwoStageExaCore(concrete = Val(true))`. -core = TwoStageExaCore(concrete = Val(true)) +# To annotate the scenario for each variable and constraint, we start with a `TwoStageExaCore` that supports scenario annotations. +core = TwoStageExaCore(ns; concrete = Val(true)) -# Now we can define the design variable and recourse variables. The `scenario` keyword argument allows us to specify which scenario(s) each variable belongs to. For the design variable `d`, we set `scenario = 0` to indicate that it is shared across all scenarios. -@add_var(core, d; start = 1.0, lvar = 0.0, uvar = Inf, scenario = 0) ## design variable d +# Design variables are shared across all scenarios — add them without `EachScenario()`. +core, d = add_var(core, nd; start = 1.0, lvar = 0.0, uvar = Inf) -# For the recourse variables `v`, we specify `scenario = [i for i=1:ns, j=1:nv]` to indicate that each variable `v[s,i]` belongs to scenario `s`. This allows us to define scenario-specific constraints and objectives that involve these recourse variables. -@add_var(core, v, ns, nv; start = 1.0, lvar = 0.0, uvar = Inf, scenario = [i for i=1:ns, j=1:nv]) ## recourse variables v +# Recourse variables are per-scenario — use `EachScenario()` to replicate them. +v = @add_var(core, EachScenario(), nv; start = 1.0, lvar = 0.0, uvar = Inf) -# Now we can define the constraints and objective function. The `scenario` keyword argument in the `constraint` and `objective` functions allows us to specify which scenario(s) each constraint or objective term belongs to. -@add_con(core, v[s,1] - v[s,2]^2 for s in 1:ns; lcon = 0.0, scenario = 1:ns) +# Per-scenario constraints use `EachScenario()`. +@add_con(core, EachScenario(), (v[(s-1)*nv+1] - v[(s-1)*nv+2]^2 for s in 1:ns); lcon = 0.0) -@add_obj(core, d^2) -@add_obj(core, weight * (v[s,i] - d)^2 for s in 1:ns, i in 1:nv) +# Objectives can mix design and recourse variables. +@add_obj(core, d[1]^2) +@add_obj(core, weight * (v[(s-1)*nv+i] - d[1])^2 for s in 1:ns, i in 1:nv) m = ExaModel(core) -# Now we can solve the model as usual. -ipopt(m) +# Now we can solve the model as usual. +ipopt(m) # If the solver knows how to exploit the scenario structure, the structure-exploiting method can be used. diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index e1f9af32..1b5bf12b 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -3,11 +3,6 @@ module ExaModelsKernelAbstractions import ExaModels: ExaModels, NLPModels import KernelAbstractions: KernelAbstractions, @kernel, @index, @Const, synchronize, CPU -function getitr(gen::UnitRange{Int64}) - return gen -end -function getitr(gen::Base.Iterators.ProductIterator{NTuple{N,UnitRange{Int64}}}) where {N} end - function ExaModels.getptr(backend, array; cmp = (x, y) -> x != y) bitarray = similar(array, Bool, length(array) + 1) @@ -33,7 +28,7 @@ function ExaModels.build_extension( c::C; prod = false, kwargs..., -) where {T,VT<:AbstractVector{T},B<:KernelAbstractions.Backend,C<:ExaModels.ExaCore{T,VT,B}} +) where {T,VT<:AbstractArray{T},B<:KernelAbstractions.Backend,C<:ExaModels.ExaCore{T,VT,B}} gsparsity = similar(c.x0, Tuple{Int,Int}, c.nnzg) @@ -58,11 +53,11 @@ function ExaModels.build_extension( jacsparsityi = similar(c.x0, Tuple{Tuple{Int,Int},Int}, c.nnzj) hesssparsityi = similar(c.x0, Tuple{Tuple{Int,Int},Int}, c.nnzh) - _jac_structure!(T, c.backend, c.cons, jacsparsityi, nothing) + ExaModels._jac_structure!(T, c.backend, c.cons, jacsparsityi, nothing) jacsparsityj = copy(jacsparsityi) - _obj_hess_structure!(T, c.backend, c.obj, hesssparsityi, nothing) - _con_hess_structure!(T, c.backend, c.cons, hesssparsityi, nothing) + ExaModels._obj_hess_structure!(T, c.backend, c.obj, hesssparsityi, nothing) + ExaModels._con_hess_structure!(T, c.backend, c.cons, hesssparsityi, nothing) hesssparsityj = copy(hesssparsityi) if !isempty(jacsparsityi) @@ -121,9 +116,14 @@ end _conaug_structure!(T, backend, ::Tuple{}, sparsity) = nothing function _conaug_structure!(T, backend, (con, cons...), sparsity) _conaug_structure!(T, backend, cons, sparsity) - con isa ExaModels.ConstraintAugmentation && !isempty(con.itr) && + _conaug_structure_item!(T, backend, con, sparsity) +end +function _conaug_structure_item!(T, backend, con::ExaModels.ConstraintAugmentation, sparsity) + if !isempty(con.itr) kers(backend)(sparsity, con.f, con.itr, con.oa, con.dims; ndrange = length(con.itr)) + end end +_conaug_structure_item!(T, backend, con, sparsity) = nothing @kernel function kers(sparsity, @Const(f), @Const(itr), @Const(oa), @Const(dims)) I = @index(Global) @inbounds sparsity[oa+I] = (ExaModels.offset0(f, itr, I, dims), oa + I) @@ -134,52 +134,59 @@ end _grad_structure!(T, backend, ::Tuple{}, gsparsity) = nothing function _grad_structure!(T, backend, (obj, objs...), gsparsity) _grad_structure!(T, backend, objs, gsparsity) - ExaModels.sgradient!(backend, gsparsity, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN)) + ExaModels.sgradient!(gsparsity, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), 1, 0, 0, length(gsparsity), backend) end +ExaModels._jac_structure!(T, backend, ::Tuple{}, rows, cols) = nothing +function ExaModels._jac_structure!(T, backend, (con, cons...), rows, cols) + ExaModels._jac_structure!(T, backend, cons, rows, cols) + ExaModels.sjacobian!(rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), 1, 0, 0, length(rows), backend) +end + +ExaModels._obj_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing +function ExaModels._obj_hess_structure!(T, backend, (obj, objs...), rows, cols) + ExaModels._obj_hess_structure!(T, backend, objs, rows, cols) + ExaModels.shessian!(rows, cols, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN), 1, 0, 0, length(rows), backend) +end +ExaModels._con_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing +function ExaModels._con_hess_structure!(T, backend, (con, cons...), rows, cols) + ExaModels._con_hess_structure!(T, backend, cons, rows, cols) + ExaModels.shessian!(rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN), 1, 0, 0, length(rows), backend) +end + + function ExaModels.jac_structure!( - m::ExaModels.AbstractExaModel{T,VT,E}, - rows::V, - cols::V, -) where {T,VT,E<:KAExtension,V<:AbstractVector} + m::ExaModels.ExaModel{T,VT,E}, + rows::AbstractVector, + cols::AbstractVector, +) where {T,VT,E<:KAExtension} if !isempty(rows) - _jac_structure!(T, m.ext.backend, m.cons, rows, cols) + ExaModels._jac_structure!(T, m.ext.backend, m.cons, rows, cols) end return rows, cols end -_jac_structure!(T, backend, ::Tuple{}, rows, cols) = nothing -function _jac_structure!(T, backend, (con, cons...), rows, cols) - _jac_structure!(T, backend, cons, rows, cols) - ExaModels.sjacobian!(backend, rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN)) -end - function ExaModels.hess_structure!( - m::ExaModels.AbstractExaModel{T,VT,E}, - rows::V, - cols::V, -) where {T,VT,E<:KAExtension,V<:AbstractVector} + m::ExaModels.ExaModel{T,VT,E}, + rows::AbstractVector, + cols::AbstractVector, +) where {T,VT,E<:KAExtension} if !isempty(rows) - _obj_hess_structure!(T, m.ext.backend, m.objs, rows, cols) - _con_hess_structure!(T, m.ext.backend, m.cons, rows, cols) - end + ExaModels._obj_hess_structure!(T, m.ext.backend, m.objs, rows, cols) + ExaModels._con_hess_structure!(T, m.ext.backend, m.cons, rows, cols) + end return rows, cols end -_obj_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing -function _obj_hess_structure!(T, backend, (obj, objs...), rows, cols) - _obj_hess_structure!(T, backend, objs, rows, cols) - ExaModels.shessian!(backend, rows, cols, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN)) -end -_con_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing -function _con_hess_structure!(T, backend, (con, cons...), rows, cols) - _con_hess_structure!(T, backend, cons, rows, cols) - ExaModels.shessian!(backend, rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN)) +function ExaModels.obj( + m::ExaModels.ExaModel{T,VT,E}, + x::AbstractVector, +) where {T,VT<:AbstractMatrix,E<:KAExtension} + ExaModels._batch_vector_error() end - function ExaModels.obj( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, ) where {T,VT,E<:KAExtension} if !isempty(m.ext.objbuffer) @@ -198,8 +205,22 @@ function _obj(backend, objbuffer, (obj, objs...), x, θ) end end +function ExaModels.obj!( + m::ExaModels.ExaModel{T,VT,E}, + bx::AbstractVecOrMat, + bf::AbstractVector, +) where {T,VT,E<:KAExtension} + fill!(bf, zero(T)) + nb = Base.size(bx, 2) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + ExaModels._obj!(bf, m.objs, vec(bx), vec(m.θ), nb, nvar, npar, m.ext.backend) + return bf +end + + function ExaModels.cons_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, y::AbstractVector, ) where {T,VT,E<:KAExtension} @@ -217,26 +238,60 @@ function ExaModels.cons_nln!( end return y end + +function ExaModels.cons_nln!( + m::ExaModels.ExaModel{T,VT,E}, + x::V, + y::V, +) where {T,VT<:AbstractMatrix,E<:KAExtension,V<:AbstractVector} + fill!(y, zero(T)) + _cons_nln!(m.ext.backend, y, m.cons, x, m.θ) + _conaugs!(m.ext.backend, m.ext.conbuffer, m.cons, x, m.θ) + if length(m.ext.conaugptr) > 1 + compress_to_dense(m.ext.backend)( + y, + m.ext.conbuffer, + m.ext.conaugptr, + m.ext.conaugsparsity; + ndrange = length(m.ext.conaugptr) - 1, + ) + end + return y +end _cons_nln!(backend, y, ::Tuple{}, x, θ) = nothing function _cons_nln!(backend, y, (con, cons...), x, θ) _cons_nln!(backend, y, cons, x, θ) - if con isa ExaModels.Constraint && !isempty(con.itr) + _cons_nln_item!(backend, y, con, x, θ) +end +function _cons_nln_item!(backend, y, con::ExaModels.Constraint, x, θ) + if !isempty(con.itr) kerf(backend)(y, con.f, con.itr, x, θ; ndrange = length(con.itr)) end end - - +_cons_nln_item!(backend, y, con, x, θ) = nothing _conaugs!(backend, y, ::Tuple{}, x, θ) = nothing function _conaugs!(backend, y, (con, cons...), x, θ) _conaugs!(backend, y, cons, x, θ) - if con isa ExaModels.ConstraintAugmentation && !isempty(con.itr) + _conaugs_item!(backend, y, con, x, θ) +end +function _conaugs_item!(backend, y, con::ExaModels.ConstraintAugmentation, x, θ) + if !isempty(con.itr) kerf2(backend)(y, con.f, con.itr, x, θ, con.oa; ndrange = length(con.itr)) end end +_conaugs_item!(backend, y, con, x, θ) = nothing function ExaModels.grad!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, + x::V, + y::V, +) where {T,VT<:AbstractMatrix,E<:KAExtension,V<:AbstractVector} + ExaModels._batch_vector_error() +end + +function ExaModels.grad!( + m::ExaModels.ExaModel{T,VT,E}, x::V, y::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} @@ -260,26 +315,11 @@ end _grad!(backend, y, ::Tuple{}, x, θ) = nothing function _grad!(backend, y, (obj, objs...), x, θ) _grad!(backend, y, objs, x, θ) - ExaModels.sgradient!(backend, y, obj, x, θ, one(eltype(y))) -end - -function ExaModels.jac_coord!( - m::ExaModels.AbstractExaModel{T,VT,E}, - x::V, - y::V, -) where {T,VT,E<:KAExtension,V<:AbstractVector} - fill!(y, zero(eltype(y))) - _jac_coord!(m.ext.backend, y, m.cons, x, m.θ) - return y -end -_jac_coord!(backend, y, ::Tuple{}, x, θ) = nothing -function _jac_coord!(backend, y, (con, cons...), x, θ) - _jac_coord!(backend, y, cons, x, θ) - ExaModels.sjacobian!(backend, y, nothing, con, x, θ, one(eltype(y))) + ExaModels.sgradient!(y, obj, x, θ, one(eltype(y)), 1, length(x), length(θ), length(y), backend) end function ExaModels.jprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jv::AbstractVector, @@ -287,7 +327,7 @@ function ExaModels.jprod_nln!( error("Prodhelper is not defined. Use ExaModels(c; prod=true) to use jprod_nln!") end function ExaModels.jtprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector, @@ -295,7 +335,7 @@ function ExaModels.jtprod_nln!( error("Prodhelper is not defined. Use ExaModels(c; prod=true) to use jtprod_nln!") end function ExaModels.jprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jv::AbstractVector, @@ -303,7 +343,7 @@ function ExaModels.jprod_nln!( fill!(Jv, zero(eltype(Jv))) fill!(m.ext.prodhelper.jacbuffer, zero(eltype(Jv))) - _jac_coord!(m.ext.backend, m.ext.prodhelper.jacbuffer, m.cons, x, m.θ) + ExaModels._jac_coord!(m.cons, x, m.θ, m.ext.prodhelper.jacbuffer, 1, length(x), length(m.θ), length(m.ext.prodhelper.jacbuffer), m.ext.backend) kerspmv(m.ext.backend)( Jv, v, @@ -315,7 +355,7 @@ function ExaModels.jprod_nln!( return Jv end function ExaModels.jtprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector, @@ -323,7 +363,7 @@ function ExaModels.jtprod_nln!( fill!(Jtv, zero(eltype(Jtv))) fill!(m.ext.prodhelper.jacbuffer, zero(eltype(Jtv))) - _jac_coord!(m.ext.backend, m.ext.prodhelper.jacbuffer, m.cons, x, m.θ) + ExaModels._jac_coord!(m.cons, x, m.θ, m.ext.prodhelper.jacbuffer, 1, length(x), length(m.θ), length(m.ext.prodhelper.jacbuffer), m.ext.backend) kerspmv2(m.ext.backend)( Jtv, v, @@ -335,7 +375,7 @@ function ExaModels.jtprod_nln!( return Jtv end function ExaModels.hprod!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, y::AbstractVector, v::AbstractVector, @@ -350,8 +390,8 @@ function ExaModels.hprod!( fill!(Hv, zero(eltype(Hv))) fill!(m.ext.prodhelper.hessbuffer, zero(eltype(Hv))) - _obj_hess_coord!(m.ext.backend, m.ext.prodhelper.hessbuffer, m.objs, x, m.θ, obj_weight) - _con_hess_coord!(m.ext.backend, m.ext.prodhelper.hessbuffer, m.cons, x, m.θ, y) + ExaModels._obj_hess_coord!(m.objs, x, m.θ, m.ext.prodhelper.hessbuffer, obj_weight, 1, length(x), length(m.θ), length(m.ext.prodhelper.hessbuffer), m.ext.backend) + ExaModels._con_hess_coord!(m.cons, x, m.θ, y, m.ext.prodhelper.hessbuffer, 1, length(x), length(m.θ), length(y), length(m.ext.prodhelper.hessbuffer), m.ext.backend) kersyspmv(m.ext.backend)( Hv, v, @@ -372,7 +412,7 @@ function ExaModels.hprod!( return Hv end function ExaModels.hprod!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Hv::AbstractVector; @@ -386,7 +426,7 @@ function ExaModels.hprod!( fill!(Hv, zero(eltype(Hv))) fill!(m.ext.prodhelper.hessbuffer, zero(eltype(Hv))) - _obj_hess_coord!(m.ext.backend, m.ext.prodhelper.hessbuffer, m.objs, x, m.θ, obj_weight) + ExaModels._obj_hess_coord!(m.objs, x, m.θ, m.ext.prodhelper.hessbuffer, obj_weight, 1, length(x), length(m.θ), length(m.ext.prodhelper.hessbuffer), m.ext.backend) kersyspmv(m.ext.backend)( Hv, v, @@ -440,161 +480,6 @@ end -function ExaModels.hess_coord!( - m::ExaModels.AbstractExaModel{T,VT,E}, - x::V, - y::V, - hess::V; - obj_weight = one(eltype(y)), -) where {T,VT,E<:KAExtension,V<:AbstractVector} - fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.ext.backend, hess, m.objs, x, m.θ, obj_weight) - _con_hess_coord!(m.ext.backend, hess, m.cons, x, m.θ, y) - return hess -end -_obj_hess_coord!(backend, hess, ::Tuple{}, x, θ, obj_weight) = nothing -function _obj_hess_coord!(backend, hess, (obj, objs...), x, θ, obj_weight) - _obj_hess_coord!(backend, hess, objs, x, θ, obj_weight) - ExaModels.shessian!(backend, hess, nothing, obj, x, θ, obj_weight, zero(eltype(hess))) -end -_con_hess_coord!(backend, hess, ::Tuple{}, x, θ, y) = nothing -function _con_hess_coord!(backend, hess, (con, cons...), x, θ, y) - _con_hess_coord!(backend, hess, cons, x, θ, y) - ExaModels.shessian!(backend, hess, nothing, con, x, θ, y, zero(eltype(hess))) -end - - -function ExaModels.sgradient!( - backend::B, - y, - f, - x, - θ, - adj, -) where {B<:KernelAbstractions.Backend} - - if !isempty(f.itr) - kerg(backend)(y, f.f, f.itr, x, θ, adj; ndrange = length(f.itr)) - end -end - -function ExaModels.sjacobian!( - backend::B, - y1, - y2, - f, - x, - θ, - adj, -) where {B<:KernelAbstractions.Backend} - if !isempty(f.itr) - kerj(backend)(y1, y2, f.f, f.itr, x, θ, adj, ExaModels._constraint_dims(f); ndrange = length(f.itr)) - end -end - -function ExaModels.shessian!( - backend::B, - y1, - y2, - f, - x, - θ, - adj, - adj2, -) where {B<:KernelAbstractions.Backend} - if !isempty(f.itr) - kerh(backend)(y1, y2, f.f, f.itr, x, θ, adj, adj2; ndrange = length(f.itr)) - end -end - -function ExaModels.shessian!( - backend::B, - y1, - y2, - f, - x, - θ, - adj::V, - adj2, -) where {B<:KernelAbstractions.Backend,V<:AbstractVector} - if !isempty(f.itr) - kerh2(backend)(y1, y2, f.f, f.itr, x, θ, adj, adj2, ExaModels._constraint_dims(f); ndrange = length(f.itr)) - end -end - -@kernel function kerh( - y1, - y2, - @Const(f), - @Const(itr), - @Const(x), - @Const(θ), - @Const(adj1), - @Const(adj2) -) - I = @index(Global) - @inbounds ExaModels.hrpass0( - f(itr[I], ExaModels.SecondAdjointNodeSource(x), θ), - f.comp2, - y1, - y2, - ExaModels.offset2(f, I), - 0, - adj1, - adj2, - ) -end - -@kernel function kerh2( - y1, - y2, - @Const(f), - @Const(itr), - @Const(x), - @Const(θ), - @Const(adjs1), - @Const(adj2), - @Const(dims) -) - I = @index(Global) - @inbounds ExaModels.hrpass0( - f(itr[I], ExaModels.SecondAdjointNodeSource(x), θ), - f.comp2, - y1, - y2, - ExaModels.offset2(f, I), - 0, - adjs1[ExaModels.offset0(f, itr, I, dims)], - adj2, - ) -end - -@kernel function kerj(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(dims)) - I = @index(Global) - @inbounds ExaModels.jrpass( - f(itr[I], ExaModels.AdjointNodeSource(x), θ), - f.comp1, - ExaModels.offset0(f, itr, I, dims), - y1, - y2, - ExaModels.offset1(f, I), - 0, - adj, - ) -end - -@kernel function kerg(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj)) - I = @index(Global) - @inbounds ExaModels.grpass( - f(itr[I], ExaModels.AdjointNodeSource(x), θ), - f.comp1, - y, - ExaModels.offset1(f, I), - 0, - adj, - ) -end - @kernel function kerf(y, @Const(f), @Const(itr), @Const(x), @Const(θ)) I = @index(Global) @inbounds y[ExaModels.offset0(f, itr, I)] = f(itr[I], x, θ) @@ -631,7 +516,7 @@ end end end -ExaModels.getbackend(m::ExaModels.AbstractExaModel{T,VT,E}) where {T,VT,E<:KAExtension} = +ExaModels.getbackend(m::ExaModels.ExaModel{T,VT,E}) where {T,VT,E<:KAExtension} = m.ext.backend function ExaModels._compress!(V, buffer, ptr, sparsity, backend) fill!(V, zero(eltype(V))) @@ -664,6 +549,204 @@ end @inbounds sparsity[i] = ((J[i], I[i]), i) end +# ============================================================================ +# Batch KA dispatch — single kernel launch for the whole batch +# ============================================================================ + +# --- Objective --- + +function ExaModels._obj_eval!(bf, obj, x, θ, nb, nvar, npar, backend::KernelAbstractions.Backend) + nitr = length(obj.itr) + if nitr > 0 + buf = similar(x, nitr, nb) + kerf_batch_fill(backend)(buf, obj.f, obj.itr, x, θ, nvar, npar, nitr; ndrange = nb * nitr) + bf .+= vec(sum(buf; dims = 1)) + end +end + +@kernel function kerf_batch_fill(buf, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + @inbounds buf[k, s] = f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) +end + +# --- Constraints --- + +function ExaModels._cons_nln_eval!(con, x, θ, g, nb, nvar, npar, ncon, backend::KernelAbstractions.Backend) + nitr = length(con.itr) + if nitr > 0 + kerf_con_batch(backend)(g, con.f, con.itr, x, θ, nvar, npar, ncon, nitr; ndrange = nb * nitr) + end +end + +function ExaModels._cons_nln_eval!(con::ExaModels.ConstraintAugmentation, x, θ, g, nb, nvar, npar, ncon, backend::KernelAbstractions.Backend) + nitr = length(con.itr) + if nitr > 0 + kerf_con_aug_batch(backend)(g, con.f, con.itr, x, θ, nvar, npar, ncon, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerf_con_batch(g, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(ncon), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + g_off = (s - 1) * ncon + @inbounds g[g_off + ExaModels.offset0(f, itr, k)] += f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) +end + +@kernel function kerf_con_aug_batch(g, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(ncon), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + g_off = (s - 1) * ncon + val = f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) + @inbounds KernelAbstractions.@atomic g[g_off + ExaModels.offset0(f, itr, k)] += val +end + +# --- Gradient --- + +function ExaModels.gradient!(y, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerg(backend)(y, f.f, f.itr, x, θ, adj, nvar, npar, nitr; ndrange = nb * nitr) + end + return y +end + +@kernel function kerg(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(nvar), @Const(npar), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nvar + @inbounds ExaModels.drpass( + f(itr[k], ExaModels.AdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), + view(y, y_off+1:y_off+nvar), + adj, + ) +end + +# --- Sparse gradient --- + +function ExaModels.sgradient!(y, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerg_sparse(backend)(y, f.f, f.itr, x, θ, adj, nvar, npar, nout, nitr; ndrange = nb * nitr) + end + return y +end + +@kernel function kerg_sparse(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + @inbounds ExaModels.grpass( + f(itr[k], ExaModels.AdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), + f.comp1, + view(y, y_off+1:y_off+nout), + ExaModels.offset1(f, k), + 0, + adj, + ) +end + +# --- Jacobian --- + +function ExaModels.sjacobian!(y1, y2, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerj(backend)(y1, y2, f.f, f.itr, x, θ, adj, ExaModels._constraint_dims(f), nvar, npar, nout, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerj(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(dims), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + @inbounds ExaModels.jrpass( + f(itr[k], ExaModels.AdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), + f.comp1, + ExaModels.offset0(f, itr, k, dims), + view(y1, y_off+1:y_off+nout), + y2, + ExaModels.offset1(f, k), + 0, + adj, + ) +end + +# --- Hessian (objective) --- + +function ExaModels.shessian!(y1, y2, f, x, θ, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerh(backend)(y1, y2, f.f, f.itr, x, θ, adj1, adj2, nvar, npar, nout, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerh(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1), @Const(adj2), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + w_s = ExaModels._get_obj_weight(adj1, s) + @inbounds ExaModels.hrpass0( + f(itr[k], ExaModels.SecondAdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), + f.comp2, + view(y1, y_off+1:y_off+nout), + y2, + ExaModels.offset2(f, k), + 0, + w_s, + adj2, + ) +end + +# --- Hessian (constraints) --- + +function ExaModels.shessian!(y1, y2, f, x, θ, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerh2(backend)(y1, y2, f.f, f.itr, x, θ, adj1s, adj2, ExaModels._constraint_dims(f), nvar, npar, ncon, nout, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerh2(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1s), @Const(adj2), @Const(dims), @Const(nvar), @Const(npar), @Const(ncon), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + a_off = (s - 1) * ncon + @inbounds ExaModels.hrpass0( + f(itr[k], ExaModels.SecondAdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), + f.comp2, + view(y1, y_off+1:y_off+nout), + y2, + ExaModels.offset2(f, k), + 0, + adj1s[a_off + ExaModels.offset0(f, itr, k, dims)], + adj2, + ) +end + end # module ExaModelsKernelAbstractions diff --git a/src/ExaModels.jl b/src/ExaModels.jl index 61ac4c85..1c4072e4 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -30,6 +30,7 @@ import Adapt: adapt import NLPModels: NLPModels, obj, + obj!, cons!, grad!, jac_coord!, @@ -77,7 +78,6 @@ export ExaModel, @add_obj, @add_con, @add_con!, - set_parameter!, solution, multipliers, multipliers_L, @@ -107,6 +107,12 @@ export ExaModel, get_lcon, set_lcon!, get_ucon, - set_ucon! + set_ucon!, + FlatNLPModel, + BatchExaCore, + BatchExaModel, + get_nbatch, + var_indices, + cons_block_indices end # module ExaModels diff --git a/src/deprecated.jl b/src/deprecated.jl index a824c19d..d9c85c2d 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -21,16 +21,23 @@ legacy mutating API (`variable`, `constraint`, etc.). `ExaCore(concrete = Val(true))` to obtain the bare immutable `ExaCore` required for AOT compilation. """ -mutable struct LegacyExaCore{T, VT <: AbstractVector{T}, B, S} <: AbstractExaCore{T,VT,B,S} +mutable struct LegacyExaCore{T, VT <: AbstractArray{T}, B, S} <: AbstractExaCore{T,VT,B,S} inner::Any # ExaCore{T,VT,B,S,...} — type erased so the tuple type params can grow end # Override the Val{false} dispatch defined in nlp.jl so ExaCore() returns a LegacyExaCore. -@inline function _make_exacore(::Val{false}, ::Type{T}, backend; kwargs...) where {T} +@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{false}, nbatch; kwargs...) where {T} @warn "`ExaCore()` is deprecated, and will be removed in v0.11. Use `ExaCore(concrete = Val(true))` for the immutable ExaCore. The default behavior for `ExaCore()` will change to return the immutable ExaCore in v0.11." inner = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) return LegacyExaCore{T, typeof(inner.x0), typeof(backend), typeof(inner.tag)}(inner) end +@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{true}, nbatch; kwargs...) where {T} + @warn "`ExaCore()` is deprecated, and will be removed in v0.11. Use `ExaCore(concrete = Val(true))` for the immutable ExaCore. The default behavior for `ExaCore()` will change to return the immutable ExaCore in v0.11." + x0 = convert_array(zeros(T, 0, nbatch), backend) + inner = _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), + y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) + return LegacyExaCore{T, typeof(inner.x0), typeof(backend), typeof(inner.tag)}(inner) +end # --------------------------------------------------------------------------- # Property forwarding for LegacyExaCore @@ -49,9 +56,6 @@ function ExaModel(c::LegacyExaCore; kwargs...) return ExaModel(c.inner; kwargs...) end -function set_parameter!(c::LegacyExaCore, param::Parameter, values::AbstractArray) - return set_parameter!(c.inner, param, values) -end # --------------------------------------------------------------------------- # Legacy named wrappers (deprecated) diff --git a/src/gradient.jl b/src/gradient.jl index 3bf806cf..512ce6ef 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -42,6 +42,17 @@ function gradient!(y, f, x, θ, adj) end return y end +function gradient!(y, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, ::Nothing = nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y_s = @view y[(s-1)*nvar+1 : s*nvar] + @simd for k in eachindex(f.itr) + gradient!(y_s, f.f, x_s, θ_s, f.itr[k], adj) + end + end + return y +end function gradient!(y, f, x, θ, p, adj) graph = f(p, AdjointNodeSource(x), θ) drpass(graph, y, adj) @@ -173,6 +184,17 @@ function sgradient!(y, f, x, θ, adj) end return y end +function sgradient!(y, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y_s = @view y[(s-1)*nout+1 : s*nout] + @simd for k in eachindex(f.itr) + sgradient!(y_s, f.f, f.itr[k], x_s, θ_s, f.itr.comp1, offset1(f, k), adj) + end + end + return y +end function sgradient!(y, f, p, x, θ, comp, o1, adj) graph = f(p, AdjointNodeSource(x), θ) diff --git a/src/hessian.jl b/src/hessian.jl index 64f92a07..2cba7ce4 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -1,3 +1,6 @@ +@inline _get_obj_weight(w::Number, s) = w +@inline _get_obj_weight(w::AbstractVector, s) = @inbounds w[s] + """ hdrpass(t1::T1, t2::T2, comp, y1, y2, o2, cnt, adj) @@ -593,13 +596,13 @@ end @inline function hrpass( t::T, comp, - y1::V, - y2::V, + y1::V1, + y2::V2, o2, cnt, adj, adj2, -) where {T<:SecondAdjointNodeVar,I<:Integer,V<:AbstractVector{I}} +) where {T<:SecondAdjointNodeVar,I<:Integer,V1<:AbstractVector{I},V2<:AbstractVector{I}} ind = o2 + comp(cnt += 1) @inbounds y1[ind] = t.i @inbounds y2[ind] = t.i @@ -623,12 +626,12 @@ end t1::T1, t2::T2, comp, - y1::V, - y2::V, + y1::V1, + y2::V2, o2, cnt, adj, -) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar,I<:Integer,V<:AbstractVector{I}} +) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar,I<:Integer,V1<:AbstractVector{I},V2<:AbstractVector{I}} i, j = t1.i, t2.i ind = o2 + comp(cnt += 1) @inbounds if i >= j @@ -694,6 +697,19 @@ function shessian!(y1, y2, f, x, θ, adj1, adj2) ) end end +function shessian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y1_s = @view y1[(s-1)*nout+1 : s*nout] + w_s = _get_obj_weight(adj1, s) + @simd for k in eachindex(f.itr) + shessian!( + y1_s, y2, f.f, f.itr[k], x_s, θ_s, f.f.comp2, offset2(f, k), w_s, adj2, + ) + end + end +end function shessian!(y1, y2, f, x, θ, adj1s::V, adj2) where {V<:AbstractVector} @simd for k in eachindex(f.itr) @inbounds shessian!( @@ -710,6 +726,20 @@ function shessian!(y1, y2, f, x, θ, adj1s::V, adj2) where {V<:AbstractVector} ) end end +function shessian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, ::Nothing = nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y1_s = @view y1[(s-1)*nout+1 : s*nout] + a_s = @view adj1s[(s-1)*ncon+1 : s*ncon] + @simd for k in eachindex(f.itr) + shessian!( + y1_s, y2, f.f, f.itr[k], x_s, θ_s, f.f.comp2, offset2(f, k), + a_s[offset0(f, k)], adj2, + ) + end + end +end function shessian!(y1, y2, f, p, x, θ, comp, o2, adj1, adj2) graph = f(p, SecondAdjointNodeSource(x), θ) diff --git a/src/jacobian.jl b/src/jacobian.jl index bc60abcb..41df9265 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -70,12 +70,12 @@ end d::D, comp, i, - y1::V, - y2::V, + y1::V1, + y2::V2, o1, cnt, adj, -) where {D<:AdjointNodeVar,I<:Integer,V<:AbstractVector{I}} +) where {D<:AdjointNodeVar,I<:Integer,V1<:AbstractVector{I},V2<:AbstractVector{I}} ind = o1 + comp(cnt += 1) @inbounds y1[ind] = i @inbounds y2[ind] = d.i @@ -125,6 +125,27 @@ function sjacobian!(y1, y2, f, x, θ, adj) ) end end +function sjacobian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y1_s = @view y1[(s-1)*nout+1 : s*nout] + @simd for i in eachindex(f.itr) + sjacobian!( + y1_s, + y2, + f.f, + f.itr[i], + x_s, + θ_s, + f.f.comp1, + offset0(f, i), + offset1(f, i), + adj, + ) + end + end +end function sjacobian!(y1, y2, f, p, x, θ, comp, o0, o1, adj) graph = f(p, AdjointNodeSource(x), θ) diff --git a/src/nlp.jl b/src/nlp.jl index 1eed2917..5f2cc6f8 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -3,9 +3,9 @@ abstract type AbstractParameter end abstract type AbstractConstraint end abstract type AbstractObjective end -struct NaNSource{T} end +struct NaNSource{T} <: AbstractVector{T} end +Base.size(::NaNSource) = (typemax(Int),) Base.getindex(::NaNSource{T}, i) where {T} = T(NaN) -Base.eltype(::NaNSource{T}) where {T} = T Base.eltype(::Type{NaNSource{T}}) where {T} = T """ @@ -73,7 +73,7 @@ end A handle to a block of model parameters added to an [`ExaCore`](@ref) via [`add_par`](@ref) / [`@add_par`](@ref). Parameter values can be updated at any -time with [`set_parameter!`](@ref) without rebuilding the model. Use indexing +time with [`set_value!`](@ref) without rebuilding the model. Use indexing (e.g. `θ[i]`) to embed parameter values in expressions. An optional `tag` field carries user-defined metadata. """ @@ -262,9 +262,12 @@ end abstract type AbstractExaCore{T,VT,B,S} end """ - ExaCore([array_eltype::Type; backend = nothing, minimize = true, name = :Generic]) + ExaCore([T::Type; backend = nothing, concrete = Val(false), batch = Val(false), nbatch = 1, minimize = true, name = :Generic]) -Creates an intermediate data object `ExaCore`, which later can be used for creating an `ExaModel` +Creates an intermediate data object `ExaCore`, which later can be used for creating an `ExaModel`. + +When `batch = Val(true)`, creates a batch core with matrix-valued storage arrays +(`nbatch` columns = instances). See [`BatchExaCore`](@ref) for a convenience alias. ## Example ```jldoctest @@ -303,7 +306,7 @@ An ExaCore number of constraint patterns: ... 0 ``` """ -struct ExaCore{T,VT<:AbstractVector{T}, B, S, V, P, O, C, R} <: AbstractExaCore{T,VT,B,S} +struct ExaCore{T,VT<:AbstractArray{T}, B, S, V, P, O, C, R} <: AbstractExaCore{T,VT,B,S} name::Symbol backend::B var::V @@ -389,15 +392,25 @@ end ) end -@inline ExaCore(::Type{T}; backend = nothing, concrete = Val(false), kwargs...) where {T<:AbstractFloat} = - _make_exacore(concrete, T, backend; kwargs...) -@inline ExaCore(; backend = nothing, concrete = Val(false), kwargs...) = ExaCore(default_T(backend); backend, concrete, kwargs...) -@inline _make_exacore(::Val{true}, ::Type{T}, backend; kwargs...) where {T} = +@inline ExaCore(::Type{T}; backend = nothing, concrete = Val(false), batch = Val(false), nbatch = 1, kwargs...) where {T<:AbstractFloat} = + _make_exacore(concrete, T, backend, batch, nbatch; kwargs...) +@inline ExaCore(; backend = nothing, concrete = Val(false), batch = Val(false), nbatch = 1, kwargs...) = ExaCore(default_T(backend); backend, concrete, batch, nbatch, kwargs...) +@inline _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{false}, nbatch; kwargs...) where {T} = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) +@inline function _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{true}, nbatch; kwargs...) where {T} + x0 = convert_array(zeros(T, 0, nbatch), backend) + _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), + y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) +end # Val{false} is overridden in deprecated.jl once LegacyExaCore is defined; # this fallback handles any other Val value by returning a concrete ExaCore. -@inline _make_exacore(::Val, ::Type{T}, backend; kwargs...) where {T} = +@inline _make_exacore(::Val, ::Type{T}, backend, ::Val{false}, nbatch; kwargs...) where {T} = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) +@inline function _make_exacore(::Val, ::Type{T}, backend, ::Val{true}, nbatch; kwargs...) where {T} + x0 = convert_array(zeros(T, 0, nbatch), backend) + _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), + y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) +end @inline ExaCore(c::C; kwargs...) where C <: ExaCore = _exa_core( ; zip(fieldnames(C), ntuple(i -> getfield(c, i), Val(fieldcount(C))))..., @@ -420,30 +433,36 @@ An ExaCore """, ) -""" - AbstractExaModel - -An abstract type for ExaModel, which is a subtype of `NLPModels.AbstractNLPModel`. -""" -abstract type AbstractExaModel{T,VT,E} <: NLPModels.AbstractNLPModel{T,VT} end - -struct ExaModel{T,VT,E,V,P,O,C,S,R} <: AbstractExaModel{T,VT,E} +struct ExaModel{T,VT,E,V,P,O,C,S,R,M} <: NLPModels.AbstractNLPModel{T,VT} name::Symbol vars::V pars::P objs::O cons::C θ::VT - meta::NLPModels.NLPModelMeta{T,VT} + meta::M counters::NLPModels.Counters ext::E tag::S refs::R end -function Base.show(io::IO, c::AbstractExaModel{T,VT}) where {T,VT} - println(io, "An ExaModel{$T, $VT, ...}\n") - Base.show(io, c.meta) +function ExaModel( + name::Symbol, vars, pars, objs, cons, θ::VT, + meta::M, + counters, ext, tag, refs, +) where {T, VT <: AbstractArray{T}, M <: NLPModels.AbstractNLPModelMeta{T}} + ExaModel{T, VT, typeof(ext), typeof(vars), typeof(pars), typeof(objs), + typeof(cons), typeof(tag), typeof(refs), M}( + name, vars, pars, objs, cons, θ, meta, counters, ext, tag, refs, + ) +end + +function Base.show(io::IO, m::ExaModel{T,VT}) where {T,VT} + nb = get_nbatch(m) + batch_str = nb > 1 ? " (batch, $nb instances)" : "" + println(io, "An ExaModel{$T, $VT, ...}$batch_str\n") + Base.show(io, m.meta) end """ @@ -488,6 +507,7 @@ julia> result = ipopt(m; print_level=0) # solve the problem ``` """ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} + nvar, ncon, nnzj, nnzh = _meta_dims(c) return ExaModel( c.name, c.var, @@ -495,20 +515,19 @@ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} c.obj, c.cons, c.θ, - NLPModels.NLPModelMeta( - c.nvar, - ncon = c.ncon, - nnzj = c.nnzj, - nnzh = c.nnzh, - x0 = (c.x0), - lvar = (c.lvar), - uvar = (c.uvar), - y0 = (c.y0), - lcon = (c.lcon), - ucon = (c.ucon), + _build_meta( + nvar, + c.x0, + c.lvar, + c.uvar, + ncon, + c.y0, + c.lcon, + c.ucon; + nnzj = nnzj, + nnzh = nnzh, minimize = c.minimize, - ) - , + ), NLPModels.Counters(), build_extension(c; prod), c.tag, @@ -516,8 +535,60 @@ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} ) end +_meta_dims(c::ExaCore) = (c.nvar, c.ncon, c.nnzj, c.nnzh) build_extension(c::ExaCore; kwargs...) = nothing +# ============================================================================ +# _build_meta: construct NLPModelMeta supporting both Vector and Matrix VT +# ============================================================================ + +_first_instance(v::AbstractVector) = v +_first_instance(m::AbstractMatrix) = @view m[:, 1] + +function _classify_bounds(lb, ub, ::Type{T}) where {T} + lb_cpu = Array(lb) + ub_cpu = Array(ub) + ifix = findall(lb_cpu .== ub_cpu) + ilow = findall((lb_cpu .> T(-Inf)) .& (ub_cpu .== T(Inf))) + iupp = findall((lb_cpu .== T(-Inf)) .& (ub_cpu .< T(Inf))) + irng = findall((lb_cpu .> T(-Inf)) .& (ub_cpu .< T(Inf)) .& (lb_cpu .< ub_cpu)) + ifree = findall((lb_cpu .== T(-Inf)) .& (ub_cpu .== T(Inf))) + iinf = findall(lb_cpu .> ub_cpu) + return ifix, ilow, iupp, irng, ifree, iinf +end + +function _build_meta( + nvar::Int, x0::VT, lvar::VT, uvar::VT, + ncon::Int, y0::VT, lcon::VT, ucon::VT; + nnzj::Int = nvar * ncon, + nnzh::Int = nvar * (nvar + 1) ÷ 2, + minimize::Bool = true, + islp::Bool = false, + name::String = "Generic", +) where {VT} + T = eltype(VT) + ifix, ilow, iupp, irng, ifree, iinf = _classify_bounds( + _first_instance(lvar), _first_instance(uvar), T) + if ncon > 0 + jfix, jlow, jupp, jrng, jfree, jinf = _classify_bounds( + _first_instance(lcon), _first_instance(ucon), T) + else + jfix = jlow = jupp = jrng = jfree = jinf = Int[] + end + nln = collect(1:ncon) + return NLPModels.NLPModelMeta{T, VT}( + nvar, x0, lvar, uvar, + ifix, ilow, iupp, irng, ifree, iinf, + nvar, nvar, nvar, + ncon, y0, lcon, ucon, + jfix, jlow, jupp, jrng, jfree, jinf, + nvar, nnzj, 0, nnzj, nnzh, + 0, ncon, Int[], nln, + minimize, islp, name, + true, true, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, + ) +end + @inline function Base.getindex(v::V, i) where {V<:AbstractVariable} _bound_check(v.size, i) _indexed_var(i, v.offset - _start(v.size[1]) + 1) @@ -589,41 +660,49 @@ function __bound_check(a::UnitRange{Int}, b::I) where {I<:Integer} end -function append!(backend, a, b::Base.Generator, lb) - if lb != 0 - b = _adapt_gen(b) - la = length(a) - resize!(a, la + lb) - map!(b.f, view(a, (la+1):(la+lb)), convert_array(b.iter, backend)) - end - return a +# Unified append! — works for both Vector and Matrix. +# For Vector: grows length by lb. +# For Matrix: grows rows by lb, broadcasting across columns. +# The trailing dimensions (columns for Matrix, nothing for Vector) are preserved. + +@inline _trailing_dims(a) = Base.size(a)[2:end] + +function _expand_to_shape(col::AbstractVector{T}, trailing::Tuple{}) where {T} + return col # Vector target — no expansion needed +end +function _expand_to_shape(col::AbstractVector{T}, trailing::Tuple) where {T} + return repeat(reshape(col, :, ntuple(_ -> 1, length(trailing))...), 1, trailing...) end -function append!(backend, a, b::Base.Generator{UnitRange{I}}, lb) where {I} - if lb != 0 - la = length(a) - resize!(a, la + lb) - map!(b.f, view(a, (la+1):(la+lb)), b.iter) - end - return a +# Reshape an array to have the given trailing dimensions. +# For vectors, delegates to _expand_to_shape. For matrices/higher-dim arrays, +# reshapes using total elements / trailing to infer the first dimension. +_reshape_to_match(arr::AbstractVector, trailing::Tuple{}) = arr +_reshape_to_match(arr::AbstractVector, trailing::Tuple) = _expand_to_shape(arr, trailing) +_reshape_to_match(arr::AbstractArray, trailing::Tuple{}) = vec(arr) +function _reshape_to_match(arr::AbstractArray, trailing::Tuple) + n = length(arr) ÷ prod(trailing) + return reshape(arr, n, trailing...) +end + +function append!(backend, a, b::Number, lb) + lb == 0 && return a + new_part = fill(eltype(a)(b), lb, _trailing_dims(a)...) + return cat(a, new_part; dims = 1) end function append!(backend, a, b::AbstractArray, lb) - if lb != 0 - la = length(a) - resize!(a, la + lb) - map!(identity, view(a, (la+1):(la+lb)), convert_array(b, backend)) - end - return a + lb == 0 && return a + arr = convert_array(b, backend) + return cat(a, _reshape_to_match(arr, _trailing_dims(a)); dims = 1) end -function append!(backend, a, b::Number, lb) - if lb != 0 - la = length(a) - resize!(a, la + lb) - fill!(view(a, (la+1):(la+lb)), eltype(a)(b)) - end - return a +function append!(backend, a, b::Base.Generator, lb) + lb == 0 && return a + b = _adapt_gen(b) + col = Vector{eltype(a)}(undef, lb) + map!(b.f, col, convert_array(b.iter, backend)) + return cat(a, _expand_to_shape(col, _trailing_dims(a)); dims = 1) end @inline total(ns) = _total(ns...) @@ -684,7 +763,6 @@ Variable end @inline function _add_var(c, tag, name, start, lvar, uvar, ns...) - o = c.nvar len = total(ns) nvar = c.nvar + len @@ -692,13 +770,17 @@ end lvar = append!(c.backend, c.lvar, lvar, len) uvar = append!(c.backend, c.uvar, uvar, len) - v = Variable(ns, len, o, _val_name(name), tag) + v = Variable(ns, len, c.nvar, _val_name(name), tag) (ExaCore(c; var = (v, c.var...), nvar=nvar, x0=x0, lvar=lvar, uvar=uvar, refs = add_refs(c.refs, name, v)), v) end @inline _val_name(::Val{N}) where {N} = N @inline _val_name(::Nothing) = :x +@inline get_nbatch(c::ExaCore{T, <:AbstractMatrix}) where {T} = Base.size(c.x0, 2) +@inline get_nbatch(::ExaCore) = 1 +@inline get_nbatch(m::ExaModel{T, <:AbstractMatrix}) where {T} = Base.size(m.meta.x0, 2) +@inline get_nbatch(::ExaModel) = 1 @inline add_refs(refs, ::Nothing, var) = refs @inline add_refs(refs, ::Val{N}, var) where {N} = (; refs..., N => var) @@ -745,46 +827,13 @@ end end @inline function _add_par(c, tag, name, start, ns...) - o = c.npar len = total(ns) npar = c.npar + len θ = append!(c.backend, c.θ, start, len) - p = Parameter(ns, len, o, tag) + p = Parameter(ns, len, c.npar, tag) (ExaCore(c; par = (p, c.par...), θ=θ, npar=npar, refs = add_refs(c.refs, name, p)), p) end -""" - set_parameter!(core, param, values) - -Updates the values of parameters in the core. - -## Example -```jldoctest -julia> using ExaModels - -julia> c = ExaCore(concrete = Val(true)); - -julia> c, p = add_par(c, ones(5)); - -julia> set_parameter!(c, p, ones(5)) -``` -""" -function set_parameter!(c::ExaCore, param::Parameter, values::AbstractArray) - if length(values) != param.length - throw( - DimensionMismatch( - "Parameter size mismatch: expected $(param.length) elements, got $(length(values))", - ), - ) - end - - start_idx = param.offset + 1 - end_idx = param.offset + param.length - - copyto!(@view(c.θ[start_idx:end_idx]), values) - - return nothing -end """ get_value(model, param) @@ -803,66 +852,58 @@ end Update all values for `param` in `model.θ` to `values`. """ -function set_value!(model::ExaModel, param::Parameter, values) +function set_value!(model::ExaModel{T, <:AbstractVector}, param::Parameter, values) where {T} if length(values) != param.length throw(DimensionMismatch( "expected $(param.length) elements, got $(length(values))" )) end - copyto!(view(model.θ, param.offset+1:param.offset+param.length), values) + rng = param.offset+1 : param.offset+param.length + copyto!(view(model.θ, rng), values) + return nothing +end + +function set_value!(model::ExaModel{T, <:AbstractMatrix}, param::Parameter, values) where {T} + if length(values) != param.length + throw(DimensionMismatch( + "expected $(param.length) elements, got $(length(values))" + )) + end + rng = param.offset+1 : param.offset+param.length + for col in axes(model.θ, 2) + copyto!(view(model.θ, rng, col), values) + end + return nothing +end + +function set_value!(model::ExaModel{T, <:AbstractMatrix{T}}, param::Parameter, values::AbstractMatrix) where {T} + nb = get_nbatch(model) + if Base.size(values, 1) != param.length || Base.size(values, 2) != nb + throw(DimensionMismatch( + "expected $(param.length) × $nb matrix, got $(Base.size(values))" + )) + end + rng = param.offset+1 : param.offset+param.length + for col in 1:nb + copyto!(view(model.θ, rng, col), view(values, :, col)) + end return nothing end @inline _var_range(v::Variable) = v.offset+1 : v.offset+v.length @inline _con_range(c::Constraint) = c.offset+1 : c.offset+total(c.size) -""" - get_start(model, var::Variable) - get_start(model, con::Constraint) - -Return a view of the initial-point values (`x0` for variables, `y0` for constraints). -""" get_start(model::ExaModel, v::Variable) = view(model.meta.x0, _var_range(v)) get_start(model::ExaModel, c::Constraint) = view(model.meta.y0, _con_range(c)) - -""" - get_lvar(model, var::Variable) - -Return a view of the lower bounds for `var`. -""" get_lvar(model::ExaModel, v::Variable) = view(model.meta.lvar, _var_range(v)) - -""" - get_uvar(model, var::Variable) - -Return a view of the upper bounds for `var`. -""" get_uvar(model::ExaModel, v::Variable) = view(model.meta.uvar, _var_range(v)) - -""" - get_lcon(model, con::Constraint) - -Return a view of the lower bounds for `con`. -""" get_lcon(model::ExaModel, c::Constraint) = view(model.meta.lcon, _con_range(c)) - -""" - get_ucon(model, con::Constraint) - -Return a view of the upper bounds for `con`. -""" get_ucon(model::ExaModel, c::Constraint) = view(model.meta.ucon, _con_range(c)) @inline function _check_len(got, expected, label) got == expected || throw(DimensionMismatch("$label: expected $expected elements, got $got")) end -""" - set_start!(model, var::Variable, values) - set_start!(model, con::Constraint, values) - -Update the initial-point values in-place (`x0` for variables, `y0` for constraints). -""" function set_start!(model::ExaModel, v::Variable, values) _check_len(length(values), v.length, "set_start!") copyto!(view(model.meta.x0, _var_range(v)), values) @@ -872,43 +913,19 @@ function set_start!(model::ExaModel, c::Constraint, values) _check_len(length(values), n, "set_start!") copyto!(view(model.meta.y0, _con_range(c)), values) end - -""" - set_lvar!(model, var::Variable, values) - -Update the lower bounds for `var` in-place. -""" function set_lvar!(model::ExaModel, v::Variable, values) _check_len(length(values), v.length, "set_lvar!") copyto!(view(model.meta.lvar, _var_range(v)), values) end - -""" - set_uvar!(model, var::Variable, values) - -Update the upper bounds for `var` in-place. -""" function set_uvar!(model::ExaModel, v::Variable, values) _check_len(length(values), v.length, "set_uvar!") copyto!(view(model.meta.uvar, _var_range(v)), values) end - -""" - set_lcon!(model, con::Constraint, values) - -Update the lower bounds for `con` in-place. -""" function set_lcon!(model::ExaModel, c::Constraint, values) n = total(c.size) _check_len(length(values), n, "set_lcon!") copyto!(view(model.meta.lcon, _con_range(c)), values) end - -""" - set_ucon!(model, con::Constraint, values) - -Update the upper bounds for `con` in-place. -""" function set_ucon!(model::ExaModel, c::Constraint, values) n = total(c.size) _check_len(length(values), n, "set_ucon!") @@ -995,7 +1012,7 @@ is intended for code that builds expression trees programmatically. @inline function add_obj(c::C, expr::N, pars = 1:1; name = nothing) where {T,C<:ExaCore{T},N<:AbstractNode} f = _simdfunction(T, expr, c.nobj, c.nnzg, c.nnzh) - _add_obj(c, f, pars, name) + _add_obj(c, f, pars, name) end @inline function _add_obj(c, f, pars, name = nothing) @@ -1075,20 +1092,19 @@ Constraint @inline function add_con( c::C, ns...; - tag = nothing, + tag = nothing, name = nothing, start = zero(T), lcon = zero(T), ucon = zero(T), kwargs... ) where {T,C<:ExaCore{T}} - gen = _get_generator(ns) dims = _get_con_dims(ns) gen = _adapt_gen(gen) f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter - + _add_con(c, f, pars, dims, start, lcon, ucon, name, tag) end @@ -1198,6 +1214,11 @@ add_con!(c, bus, gen.bus => -pg[gen.i] for gen in data.gen) # subtrac _add_con!(c, f, pars, _constraint_dims(c1), tag) end +_probe_conaug(probe::ConAugPair) = probe.con +_probe_conaug(probe) = error( + "add_con! two-argument form requires `constraint[idx] += expr` syntax in the generator body" +) + """ add_con!(core::ExaCore, gen::Base.Generator; tag = nothing) @@ -1221,11 +1242,7 @@ c, _ = add_con!(c, g[i] += x[i] + x[i+1] for i = 1:9) # Probe the generator: the result is a ConAugPair which carries the target # constraint alongside the index/expression pair. probe = gen.f(DataSource()) - probe isa ConAugPair || error( - "add_con! two-argument form requires `constraint[idx] += expr` syntax " * - "in the generator body" - ) - con = probe.con + con = _probe_conaug(probe) # The generator yields ConAugPair values; replace_T unwraps them to plain # Pairs before SIMDFunction stores them, so offset0 dispatch is unchanged. @@ -1305,8 +1322,10 @@ c, s = add_expr(c, x[i, k]^2 for (i, k) in itr) return (ExaCore(c; refs = add_refs(c.refs, name, ex)), ex) end -function jac_structure!(m::AbstractExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T - _jac_structure!(T, m.cons, rows, cols) +function jac_structure!(m::ExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T + if !isempty(rows) + _jac_structure!(T, m.cons, rows, cols) + end return rows, cols end @@ -1316,9 +1335,11 @@ _jac_structure!(T, cons::Tuple{}, rows, cols) = nothing sjacobian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN)) end -function hess_structure!(m::AbstractExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T - _obj_hess_structure!(T, m.objs, rows, cols) - _con_hess_structure!(T, m.cons, rows, cols) +function hess_structure!(m::ExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T + if !isempty(rows) + _obj_hess_structure!(T, m.objs, rows, cols) + _con_hess_structure!(T, m.cons, rows, cols) + end return rows, cols end @@ -1334,62 +1355,140 @@ _con_hess_structure!(T, cons::Tuple{}, rows, cols) = nothing shessian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN)) end -function obj(m::AbstractExaModel, x::AbstractVector) - return _obj(m.objs, x, m.θ) +# ============================================================================ +# Batch-aware evaluation — all low-level functions loop over 1:nb, +# striding x/θ/g/etc. by per-instance sizes. For nb=1, the loop runs +# once with views of the full arrays (no overhead). +# ============================================================================ + +function NLPModels.obj(m::ExaModel{T, <:AbstractVector, Nothing}, x::AbstractVector) where {T} + return _obj_scalar(m.objs, x, m.θ) +end +function obj(m::ExaModel{T}, x::AbstractVector) where {T} + bf = similar(x, T, 1) + obj!(m, x, bf) + return @inbounds bf[1] end -@inline function _obj((obj, objs...), x, θ) - s = _obj(objs, x, θ) - for i in obj.itr +@inline function _obj_scalar((obj, objs...), x, θ) + s = _obj_scalar(objs, x, θ) + @inbounds for i in obj.itr s += obj.f(i, x, θ) end return s end +@inline _obj_scalar(::Tuple{}, x, θ) = zero(eltype(x)) + +@inline function _obj!(bf, (obj, objs...), x, θ, nb, nvar, npar, backend = nothing) + _obj!(bf, objs, x, θ, nb, nvar, npar, backend) + _obj_eval!(bf, obj, x, θ, nb, nvar, npar, backend) +end +@inline _obj!(bf, ::Tuple{}, x, θ, nb, nvar, npar, backend = nothing) = nothing -@inline _obj(obj::Tuple{}, x, θ) = zero(eltype(x)) +@inline function _obj_eval!(bf, obj, x, θ, nb, nvar, npar, ::Nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + for i in obj.itr + bf[s] += obj.f(i, x_s, θ_s) + end + end +end -function cons_nln!(m::AbstractExaModel, x::AbstractVector, g::AbstractVector) - fill!(g, zero(eltype(g))) +function NLPModels.cons_nln!(m::ExaModel{T, <:AbstractVector, Nothing}, x::AbstractVector, g::AbstractVector) where {T} + fill!(g, zero(T)) _cons_nln!(m.cons, x, m.θ, g) return g end +function NLPModels.cons_nln!(m::ExaModel{T}, bx::AbstractVecOrMat, bc::AbstractVecOrMat) where {T} + fill!(vec(bc), zero(T)) + nb = Base.size(bx, 2) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + npar = Base.size(m.θ, 1) + _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) + return bc +end +_cons_nln!(cons::Tuple{}, x, θ, g) = nothing @inline function _cons_nln!(cons::Tuple, x, θ, g) con = first(cons) _cons_nln!(Base.tail(cons), x, θ, g) @simd for i in eachindex(con.itr) - g[offset0(con, i)] += con.f(con.itr[i], x, θ) + @inbounds g[offset0(con, i)] += con.f(con.itr[i], x, θ) end end -_cons_nln!(cons::Tuple{}, x, θ, g) = nothing - +@inline function _cons_nln!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) + _cons_nln!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) + _cons_nln_eval!(first(cons), x, θ, g, nb, nvar, npar, ncon, backend) +end +_cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing +@inline function _cons_nln_eval!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + g_s = @view g[(s-1)*ncon+1 : s*ncon] + @simd for i in eachindex(con.itr) + g_s[offset0(con, i)] += con.f(con.itr[i], x_s, θ_s) + end + end +end -function grad!(m::AbstractExaModel, x::AbstractVector, f::AbstractVector) - fill!(f, zero(eltype(f))) +function NLPModels.grad!(m::ExaModel{T, <:AbstractVector, Nothing}, x::AbstractVector, f::AbstractVector) where {T} + fill!(f, zero(T)) _grad!(m.objs, x, m.θ, f) return f end +function NLPModels.grad!(m::ExaModel{T}, bx::AbstractVecOrMat, bg::AbstractVecOrMat) where {T} + fill!(vec(bg), zero(T)) + nb = Base.size(bx, 2) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + _grad!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) + return bg +end +_grad!(objs::Tuple{}, x, θ, f) = nothing @inline function _grad!(objs::Tuple, x, θ, f) _grad!(Base.tail(objs), x, θ, f) gradient!(f, first(objs), x, θ, one(eltype(f))) end -_grad!(objs::Tuple{}, x, θ, f) = nothing +@inline function _grad!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) + _grad!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) + gradient!(f, first(objs), x, θ, one(eltype(f)), nb, nvar, npar, backend) +end +_grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing -function jac_coord!(m::AbstractExaModel, x::AbstractVector, jac::AbstractVector) - fill!(jac, zero(eltype(jac))) - _jac_coord!(m.cons, x, m.θ, jac) - return jac +@inline function _jac_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, jvals) where {T} + fill!(jvals, zero(T)) + _jac_coord!(m.cons, x, m.θ, jvals) + return jvals end +@inline function _jac_coord_impl!(m::ExaModel{T}, bx, jvals) where {T} + fill!(vec(jvals), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + nnzj = NLPModels.get_nnzj(m) + _jac_coord!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) + return jvals +end +NLPModels.jac_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, jvals::AbstractVecOrMat) where {T} = _jac_coord_impl!(m, bx, jvals) +NLPModels.jac_coord!(m::ExaModel{T}, x::AbstractVector, jac::AbstractVector) where {T} = _jac_coord_impl!(m, x, jac) _jac_coord!(cons::Tuple{}, x, θ, jac) = nothing @inline function _jac_coord!(cons::Tuple, x, θ, jac) _jac_coord!(Base.tail(cons), x, θ, jac) sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac))) end +_jac_coord!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing +@inline function _jac_coord!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) + _jac_coord!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) + sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac)), nb, nvar, npar, nnzj, backend) +end -function jprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) +function jprod_nln!(m::ExaModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) fill!(Jv, zero(eltype(Jv))) _jprod_nln!(m.cons, x, m.θ, v, Jv) return Jv @@ -1401,7 +1500,7 @@ _jprod_nln!(cons::Tuple{}, x, θ, v, Jv) = nothing sjacobian!((Jv, v), nothing, first(cons), x, θ, one(eltype(Jv))) end -function jtprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) +function jtprod_nln!(m::ExaModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) fill!(Jtv, zero(eltype(Jtv))) _jtprod_nln!(m.cons, x, m.θ, v, Jtv) return Jtv @@ -1413,44 +1512,71 @@ _jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv) = nothing sjacobian!(nothing, (Jtv, v), first(cons), x, θ, one(eltype(Jtv))) end -function hess_coord!( - m::AbstractExaModel, - x::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) - return hess -end - -function hess_coord!( - m::AbstractExaModel, - x::AbstractVector, - y::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) - _con_hess_coord!(m.cons, x, m.θ, y, hess, obj_weight) - return hess -end +_as_weight(::Type{T}, w::Number) where {T} = T(w) +_as_weight(::Type{T}, w) where {T} = w + +@inline function _obj_hess_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, hvals, obj_weight) where {T} + fill!(hvals, zero(T)) + _obj_hess_coord!(m.objs, x, m.θ, hvals, _as_weight(T, obj_weight)) + return hvals +end +@inline function _obj_hess_coord_impl!(m::ExaModel{T}, bx, hvals, obj_weight) where {T} + fill!(vec(hvals), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + nnzh = NLPModels.get_nnzh(m) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), _as_weight(T, obj_weight), nb, nvar, npar, nnzh, getbackend(m)) + return hvals +end +NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, bx, hvals, obj_weight) +NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, x, hess, obj_weight) + +@inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, y, hvals, obj_weight) where {T} + fill!(hvals, zero(T)) + _obj_hess_coord!(m.objs, x, m.θ, hvals, _as_weight(T, obj_weight)) + _con_hess_coord!(m.cons, x, m.θ, y, hvals) + return hvals +end +@inline function _hess_coord_impl!(m::ExaModel{T}, bx, by, hvals, obj_weight) where {T} + fill!(vec(hvals), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + npar = Base.size(m.θ, 1) + nnzh = NLPModels.get_nnzh(m) + backend = getbackend(m) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), _as_weight(T, obj_weight), nb, nvar, npar, nnzh, backend) + _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) + return hvals +end +NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, by::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, bx, by, hvals, obj_weight) +NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, y::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, x, y, hess, obj_weight) _obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing @inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight) shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess))) end +_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing +@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) + _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) + shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess)), nb, nvar, npar, nnzh, backend) +end -_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, obj_weight) = nothing -@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, obj_weight) - _con_hess_coord!(Base.tail(cons), x, θ, y, hess, obj_weight) +_con_hess_coord!(cons::Tuple{}, x, θ, y, hess) = nothing +@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess) + _con_hess_coord!(Base.tail(cons), x, θ, y, hess) shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess))) end +_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing +@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) + _con_hess_coord!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) + shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess)), nb, nvar, npar, ncon, nnzh, backend) +end function hprod!( - m::AbstractExaModel, + m::ExaModel, x::AbstractVector, v::AbstractVector, Hv::AbstractVector; @@ -1462,7 +1588,7 @@ function hprod!( end function hprod!( - m::AbstractExaModel, + m::ExaModel, x::AbstractVector, y::AbstractVector, v::AbstractVector, @@ -1675,6 +1801,174 @@ end _adapt_gen(gen) = Base.Generator(gen.f, collect(gen.iter)) _adapt_gen(gen::Base.Generator{P}) where {P<:Union{AbstractArray,AbstractRange}} = gen +# ============================================================================ +# Batch ExaModel — dispatch on VT <: AbstractMatrix +# ============================================================================ +# +# BatchExaCore / BatchExaModel are determined by VT <: AbstractMatrix. +# nbatch is derived from size(x0, 2). +# +# All evaluation is handled by the batch-aware low-level functions above +# (_obj, _grad!, _cons_nln!, etc.) which loop over 1:nb with strided views. +# This section provides type aliases, the BatchExaCore constructor, +# and thin batch API wrappers (matrix-argument dispatch). + +# ============================================================================ +# Type aliases — determined by VT <: AbstractMatrix +# ============================================================================ + +""" + BatchExaCore{T,VT,B} + +Type alias for an [`ExaCore`](@ref) whose storage arrays are matrices +(columns = instances). +""" +const BatchExaCore{T,VT<:AbstractMatrix{T},B} = ExaCore{T,VT,B} + +""" + BatchExaModel{T,VT,E,V,P,O,C,S,R,M} + +Type alias for an [`ExaModel`](@ref) built from a [`BatchExaCore`](@ref). +""" +const BatchExaModel{T,VT<:AbstractMatrix{T},E,V,P,O,C,S,R,M} = ExaModel{T,VT,E,V,P,O,C,S,R,M} + +# ============================================================================ +# BatchExaCore constructor — alias for ExaCore with nbatch +# ============================================================================ + +""" + BatchExaCore(nbatch; kwargs...) + +Alias for `ExaCore(; concrete = Val(true), batch = Val(true), nbatch, kwargs...)`. + +Creates an [`ExaCore`](@ref) for building batch optimization models with +`nbatch` independent instances. Generators should iterate over per-instance +dimensions only — the batch dimension is handled automatically at evaluation +time by striding through the data. + +## Example +```julia +core = BatchExaCore(3) +c, v = add_var(core, 10; start = 1.0, lvar = 0.0, uvar = 10.0) +c, _ = add_obj(c, v[i]^2 for i in 1:10) +model = ExaModel(c) +``` +""" +BatchExaCore(nbatch::Integer; kwargs...) = ExaCore(; concrete = Val(true), batch = Val(true), nbatch, kwargs...) + + +""" + var_indices(model, i) -> UnitRange + +Variable index range for instance `i` in the fused model's global variable vector. +""" +var_indices(model::BatchExaModel, i::Int) = + ((i - 1) * NLPModels.get_nvar(model) + 1):(i * NLPModels.get_nvar(model)) + +""" + cons_block_indices(model, i) -> UnitRange + +Constraint index range for instance `i` in the fused model's global constraint vector. +""" +cons_block_indices(model::BatchExaModel, i::Int) = + ((i - 1) * NLPModels.get_ncon(model) + 1):(i * NLPModels.get_ncon(model)) + +# ============================================================================ +# Batch getters / setters +# ============================================================================ + +get_start(model::BatchExaModel, v::Variable) = view(model.meta.x0, _var_range(v), :) +get_start(model::BatchExaModel, c::Constraint) = view(model.meta.y0, _con_range(c), :) +get_start(model::BatchExaModel, v::Variable, i::Int) = view(model.meta.x0, _var_range(v), i) +get_start(model::BatchExaModel, c::Constraint, i::Int) = view(model.meta.y0, _con_range(c), i) +get_lvar(model::BatchExaModel, v::Variable) = view(model.meta.lvar, _var_range(v), :) +get_lvar(model::BatchExaModel, v::Variable, i::Int) = view(model.meta.lvar, _var_range(v), i) +get_uvar(model::BatchExaModel, v::Variable) = view(model.meta.uvar, _var_range(v), :) +get_uvar(model::BatchExaModel, v::Variable, i::Int) = view(model.meta.uvar, _var_range(v), i) +get_lcon(model::BatchExaModel, c::Constraint) = view(model.meta.lcon, _con_range(c), :) +get_lcon(model::BatchExaModel, c::Constraint, i::Int) = view(model.meta.lcon, _con_range(c), i) +get_ucon(model::BatchExaModel, c::Constraint) = view(model.meta.ucon, _con_range(c), :) +get_ucon(model::BatchExaModel, c::Constraint, i::Int) = view(model.meta.ucon, _con_range(c), i) + +function set_start!(model::BatchExaModel, v::Variable, values) + copyto!(view(model.meta.x0, _var_range(v), :), values) +end +function set_start!(model::BatchExaModel, c::Constraint, values) + copyto!(view(model.meta.y0, _con_range(c), :), values) +end +function set_start!(model::BatchExaModel, v::Variable, values, i::Int) + _check_len(length(values), v.length, "set_start!") + copyto!(view(model.meta.x0, _var_range(v), i), values) +end +function set_start!(model::BatchExaModel, c::Constraint, values, i::Int) + n = total(c.size) + _check_len(length(values), n, "set_start!") + copyto!(view(model.meta.y0, _con_range(c), i), values) +end +function set_lvar!(model::BatchExaModel, v::Variable, values) + copyto!(view(model.meta.lvar, _var_range(v), :), values) +end +function set_lvar!(model::BatchExaModel, v::Variable, values, i::Int) + _check_len(length(values), v.length, "set_lvar!") + copyto!(view(model.meta.lvar, _var_range(v), i), values) +end +function set_uvar!(model::BatchExaModel, v::Variable, values) + copyto!(view(model.meta.uvar, _var_range(v), :), values) +end +function set_uvar!(model::BatchExaModel, v::Variable, values, i::Int) + _check_len(length(values), v.length, "set_uvar!") + copyto!(view(model.meta.uvar, _var_range(v), i), values) +end +function set_lcon!(model::BatchExaModel, c::Constraint, values) + copyto!(view(model.meta.lcon, _con_range(c), :), values) +end +function set_lcon!(model::BatchExaModel, c::Constraint, values, i::Int) + n = total(c.size) + _check_len(length(values), n, "set_lcon!") + copyto!(view(model.meta.lcon, _con_range(c), i), values) +end +function set_ucon!(model::BatchExaModel, c::Constraint, values) + copyto!(view(model.meta.ucon, _con_range(c), :), values) +end +function set_ucon!(model::BatchExaModel, c::Constraint, values, i::Int) + n = total(c.size) + _check_len(length(values), n, "set_ucon!") + copyto!(view(model.meta.ucon, _con_range(c), i), values) +end + +# ============================================================================ +# Batch API: matrix-argument wrappers +# These delegate to the unified batch-aware functions above. +# ============================================================================ + +function obj!(m::ExaModel{T}, bx::AbstractVecOrMat, bf::AbstractVector) where {T} + fill!(bf, zero(T)) + nb = Base.size(bx, 2) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + _obj!(bf, m.objs, vec(bx), vec(m.θ), nb, nvar, npar, getbackend(m)) + return bf +end + +function obj(m::ExaModel{T}, bx::AbstractMatrix) where {T} + bf = similar(bx, T, Base.size(bx, 2)) + obj!(m, bx, bf) + return bf +end + +NLPModels.cons!(m::ExaModel{T}, bx::AbstractMatrix, bc::AbstractMatrix) where {T} = + NLPModels.cons_nln!(m, bx, bc) + +_batch_vector_error() = throw(ArgumentError("batch model requires matrix input; got vector")) + +obj(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector) where {T} = _batch_vector_error() +NLPModels.grad!(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector, g::AbstractVector) where {T} = _batch_vector_error() +NLPModels.cons!(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector, c::AbstractVector) where {T} = _batch_vector_error() +NLPModels.cons_nln!(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector, c::AbstractVector) where {T} = _batch_vector_error() + + + + function Base.getproperty(core::E, name::Symbol) where {E <: Union{ExaCore, ExaModel}} if hasfield(E, name) getfield(core,name) diff --git a/src/two_stage.jl b/src/two_stage.jl index 03bed951..40090022 100644 --- a/src/two_stage.jl +++ b/src/two_stage.jl @@ -60,7 +60,7 @@ const TwoStageExaModel{T,VT,E,V,P,O,C,R} = ExaModel{T,VT,E,V,P,O,C,<:TwoStageExa """ - TwoStageExaCore(nscen; backend = nothing, concrete = Val(false), kwargs...) + TwoStageExaCore(nscen; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) Create an [`ExaCore`](@ref) for building two-stage stochastic programs with `nscen` scenarios. @@ -69,6 +69,9 @@ Use [`add_var`](@ref), [`add_par`](@ref), and [`add_con`](@ref) with [`EachScenario()`](@ref) to declare per-scenario components, or without it for first-stage (design) components. +When `nbatch = Val(N)` with `N > 1`, creates a batched two-stage core that +combines batch optimization with two-stage structure. + ## Example ```julia core = TwoStageExaCore(5) # 5 scenarios @@ -77,10 +80,16 @@ c, v = add_var(c, EachScenario(), 2) # 2 recourse variables per scenario model = ExaModel(c) ``` """ -function TwoStageExaCore(ns::Integer; backend = nothing, concrete = Val(false), kwargs...) +_ts_nb(::Val{N}) where {N} = N +_ts_nb(n::Integer) = Int(n) + +function TwoStageExaCore(ns::Integer; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) + nb = _ts_nb(nbatch) return ExaCore(; backend, concrete, + batch = Val(nb > 1), + nbatch = nb, tag = TwoStageExaModelTag( ns, convert_array(zeros(Int, 0), backend), @@ -103,13 +112,14 @@ function add_var( start = zero(T), lvar = T(-Inf), uvar = T(Inf), - ) where {T,VT<:AbstractVector{T},B} - + ) where {T,VT<:AbstractArray{T},B} + len = total(ns) - append!(c.backend, c.tag.var_scen, 0, len) + new_var_scen = append!(c.backend, c.tag.var_scen, 0, len) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, new_var_scen, c.tag.con_scen)) return _add_var( c, FirstStageTag(), name, start, lvar, uvar, ns... - ) + ) end """ add_var(core::TwoStageExaCore, ::EachScenario, dims...; start = 0, lvar = -Inf, uvar = Inf, name = nothing) @@ -127,13 +137,14 @@ function add_var( start = zero(T), lvar = T(-Inf), uvar = T(Inf), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} nscen = c.tag.nscen len = total(ns) - append!(c.backend, c.tag.var_scen, _scen_each_tag(nscen, len), len * nscen) + new_var_scen = append!(c.backend, c.tag.var_scen, _scen_each_tag(nscen, len), len * nscen) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, new_var_scen, c.tag.con_scen)) return _add_var( c, SecondStageTag(), name, start, lvar, uvar, ns..., nscen - ) + ) end """ @@ -153,7 +164,7 @@ function add_par( c::TwoStageExaCore{T,VT,B}, value::AbstractArray; name = nothing, - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} return _add_par(c, FirstStageTag(), name, value, Base.size(value)...) end function add_par( @@ -161,7 +172,7 @@ function add_par( n::AbstractRange; name = nothing, value = zero(T), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} return _add_par(c, FirstStageTag(), name, value, n) end function add_par( @@ -169,7 +180,7 @@ function add_par( ns...; name = nothing, value = zero(T), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} return _add_par(c, FirstStageTag(), name, value, ns...) end @@ -184,7 +195,7 @@ function add_par( ::EachScenario, value::AbstractVector; name = nothing, - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} combined = cat((value for _ in 1:c.tag.nscen)...; dims = ndims(value) + 1) return _add_par(c, SecondStageTag(), name, combined, Base.size(combined)...) end @@ -203,7 +214,7 @@ function add_con( start = zero(T), lcon = zero(T), ucon = zero(T), - ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} + ) where {T,VT<:AbstractArray{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} gen = _get_generator(ns) dims = _get_con_dims(ns) @@ -211,7 +222,8 @@ function add_con( f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter - append!(c.backend, c.tag.con_scen, 0, length(pars)) + new_con_scen = append!(c.backend, c.tag.con_scen, 0, length(pars)) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) return _add_con(c, f, pars, dims, start, lcon, ucon, name, FirstStageConstraintTag()) end @@ -232,7 +244,7 @@ function add_con( start = zero(T), lcon = zero(T), ucon = zero(T), - ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} + ) where {T,VT<:AbstractArray{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} gen = _get_generator(ns) dims = _get_con_dims(ns) @@ -242,7 +254,8 @@ function add_con( nscen = c.tag.nscen len = length(pars) - append!(c.backend, c.tag.con_scen, _scen_each_tag(nscen, div(len, nscen)), len) + new_con_scen = append!(c.backend, c.tag.con_scen, _scen_each_tag(nscen, div(len, nscen)), len) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) return _add_con(c, f, pars, dims, start, lcon, ucon, name, SecondStageConstraintTag()) end diff --git a/src/utils.jl b/src/utils.jl index 54abf4d3..f0f8226f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -579,3 +579,168 @@ function _structure!(I, J, ptr, sparsity, backend::Nothing) end export WrapperNLPModel, TimedNLPModel, CompressedNLPModel + +# ============================================================================ +# get_nbatch for generic AbstractNLPModel (used by FlatNLPModel) +# ============================================================================ + +get_nbatch(meta::NLPModels.NLPModelMeta{T, <:AbstractMatrix}) where {T} = Base.size(meta.x0, 2) +get_nbatch(meta::NLPModels.NLPModelMeta) = 1 +get_nbatch(m::NLPModels.AbstractNLPModel) = get_nbatch(m.meta) + +# ============================================================================ +# FlatNLPModel +# ============================================================================ + +""" + FlatNLPModel{T, VT, M} <: AbstractNLPModel{T, VT} + +Wrapper that presents a batch NLP model as a flat (Vector-based) NLP model. +All NLPModels callbacks delegate to the underlying batch model's matrix API. + + FlatNLPModel(model::AbstractNLPModel) + +Construct a flat model from a batch model whose `meta.x0` is a matrix. +""" +struct FlatNLPModel{T, VT <: AbstractVector{T}, M <: NLPModels.AbstractNLPModel{T}} <: NLPModels.AbstractNLPModel{T, VT} + batch::M + meta::NLPModels.NLPModelMeta{T, VT} + counters::NLPModels.Counters +end + +function FlatNLPModel(model::NLPModels.AbstractNLPModel{T}) where {T} + nb = get_nbatch(model) + nvar_s = NLPModels.get_nvar(model) + ncon_s = NLPModels.get_ncon(model) + nvar = nvar_s * nb + ncon = ncon_s * nb + nnzj = NLPModels.get_nnzj(model) * nb + nnzh = NLPModels.get_nnzh(model) * nb + x0 = vec(model.meta.x0) + lvar = vec(model.meta.lvar) + uvar = vec(model.meta.uvar) + y0 = vec(model.meta.y0) + lcon = vec(model.meta.lcon) + ucon = vec(model.meta.ucon) + + meta = _build_meta( + nvar, x0, lvar, uvar, + ncon, y0, lcon, ucon; + nnzj = nnzj, + nnzh = nnzh, + minimize = model.meta.minimize, + name = String(model.meta.name), + ) + return FlatNLPModel(model, meta, NLPModels.Counters()) +end + +function NLPModels.obj(m::FlatNLPModel{T}, x::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + bx = reshape(x, nvar, nb) + bf = similar(x, T, nb) + obj!(m.batch, bx, bf) + return sum(bf) +end + +function NLPModels.grad!(m::FlatNLPModel{T}, x::AbstractVector, g::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + NLPModels.grad!(m.batch, reshape(x, nvar, nb), reshape(g, nvar, nb)) + return g +end + +function NLPModels.cons_nln!(m::FlatNLPModel{T}, x::AbstractVector, c::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + NLPModels.cons_nln!(m.batch, reshape(x, nvar, nb), reshape(c, ncon, nb)) + return c +end + +function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + nnzj = NLPModels.get_nnzj(m.batch) + + r1_dev = similar(m.batch.meta.x0, Int, nnzj) + c1_dev = similar(m.batch.meta.x0, Int, nnzj) + NLPModels.jac_structure!(m.batch, r1_dev, c1_dev) + + r1 = Vector{Int}(r1_dev) + c1 = Vector{Int}(c1_dev) + r_cpu = Vector{Int}(undef, nnzj * nb) + c_cpu = Vector{Int}(undef, nnzj * nb) + copyto!(r_cpu, 1, r1, 1, nnzj) + copyto!(c_cpu, 1, c1, 1, nnzj) + + @inbounds for s in 2:nb + offset = (s - 1) * nnzj + row_shift = (s - 1) * ncon + col_shift = (s - 1) * nvar + for k in 1:nnzj + r_cpu[offset + k] = r1[k] + row_shift + c_cpu[offset + k] = c1[k] + col_shift + end + end + copyto!(rows, r_cpu) + copyto!(cols, c_cpu) + return rows, cols +end + +function NLPModels.jac_nln_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + nnzj = NLPModels.get_nnzj(m.batch) + NLPModels.jac_coord!(m.batch, reshape(x, nvar, nb), reshape(jvals, nnzj, nb)) + return jvals +end + +function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + nnzh = NLPModels.get_nnzh(m.batch) + + r1_dev = similar(m.batch.meta.x0, Int, nnzh) + c1_dev = similar(m.batch.meta.x0, Int, nnzh) + NLPModels.hess_structure!(m.batch, r1_dev, c1_dev) + + r1 = Vector{Int}(r1_dev) + c1 = Vector{Int}(c1_dev) + r_cpu = Vector{Int}(undef, nnzh * nb) + c_cpu = Vector{Int}(undef, nnzh * nb) + copyto!(r_cpu, 1, r1, 1, nnzh) + copyto!(c_cpu, 1, c1, 1, nnzh) + + @inbounds for s in 2:nb + offset = (s - 1) * nnzh + shift = (s - 1) * nvar + for k in 1:nnzh + r_cpu[offset + k] = r1[k] + shift + c_cpu[offset + k] = c1[k] + shift + end + end + copyto!(rows, r_cpu) + copyto!(cols, c_cpu) + return rows, cols +end + +function NLPModels.hess_coord!( + m::FlatNLPModel{T}, x::AbstractVector, y::AbstractVector, + hvals::AbstractVector; obj_weight = one(T), +) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + nnzh = NLPModels.get_nnzh(m.batch) + NLPModels.hess_coord!(m.batch, reshape(x, nvar, nb), reshape(y, ncon, nb), reshape(hvals, nnzh, nb); obj_weight) + return hvals +end + +function NLPModels.jtprod_nln!(m::FlatNLPModel{T}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) where {T} + NLPModels.jtprod_nln!(m.batch, x, v, Jtv) + return Jtv +end + +export FlatNLPModel diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl new file mode 100644 index 00000000..ee67fe59 --- /dev/null +++ b/test/BatchTest/BatchTest.jl @@ -0,0 +1,626 @@ +module BatchTest + +using Test +using ExaModels +import NLPModels +import NLPModels: + obj, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, + hess_structure! +import NLPModels: obj! +import ExaModels: var_indices, cons_block_indices, get_nbatch, + get_start, get_lvar, get_uvar, get_lcon, get_ucon, WrapperNLPModel, FlatNLPModel + +import NLPModelsIpopt: ipopt +import MadNLP +import MadNLP: madnlp, ERROR as MADNLP_ERROR +import PowerModels +import Downloads + +import ..BACKENDS +using Adapt + +# ============================================================================ +# Helper: build a standard test problem +# ============================================================================ + +function build_batch_model(; ns=2, nv=2, θ_val=[2.0], backend = nothing) + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, θ_val) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j] - θ[1] for j in 1:nv; lcon = 0.0) + return ExaModel(c) +end + +# ============================================================================ +# Helper utilities for backend-aware allocation +# ============================================================================ + +_ca(x, ::Nothing) = x +_ca(x, backend) = ExaModels.convert_array(x, backend) +_to_cpu(x) = Array(x) +_to_cpu(x::Array) = x + +# ============================================================================ +# Test functions — all accept backend parameter +# ============================================================================ + +function test_construction(; backend = nothing) + model = build_batch_model(ns=3; backend) + @test get_nbatch(model) == 3 + @test NLPModels.get_nvar(model) == 2 + @test NLPModels.get_ncon(model) == 2 + @test model isa NLPModels.AbstractNLPModel + @test size(model.meta.x0) == (2, 3) +end + +function test_obj(; backend = nothing) + model = build_batch_model(; backend) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + bf = _ca(zeros(2), backend) + obj!(model, bx, bf) + @test _to_cpu(bf)[1] ≈ 10.0 + @test _to_cpu(bf)[2] ≈ 50.0 + @test _to_cpu(obj(model, bx)) ≈ _to_cpu(bf) + flat = FlatNLPModel(model) + @test sum(_to_cpu(bf)) ≈ obj(flat, vec(bx)) +end + +function test_grad(; backend = nothing) + model = build_batch_model(; backend) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + bg = _ca(zeros(2, 2), backend) + grad!(model, bx, bg) + bg_cpu = _to_cpu(bg) + @test bg_cpu[:, 1] ≈ [4.0, 8.0] + @test bg_cpu[:, 2] ≈ [12.0, 16.0] + flat = FlatNLPModel(model) + g_flat = _ca(zeros(4), backend) + grad!(flat, vec(bx), g_flat) + @test vec(bg_cpu) ≈ _to_cpu(g_flat) +end + +function test_cons(; backend = nothing) + model = build_batch_model(; backend) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + bc = _ca(zeros(2, 2), backend) + cons!(model, bx, bc) + bc_cpu = _to_cpu(bc) + @test bc_cpu[:, 1] ≈ [-1.0, 0.0] + @test bc_cpu[:, 2] ≈ [1.0, 2.0] + flat = FlatNLPModel(model) + c_flat = _ca(zeros(4), backend) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc_cpu) ≈ _to_cpu(c_flat) +end + +function test_jac_hess(; backend = nothing) + model = build_batch_model(; backend) + ns, nv = 2, 2 + flat = FlatNLPModel(model) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + + # --- Jacobian values --- + nnzj = NLPModels.get_nnzj(model) + jvals = _ca(zeros(nnzj, ns), backend) + jac_coord!(model, bx, jvals) + jvals_flat = _ca(zeros(NLPModels.get_nnzj(flat)), backend) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(_to_cpu(jvals)) ≈ _to_cpu(jvals_flat) + + # --- Hessian values --- + nnzh = NLPModels.get_nnzh(model) + by = _ca(ones(nv, ns), backend) + hvals = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals) + hvals_flat = _ca(zeros(NLPModels.get_nnzh(flat)), backend) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(_to_cpu(hvals)) ≈ _to_cpu(hvals_flat) +end + +function test_hess_obj_weight(; backend = nothing) + ns, nv = 2, 2 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j]^2 for j in 1:nv) + model = ExaModel(c) + + nc = NLPModels.get_ncon(model) + nnzh = NLPModels.get_nnzh(model) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + by = _ca(ones(nc, ns), backend) + flat = FlatNLPModel(model) + + hvals_w1 = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_w1; obj_weight = 1.0) + hvals_flat_w1 = _ca(zeros(NLPModels.get_nnzh(flat)), backend) + hess_coord!(flat, vec(bx), vec(by), hvals_flat_w1; obj_weight = 1.0) + @test vec(_to_cpu(hvals_w1)) ≈ _to_cpu(hvals_flat_w1) + + hvals_w2 = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_w2; obj_weight = 2.0) + hvals_flat_w2 = _ca(zeros(NLPModels.get_nnzh(flat)), backend) + hess_coord!(flat, vec(bx), vec(by), hvals_flat_w2; obj_weight = 2.0) + @test vec(_to_cpu(hvals_w2)) ≈ _to_cpu(hvals_flat_w2) + + @test _to_cpu(hvals_w1) != _to_cpu(hvals_w2) +end + +function test_hess_vector_obj_weight(; backend = nothing) + ns, nv = 2, 2 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j]^2 for j in 1:nv) + model = ExaModel(c) + + nc = NLPModels.get_ncon(model) + nnzh = NLPModels.get_nnzh(model) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + by = _ca(ones(nc, ns), backend) + + # Vector obj_weight = [w1, w2] + wvec = _ca([1.5, 3.0], backend) + hvals_vec = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_vec; obj_weight = wvec) + + # Uniform scalar weights for comparison + hvals_w1 = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_w1; obj_weight = 1.5) + hvals_w2 = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_w2; obj_weight = 3.0) + + @test _to_cpu(hvals_vec) != _to_cpu(hvals_w1) + @test _to_cpu(hvals_vec) != _to_cpu(hvals_w2) + + # Verify consistency: uniform weight = special case of vector weight + hvals_uniform = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_uniform; obj_weight = _ca([2.0, 2.0], backend)) + hvals_scalar = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_scalar; obj_weight = 2.0) + @test _to_cpu(hvals_uniform) ≈ _to_cpu(hvals_scalar) +end + +function test_multiple_constraints(; backend = nothing) + ns, nv = 2, 3 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [1.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j] - θ[1] for j in 1:nv) + c, _ = add_con(c, v[1] + v[2] + v[3] for _ in 1:1; ucon = 10.0) + model = ExaModel(c) + flat = FlatNLPModel(model) + + nc = NLPModels.get_ncon(model) + @test nc == nv + 1 + @test get_nbatch(model) == ns + + bx = _ca(reshape(Float64[1, 2, 3, 4, 5, 6], nv, ns), backend) + + # cons! + bc = _ca(zeros(nc, ns), backend) + cons!(model, bx, bc) + c_flat = _ca(zeros(nc * ns), backend) + cons_nln!(flat, vec(bx), c_flat) + @test vec(_to_cpu(bc)) ≈ _to_cpu(c_flat) + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = _ca(zeros(nnzj, ns), backend) + jac_coord!(model, bx, jvals) + jvals_flat = _ca(zeros(NLPModels.get_nnzj(flat)), backend) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(_to_cpu(jvals)) ≈ _to_cpu(jvals_flat) + + # hess + nnzh = NLPModels.get_nnzh(model) + by = _ca(ones(nc, ns), backend) + hvals = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals) + hvals_flat = _ca(zeros(NLPModels.get_nnzh(flat)), backend) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(_to_cpu(hvals)) ≈ _to_cpu(hvals_flat) +end + +function test_error_guards(; backend = nothing) + model = build_batch_model(; backend) + x_vec = _ca(ones(2), backend) + @test_throws ArgumentError obj(model, x_vec) + @test_throws ArgumentError cons!(model, x_vec, _ca(zeros(2), backend)) + @test_throws ArgumentError grad!(model, x_vec, _ca(zeros(2), backend)) +end + +function test_bounds(; backend = nothing) + ns, nv = 2, 2 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv; start = 0.5, lvar = 0.0, uvar = 10.0) + c, _ = add_obj(c, v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 100.0) + model = ExaModel(c) + + @test size(model.meta.x0) == (nv, ns) + @test _to_cpu(model.meta.x0) ≈ fill(0.5, nv, ns) + @test _to_cpu(model.meta.lvar) ≈ fill(0.0, nv, ns) + @test _to_cpu(model.meta.uvar) ≈ fill(10.0, nv, ns) + + flat = FlatNLPModel(model) + @test NLPModels.get_nvar(flat) == nv * ns + @test _to_cpu(flat.meta.x0) ≈ fill(0.5, nv * ns) + @test _to_cpu(flat.meta.lvar) ≈ fill(0.0, nv * ns) + @test _to_cpu(flat.meta.uvar) ≈ fill(10.0, nv * ns) +end + +function test_flatten_model(; backend = nothing) + model = build_batch_model(; backend) + flat = FlatNLPModel(model) + @test flat isa FlatNLPModel + @test NLPModels.get_nvar(flat) == 2 * 2 + @test NLPModels.get_ncon(flat) == 2 * 2 +end + +function test_ipopt_simple(; backend = nothing) + ns, nv = 3, 1 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, (v[1] - θ[1])^2 for _ in 1:1) + c, _ = add_con(c, v[1] for _ in 1:1; lcon = 0.0, ucon = Inf) + model = ExaModel(c) + + result = ipopt(WrapperNLPModel(FlatNLPModel(model)); print_level = 0) + @test result.status == :first_order + for i in 1:ns + @test result.solution[var_indices(model, i)] ≈ [2.0] atol = 1e-5 + end + @test isapprox(result.objective, 0.0; atol = 1e-8) +end + +function test_ipopt_multi(; backend = nothing) + ns, nv = 2, 2 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [1.0, 3.0]) + c, _ = add_obj(c, (v[j] - θ[j])^2 for j in 1:nv) + c, _ = add_con(c, v[1] + v[2] for _ in 1:1; ucon = 10.0) + model = ExaModel(c) + + result = ipopt(WrapperNLPModel(FlatNLPModel(model)); print_level = 0) + @test result.status == :first_order + @test result.solution[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1e-5 + @test result.solution[var_indices(model, 2)] ≈ [1.0, 3.0] atol = 1e-5 + @test isapprox(result.objective, 0.0; atol = 1e-8) +end + +function test_set_parameter(; backend = nothing) + ns, nv = 2, 1 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, (v[1] - θ[1])^2 for _ in 1:1) + model = ExaModel(c) + + # Default: both instances share θ = [2.0] + bx = _ca(reshape([1.0, 3.0], nv, ns), backend) + bf = _ca(zeros(ns), backend) + obj!(model, bx, bf) + bf_cpu = _to_cpu(bf) + @test bf_cpu[1] ≈ 1.0 # (1-2)^2 + @test bf_cpu[2] ≈ 1.0 # (3-2)^2 + + # Update parameters via set_value! on the model + set_value!(model, θ, [10.0]) + bf2 = _ca(zeros(ns), backend) + obj!(model, bx, bf2) + bf2_cpu = _to_cpu(bf2) + @test bf2_cpu[1] ≈ (1.0 - 10.0)^2 + @test bf2_cpu[2] ≈ (3.0 - 10.0)^2 +end + +function test_multidim_vars(; backend = nothing) + ns = 2 + nh, nc = 3, 2 # multi-dimensional variable + c = BatchExaCore(ns; backend) + @add_var(c, w, nh, nc; start = zeros(nh, nc)) + @add_par(c, θ, [1.0]) + c, _ = add_obj(c, w[i, j]^2 for i in 1:nh, j in 1:nc) + model = ExaModel(c) + + @test NLPModels.get_nvar(model) == nh * nc + @test get_nbatch(model) == ns + + bx = _ca(reshape(Float64.(1:(nh*nc*ns)), nh * nc, ns), backend) + bf = _ca(zeros(ns), backend) + obj!(model, bx, bf) + bx_cpu = _to_cpu(bx) + bf_cpu = _to_cpu(bf) + @test bf_cpu[1] ≈ sum(bx_cpu[:, 1] .^ 2) + @test bf_cpu[2] ≈ sum(bx_cpu[:, 2] .^ 2) +end + +function test_add_con_aug(; backend = nothing) + ns, nv = 2, 3 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [1.0]) + c, _ = add_obj(c, v[j]^2 for j in 1:nv) + # Create a constraint, then augment it + @add_con(c, g, v[j] for j in 1:nv; lcon = -10.0, ucon = 10.0) + @add_con!(c, g, j => θ[1] * v[j] for j in 1:nv) + model = ExaModel(c) + flat = FlatNLPModel(model) + + nc = NLPModels.get_ncon(model) + @test nc == nv + + bx = _ca(reshape(Float64.(1:(nv*ns)), nv, ns), backend) + + # cons! — augmented constraint = v[j] + θ[1]*v[j] = 2*v[j] + bc = _ca(zeros(nc, ns), backend) + cons!(model, bx, bc) + c_flat = _ca(zeros(nc * ns), backend) + cons_nln!(flat, vec(bx), c_flat) + bc_cpu = _to_cpu(bc) + bx_cpu = _to_cpu(bx) + @test vec(bc_cpu) ≈ _to_cpu(c_flat) + + # Each constraint value should be v[j] + 1.0*v[j] = 2*v[j] + @test bc_cpu[:, 1] ≈ 2.0 .* bx_cpu[:, 1] + @test bc_cpu[:, 2] ≈ 2.0 .* bx_cpu[:, 2] + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = _ca(zeros(nnzj, ns), backend) + jac_coord!(model, bx, jvals) + jvals_flat = _ca(zeros(NLPModels.get_nnzj(flat)), backend) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(_to_cpu(jvals)) ≈ _to_cpu(jvals_flat) + + # hess + nnzh = NLPModels.get_nnzh(model) + by = _ca(ones(nc, ns), backend) + hvals = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals) + hvals_flat = _ca(zeros(NLPModels.get_nnzh(flat)), backend) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(_to_cpu(hvals)) ≈ _to_cpu(hvals_flat) +end + +function test_add_expr(; backend = nothing) + ns, nv = 2, 3 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + + # Create a subexpression and use it in objective and constraint + @add_expr(c, s, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_obj(c, s[j] for j in 1:nv) + c, _ = add_con(c, s[j] - v[j] for j in 1:nv; lcon = 0.0) + model = ExaModel(c) + flat = FlatNLPModel(model) + + nc = NLPModels.get_ncon(model) + @test nc == nv + + bx = _ca(reshape(Float64.(1:(nv*ns)), nv, ns), backend) + + # obj — should be sum of θ[1]*v[j]^2 = 2*v[j]^2 + bf = _ca(zeros(ns), backend) + obj!(model, bx, bf) + bx_cpu = _to_cpu(bx) + bf_cpu = _to_cpu(bf) + @test bf_cpu[1] ≈ 2.0 * sum(bx_cpu[:, 1] .^ 2) + @test bf_cpu[2] ≈ 2.0 * sum(bx_cpu[:, 2] .^ 2) + @test sum(bf_cpu) ≈ obj(flat, vec(bx)) + + # cons — s[j] - v[j] = 2*v[j]^2 - v[j] + bc = _ca(zeros(nc, ns), backend) + cons!(model, bx, bc) + c_flat = _ca(zeros(nc * ns), backend) + cons_nln!(flat, vec(bx), c_flat) + bc_cpu = _to_cpu(bc) + @test vec(bc_cpu) ≈ _to_cpu(c_flat) + @test bc_cpu[:, 1] ≈ 2.0 .* bx_cpu[:, 1] .^ 2 .- bx_cpu[:, 1] + + # grad + bg = _ca(zeros(nv, ns), backend) + grad!(model, bx, bg) + g_flat = _ca(zeros(nv * ns), backend) + grad!(flat, vec(bx), g_flat) + @test vec(_to_cpu(bg)) ≈ _to_cpu(g_flat) +end + +function _get_opf_case(filename) + isfile(filename) && return filename + tmpdir = tempname() + mkdir(tmpdir) + ff = joinpath(tmpdir, filename) + Downloads.download( + "https://raw.githubusercontent.com/power-grid-lib/pglib-opf/dc6be4b2f85ca0e776952ec22cbd4c22396ea5a3/$filename", + ff, + ) + return ff +end + +function _parse_opf_data(filename) + data = PowerModels.parse_file(_get_opf_case(filename)) + PowerModels.standardize_cost_terms!(data, order = 2) + PowerModels.calc_thermal_limits!(data) + ref = PowerModels.build_ref(data)[:it][:pm][:nw][0] + + arcdict = Dict(a => k for (k, a) in enumerate(ref[:arcs])) + busdict = Dict(k => i for (i, (k, v)) in enumerate(ref[:bus])) + branchdict = Dict(k => i for (i, (k, v)) in enumerate(ref[:branch])) + + bus = [ + begin + bus_loads = [ref[:load][l] for l in ref[:bus_loads][k]] + bus_shunts = [ref[:shunt][s] for s in ref[:bus_shunts][k]] + pd = sum(load["pd"] for load in bus_loads; init = 0.0) + gs = sum(sh["gs"] for sh in bus_shunts; init = 0.0) + qd = sum(load["qd"] for load in bus_loads; init = 0.0) + bs = sum(sh["bs"] for sh in bus_shunts; init = 0.0) + (i = busdict[k], pd = pd, gs = gs, qd = qd, bs = bs) + end for (k, v) in ref[:bus] + ] + gen = [ + (i = gendict_i, cost1 = v["cost"][1], cost2 = v["cost"][2], cost3 = v["cost"][3], + bus = busdict[v["gen_bus"]]) + for (gendict_i, (k, v)) in enumerate(ref[:gen]) + ] + arc = [ + (i = k, rate_a = ref[:branch][arc_l]["rate_a"], bus = busdict[arc_i]) + for (k, (arc_l, arc_i, arc_j)) in enumerate(ref[:arcs]) + ] + branch = [ + begin + g, b = PowerModels.calc_branch_y(br) + tr, ti = PowerModels.calc_branch_t(br) + ttm = tr^2 + ti^2 + ( + i = branchdict[bi], j = 1, + f_idx = arcdict[bi, br["f_bus"], br["t_bus"]], + t_idx = arcdict[bi, br["t_bus"], br["f_bus"]], + f_bus = busdict[br["f_bus"]], t_bus = busdict[br["t_bus"]], + c1 = (-g*tr - b*ti)/ttm, c2 = (-b*tr + g*ti)/ttm, + c3 = (-g*tr + b*ti)/ttm, c4 = (-b*tr - g*ti)/ttm, + c5 = (g + br["g_fr"])/ttm, c6 = (b + br["b_fr"])/ttm, + c7 = (g + br["g_to"]), c8 = (b + br["b_to"]), + rate_a_sq = br["rate_a"]^2, + ) + end for (bi, br) in ref[:branch] + ] + return ( + bus = bus, + gen = gen, + arc = arc, + branch = branch, + ref_buses = [busdict[i] for (i, _) in ref[:ref_buses]], + vmax = [v["vmax"] for (k, v) in ref[:bus]], + vmin = [v["vmin"] for (k, v) in ref[:bus]], + pmax = [v["pmax"] for (k, v) in ref[:gen]], + pmin = [v["pmin"] for (k, v) in ref[:gen]], + qmax = [v["qmax"] for (k, v) in ref[:gen]], + qmin = [v["qmin"] for (k, v) in ref[:gen]], + rate_a = [ref[:branch][arc_l]["rate_a"] for (arc_l, arc_i, arc_j) in ref[:arcs]], + angmax = [b["angmax"] for (k, b) in ref[:branch]], + angmin = [b["angmin"] for (k, b) in ref[:branch]], + ) +end + +function _build_batch_opf(data; backend, nbatch) + core = BatchExaCore(nbatch; backend) + @add_var(core, va, length(data.bus)) + @add_var(core, vm, length(data.bus); + start = fill!(similar(data.bus, Float64), 1.0), + lvar = data.vmin, uvar = data.vmax) + @add_var(core, pg, length(data.gen); lvar = data.pmin, uvar = data.pmax) + @add_var(core, qg, length(data.gen); lvar = data.qmin, uvar = data.qmax) + @add_var(core, p, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + @add_var(core, q, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + + @add_obj(core, g.cost1 * pg[g.i]^2 + g.cost2 * pg[g.i] + g.cost3 for g in data.gen) + + @add_con(core, c_ref_angle, va[i] for i in data.ref_buses) + @add_con(core, c_from_p, + p[b.f_idx] - b.c5*vm[b.f_bus]^2 - + b.c3*(vm[b.f_bus]*vm[b.t_bus]*cos(va[b.f_bus]-va[b.t_bus])) - + b.c4*(vm[b.f_bus]*vm[b.t_bus]*sin(va[b.f_bus]-va[b.t_bus])) + for b in data.branch) + @add_con(core, c_from_q, + q[b.f_idx] + b.c6*vm[b.f_bus]^2 + + b.c4*(vm[b.f_bus]*vm[b.t_bus]*cos(va[b.f_bus]-va[b.t_bus])) - + b.c3*(vm[b.f_bus]*vm[b.t_bus]*sin(va[b.f_bus]-va[b.t_bus])) + for b in data.branch) + @add_con(core, c_to_p, + p[b.t_idx] - b.c7*vm[b.t_bus]^2 - + b.c1*(vm[b.t_bus]*vm[b.f_bus]*cos(va[b.t_bus]-va[b.f_bus])) - + b.c2*(vm[b.t_bus]*vm[b.f_bus]*sin(va[b.t_bus]-va[b.f_bus])) + for b in data.branch) + @add_con(core, c_to_q, + q[b.t_idx] + b.c8*vm[b.t_bus]^2 + + b.c2*(vm[b.t_bus]*vm[b.f_bus]*cos(va[b.t_bus]-va[b.f_bus])) - + b.c1*(vm[b.t_bus]*vm[b.f_bus]*sin(va[b.t_bus]-va[b.f_bus])) + for b in data.branch) + @add_con(core, c_angle_diff, + va[b.f_bus] - va[b.t_bus] for b in data.branch; + lcon = data.angmin, ucon = data.angmax) + @add_con(core, c_thermal_f, + p[b.f_idx]^2 + q[b.f_idx]^2 - b.rate_a_sq for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf)) + @add_con(core, c_thermal_t, + p[b.t_idx]^2 + q[b.t_idx]^2 - b.rate_a_sq for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf)) + @add_con(core, c_p_balance, b.pd + b.gs*vm[b.i]^2 for b in data.bus) + @add_con(core, c_q_balance, b.qd - b.bs*vm[b.i]^2 for b in data.bus) + @add_con!(core, c_p_balance, a.bus => p[a.i] for a in data.arc) + @add_con!(core, c_q_balance, a.bus => q[a.i] for a in data.arc) + @add_con!(core, c_p_balance, g.bus => -pg[g.i] for g in data.gen) + @add_con!(core, c_q_balance, g.bus => -qg[g.i] for g in data.gen) + + return FlatNLPModel(ExaModel(core; prod = true)) +end + +function test_batch_opf_flat(; backend = nothing) + data = _parse_opf_data("pglib_opf_case3_lmbd.m") + m = _build_batch_opf(data; backend, nbatch = 3) + result = madnlp(WrapperNLPModel(m); print_level = MADNLP_ERROR) + @test result.status == MadNLP.SOLVE_SUCCEEDED +end + +function test_per_instance_accessors(; backend = nothing) + ns, nv = 2, 3 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv; start = 1.0, lvar = -5.0, uvar = 5.0) + c, _ = add_obj(c, v[j]^2 for j in 1:nv) + c, con = add_con(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 10.0) + model = ExaModel(c) + + # Test per-instance variable accessors + @test _to_cpu(get_start(model, v, 1)) ≈ fill(1.0, nv) + @test _to_cpu(get_lvar(model, v, 1)) ≈ fill(-5.0, nv) + @test _to_cpu(get_uvar(model, v, 2)) ≈ fill(5.0, nv) + + # Test per-instance constraint accessors + @test _to_cpu(get_lcon(model, con, 1)) ≈ fill(0.0, nv) + @test _to_cpu(get_ucon(model, con, 2)) ≈ fill(10.0, nv) + + # Test cons_block_indices + @test cons_block_indices(model, 1) == 1:nv + @test cons_block_indices(model, 2) == (nv+1):(2*nv) +end + +# ============================================================================ + +function runtests() + return @testset "Batch ExaModel" begin + for backend in BACKENDS + @testset "Backend: $(something(backend, :CPU))" begin + @testset "Construction" test_construction(; backend) + @testset "obj!" test_obj(; backend) + @testset "grad!" test_grad(; backend) + @testset "cons!" test_cons(; backend) + @testset "jac and hess" test_jac_hess(; backend) + @testset "hess obj_weight" test_hess_obj_weight(; backend) + @testset "hess vector obj_weight" test_hess_vector_obj_weight(; backend) + @testset "Multiple constraints" test_multiple_constraints(; backend) + @testset "Error guards" test_error_guards(; backend) + @testset "Bounds" test_bounds(; backend) + @testset "flatten_model" test_flatten_model(; backend) + @testset "Ipopt simple" test_ipopt_simple(; backend) + @testset "Ipopt multi" test_ipopt_multi(; backend) + @testset "set_value!" test_set_parameter(; backend) + @testset "Multidim vars" test_multidim_vars(; backend) + @testset "add_con!" test_add_con_aug(; backend) + @testset "add_expr" test_add_expr(; backend) + @testset "Per-instance accessors" test_per_instance_accessors(; backend) + @testset "Batch OPF FlatNLPModel" test_batch_opf_flat(; backend) + end + end + end +end + +end # module diff --git a/test/LinAlgTest/LinAlgTest.jl b/test/LinAlgTest/LinAlgTest.jl deleted file mode 100644 index f1bd0841..00000000 --- a/test/LinAlgTest/LinAlgTest.jl +++ /dev/null @@ -1,1560 +0,0 @@ -module LinAlgTest - -using ExaModels -using Test, ForwardDiff, LinearAlgebra - -# --- AD correctness helpers (from original LinAlgTest) --- - -function gradient(f, x) - T = eltype(x) - y = fill!(similar(x), zero(T)) - ExaModels.gradient!(y, (p, x, θ) -> f(x), x, nothing, nothing, one(T)) - return y -end - -function sgradient(f, x) - T = eltype(x) - - ff = f(ExaModels.VarSource()) - d = ff(ExaModels.Identity(), ExaModels.AdjointNodeSource(nothing), nothing) - y1 = [] - ExaModels.grpass(d, nothing, y1, nothing, 0, NaN) - - a1 = unique(y1) - comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a1) for i in y1)) - - n = length(a1) - buffer = fill!(similar(x, n), zero(T)) - buffer_I = similar(x, Tuple{Int, Int}, n) - - ExaModels.sgradient!(buffer_I, ff, nothing, nothing, nothing, comp, 0, NaN) - ExaModels.sgradient!(buffer, ff, nothing, x, nothing, comp, 0, one(T)) - - y = zeros(length(x)) - y[collect(i for (i, j) in buffer_I)] += buffer - - return y -end - -function shessian(f, x) - T = eltype(x) - - ff = f(ExaModels.VarSource()) - t = ff(ExaModels.Identity(), ExaModels.SecondAdjointNodeSource(nothing), nothing) - y2 = [] - ExaModels.hrpass0(t, nothing, y2, nothing, nothing, 0, NaN, NaN) - - a2 = unique(y2) - comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a2) for i in y2)) - - n = length(a2) - buffer = fill!(similar(x, n), zero(T)) - buffer_I = similar(x, Int, n) - buffer_J = similar(x, Int, n) - - ExaModels.shessian!( - buffer_I, - buffer_J, - ff, - nothing, - nothing, - nothing, - comp, - 0, - NaN, - NaN, - ) - ExaModels.shessian!(buffer, nothing, ff, nothing, x, nothing, comp, 0, one(T), zero(T)) - - y = zeros(length(x), length(x)) - for (k, (i, j)) in enumerate(zip(buffer_I, buffer_J)) - if i == j - y[i, j] += buffer[k] - else - y[i, j] += buffer[k] - y[j, i] += buffer[k] - end - end - return y -end - -_vec(x, inds) = [x[i] for i in inds] -_mat(x, rows, cols) = [x[rows[i] + cols[j] - 1] for i in eachindex(rows), j in eachindex(cols)] - -const LINALG_FUNCTIONS = [ - ("linalg-dot-node-node", x -> dot(_vec(x, 1:3), _vec(x, 4:6))), - ("linalg-dot-real-node", x -> dot([1.0, 2.0, 3.0], _vec(x, 1:3))), - ("linalg-dot-node-real", x -> dot(_vec(x, 1:3), [1.0, 2.0, 3.0])), - ("linalg-sum", x -> sum(_vec(x, 1:4))), - ("linalg-norm2", x -> norm(_vec(x, 1:3))), - ("linalg-norm3", x -> norm(_vec(x, 1:3), 3)), - ( - "linalg-matvec-node-node", - x -> begin - A = [x[1] x[3]; x[2] x[4]] - v = [x[5], x[6]] - r = A * v - r[1] + r[2] - end, - ), - ( - "linalg-matvec-real-node", - x -> begin - A = [1.0 2.0; 3.0 4.0] - v = [x[1], x[2]] - r = A * v - r[1] + r[2] - end, - ), - ( - "linalg-tr", - x -> begin - A = [x[1] x[3]; x[2] x[4]] - tr(A) - end, - ), - ( - "linalg-det-2x2", - x -> begin - A = [x[1] x[2]; x[3] x[4]] - det(A) - end, - ), - ( - "linalg-det-3x3", - x -> begin - A = [x[1] x[2] x[3]; x[4] x[5] x[6]; x[7] x[8] x[9]] - det(A) - end, - ), - ( - "linalg-cross", - x -> begin - c = cross(_vec(x, 1:3), _vec(x, 4:6)) - c[1] + c[2] + c[3] - end, - ), - ( - "linalg-vec-add", - x -> begin - r = _vec(x, 1:3) + _vec(x, 4:6) - r[1] + r[2] + r[3] - end, - ), - ( - "linalg-vec-sub", - x -> begin - r = _vec(x, 1:3) - _vec(x, 4:6) - r[1] + r[2] + r[3] - end, - ), - ( - "linalg-scalar-vec", - x -> begin - r = x[1] * _vec(x, 2:4) - r[1] + r[2] + r[3] - end, - ), - ( - "linalg-matmul", - x -> begin - A = [x[1] x[2]; x[3] x[4]] - B = [x[5] x[6]; x[7] x[8]] - C = A * B - C[1, 1] + C[1, 2] + C[2, 1] + C[2, 2] - end, - ), - ( - "linalg-composite-norm-matvec", - x -> begin - A = [1.0 2.0; 3.0 4.0] - v = [x[1], x[2]] - norm(A * v) - end, - ), - ( - "linalg-composite-det-dot", - x -> begin - A = [x[1] x[2]; x[3] x[4]] - det(A) * dot(_vec(x, 5:6), _vec(x, 7:8)) - end, - ), -] - -# --- Type dispatch helpers (from original LinAlgTest2) --- - -is_null_zero(x::ExaModels.Null) = iszero(x.value) -is_null_zero(x::ExaModels.AbstractNode) = false - -function create_nodes() - x = ExaModels.Null(1.0) - y = ExaModels.Null(2.0) - z = ExaModels.Null(3.0) - w = ExaModels.Null(4.0) - return x, y, z, w -end - -# --- Main test runner --- - -function runtests() - return @testset "Linear Algebra test" begin - - # ===================================================================== - # AD correctness: gradient, sparse gradient, sparse Hessian vs ForwardDiff - # ===================================================================== - @testset "AD correctness" begin - for (name, f) in LINALG_FUNCTIONS - x0 = 0.5 .+ rand(10) # avoid zero for norm derivatives - @testset "$name" begin - g_fd = ForwardDiff.gradient(f, x0) - h_fd = ForwardDiff.hessian(f, x0) - @test gradient(f, x0) ≈ g_fd atol = 1.0e-6 - @test sgradient(f, x0) ≈ g_fd atol = 1.0e-6 - @test shessian(f, x0) ≈ h_fd atol = 1.0e-6 - end - end - end - - # ===================================================================== - # Type dispatch and structure: return types, sizes, zero optimizations - # ===================================================================== - @testset "Type dispatch and structure" begin - x, y, z, w = create_nodes() - - @testset "Type conversions and promotions" begin - node = convert(ExaModels.AbstractNode, 5) - @test node isa ExaModels.Null - @test node isa ExaModels.AbstractNode - @test node.value == 5 - - zero_int = convert(ExaModels.AbstractNode, 0) - @test zero_int isa ExaModels.Null - @test iszero(zero_int.value) - @test zero_int === zero(ExaModels.AbstractNode) - - zero_float = convert(ExaModels.AbstractNode, 0.0) - @test zero_float isa ExaModels.Null - @test iszero(zero_float.value) - @test zero_float === zero(ExaModels.AbstractNode) - - arr = [x, 2.0, 3.0] - @test eltype(arr) == ExaModels.AbstractNode - end - - @testset "Scalar × Vector multiplication" begin - v_num = [1.0, 2.0, 3.0] - - result1 = x * v_num - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - vec_nodes = [x, y, z] - result2 = 2.0 * vec_nodes - @test length(result2) == 3 - @test result2 isa Vector - @test result2[1] isa ExaModels.AbstractNode - end - - @testset "Vector × Scalar multiplication" begin - v_num = [1.0, 2.0, 3.0] - vec_nodes = [x, y, z] - - result1 = v_num * x - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = vec_nodes .* 2.5 - @test length(result2) == 3 - @test result2 isa Vector - @test result2[1] isa ExaModels.AbstractNode - end - - @testset "Scalar × Matrix multiplication" begin - A_num = [1.0 2.0; 3.0 4.0] - mat_nodes = [x y; z w] - - result1 = x * A_num - @test size(result1) == (2, 2) - @test result1 isa Matrix - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = 3.0 * mat_nodes - @test size(result2) == (2, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Scalar multiplication" begin - A_num = [1.0 2.0; 3.0 4.0] - mat_nodes = [x y; z w] - - result1 = A_num * x - @test size(result1) == (2, 2) - @test result1 isa Matrix - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = mat_nodes .* 1.5 - @test size(result2) == (2, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Dot product" begin - v_num = [1.0, 2.0, 3.0] - vec_nodes = [x, y, z] - - result1 = dot(v_num, vec_nodes) - @test result1 isa ExaModels.AbstractNode - - result2 = dot(vec_nodes, v_num) - @test result2 isa ExaModels.AbstractNode - end - - @testset "Dot product (Real × Real fallback)" begin - v1 = [1.0, 2.0, 3.0] - v2 = [4.0, 5.0, 6.0] - - result1 = dot(v1, v2) - @test result1 isa Real - @test result1 ≈ 32.0 - - v1_view = @view v1[1:3] - v2_view = @view v2[1:3] - result2 = dot(v1_view, v2_view) - @test result2 isa Real - @test result2 ≈ 32.0 - - v1_reshaped = reshape([1.0, 2.0, 3.0], 3) - v2_reshaped = reshape([4.0, 5.0, 6.0], 3) - result3 = dot(v1_reshaped, v2_reshaped) - @test result3 isa Real - @test result3 ≈ 32.0 - - complex_vec1 = ComplexF64[1.0 + 2.0im, 3.0 + 4.0im] - complex_vec2 = ComplexF64[5.0 + 6.0im, 7.0 + 8.0im] - real_reinterp1 = reinterpret(Float64, complex_vec1) - real_reinterp2 = reinterpret(Float64, complex_vec2) - result4 = dot(real_reinterp1, real_reinterp2) - @test result4 isa Real - @test result4 ≈ 1*5 + 2*6 + 3*7 + 4*8 - - result5 = dot(v1_view, v2_reshaped) - @test result5 isa Real - @test result5 ≈ 32.0 - end - - @testset "Matrix × Vector product" begin - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0] - vec_nodes = [x, y, z] - - result = A_num * vec_nodes - @test length(result) == 2 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Matrix product" begin - A_num = [1.0 2.0; 3.0 4.0] - B_nodes = [x y; z w] - - result1 = A_num * B_nodes - @test size(result1) == (2, 2) - @test result1 isa Matrix - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = B_nodes * A_num - @test size(result2) == (2, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Adjoint Vector × Vector product" begin - vec_nodes = [x, y, z] - v_num = [1.0, 2.0, 3.0] - - result1 = vec_nodes' * v_num - @test result1 isa ExaModels.AbstractNode - - result2 = v_num' * vec_nodes - @test result2 isa ExaModels.AbstractNode - - vec_nodes2 = [y, z, x] - result3 = vec_nodes' * vec_nodes2 - @test result3 isa ExaModels.AbstractNode - end - - @testset "Adjoint Vector × Matrix product" begin - vec_nodes = [x, y, z] - A_num = [1.0 2.0; 3.0 4.0; 5.0 6.0] - - result = vec_nodes' * A_num - @test size(result) == (1, 2) - @test result isa LinearAlgebra.Adjoint - end - - @testset "Matrix adjoint" begin - mat_nodes = [x y z; y z x] - - result = adjoint(mat_nodes) - @test size(result) == (3, 2) - @test result isa Matrix - end - - @testset "Determinant" begin - A1 = reshape([x], 1, 1) - result1 = det(A1) - @test result1 isa ExaModels.AbstractNode - - A2 = [x y; z w] - result2 = det(A2) - @test result2 isa ExaModels.AbstractNode - - A3 = [x y z; y z x; z x y] - result3 = det(A3) - @test result3 isa ExaModels.AbstractNode - end - - @testset "Broadcasting operations" begin - vec_nodes = [x, y, z] - mat_nodes = [x y; z w] - - result1 = cos.(vec_nodes) - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = sin.(vec_nodes) - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - result3 = exp.(vec_nodes) - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - result4 = cos.(mat_nodes) - @test size(result4) == (2, 2) - @test result4 isa Matrix - @test result4[1, 1] isa ExaModels.AbstractNode - - result5 = vec_nodes .+ 1.0 - @test length(result5) == 3 - @test result5[1] isa ExaModels.AbstractNode - - result6 = vec_nodes .* 2.0 - @test length(result6) == 3 - @test result6[1] isa ExaModels.AbstractNode - - v_num = [1.0, 2.0, 3.0] - result7 = vec_nodes .+ v_num - @test length(result7) == 3 - @test result7[1] isa ExaModels.AbstractNode - - result8 = vec_nodes .* v_num - @test length(result8) == 3 - @test result8[1] isa ExaModels.AbstractNode - end - - @testset "Trace" begin - A2 = [x y; z w] - result1 = tr(A2) - @test result1 isa ExaModels.AbstractNode - - A3 = [x y z; y z w; z w x] - result2 = tr(A3) - @test result2 isa ExaModels.AbstractNode - - A_rect = [x y z; y z w] - @test_throws AssertionError tr(A_rect) - end - - @testset "Norms" begin - vec_nodes = [x, y, z] - mat_nodes = [x y; z w] - - result1 = norm(vec_nodes) - @test result1 isa ExaModels.AbstractNode - - result2 = norm(vec_nodes, 1) - @test result2 isa ExaModels.AbstractNode - - result3 = norm(vec_nodes, 2) - @test result3 isa ExaModels.AbstractNode - - result4 = norm(vec_nodes, 3) - @test result4 isa ExaModels.AbstractNode - - result5 = norm(mat_nodes) - @test result5 isa ExaModels.AbstractNode - - @test_throws ErrorException norm(vec_nodes, Inf) - end - - @testset "Array addition" begin - vec_nodes1 = [x, y, z] - vec_nodes2 = [y, z, w] - v_num = [1.0, 2.0, 3.0] - - result1 = vec_nodes1 + v_num - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = v_num + vec_nodes1 - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - result3 = vec_nodes1 + vec_nodes2 - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - mat_nodes1 = [x y; z w] - mat_nodes2 = [y z; w x] - A_num = [1.0 2.0; 3.0 4.0] - - result4 = mat_nodes1 + A_num - @test size(result4) == (2, 2) - @test result4 isa Matrix - @test result4[1, 1] isa ExaModels.AbstractNode - - result5 = A_num + mat_nodes1 - @test size(result5) == (2, 2) - @test result5[1, 1] isa ExaModels.AbstractNode - - result6 = mat_nodes1 + mat_nodes2 - @test size(result6) == (2, 2) - @test result6[1, 1] isa ExaModels.AbstractNode - end - - @testset "Array subtraction" begin - vec_nodes1 = [x, y, z] - vec_nodes2 = [y, z, w] - v_num = [1.0, 2.0, 3.0] - - result1 = vec_nodes1 - v_num - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = v_num - vec_nodes1 - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - result3 = vec_nodes1 - vec_nodes2 - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - mat_nodes1 = [x y; z w] - mat_nodes2 = [y z; w x] - A_num = [1.0 2.0; 3.0 4.0] - - result4 = mat_nodes1 - A_num - @test size(result4) == (2, 2) - @test result4 isa Matrix - @test result4[1, 1] isa ExaModels.AbstractNode - - result5 = A_num - mat_nodes1 - @test size(result5) == (2, 2) - @test result5[1, 1] isa ExaModels.AbstractNode - - result6 = mat_nodes1 - mat_nodes2 - @test size(result6) == (2, 2) - @test result6[1, 1] isa ExaModels.AbstractNode - end - - @testset "Diagonal operations" begin - mat_nodes = [x y z; w x y; z w x] - result1 = diag(mat_nodes) - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - mat_rect = [x y z; w x y] - result2 = diag(mat_rect) - @test length(result2) == 2 - @test result2[1] isa ExaModels.AbstractNode - - vec_nodes = [x, y, z] - result3 = diagm(vec_nodes) - @test size(result3) == (3, 3) - @test result3 isa Matrix - @test result3[1, 1] isa ExaModels.AbstractNode - @test is_null_zero(result3[1, 2]) - @test is_null_zero(result3[2, 1]) - - result4 = diagm(1 => vec_nodes) - @test size(result4) == (4, 4) - @test result4[1, 2] isa ExaModels.AbstractNode - @test is_null_zero(result4[1, 1]) - - result5 = diagm(-1 => vec_nodes) - @test size(result5) == (4, 4) - @test result5[2, 1] isa ExaModels.AbstractNode - @test is_null_zero(result5[1, 1]) - end - - @testset "Transpose operations" begin - result1 = transpose(x) - @test result1 isa ExaModels.AbstractNode - - mat_nodes = [x y z; w x y] - result2 = transpose(mat_nodes) - @test size(result2) == (3, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Dimension mismatch errors" begin - v1 = [1.0, 2.0] - vec_nodes = [x, y, z] - @test_throws AssertionError dot(v1, vec_nodes) - - A = [1.0 2.0; 3.0 4.0] - v = [x, y, z] - @test_throws AssertionError A * v - - A1 = [1.0 2.0; 3.0 4.0] - A2_nodes = [x y; z w; y x] - @test_throws AssertionError A1 * A2_nodes - - A_rect = [x y z; y z x] - @test_throws AssertionError det(A_rect) - - v1 = [x, y] - v2 = [x, y, z] - @test_throws AssertionError v1 + v2 - @test_throws AssertionError v1 - v2 - - A1 = [x y; z w] - A2 = [x y z; w x y] - @test_throws AssertionError A1 + A2 - @test_throws AssertionError A1 - A2 - end - - @testset "ExaCore variable arrays" begin - c = ExaModels.ExaCore(concrete = Val(true)) - @add_var(c, xvar, 2, 0:10, lvar=0, uvar=1) - - v = [xvar[i, 1] for i in 1:2] - @test length(v) == 2 - @test v isa Vector - @test v[1] isa ExaModels.AbstractNode - - A = [xvar[i, j] for (i, j) ∈ Base.product(1:2, 0:3)] - @test size(A) == (2, 4) - @test A isa Matrix - @test A[1, 1] isa ExaModels.AbstractNode - - v_num = [1.0, 2.0] - result1 = v + v_num - @test length(result1) == 2 - @test result1[1] isa ExaModels.AbstractNode - - result2 = v - v_num - @test length(result2) == 2 - @test result2[1] isa ExaModels.AbstractNode - - result3 = 2.0 * v - @test length(result3) == 2 - @test result3[1] isa ExaModels.AbstractNode - - result4 = dot(v, v_num) - @test result4 isa ExaModels.AbstractNode - - result5 = norm(v) - @test result5 isa ExaModels.AbstractNode - - v2 = [xvar[i, 2] for i in 1:2] - result6 = A * [1.0, 2.0, 3.0, 4.0] - @test length(result6) == 2 - @test result6[1] isa ExaModels.AbstractNode - - A_square = [xvar[i, j] for (i, j) ∈ Base.product(1:2, 1:2)] - @test size(A_square) == (2, 2) - - result7 = det(A_square) - @test result7 isa ExaModels.AbstractNode - - result8 = tr(A_square) - @test result8 isa ExaModels.AbstractNode - - result9 = diag(A_square) - @test length(result9) == 2 - @test result9[1] isa ExaModels.AbstractNode - - A_num = [1.0 2.0 3.0 4.0; 5.0 6.0 7.0 8.0] - result10 = A + A_num - @test size(result10) == (2, 4) - @test result10[1, 1] isa ExaModels.AbstractNode - - result11 = A - A_num - @test size(result11) == (2, 4) - @test result11[1, 1] isa ExaModels.AbstractNode - - result12 = cos.(v) - @test length(result12) == 2 - @test result12[1] isa ExaModels.AbstractNode - - result13 = sin.(A_square) - @test size(result13) == (2, 2) - @test result13[1, 1] isa ExaModels.AbstractNode - - result14 = v' - @test size(result14) == (1, 2) - @test result14 isa LinearAlgebra.Adjoint - - result15 = transpose(A) - @test size(result15) == (4, 2) - @test result15 isa Matrix - - result16 = diagm(v) - @test size(result16) == (2, 2) - @test result16[1, 1] isa ExaModels.AbstractNode - @test is_null_zero(result16[1, 2]) - end - - @testset "Method ambiguity fixes" begin - x, y, z, w = create_nodes() - vec_nodes1 = [x, y, z] - vec_nodes2 = [y, z, w] - mat_nodes1 = [x y; z w] - mat_nodes2 = [y z; w x] - - @testset "Scalar × Vector (both AbstractNode)" begin - result = x * vec_nodes1 - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Vector × Scalar (both AbstractNode)" begin - result = vec_nodes1 * x - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Scalar × Matrix (both AbstractNode)" begin - result = x * mat_nodes1 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Scalar (both AbstractNode)" begin - result = mat_nodes1 * x - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "dot product (both AbstractNode)" begin - result = dot(vec_nodes1, vec_nodes2) - @test result isa ExaModels.AbstractNode - end - - @testset "Matrix × Vector (both AbstractNode)" begin - A = [x y z; w x y] - v = [x, y, z] - - result = A * v - @test length(result) == 2 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Matrix (both AbstractNode)" begin - result = mat_nodes1 * mat_nodes2 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Adjoint Vector × Matrix (both AbstractNode)" begin - A = [x y; z w; y x] - v = [x, y, z] - - result = v' * A - @test size(result) == (1, 2) - @test result isa LinearAlgebra.Adjoint - end - - @testset "Vector + Vector (both AbstractNode)" begin - result = vec_nodes1 + vec_nodes2 - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Vector - Vector (both AbstractNode)" begin - result = vec_nodes1 - vec_nodes2 - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Matrix + Matrix (both AbstractNode)" begin - result = mat_nodes1 + mat_nodes2 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Matrix - Matrix (both AbstractNode)" begin - result = mat_nodes1 - mat_nodes2 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Mixed operations (no standard library conflicts)" begin - v_num = [1.0, 2.0, 3.0] - v_num_2 = [1.0, 2.0] - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0] - - @test (vec_nodes1 * 2.0) isa Vector - @test (2.0 * vec_nodes1) isa Vector - @test (mat_nodes1 * 2.0) isa Matrix - @test (2.0 * mat_nodes1) isa Matrix - @test dot(v_num, vec_nodes1) isa ExaModels.AbstractNode - @test dot(vec_nodes1, v_num) isa ExaModels.AbstractNode - @test (A_num * vec_nodes1) isa Vector - @test (mat_nodes1 * v_num_2) isa Vector - @test (A_num * [x y; z w; y x]) isa Matrix - @test (mat_nodes1 * mat_nodes2) isa Matrix - end - end - - @testset "Canonical nodes" begin - @testset "zero and one helpers" begin - z = zero(ExaModels.AbstractNode) - @test z isa ExaModels.Null - @test iszero(z.value) - @test z.value == 0 - - o = one(ExaModels.AbstractNode) - @test o isa ExaModels.Null - @test isone(o.value) - @test o.value == 1 - end - - @testset "zeros and ones array creation" begin - z1 = zeros(ExaModels.AbstractNode, 3) - @test length(z1) == 3 - @test z1 isa Vector{<:ExaModels.AbstractNode} - @test all(is_null_zero.(z1)) - @test all(x -> x isa ExaModels.Null, z1) - - z2 = zeros(ExaModels.AbstractNode, 2, 3) - @test size(z2) == (2, 3) - @test z2 isa Matrix{<:ExaModels.AbstractNode} - @test all(is_null_zero.(z2)) - - z3 = zeros(ExaModels.AbstractNode, 2, 2, 2) - @test size(z3) == (2, 2, 2) - @test z3 isa Array{<:ExaModels.AbstractNode, 3} - @test all(is_null_zero.(z3)) - - o1 = ones(ExaModels.AbstractNode, 3) - @test length(o1) == 3 - @test o1 isa Vector{<:ExaModels.AbstractNode} - @test all(x -> x isa ExaModels.Null && isone(x.value), o1) - - o2 = ones(ExaModels.AbstractNode, 2, 3) - @test size(o2) == (2, 3) - @test o2 isa Matrix{<:ExaModels.AbstractNode} - @test all(x -> x isa ExaModels.Null && isone(x.value), o2) - - o3 = ones(ExaModels.AbstractNode, 2, 2, 2) - @test size(o3) == (2, 2, 2) - @test o3 isa Array{<:ExaModels.AbstractNode, 3} - @test all(x -> x isa ExaModels.Null && isone(x.value), o3) - - z_var = zeros(ExaModels.AbstractNode, 3) - @test length(z_var) == 3 - @test eltype(z_var) <: ExaModels.AbstractNode - @test all(is_null_zero.(z_var)) - - o_var = ones(ExaModels.AbstractNode, 2, 2) - @test size(o_var) == (2, 2) - @test eltype(o_var) <: ExaModels.AbstractNode - @test all(x -> x isa ExaModels.Null && isone(x.value), o_var) - end - - @testset "zeros and ones in operations" begin - x, y, z, w = create_nodes() - - z_vec = zeros(ExaModels.AbstractNode, 3) - vec_nodes = [x, y, z] - - result1 = vec_nodes + z_vec - @test result1[1].value == x.value - @test result1[2].value == y.value - @test result1[3].value == z.value - - result2 = [ExaModels.Null(2), ExaModels.Null(3), ExaModels.Null(4)] .* z_vec - @test all(is_null_zero.(result2)) - - o_vec = ones(ExaModels.AbstractNode, 3) - - result3 = vec_nodes .* o_vec - @test result3[1].value == x.value - @test result3[2].value == y.value - @test result3[3].value == z.value - - I_like = diagm(ones(ExaModels.AbstractNode, 2)) - vec2 = [x, y] - result4 = I_like * vec2 - @test result4[1].value == x.value - @test result4[2].value == y.value - end - end - - @testset "Zero multiplication optimizations" begin - x, y, z, w = create_nodes() - - @testset "Scalar × Vector with zero scalar" begin - result1 = ExaModels.Null(0) * [1.0, 2.0, 3.0] - @test all(is_null_zero.(result1)) - @test result1 isa Vector{<:ExaModels.AbstractNode} - @test length(result1) == 3 - - vec_nodes = [x, y, z] - result2 = 0 * vec_nodes - @test all(is_null_zero.(result2)) - @test result2 isa Vector{<:ExaModels.AbstractNode} - - result3 = 0.0 * vec_nodes - @test all(is_null_zero.(result3)) - end - - @testset "Scalar × Vector with zero elements" begin - result = x * [0, 1.0, 0, 2.0] - @test is_null_zero(result[1]) - @test !is_null_zero(result[2]) - @test is_null_zero(result[3]) - @test !is_null_zero(result[4]) - end - - @testset "Vector × Scalar with zero scalar" begin - vec_nodes = [x, y, z] - - result1 = vec_nodes * 0 - @test all(is_null_zero.(result1)) - @test result1 isa Vector{<:ExaModels.AbstractNode} - - result2 = vec_nodes * 0.0 - @test all(is_null_zero.(result2)) - - result3 = [1.0, 2.0, 3.0] * ExaModels.Null(0) - @test all(is_null_zero.(result3)) - end - - @testset "Vector × Scalar with zero elements" begin - result = [0, 1.0, ExaModels.Null(0), 2.0] * x - @test is_null_zero(result[1]) - @test !is_null_zero(result[2]) - @test is_null_zero(result[3]) - @test !is_null_zero(result[4]) - end - - @testset "Scalar × Matrix with zero scalar" begin - mat_nodes = [x y; z w] - - result1 = 0 * mat_nodes - @test all(is_null_zero.(result1)) - @test result1 isa Matrix{<:ExaModels.AbstractNode} - @test size(result1) == (2, 2) - - result2 = ExaModels.Null(0) * [1.0 2.0; 3.0 4.0] - @test all(is_null_zero.(result2)) - end - - @testset "Matrix × Scalar with zero scalar" begin - mat_nodes = [x y; z w] - - result = mat_nodes * 0.0 - @test all(is_null_zero.(result)) - @test result isa Matrix{<:ExaModels.AbstractNode} - end - end - - @testset "Zero addition optimizations" begin - x, y, z, w = create_nodes() - - @testset "Vector + Vector with zero elements" begin - vec_nodes = [x, y, z] - - result1 = vec_nodes + [0, 0, 0] - @test result1[1].value == x.value - @test result1[2].value == y.value - @test result1[3].value == z.value - @test result1 isa Vector{<:ExaModels.AbstractNode} - - result2 = [0.0, 0.0, 0.0] + vec_nodes - @test result2[1].value == x.value - @test result2[2].value == y.value - @test result2[3].value == z.value - - result3 = vec_nodes + [0, 1.0, 0] - @test result3[1].value == x.value - @test result3[2].value == y.value + 1.0 - @test result3[3].value == z.value - - zero_vec = [ExaModels.Null(0), ExaModels.Null(0), ExaModels.Null(0)] - result4 = vec_nodes + zero_vec - @test result4[1].value == x.value - @test result4[2].value == y.value - @test result4[3].value == z.value - end - - @testset "Matrix + Matrix with zero elements" begin - mat_nodes = [x y; z w] - - result1 = mat_nodes + [0 0; 0 0] - @test result1[1,1].value == x.value - @test result1[1,2].value == y.value - @test result1[2,1].value == z.value - @test result1[2,2].value == w.value - @test result1 isa Matrix{<:ExaModels.AbstractNode} - - result2 = [0.0 0.0; 0.0 0.0] + mat_nodes - @test result2[1,1].value == x.value - @test result2[1,2].value == y.value - - result3 = mat_nodes + [0 1.0; 0 0] - @test result3[1,1].value == x.value - @test result3[1,2].value == y.value + 1.0 - @test result3[2,1].value == z.value - @test result3[2,2].value == w.value - end - end - - @testset "Zero subtraction optimizations" begin - x, y, z, w = create_nodes() - - @testset "Vector - Vector with zero elements" begin - vec_nodes = [x, y, z] - - result1 = vec_nodes - [0, 0, 0] - @test result1[1].value == x.value - @test result1[2].value == y.value - @test result1[3].value == z.value - @test result1 isa Vector{<:ExaModels.AbstractNode} - - result2 = vec_nodes - [0, 1.0, 0] - @test result2[1].value == x.value - @test result2[2].value == y.value - 1.0 - @test result2[3].value == z.value - - zero_vec = [ExaModels.Null(0), ExaModels.Null(0), ExaModels.Null(0)] - result3 = vec_nodes - zero_vec - @test result3[1].value == x.value - @test result3[2].value == y.value - @test result3[3].value == z.value - end - - @testset "Matrix - Matrix with zero elements" begin - mat_nodes = [x y; z w] - - result1 = mat_nodes - [0 0; 0 0] - @test result1[1,1].value == x.value - @test result1[1,2].value == y.value - @test result1[2,1].value == z.value - @test result1[2,2].value == w.value - @test result1 isa Matrix{<:ExaModels.AbstractNode} - - result2 = mat_nodes - [0 1.0; 0 0] - @test result2[1,1].value == x.value - @test result2[1,2].value == y.value - 1.0 - @test result2[2,1].value == z.value - @test result2[2,2].value == w.value - end - end - - @testset "Scalar operations on Null nodes (+, -, *)" begin - x, y, z, w = create_nodes() - - e = ExaModels.Node2(+, x, y) - f = ExaModels.Node2(+, z, w) - - @testset "+ operator rules" begin - result1 = ExaModels.Null(3) + ExaModels.Null(5) - @test result1 isa ExaModels.Null - @test result1.value == 8 - - result2 = ExaModels.Null(3) + e - @test result2 isa ExaModels.AbstractNode - @test !(result2 isa ExaModels.Null) - - result3 = e + ExaModels.Null(3) - @test result3 isa ExaModels.AbstractNode - @test !(result3 isa ExaModels.Null) - - result4 = e + f - @test result4 isa ExaModels.AbstractNode - end - - @testset "- operator rules" begin - result1 = ExaModels.Null(5) - ExaModels.Null(3) - @test result1 isa ExaModels.Null - @test result1.value == 2 - - result2 = ExaModels.Null(0) - e - @test result2 isa ExaModels.Node1 - - result3 = ExaModels.Null(3) - e - @test result3 isa ExaModels.AbstractNode - @test !(result3 isa ExaModels.Null) - - result4 = e - ExaModels.Null(3) - @test result4 isa ExaModels.AbstractNode - @test !(result4 isa ExaModels.Null) - - result5 = e - f - @test result5 isa ExaModels.AbstractNode - end - - @testset "* operator rules" begin - result1 = ExaModels.Null(3) * ExaModels.Null(5) - @test result1 isa ExaModels.Null - @test result1.value == 15 - - result2 = ExaModels.Null(0) * e - @test result2 isa ExaModels.Null - @test iszero(result2.value) - - result3 = e * ExaModels.Null(0) - @test result3 isa ExaModels.Null - @test iszero(result3.value) - - result4 = ExaModels.Null(3) * e - @test result4 isa ExaModels.AbstractNode - @test !(result4 isa ExaModels.Null) - - result5 = e * ExaModels.Null(3) - @test result5 isa ExaModels.AbstractNode - @test !(result5 isa ExaModels.Null) - - result6 = e * f - @test result6 isa ExaModels.AbstractNode - end - end - - @testset "sum function" begin - x, y, z, w = create_nodes() - - @testset "sum with zeros" begin - result1 = sum([zero(ExaModels.AbstractNode), zero(ExaModels.AbstractNode), zero(ExaModels.AbstractNode)]) - @test result1 isa ExaModels.Null - @test iszero(result1.value) - - result2 = sum([zero(ExaModels.AbstractNode), zero(ExaModels.AbstractNode), x, zero(ExaModels.AbstractNode)]) - @test result2 isa ExaModels.Null - @test result2.value == x.value - - result3 = sum([zero(ExaModels.AbstractNode), x, zero(ExaModels.AbstractNode), y, zero(ExaModels.AbstractNode)]) - @test result3 isa ExaModels.Null - @test result3.value == x.value + y.value - end - - @testset "sum with single element" begin - result1 = sum([x]) - @test result1 isa ExaModels.Null - @test result1.value == x.value - - result2 = sum([zero(ExaModels.AbstractNode)]) - @test result2 isa ExaModels.Null - @test iszero(result2.value) - end - - @testset "sum with all non-zeros" begin - result1 = sum([x, y, z]) - @test result1 isa ExaModels.Null - @test result1.value == x.value + y.value + z.value - end - - @testset "sum on matrices" begin - mat = [x zero(ExaModels.AbstractNode); zero(ExaModels.AbstractNode) y] - result = sum(mat) - @test result isa ExaModels.Null - @test result.value == x.value + y.value - end - end - - @testset "Optimized dot product" begin - x, y, z, t = create_nodes() - - @testset "dot([1, 0, 1, 0], [x, y, z, t]) = x + z" begin - result = dot([1, 0, 1, 0], [x, y, z, t]) - @test result isa ExaModels.Null - @test result.value == x.value + z.value - end - - @testset "dot with all zeros" begin - result = dot([0, 0, 0], [x, y, z]) - @test result isa ExaModels.Null - @test iszero(result.value) - end - - @testset "dot with all ones" begin - result = dot([1, 1, 1], [x, y, z]) - @test result isa ExaModels.Null - @test result.value == x.value + y.value + z.value - end - - @testset "dot with single non-zero" begin - result = dot([0, 1, 0], [x, y, z]) - @test result isa ExaModels.Null - @test result.value == y.value - end - end - - @testset "Matrix operations with optimization" begin - x, y, z, w = create_nodes() - - @testset "Identity matrix multiplication" begin - I2 = [1 0; 0 1] - vec = [x, y] - result = I2 * vec - - @test result[1] isa ExaModels.Null - @test result[1].value == x.value - @test result[2] isa ExaModels.Null - @test result[2].value == y.value - end - - @testset "Zero matrix multiplication" begin - Z2 = [0 0; 0 0] - vec = [x, y] - result = Z2 * vec - - @test result[1] isa ExaModels.Null - @test iszero(result[1].value) - @test result[2] isa ExaModels.Null - @test iszero(result[2].value) - end - - @testset "Sparse matrix multiplication" begin - A = [1 0 0; 0 0 1] - vec = [x, y, z] - result = A * vec - - @test result[1] isa ExaModels.Null - @test result[1].value == x.value - @test result[2] isa ExaModels.Null - @test result[2].value == z.value - end - end - - @testset "SubArray, ReshapedArray, ReinterpretArray" begin - x, y, z, w = create_nodes() - - @testset "SubArray (views) - vectors of Real" begin - v_num = [1.0, 2.0, 3.0, 4.0] - vec_nodes = [x, y, z, w] - - v_view = @view v_num[1:3] - @test v_view isa SubArray - - result1 = dot(v_view, vec_nodes[1:3]) - @test result1 isa ExaModels.AbstractNode - - result2 = dot(vec_nodes[1:3], v_view) - @test result2 isa ExaModels.AbstractNode - - result3 = x * v_view - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - result4 = 2.0 * vec_nodes[1:3] - @test length(result4) == 3 - - result5 = v_view * x - @test length(result5) == 3 - @test result5[1] isa ExaModels.AbstractNode - - result6 = vec_nodes[1:3] + v_view - @test length(result6) == 3 - @test result6[1] isa ExaModels.AbstractNode - - result7 = v_view + vec_nodes[1:3] - @test length(result7) == 3 - - result8 = vec_nodes[1:3] - v_view - @test length(result8) == 3 - @test result8[1] isa ExaModels.AbstractNode - - result9 = v_view - vec_nodes[1:3] - @test length(result9) == 3 - - result10 = norm(@view vec_nodes[1:3]) - @test result10 isa ExaModels.AbstractNode - end - - @testset "SubArray (views) - vectors of AbstractNode" begin - vec_nodes = [x, y, z, w] - - v_view = @view vec_nodes[1:3] - @test v_view isa SubArray - - v_view2 = @view vec_nodes[2:4] - result1 = v_view + v_view2 - @test length(result1) == 3 - @test result1[1] isa ExaModels.AbstractNode - - result2 = v_view - v_view2 - @test length(result2) == 3 - - result3 = dot(v_view, v_view2) - @test result3 isa ExaModels.AbstractNode - - result4 = 2.5 * v_view - @test length(result4) == 3 - @test result4[1] isa ExaModels.AbstractNode - - result5 = x * v_view - @test length(result5) == 3 - end - - @testset "SubArray (views) - matrices of Real" begin - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0] - mat_nodes = [x y z; w x y; z w x] - - A_view = @view A_num[1:2, 1:2] - @test A_view isa SubArray - @test size(A_view) == (2, 2) - - result1 = x * A_view - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = 3.0 * (@view mat_nodes[1:2, 1:2]) - @test size(result2) == (2, 2) - - result3 = A_view * x - @test size(result3) == (2, 2) - @test result3[1, 1] isa ExaModels.AbstractNode - - vec_nodes = [x, y] - result4 = A_view * vec_nodes - @test length(result4) == 2 - @test result4[1] isa ExaModels.AbstractNode - - result5 = (@view mat_nodes[1:2, 1:2]) + A_view - @test size(result5) == (2, 2) - @test result5[1, 1] isa ExaModels.AbstractNode - end - - @testset "SubArray (views) - matrices of AbstractNode" begin - mat_nodes = [x y z; w x y; z w x] - - A_view = @view mat_nodes[1:2, 1:2] - @test A_view isa SubArray - @test size(A_view) == (2, 2) - - B_view = @view mat_nodes[2:3, 2:3] - result1 = A_view * B_view - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = A_view + B_view - @test size(result2) == (2, 2) - - result3 = A_view - B_view - @test size(result3) == (2, 2) - - result4 = adjoint(A_view) - @test size(result4) == (2, 2) - - result5 = transpose(A_view) - @test size(result5) == (2, 2) - end - - @testset "ReshapedArray - vectors to matrices" begin - v_num = [1.0, 2.0, 3.0, 4.0] - vec_nodes = [x, y, z, w] - - A_reshaped = reshape(v_num, 2, 2) - @test A_reshaped isa Union{Matrix, Base.ReshapedArray} - - result1 = x * A_reshaped - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - - mat_reshaped = reshape(vec_nodes, 2, 2) - @test mat_reshaped isa Union{Matrix, Base.ReshapedArray} - - result2 = 2.0 * mat_reshaped - @test size(result2) == (2, 2) - @test result2[1, 1] isa ExaModels.AbstractNode - - result3 = det(mat_reshaped) - @test result3 isa ExaModels.AbstractNode - - result4 = tr(mat_reshaped) - @test result4 isa ExaModels.AbstractNode - end - - @testset "ReshapedArray - matrices to vectors" begin - A_num = [1.0 2.0; 3.0 4.0] - mat_nodes = [x y; z w] - - v_reshaped = reshape(A_num, 4) - @test v_reshaped isa Union{Vector, Base.ReshapedArray} - - result1 = x * v_reshaped - @test length(result1) == 4 - @test result1[1] isa ExaModels.AbstractNode - - vec_reshaped = reshape(mat_nodes, 4) - @test vec_reshaped isa Union{Vector, Base.ReshapedArray} - - result2 = 2.0 * vec_reshaped - @test length(result2) == 4 - @test result2[1] isa ExaModels.AbstractNode - - result3 = norm(vec_reshaped) - @test result3 isa ExaModels.AbstractNode - - result4 = dot(v_reshaped, vec_reshaped) - @test result4 isa ExaModels.AbstractNode - end - - @testset "ReinterpretArray - Complex to Real" begin - complex_vec = ComplexF64[1.0 + 2.0im, 3.0 + 4.0im, 5.0 + 6.0im] - - real_reinterp = reinterpret(Float64, complex_vec) - @test real_reinterp isa Base.ReinterpretArray - @test length(real_reinterp) == 6 - - vec_nodes = [x, y, z, w, ExaModels.Null(5), ExaModels.Null(6)] - - result1 = dot(real_reinterp, vec_nodes) - @test result1 isa ExaModels.AbstractNode - - result2 = dot(vec_nodes, real_reinterp) - @test result2 isa ExaModels.AbstractNode - - result3 = x * real_reinterp - @test length(result3) == 6 - @test result3[1] isa ExaModels.AbstractNode - - result4 = real_reinterp * x - @test length(result4) == 6 - @test result4[1] isa ExaModels.AbstractNode - - result5 = vec_nodes + real_reinterp - @test length(result5) == 6 - @test result5[1] isa ExaModels.AbstractNode - - result6 = real_reinterp + vec_nodes - @test length(result6) == 6 - - result7 = vec_nodes - real_reinterp - @test length(result7) == 6 - @test result7[1] isa ExaModels.AbstractNode - - result8 = real_reinterp - vec_nodes - @test length(result8) == 6 - end - - @testset "Mixed wrapper combinations" begin - v_num = [1.0, 2.0, 3.0, 4.0] - vec_nodes = [x, y, z, w] - - v_view = @view v_num[1:3] - v_reshaped = reshape([x, y, z], 3) - result1 = v_view + v_reshaped - @test length(result1) == 3 - @test result1[1] isa ExaModels.AbstractNode - - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0] - A_view = @view A_num[1:3, 1:3] - vec_reshaped = reshape([x, y, z], 3) - result2 = A_view * vec_reshaped - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - complex_vec = ComplexF64[1.0 + 2.0im, 3.0 + 4.0im] - real_reinterp = reinterpret(Float64, complex_vec) - vec_view = @view vec_nodes[1:4] - result3 = real_reinterp + vec_view - @test length(result3) == 4 - @test result3[1] isa ExaModels.AbstractNode - end - - @testset "diagm with views and reshaped arrays" begin - vec_nodes = [x, y, z] - - v_view = @view vec_nodes[1:2] - result1 = diagm(v_view) - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - @test is_null_zero(result1[1, 2]) - - v_reshaped = reshape([x, y], 2) - result2 = diagm(v_reshaped) - @test size(result2) == (2, 2) - @test result2[1, 1] isa ExaModels.AbstractNode - - result3 = diagm(1 => v_view) - @test size(result3) == (3, 3) - @test result3[1, 2] isa ExaModels.AbstractNode - end - - @testset "diag with views and reshaped arrays" begin - mat_nodes = [x y z; w x y; z w x] - - A_view = @view mat_nodes[1:2, 1:2] - result1 = diag(A_view) - @test length(result1) == 2 - @test result1[1] isa ExaModels.AbstractNode - - vec = [x, y, z, w] - A_reshaped = reshape(vec, 2, 2) - result2 = diag(A_reshaped) - @test length(result2) == 2 - @test result2[1] isa ExaModels.AbstractNode - end - - @testset "Adjoint operations with views" begin - vec_nodes = [x, y, z] - A_num = [1.0 2.0; 3.0 4.0; 5.0 6.0] - - v_view = @view vec_nodes[1:3] - v_num = [1.0, 2.0, 3.0] - result0 = v_view' * v_num - @test result0 isa ExaModels.AbstractNode - - result0b = v_num' * v_view - @test result0b isa ExaModels.AbstractNode - - v_view2 = @view vec_nodes[1:3] - result0c = v_view' * v_view2 - @test result0c isa ExaModels.AbstractNode - - result1 = v_view' * A_num - @test size(result1) == (1, 2) - @test result1 isa LinearAlgebra.Adjoint - - mat_nodes = [x y z; w x y; z w x] - A_view = @view mat_nodes[1:2, 1:2] - v_num2 = [1.0, 2.0] - result2 = v_num2' * A_view - @test size(result2) == (1, 2) - - vec_reshaped = reshape([x, y, z], 3) - result3 = vec_reshaped' * v_num - @test result3 isa ExaModels.AbstractNode - - result4 = v_num' * vec_reshaped - @test result4 isa ExaModels.AbstractNode - end - end - end - end -end - -end # module diff --git a/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl b/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl index f63d1c19..79fa45a0 100644 --- a/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl +++ b/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl @@ -37,11 +37,11 @@ function (@main)(ARGS) name = ARGS[1] N = parse(Int, ARGS[2]) println(Core.stdout, "Solving $name (N=$N) with Ipopt...") - + if name == "rosenrock" m = LV.rosenrock_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "wood" m = LV.wood_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) @@ -53,31 +53,31 @@ function (@main)(ARGS) elseif name == "broyden_banded" m = LV.broyden_banded_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "broyden_tridiagonal" m = LV.broyden_tridiagonal_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "chained_powell" m = LV.chained_powell_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "cragg_levy" m = LV.cragg_levy_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "generalized_brown" m = LV.generalized_brown_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "modified_brown" m = LV.modified_brown_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "trigo_tridiagonal" m = LV.trigo_tridiagonal_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "Chained_HS46" m = LV.Chained_HS46_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) @@ -89,7 +89,7 @@ function (@main)(ARGS) elseif name == "Chained_HS48" m = LV.Chained_HS48_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "Chained_HS49" m = LV.Chained_HS49_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) @@ -113,7 +113,7 @@ function (@main)(ARGS) println(Core.stdout, "Unknown model: $name") return 1 end - + println(Core.stdout, "Ipopt status : ", result.status) return result.status == 0 ? 0 : 1 end diff --git a/test/NLPTest/feature_test.jl b/test/NLPTest/feature_test.jl index c19b9701..ef5b7195 100644 --- a/test/NLPTest/feature_test.jl +++ b/test/NLPTest/feature_test.jl @@ -125,10 +125,9 @@ function test_add_par_dims(backend) @test g_vals ≈ [1.0, 5.0, 9.0] end - @testset "set_parameter! with range-sized parameter" begin + @testset "set_value! with range-sized parameter" begin c = ExaCore(; backend, concrete = Val(true)) - c, θ = add_par(c, 2:4; value = ones(3)) - set_parameter!(c, θ, [5.0, 6.0, 7.0]) + c, θ = add_par(c, 2:4; value = [5.0, 6.0, 7.0]) c, x = add_var(c, 1) c, g = add_con(c, θ[j] * x[1] for j in 2:4; lcon = 0.0, ucon = 0.0) m = ExaModel(c) diff --git a/test/NLPTest/parameter_test.jl b/test/NLPTest/parameter_test.jl index 7aacc715..de5f9501 100644 --- a/test/NLPTest/parameter_test.jl +++ b/test/NLPTest/parameter_test.jl @@ -29,12 +29,8 @@ function exa_luksan_vlcek_parametric( @add_var(c, x, N, M; start = [luksan_vlcek_x0(i) for i = 1:N, j = 1:M]) if use_parameters - @add_par(c, θ, zeros(7)) - if !isnothing(param_values) - set_parameter!(c, θ, param_values) - else - set_parameter!(c, θ, [100.0, 1.0, 3.0, 2.0, 5.0, 4.0, 3.0]) - end + default_params = isnothing(param_values) ? [100.0, 1.0, 3.0, 2.0, 5.0, 4.0, 3.0] : param_values + @add_par(c, θ, default_params) @add_con(c, s, luksan_vlcek_con1_param(x, θ, i, j) for i = 1:(N-2), j = 1:M) @add_con!( @@ -349,7 +345,7 @@ function test_parametric_vs_nonparametric(backend) m_param, c_param, (_, θ_param), _ = exa_luksan_vlcek_parametric(backend, 3, M = 2, use_parameters = true) new_params = [75.0, 1.5, 4.0, 3.0, 6.0, 5.0, 2.0] - set_parameter!(c_param, θ_param, new_params) + set_value!(m_param, θ_param, new_params) m_nonparam, _, _, _ = exa_luksan_vlcek_parametric( backend, 3, diff --git a/test/Project.toml b/test/Project.toml index c4761037..bbd6deca 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -24,3 +25,6 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" + +[sources] +ExaModels = {path = ".."} diff --git a/test/TwoStageTest/TwoStageTest.jl b/test/TwoStageTest/TwoStageTest.jl index bba7501e..aee0efa9 100644 --- a/test/TwoStageTest/TwoStageTest.jl +++ b/test/TwoStageTest/TwoStageTest.jl @@ -516,6 +516,113 @@ function runtests() end end + + @testset "Batched two-stage" begin + + @testset "Construction" begin + ns_scen, nb = 3, 2 + core = TwoStageExaCore(ns_scen; nbatch = Val(nb)) + v = @add_var(core, EachScenario(), 2) + core, d = add_var(core, 1) + + model = ExaModel(core) + @test get_nscen(model) == ns_scen + @test ExaModels.get_nbatch(model) == nb + @test NLPModels.get_nvar(model) == ns_scen * 2 + 1 + @test size(model.meta.x0) == (ns_scen * 2 + 1, nb) + end + + @testset "Evaluation" begin + ns_scen, nb = 2, 3 + nv, nd = 1, 1 + θ_vals = [2.0, 4.0] + + core = TwoStageExaCore(ns_scen; nbatch = Val(nb)) + v = @add_var(core, EachScenario(), nv) + core, d = add_var(core, nd) + core, θ = add_par(core, θ_vals) + + obj_data = [(i, (i-1)*nv+1, i) for i in 1:ns_scen] + @add_obj(core, (v[v_idx] - θ[θi])^2 + d[1]^2 for (i, v_idx, θi) in obj_data) + + con_data = [(i, (i-1)*nv+1) for i in 1:ns_scen] + @add_con(core, EachScenario(), (v[v_idx] - d[1] for (i, v_idx) in con_data); lcon = 0.0) + + model = ExaModel(core) + flat = ExaModels.FlatNLPModel(model) + + nvar = NLPModels.get_nvar(model) + ncon = NLPModels.get_ncon(model) + + # obj + bx = reshape(Float64.(1:(nvar*nb)), nvar, nb) + bf = zeros(nb) + NLPModels.obj!(model, bx, bf) + @test sum(bf) ≈ NLPModels.obj(flat, vec(bx)) + + # grad + bg = zeros(nvar, nb) + NLPModels.grad!(model, bx, bg) + g_flat = zeros(nvar * nb) + grad!(flat, vec(bx), g_flat) + @test vec(bg) ≈ g_flat + + # cons + bc = zeros(ncon, nb) + NLPModels.cons!(model, bx, bc) + c_flat = zeros(ncon * nb) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc) ≈ c_flat + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = zeros(nnzj, nb) + jac_coord!(model, bx, jvals) + jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(jvals) ≈ jvals_flat + + # hess + nnzh = NLPModels.get_nnzh(model) + by = ones(ncon, nb) + hvals = zeros(nnzh, nb) + hess_coord!(model, bx, by, hvals) + hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(hvals) ≈ hvals_flat + end + + @testset "Ipopt" begin + ns_scen, nb = 2, 2 + nv, nd = 1, 1 + θ_vals = [2.0, 4.0] + + core = TwoStageExaCore(ns_scen; nbatch = Val(nb)) + v = @add_var(core, EachScenario(), nv) + core, d = add_var(core, nd) + core, θ = add_par(core, θ_vals) + + @add_obj(core, d[1]^2) + @add_obj(core, (v[i] - θ[i])^2 for i in 1:ns_scen) + @add_con(core, EachScenario(), (v[i] - d[1] for i in 1:ns_scen); lcon = 0.0) + + model = ExaModel(core) + flat = ExaModels.FlatNLPModel(model) + + result = ipopt(flat; print_level = 0) + @test result.status == :first_order + + # Each batch instance should have the same solution + nvar = NLPModels.get_nvar(model) + for b in 1:nb + x_b = result.solution[ExaModels.var_indices(model, b)] + # d² + Σ(v_i - θ_i)² s.t. v_i = d → d* = Σθ / (1 + ns) + d_expected = sum(θ_vals) / (1 + ns_scen) + @test x_b[end] ≈ d_expected atol = 1e-4 + end + end + + end end end # module TwoStageTest diff --git a/test/runtests.jl b/test/runtests.jl index f115393a..f9980154 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -20,8 +20,9 @@ include("ADTest/ADTest.jl") include("DeprecatedTest/DeprecatedTest.jl") include("JuMPTest/JuMPTest.jl") include("UtilsTest/UtilsTest.jl") -include("JuliaCTest/JuliaCTest.jl") include("TwoStageTest/TwoStageTest.jl") +include("BatchTest/BatchTest.jl") +include("JuliaCTest/JuliaCTest.jl") include("GetterSetterTest/GetterSetterTest.jl") include("PrettyPrintTest.jl") # include("OptimalControlTest/OptimalControlTest.jl") @@ -57,6 +58,8 @@ include("PrettyPrintTest.jl") # @info "Running OptimalControl Test" # OptimalControlTest.runtests() + @info "Running Batch Test" + BatchTest.runtests() end # Force full GC before Julia exits so that OpenCL/PoCL objects are finalized