BatchExaModel: batch optimization with GPU support

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"
+oneAPI = "2"

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

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.θ)

docs/src/core.md

@@ -2,3 +2,4 @@
 ```@autodocs
 Modules = [ExaModels]
 ```
+

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)")

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)")
 ````

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.