From 8f06f5206856ec26ded211e780764a95efc8e886 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 20:39:59 -0400 Subject: [PATCH 01/10] rename --- README.md | 18 ++++++++--------- ext/BNKChainRulesCore.jl | 24 +++++++++++------------ src/BatchNLPKernels.jl | 5 +++-- src/api/cons.jl | 26 ++++++++++++------------- src/api/grad.jl | 24 +++++++++++------------ src/api/hess.jl | 30 ++++++++++++++-------------- src/api/hprod.jl | 16 +++++++-------- src/api/jac.jl | 22 ++++++++++----------- src/api/jprod.jl | 36 +++++++++++++++++----------------- src/api/obj.jl | 14 +++++++------- src/api/viols.jl | 22 +++++++++++++++++---- test/api.jl | 42 ++++++++++++++++++++-------------------- test/config.jl | 20 +++++++++---------- test/test_diff.jl | 28 +++++++++++++-------------- test/test_viols.jl | 4 ++-- 15 files changed, 173 insertions(+), 158 deletions(-) diff --git a/README.md b/README.md index 33f0c96..07d4c45 100644 --- a/README.md +++ b/README.md @@ -5,14 +5,14 @@ `BatchNLPKernels.jl` provides [`KernelAbstractions.jl`](https://github.com/JuliaGPU/KernelAbstractions.jl) kernels for evaluating problem data from a (parametric) [`ExaModel`](https://github.com/exanauts/ExaModels.jl) for batches of solutions (and parameters). Currently the following functions (as well as their non-parametric variants) are exported: -- `obj_batch!(::BatchModel, X, Θ)` -- `grad_batch!(::BatchModel, X, Θ)` -- `cons_nln_batch!(::BatchModel, X, Θ)` -- `jac_coord_batch!(::BatchModel, X, Θ)` -- `hess_coord_batch!(::BatchModel, X, Θ, Y; obj_weight=1.0)` -- `jprod_nln_batch!(::BatchModel, X, Θ, V)` -- `jtprod_nln_batch!(::BatchModel, X, Θ, V)` -- `hprod_batch!(::BatchModel, X, Θ, Y, V; obj_weight=1.0)` +- `objective!(::BatchModel, X, Θ)` +- `objective_gradient!(::BatchModel, X, Θ)` +- `constraints!(::BatchModel, X, Θ)` +- `constraints_jacobian!(::BatchModel, X, Θ)` +- `lagrangian_hessian!(::BatchModel, X, Θ, Y; obj_weight=1.0)` +- `constraints_jprod!(::BatchModel, X, Θ, V)` +- `constraints_jtprod!(::BatchModel, X, Θ, V)` +- `lagrangian_hprod!(::BatchModel, X, Θ, Y, V; obj_weight=1.0)` To use these functions, first wrap your `ExaModel` in a `BatchModel`: @@ -30,7 +30,7 @@ This pre-allocates work and output buffers. By default, only the buffers to supp Then, you can call the batch functions as follows: ```julia -objs = obj_batch!(bm, X, Θ) +objs = objective!(bm, X, Θ) ``` where `X` and `Θ` are (device) matrices with dimensions `(nvar, batch_size)` and `(nθ, batch_size)` respectively. diff --git a/ext/BNKChainRulesCore.jl b/ext/BNKChainRulesCore.jl index a9791c9..b3fb93f 100644 --- a/ext/BNKChainRulesCore.jl +++ b/ext/BNKChainRulesCore.jl @@ -3,12 +3,12 @@ module BNKChainRulesCore using BatchNLPKernels using ChainRulesCore -function ChainRulesCore.rrule(::typeof(BatchNLPKernels.obj_batch!), bm::BatchModel, X, Θ) - y = BatchNLPKernels.obj_batch!(bm, X, Θ) +function ChainRulesCore.rrule(::typeof(BatchNLPKernels.objective!), bm::BatchModel, X, Θ) + y = BatchNLPKernels.objective!(bm, X, Θ) function obj_batch_pullback(Ȳ) Ȳ = ChainRulesCore.unthunk(Ȳ) - gradients = BatchNLPKernels.grad_batch!(bm, X, Θ) + gradients = BatchNLPKernels.objective_gradient!(bm, X, Θ) X̄ = gradients .* Ȳ' @@ -17,12 +17,12 @@ function ChainRulesCore.rrule(::typeof(BatchNLPKernels.obj_batch!), bm::BatchMod return y, obj_batch_pullback end -function ChainRulesCore.rrule(::typeof(BatchNLPKernels.obj_batch!), bm::BatchModel, X) - y = BatchNLPKernels.obj_batch!(bm, X) +function ChainRulesCore.rrule(::typeof(BatchNLPKernels.objective!), bm::BatchModel, X) + y = BatchNLPKernels.objective!(bm, X) function obj_batch_pullback(Ȳ) Ȳ = ChainRulesCore.unthunk(Ȳ) - gradients = BatchNLPKernels.grad_batch!(bm, X) + gradients = BatchNLPKernels.objective_gradient!(bm, X) X̄ = gradients .* Ȳ' @@ -33,23 +33,23 @@ function ChainRulesCore.rrule(::typeof(BatchNLPKernels.obj_batch!), bm::BatchMod end -function ChainRulesCore.rrule(::typeof(BatchNLPKernels.cons_nln_batch!), bm::BatchModel, X, Θ) - y = BatchNLPKernels.cons_nln_batch!(bm, X, Θ) +function ChainRulesCore.rrule(::typeof(BatchNLPKernels.constraints!), bm::BatchModel, X, Θ) + y = BatchNLPKernels.constraints!(bm, X, Θ) function cons_nln_batch_pullback(Ȳ) Ȳ = ChainRulesCore.unthunk(Ȳ) - X̄ = BatchNLPKernels.jtprod_nln_batch!(bm, X, Θ, Ȳ) + X̄ = BatchNLPKernels.constraints_jtprod!(bm, X, Θ, Ȳ) return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄, ChainRulesCore.NoTangent() end return y, cons_nln_batch_pullback end -function ChainRulesCore.rrule(::typeof(BatchNLPKernels.cons_nln_batch!), bm::BatchModel, X) - y = BatchNLPKernels.cons_nln_batch!(bm, X) +function ChainRulesCore.rrule(::typeof(BatchNLPKernels.constraints!), bm::BatchModel, X) + y = BatchNLPKernels.constraints!(bm, X) function cons_nln_batch_pullback(Ȳ) Ȳ = ChainRulesCore.unthunk(Ȳ) - X̄ = BatchNLPKernels.jtprod_nln_batch!(bm, X, Ȳ) + X̄ = BatchNLPKernels.constraints_jtprod!(bm, X, Ȳ) return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄ end diff --git a/src/BatchNLPKernels.jl b/src/BatchNLPKernels.jl index d09c930..16602e3 100644 --- a/src/BatchNLPKernels.jl +++ b/src/BatchNLPKernels.jl @@ -11,8 +11,9 @@ include("batch_model.jl") const BOI = BatchNLPKernels export BOI, BatchModel, BatchModelConfig -export obj_batch!, grad_batch!, cons_nln_batch!, jac_coord_batch!, hess_coord_batch! -export jprod_nln_batch!, jtprod_nln_batch!, hprod_batch! +export objective!, objective_gradient!, constraints!, constraints_jacobian!, lagrangian_hessian! +export constraints_jprod!, constraints_jtprod!, lagrangian_hprod! +export all_violations!, constraint_violations!, bound_violations! include("utils.jl") include("kernels.jl") diff --git a/src/api/cons.jl b/src/api/cons.jl index 46e2002..57918e7 100644 --- a/src/api/cons.jl +++ b/src/api/cons.jl @@ -1,26 +1,26 @@ """ - cons_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) + constraints!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) Evaluate constraints for a batch of solutions and parameters. """ -function cons_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) +function constraints!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) C = _maybe_view(bm, :cons_out, X) - cons_nln_batch!(bm, X, Θ, C) + constraints!(bm, X, Θ, C) return C end """ - cons_nln_batch!(bm::BatchModel, X::AbstractMatrix) + constraints!(bm::BatchModel, X::AbstractMatrix) Evaluate constraints for a batch of solutions. """ -function cons_nln_batch!(bm::BatchModel, X::AbstractMatrix) +function constraints!(bm::BatchModel, X::AbstractMatrix) Θ = _repeat_params(bm, X) - cons_nln_batch!(bm, X, Θ) + constraints!(bm, X, Θ) end -function cons_nln_batch!( +function constraints!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -34,7 +34,7 @@ function cons_nln_batch!( _assert_batch_size(batch_size, bm.batch_size) backend = _get_backend(bm.model) - _cons_nln_batch!(backend, C, bm.model.cons, X, Θ) + _constraints!(backend, C, bm.model.cons, X, Θ) conbuffers_batch = _maybe_view(bm, :cons_work, X) @@ -53,17 +53,17 @@ function cons_nln_batch!( return C end -function _cons_nln_batch!(backend, C, con::ExaModels.Constraint, X, Θ) +function _constraints!(backend, C, con::ExaModels.Constraint, X, Θ) if !isempty(con.itr) batch_size = size(X, 2) kerf_batch(backend)(C, con.f, con.itr, X, Θ; ndrange = (length(con.itr), batch_size)) end - _cons_nln_batch!(backend, C, con.inner, X, Θ) + _constraints!(backend, C, con.inner, X, Θ) synchronize(backend) end -function _cons_nln_batch!(backend, C, con::ExaModels.ConstraintNull, X, Θ) end -function _cons_nln_batch!(backend, C, con::ExaModels.ConstraintAug, X, Θ) - _cons_nln_batch!(backend, C, con.inner, X, Θ) +function _constraints!(backend, C, con::ExaModels.ConstraintNull, X, Θ) end +function _constraints!(backend, C, con::ExaModels.ConstraintAug, X, Θ) + _constraints!(backend, C, con.inner, X, Θ) end function _conaugs_batch!(backend, conbuffers, con::ExaModels.ConstraintAug, X, Θ) diff --git a/src/api/grad.jl b/src/api/grad.jl index 0021eda..b363130 100644 --- a/src/api/grad.jl +++ b/src/api/grad.jl @@ -1,30 +1,30 @@ """ - grad_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) + objective_gradient!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) Evaluate objective gradient for a batch of points. """ -function grad_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) +function objective_gradient!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) G = _maybe_view(bm, :grad_out, X) - grad_batch!(bm, X, Θ, G) + objective_gradient!(bm, X, Θ, G) return G end """ - grad_batch!(bm::BatchModel, X::AbstractMatrix) + objective_gradient!(bm::BatchModel, X::AbstractMatrix) Evaluate objective gradient for a batch of points. """ -function grad_batch!(bm::BatchModel, X::AbstractMatrix) +function objective_gradient!(bm::BatchModel, X::AbstractMatrix) Θ = _repeat_params(bm, X) - grad_batch!(bm, X, Θ) + objective_gradient!(bm, X, Θ) end -function _grad_batch!(backend, grad_work, objs, X, Θ) +function _objective_gradient!(backend, grad_work, objs, X, Θ) sgradient_batch!(backend, grad_work, objs, X, Θ, one(eltype(grad_work))) - _grad_batch!(backend, grad_work, objs.inner, X, Θ) + _objective_gradient!(backend, grad_work, objs.inner, X, Θ) synchronize(backend) end -function _grad_batch!(backend, grad_work, objs::ExaModels.ObjectiveNull, X, Θ) end +function _objective_gradient!(backend, grad_work, objs::ExaModels.ObjectiveNull, X, Θ) end function sgradient_batch!( backend::B, @@ -41,11 +41,11 @@ function sgradient_batch!( end """ - grad_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, G::AbstractMatrix) + objective_gradient!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, G::AbstractMatrix) Evaluate gradients for a batch of points with different parameters. """ -function grad_batch!( +function objective_gradient!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -63,7 +63,7 @@ function grad_batch!( if !isempty(grad_work) fill!(grad_work, zero(eltype(grad_work))) - _grad_batch!(backend, grad_work, bm.model.objs, X, Θ) + _objective_gradient!(backend, grad_work, bm.model.objs, X, Θ) fill!(G, zero(eltype(G))) compress_to_dense_batch(backend)( diff --git a/src/api/hess.jl b/src/api/hess.jl index 8b52565..4a43f39 100644 --- a/src/api/hess.jl +++ b/src/api/hess.jl @@ -1,25 +1,25 @@ """ - hess_coord_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) + lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) Evaluate Hessian coordinates for a batch of points. """ -function hess_coord_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) +function lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) H_view = _maybe_view(bm, :hprod_work, X) - hess_coord_batch!(bm, X, Θ, Y, H_view; obj_weight=obj_weight) + lagrangian_hessian!(bm, X, Θ, Y, H_view; obj_weight=obj_weight) return H_view end """ - hess_coord_batch!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) + lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) Evaluate Hessian coordinates for a batch of points. """ -function hess_coord_batch!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) +function lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0) Θ = _repeat_params(bm, X) - hess_coord_batch!(bm, X, Θ, Y; obj_weight=obj_weight) + lagrangian_hessian!(bm, X, Θ, Y; obj_weight=obj_weight) end -function hess_coord_batch!( +function lagrangian_hessian!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -37,24 +37,24 @@ function hess_coord_batch!( backend = _get_backend(bm.model) fill!(H, zero(eltype(H))) - _obj_hess_coord_batch!(backend, H, bm.model.objs, X, Θ, obj_weight) - _con_hess_coord_batch!(backend, H, bm.model.cons, X, Θ, Y) + _obj_lagrangian_hessian!(backend, H, bm.model.objs, X, Θ, obj_weight) + _con_lagrangian_hessian!(backend, H, bm.model.cons, X, Θ, Y) return H end -function _obj_hess_coord_batch!(backend, H, objs, X, Θ, obj_weight) +function _obj_lagrangian_hessian!(backend, H, objs, X, Θ, obj_weight) shessian_batch!(backend, H, nothing, objs, X, Θ, obj_weight, zero(eltype(H))) - _obj_hess_coord_batch!(backend, H, objs.inner, X, Θ, obj_weight) + _obj_lagrangian_hessian!(backend, H, objs.inner, X, Θ, obj_weight) synchronize(backend) end -function _obj_hess_coord_batch!(backend, H, objs::ExaModels.ObjectiveNull, X, Θ, obj_weight) end +function _obj_lagrangian_hessian!(backend, H, objs::ExaModels.ObjectiveNull, X, Θ, obj_weight) end -function _con_hess_coord_batch!(backend, H, cons, X, Θ, Y) +function _con_lagrangian_hessian!(backend, H, cons, X, Θ, Y) shessian_batch!(backend, H, nothing, cons, X, Θ, Y, zero(eltype(H))) - _con_hess_coord_batch!(backend, H, cons.inner, X, Θ, Y) + _con_lagrangian_hessian!(backend, H, cons.inner, X, Θ, Y) synchronize(backend) end -function _con_hess_coord_batch!(backend, H, cons::ExaModels.ConstraintNull, X, Θ, Y) end +function _con_lagrangian_hessian!(backend, H, cons::ExaModels.ConstraintNull, X, Θ, Y) end function shessian_batch!( backend::B, diff --git a/src/api/hprod.jl b/src/api/hprod.jl index 36a9d67..bdbd372 100644 --- a/src/api/hprod.jl +++ b/src/api/hprod.jl @@ -1,26 +1,26 @@ """ - hprod_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) + lagrangian_hprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) Evaluate Hessian-vector products for a batch of points. """ -function hprod_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) +function lagrangian_hprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) Hv = _maybe_view(bm, :hprod_out, X) - hprod_batch!(bm, X, Θ, Y, V, Hv; obj_weight=obj_weight) + lagrangian_hprod!(bm, X, Θ, Y, V, Hv; obj_weight=obj_weight) return Hv end """ - hprod_batch!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) + lagrangian_hprod!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) Evaluate Hessian-vector products for a batch of points. """ -function hprod_batch!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) +function lagrangian_hprod!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix, V::AbstractMatrix; obj_weight=1.0) Θ = _repeat_params(bm, X) - hprod_batch!(bm, X, Θ, Y, V; obj_weight=obj_weight) + lagrangian_hprod!(bm, X, Θ, Y, V; obj_weight=obj_weight) return Hv end -function hprod_batch!( +function lagrangian_hprod!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -40,7 +40,7 @@ function hprod_batch!( H_batch = _maybe_view(bm, :hprod_work, X) - hess_coord_batch!(bm, X, Θ, Y, H_batch; obj_weight=obj_weight) + lagrangian_hessian!(bm, X, Θ, Y, H_batch; obj_weight=obj_weight) fill!(Hv, zero(eltype(Hv))) kersyspmv_batch(backend)( diff --git a/src/api/jac.jl b/src/api/jac.jl index 72bffc9..0e970a6 100644 --- a/src/api/jac.jl +++ b/src/api/jac.jl @@ -1,25 +1,25 @@ """ - jac_coord_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) + constraints_jacobian!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) Evaluate Jacobian coordinates for a batch of points. """ -function jac_coord_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) +function constraints_jacobian!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) J_view = _maybe_view(bm, :jprod_work, X) - jac_coord_batch!(bm, X, Θ, J_view) + constraints_jacobian!(bm, X, Θ, J_view) return J_view end """ - jac_coord_batch!(bm::BatchModel, X::AbstractMatrix) + constraints_jacobian!(bm::BatchModel, X::AbstractMatrix) Evaluate Jacobian coordinates for a batch of points. """ -function jac_coord_batch!(bm::BatchModel, X::AbstractMatrix) +function constraints_jacobian!(bm::BatchModel, X::AbstractMatrix) Θ = _repeat_params(bm, X) - jac_coord_batch!(bm, X, Θ) + constraints_jacobian!(bm, X, Θ) end -function jac_coord_batch!( +function constraints_jacobian!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -34,16 +34,16 @@ function jac_coord_batch!( backend = _get_backend(bm.model) fill!(J, zero(eltype(J))) - _jac_coord_batch!(backend, J, bm.model.cons, X, Θ) + _constraints_jacobian!(backend, J, bm.model.cons, X, Θ) return J end -function _jac_coord_batch!(backend, J, cons, X, Θ) +function _constraints_jacobian!(backend, J, cons, X, Θ) sjacobian_batch!(backend, J, nothing, cons, X, Θ, one(eltype(J))) - _jac_coord_batch!(backend, J, cons.inner, X, Θ) + _constraints_jacobian!(backend, J, cons.inner, X, Θ) synchronize(backend) end -function _jac_coord_batch!(backend, J, cons::ExaModels.ConstraintNull, X, Θ) end +function _constraints_jacobian!(backend, J, cons::ExaModels.ConstraintNull, X, Θ) end function sjacobian_batch!( backend::B, diff --git a/src/api/jprod.jl b/src/api/jprod.jl index 60dc0de..b298cd6 100644 --- a/src/api/jprod.jl +++ b/src/api/jprod.jl @@ -1,31 +1,31 @@ """ - jprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) + constraints_jprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) Evaluate Jacobian-vector products for a batch of points. """ -function jprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) +function constraints_jprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) Jv = _maybe_view(bm, :jprod_out, X) - jprod_nln_batch!(bm, X, Θ, V, Jv) + constraints_jprod!(bm, X, Θ, V, Jv) return Jv end """ - jprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) + constraints_jprod!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) Evaluate Jacobian-vector products for a batch of points. """ -function jprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) +function constraints_jprod!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) Θ = _repeat_params(bm, X) - jprod_nln_batch!(bm, X, Θ, V) + constraints_jprod!(bm, X, Θ, V) end """ - jprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix, Jv::AbstractMatrix) + constraints_jprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix, Jv::AbstractMatrix) Evaluate Jacobian-vector products for a batch of points. """ -function jprod_nln_batch!( +function constraints_jprod!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -42,7 +42,7 @@ function jprod_nln_batch!( J_batch = _maybe_view(bm, :jprod_work, X) - jac_coord_batch!(bm, X, Θ, J_batch) + constraints_jacobian!(bm, X, Θ, J_batch) fill!(Jv, zero(eltype(Jv))) kerspmv_batch(bm.model.ext.backend)( @@ -60,32 +60,32 @@ end """ - jtprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) + constraints_jtprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) Evaluate Jacobian-transpose-vector products for a batch of points. """ -function jtprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) +function constraints_jtprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix) Jtv = _maybe_view(bm, :jtprod_out, X) - jtprod_nln_batch!(bm, X, Θ, V, Jtv) + constraints_jtprod!(bm, X, Θ, V, Jtv) return Jtv end """ - jtprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) + constraints_jtprod!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) Evaluate Jacobian-transpose-vector products for a batch of points. """ -function jtprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) +function constraints_jtprod!(bm::BatchModel, X::AbstractMatrix, V::AbstractMatrix) Θ = _repeat_params(bm, X) - jtprod_nln_batch!(bm, X, Θ, V) + constraints_jtprod!(bm, X, Θ, V) end """ - jtprod_nln_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix, Jtv::AbstractMatrix) + constraints_jtprod!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, V::AbstractMatrix, Jtv::AbstractMatrix) Evaluate Jacobian-transpose-vector products for a batch of points. """ -function jtprod_nln_batch!( +function constraints_jtprod!( bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, @@ -104,7 +104,7 @@ function jtprod_nln_batch!( J_batch = _maybe_view(bm, :jprod_work, X) - jac_coord_batch!(bm, X, Θ, J_batch) + constraints_jacobian!(bm, X, Θ, J_batch) fill!(Jtv, zero(eltype(Jtv))) kerspmv2_batch(backend)( diff --git a/src/api/obj.jl b/src/api/obj.jl index 1b2ce21..15a7c15 100644 --- a/src/api/obj.jl +++ b/src/api/obj.jl @@ -1,25 +1,25 @@ """ - obj_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) + objective!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) Evaluate objective function for a batch of points. """ -function obj_batch!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) +function objective!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) obj_work = _maybe_view(bm, :obj_work, X) - return obj_batch!(bm, obj_work, X, Θ) + return objective!(bm, obj_work, X, Θ) end """ - obj_batch!(bm::BatchModel, X::AbstractMatrix) + objective!(bm::BatchModel, X::AbstractMatrix) Evaluate objective function for a batch of points. """ -function obj_batch!(bm::BatchModel, X::AbstractMatrix) +function objective!(bm::BatchModel, X::AbstractMatrix) Θ = _repeat_params(bm, X) - return obj_batch!(bm, X, Θ) + return objective!(bm, X, Θ) end -function obj_batch!( +function objective!( bm::BatchModel, obj_work, X::AbstractMatrix, diff --git a/src/api/viols.jl b/src/api/viols.jl index 360d3e5..18c1799 100644 --- a/src/api/viols.jl +++ b/src/api/viols.jl @@ -4,7 +4,7 @@ Compute all constraint and variable violations for a batch of solutions. """ function all_violations!(bm::BatchModel, X::AbstractMatrix) - V = cons_nln_batch!(bm, X) + V = constraints!(bm, X) Vc = constraint_violations!(bm, V) Vb = bound_violations!(bm, X) @@ -18,7 +18,7 @@ end Compute all constraint and variable violations for a batch of solutions and parameters. """ function all_violations!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) - V = cons_nln_batch!(bm, X, Θ) + V = constraints!(bm, X, Θ) Vc = constraint_violations!(bm, V) Vb = bound_violations!(bm, X) @@ -27,11 +27,25 @@ function all_violations!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) end """ - constraint_violations!(bm::BatchModel, V::AbstractMatrix) + constraint_violations!(bm::BatchModel, X::AbstractMatrix) Compute constraint violations for a batch of constraint primal values. """ -function constraint_violations!(bm::BatchModel, V::AbstractMatrix) +function constraint_violations!(bm::BatchModel, X::AbstractMatrix) + V = constraints!(bm, X) + return _constraint_violations!(bm, V) +end +""" + constraint_violations!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) + +Compute constraint violations for a batch of constraint primal values. +""" +function constraint_violations!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) + V = constraints!(bm, X, Θ) + return _constraint_violations!(bm, V) +end + +function _constraint_violations!(bm::BatchModel, V::AbstractMatrix) viols_cons_out = _maybe_view(bm, :viols_cons_out, V) _violation!.(eachcol(viols_cons_out), eachcol(V), bm.viols_cons) diff --git a/test/api.jl b/test/api.jl index c1d5f33..6686255 100644 --- a/test/api.jl +++ b/test/api.jl @@ -12,7 +12,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Model Info: $(nvar) vars, $(ncon) cons, $(nθ) params" begin @testset "Objective" begin - obj_vals = BOI.obj_batch!(bm, X, Θ) + obj_vals = BOI.objective!(bm, X, Θ) @test length(obj_vals) == batch_size @test all(isfinite, obj_vals) for i in 1:batch_size @@ -23,7 +23,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Constraint" begin if ncon > 0 - cons_vals = BOI.cons_nln_batch!(bm, X, Θ) + cons_vals = BOI.constraints!(bm, X, Θ) @test size(cons_vals) == (ncon, batch_size) @test all(isfinite, cons_vals) for i in 1:batch_size @@ -36,7 +36,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; end @testset "Gradient" begin - grad_vals = BOI.grad_batch!(bm, X, Θ) + grad_vals = BOI.objective_gradient!(bm, X, Θ) @test size(grad_vals) == (nvar, batch_size) @test all(isfinite, grad_vals) for i in 1:batch_size @@ -50,7 +50,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Jacobian-Vector Product" begin if ncon > 0 V = MOD.randn(nvar, batch_size) - jprod_vals = BOI.jprod_nln_batch!(bm, X, Θ, V) + jprod_vals = BOI.constraints_jprod!(bm, X, Θ, V) @test size(jprod_vals) == (ncon, batch_size) @test all(isfinite, jprod_vals) for i in 1:batch_size @@ -65,7 +65,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Jacobian-Transpose-Vector Product" begin if ncon > 0 V = MOD.randn(ncon, batch_size) - jtprod_vals = BOI.jtprod_nln_batch!(bm, X, Θ, V) + jtprod_vals = BOI.constraints_jtprod!(bm, X, Θ, V) @test size(jtprod_vals) == (nvar, batch_size) @test all(isfinite, jtprod_vals) for i in 1:batch_size @@ -81,7 +81,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; V = MOD.randn(nvar, batch_size) if ncon > 0 Y = MOD.randn(ncon, batch_size) - hprod_vals = BOI.hprod_batch!(bm, X, Θ, Y, V) + hprod_vals = BOI.lagrangian_hprod!(bm, X, Θ, Y, V) @test size(hprod_vals) == (nvar, batch_size) @test all(isfinite, hprod_vals) for i in 1:batch_size @@ -92,7 +92,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; end else Y = MOD.zeros(ncon, batch_size) - hprod_vals = BOI.hprod_batch!(bm, X, Θ, Y, V) + hprod_vals = BOI.lagrangian_hprod!(bm, X, Θ, Y, V) @test size(hprod_vals) == (nvar, batch_size) @test all(isfinite, hprod_vals) for i in 1:batch_size @@ -106,60 +106,60 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Batch Size Validation" begin X_large = MOD.randn(nvar, batch_size + 1) - @test_throws AssertionError BOI.obj_batch!(bm, X_large) + @test_throws AssertionError BOI.objective!(bm, X_large) if ncon > 0 - @test_throws AssertionError BOI.cons_nln_batch!(bm, X_large) + @test_throws AssertionError BOI.constraints!(bm, X_large) end - @test_throws AssertionError BOI.grad_batch!(bm, X_large) + @test_throws AssertionError BOI.objective_gradient!(bm, X_large) if ncon > 0 V_jprod = MOD.randn(nvar, batch_size + 1) - @test_throws AssertionError BOI.jprod_nln_batch!(bm, X_large, V_jprod) + @test_throws AssertionError BOI.constraints_jprod!(bm, X_large, V_jprod) V_jtprod = MOD.randn(ncon, batch_size + 1) - @test_throws AssertionError BOI.jtprod_nln_batch!(bm, X_large, V_jtprod) + @test_throws AssertionError BOI.constraints_jtprod!(bm, X_large, V_jtprod) end V_hprod = MOD.randn(nvar, batch_size + 1) if ncon > 0 Y_large = MOD.randn(ncon, batch_size + 1) - @test_throws AssertionError BOI.hprod_batch!(bm, X_large, Y_large, V_hprod) + @test_throws AssertionError BOI.lagrangian_hprod!(bm, X_large, Y_large, V_hprod) else Y_large = MOD.zeros(ncon, batch_size + 1) - @test_throws AssertionError BOI.hprod_batch!(bm, X_large, Y_large, V_hprod) + @test_throws AssertionError BOI.lagrangian_hprod!(bm, X_large, Y_large, V_hprod) end end @testset "Dimension Validation" begin X_wrong = MOD.randn(nvar + 1, batch_size) - @test_throws DimensionMismatch BOI.obj_batch!(bm, X_wrong) + @test_throws DimensionMismatch BOI.objective!(bm, X_wrong) if nθ > 0 Θ_wrong = MOD.randn(nθ + 1, batch_size) - @test_throws DimensionMismatch BOI.obj_batch!(bm, X, Θ_wrong) + @test_throws DimensionMismatch BOI.objective!(bm, X, Θ_wrong) end if ncon > 0 V_jprod_wrong = MOD.randn(nvar + 1, batch_size) - @test_throws DimensionMismatch BOI.jprod_nln_batch!(bm, X, V_jprod_wrong) + @test_throws DimensionMismatch BOI.constraints_jprod!(bm, X, V_jprod_wrong) V_jtprod_wrong = MOD.randn(ncon + 1, batch_size) - @test_throws DimensionMismatch BOI.jtprod_nln_batch!(bm, X, V_jtprod_wrong) + @test_throws DimensionMismatch BOI.constraints_jtprod!(bm, X, V_jtprod_wrong) Y_wrong = MOD.randn(ncon + 1, batch_size) V_hprod = MOD.randn(nvar, batch_size) - @test_throws DimensionMismatch BOI.hprod_batch!(bm, X, Y_wrong, V_hprod) + @test_throws DimensionMismatch BOI.lagrangian_hprod!(bm, X, Y_wrong, V_hprod) end V_hprod_wrong = MOD.randn(nvar + 1, batch_size) if ncon > 0 Y = MOD.randn(ncon, batch_size) - @test_throws DimensionMismatch BOI.hprod_batch!(bm, X, Y, V_hprod_wrong) + @test_throws DimensionMismatch BOI.lagrangian_hprod!(bm, X, Y, V_hprod_wrong) else Y = MOD.zeros(ncon, batch_size) - @test_throws DimensionMismatch BOI.hprod_batch!(bm, X, Y, V_hprod_wrong) + @test_throws DimensionMismatch BOI.lagrangian_hprod!(bm, X, Y, V_hprod_wrong) end end end diff --git a/test/config.jl b/test/config.jl index 163c0bc..f8eb7fc 100644 --- a/test/config.jl +++ b/test/config.jl @@ -10,11 +10,11 @@ @testset "Minimal" begin bm_minimal = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:minimal)) - @test_throws ArgumentError BOI.grad_batch!(bm_minimal, X, Θ) - @test_throws ArgumentError BOI.jac_coord_batch!(bm_minimal, X, Θ) + @test_throws ArgumentError BOI.objective_gradient!(bm_minimal, X, Θ) + @test_throws ArgumentError BOI.constraints_jacobian!(bm_minimal, X, Θ) if ncon > 0 Y = OpenCL.randn(ncon, batch_size) - @test_throws ArgumentError BOI.hess_coord_batch!(bm_minimal, X, Θ, Y) + @test_throws ArgumentError BOI.lagrangian_hessian!(bm_minimal, X, Θ, Y) end end @@ -24,29 +24,29 @@ if ncon > 0 V = OpenCL.randn(nvar, batch_size) - @test_throws ArgumentError BOI.jprod_nln_batch!(bm_partial, X, Θ, V) + @test_throws ArgumentError BOI.constraints_jprod!(bm_partial, X, Θ, V) V_t = OpenCL.randn(ncon, batch_size) - @test_throws ArgumentError BOI.jtprod_nln_batch!(bm_partial, X, Θ, V_t) + @test_throws ArgumentError BOI.constraints_jtprod!(bm_partial, X, Θ, V_t) end V_h = OpenCL.randn(nvar, batch_size) if ncon > 0 Y = OpenCL.randn(ncon, batch_size) - @test_throws ArgumentError BOI.hprod_batch!(bm_partial, X, Θ, Y, V_h) + @test_throws ArgumentError BOI.lagrangian_hprod!(bm_partial, X, Θ, Y, V_h) else Y = OpenCL.zeros(ncon, batch_size) - @test_throws ArgumentError BOI.hprod_batch!(bm_partial, X, Θ, Y, V_h) + @test_throws ArgumentError BOI.lagrangian_hprod!(bm_partial, X, Θ, Y, V_h) end end @testset "Output" begin bm_no_gradout = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(obj=true, cons=true, grad=false, jac=false, hess=false, jprod=false, jtprod=false, hprod=false)) G = OpenCL.randn(nvar, batch_size) - @test_throws ArgumentError BOI.grad_batch!(bm_no_gradout, X, Θ, G) - @test_throws ArgumentError BOI.grad_batch!(bm_no_gradout, X, Θ) + @test_throws ArgumentError BOI.objective_gradient!(bm_no_gradout, X, Θ, G) + @test_throws ArgumentError BOI.objective_gradient!(bm_no_gradout, X, Θ) bm_no_cons = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(obj=true, cons=false, grad=false, jac=false, hess=false, jprod=false, jtprod=false, hprod=false)) - @test_throws ArgumentError BOI.cons_nln_batch!(bm_no_cons, X, Θ) + @test_throws ArgumentError BOI.constraints!(bm_no_cons, X, Θ) end end diff --git a/test/test_diff.jl b/test/test_diff.jl index 6298fb3..af70d39 100644 --- a/test/test_diff.jl +++ b/test/test_diff.jl @@ -8,14 +8,14 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) X_gpu = MOD.randn(nvar, batch_size) Θ_gpu = MOD.randn(nθ, batch_size) - @testset "obj_batch!" begin - y = BOI.obj_batch!(bm, X_gpu, Θ_gpu) + @testset "objective!" begin + y = BOI.objective!(bm, X_gpu, Θ_gpu) @test size(y) == (batch_size,) function f_gpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.obj_batch!(bm, X, Θ)) + return sum(BOI.objective!(bm, X, Θ)) end params = vcat(X_gpu, Θ_gpu) @@ -26,14 +26,14 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) ncon == 0 && return - @testset "cons_nln_batch!" begin - y = BOI.cons_nln_batch!(bm, X_gpu, Θ_gpu) + @testset "constraints!" begin + y = BOI.constraints!(bm, X_gpu, Θ_gpu) @test size(y) == (ncon, batch_size) function f_gpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.cons_nln_batch!(bm, X, Θ)) + return sum(BOI.constraints!(bm, X, Θ)) end params = vcat(X_gpu, Θ_gpu) @@ -53,14 +53,14 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) X_cpu = randn(nvar, batch_size) Θ_cpu = randn(nθ, batch_size) - @testset "obj_batch! CPU" begin - y = BOI.obj_batch!(bm, X_cpu, Θ_cpu) + @testset "objective! CPU" begin + y = BOI.objective!(bm, X_cpu, Θ_cpu) @test size(y) == (batch_size,) function f_cpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.obj_batch!(bm, X, Θ)) + return sum(BOI.objective!(bm, X, Θ)) end params = vcat(X_cpu, Θ_cpu) @@ -68,7 +68,7 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) @test grad isa AbstractMatrix @test size(grad) == size(params) - @testset "FiniteDifferences obj_batch!" begin + @testset "FiniteDifferences objective!" begin gradfd = DI.gradient(f_cpu, AutoFiniteDifferences(fdm=FiniteDifferences.central_fdm(3,1)), params) @test gradfd[1:nvar,:] ≈ grad[1:nvar,:] atol=1e-4 rtol=1e-4 end @@ -76,14 +76,14 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) ncon == 0 && return - @testset "cons_nln_batch! CPU" begin - y = BOI.cons_nln_batch!(bm, X_cpu, Θ_cpu) + @testset "constraints! CPU" begin + y = BOI.constraints!(bm, X_cpu, Θ_cpu) @test size(y) == (ncon, batch_size) function f_cpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.cons_nln_batch!(bm, X, Θ)) + return sum(BOI.constraints!(bm, X, Θ)) end params = vcat(X_cpu, Θ_cpu) @@ -91,7 +91,7 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) @test grad isa AbstractMatrix @test size(grad) == size(params) - @testset "FiniteDifferences cons_nln_batch!" begin + @testset "FiniteDifferences constraints!" begin gradfd = DI.gradient(f_cpu, AutoFiniteDifferences(fdm=FiniteDifferences.central_fdm(3,1)), params) @test gradfd[1:nvar,:] ≈ grad[1:nvar,:] atol=1e-4 rtol=1e-4 end diff --git a/test/test_viols.jl b/test/test_viols.jl index e678db8..b80eca2 100644 --- a/test/test_viols.jl +++ b/test/test_viols.jl @@ -52,7 +52,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; @testset "Constraint Violations" begin if ncon > 0 - V_cons = BOI.cons_nln_batch!(bm, X, Θ) + V_cons = BOI.constraints!(bm, X, Θ) Vc = BOI.constraint_violations!(bm, V_cons) @test size(Vc) == (ncon, batch_size) @@ -178,7 +178,7 @@ function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; function f_cons_viols(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - V_cons = BOI.cons_nln_batch!(bm, X, Θ) + V_cons = BOI.constraints!(bm, X, Θ) Vc = BOI.constraint_violations!(bm, V_cons) return sum(Vc) end From 46a0b2f39b42660d967d5cd3cef1e4cc69777c6f Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 20:41:30 -0400 Subject: [PATCH 02/10] readme viols --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 07d4c45..84514af 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,9 @@ - `constraints_jprod!(::BatchModel, X, Θ, V)` - `constraints_jtprod!(::BatchModel, X, Θ, V)` - `lagrangian_hprod!(::BatchModel, X, Θ, Y, V; obj_weight=1.0)` - +- `all_violations!(::BatchModel, X, Θ)` +- `constraint_violations!(::BatchModel, X, Θ)` +- `bound_violations!(::BatchModel, X)` To use these functions, first wrap your `ExaModel` in a `BatchModel`: From 9d2a22a62d80d7d27343bb80f645b1e6a11580c9 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 20:42:11 -0400 Subject: [PATCH 03/10] BOI->BNK --- src/BatchNLPKernels.jl | 4 ++-- test/api.jl | 44 +++++++++++++++++------------------ test/config.jl | 28 +++++++++++------------ test/test_diff.jl | 20 ++++++++-------- test/test_viols.jl | 52 +++++++++++++++++++++--------------------- 5 files changed, 74 insertions(+), 74 deletions(-) diff --git a/src/BatchNLPKernels.jl b/src/BatchNLPKernels.jl index 16602e3..1d8091e 100644 --- a/src/BatchNLPKernels.jl +++ b/src/BatchNLPKernels.jl @@ -9,8 +9,8 @@ const KAExtension = ExaKA.KAExtension include("interval.jl") include("batch_model.jl") -const BOI = BatchNLPKernels -export BOI, BatchModel, BatchModelConfig +const BNK = BatchNLPKernels +export BNK, BatchModel, BatchModelConfig export objective!, objective_gradient!, constraints!, constraints_jacobian!, lagrangian_hessian! export constraints_jprod!, constraints_jtprod!, lagrangian_hprod! export all_violations!, constraint_violations!, bound_violations! diff --git a/test/api.jl b/test/api.jl index 6686255..d0edb53 100644 --- a/test/api.jl +++ b/test/api.jl @@ -1,7 +1,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; atol::Float64=1e-10, rtol::Float64=1e-10, MOD=OpenCL) - bm = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:full)) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) nvar = model.meta.nvar ncon = model.meta.ncon @@ -12,7 +12,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Model Info: $(nvar) vars, $(ncon) cons, $(nθ) params" begin @testset "Objective" begin - obj_vals = BOI.objective!(bm, X, Θ) + obj_vals = BNK.objective!(bm, X, Θ) @test length(obj_vals) == batch_size @test all(isfinite, obj_vals) for i in 1:batch_size @@ -23,7 +23,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Constraint" begin if ncon > 0 - cons_vals = BOI.constraints!(bm, X, Θ) + cons_vals = BNK.constraints!(bm, X, Θ) @test size(cons_vals) == (ncon, batch_size) @test all(isfinite, cons_vals) for i in 1:batch_size @@ -36,7 +36,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; end @testset "Gradient" begin - grad_vals = BOI.objective_gradient!(bm, X, Θ) + grad_vals = BNK.objective_gradient!(bm, X, Θ) @test size(grad_vals) == (nvar, batch_size) @test all(isfinite, grad_vals) for i in 1:batch_size @@ -50,7 +50,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Jacobian-Vector Product" begin if ncon > 0 V = MOD.randn(nvar, batch_size) - jprod_vals = BOI.constraints_jprod!(bm, X, Θ, V) + jprod_vals = BNK.constraints_jprod!(bm, X, Θ, V) @test size(jprod_vals) == (ncon, batch_size) @test all(isfinite, jprod_vals) for i in 1:batch_size @@ -65,7 +65,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Jacobian-Transpose-Vector Product" begin if ncon > 0 V = MOD.randn(ncon, batch_size) - jtprod_vals = BOI.constraints_jtprod!(bm, X, Θ, V) + jtprod_vals = BNK.constraints_jtprod!(bm, X, Θ, V) @test size(jtprod_vals) == (nvar, batch_size) @test all(isfinite, jtprod_vals) for i in 1:batch_size @@ -81,7 +81,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; V = MOD.randn(nvar, batch_size) if ncon > 0 Y = MOD.randn(ncon, batch_size) - hprod_vals = BOI.lagrangian_hprod!(bm, X, Θ, Y, V) + hprod_vals = BNK.lagrangian_hprod!(bm, X, Θ, Y, V) @test size(hprod_vals) == (nvar, batch_size) @test all(isfinite, hprod_vals) for i in 1:batch_size @@ -92,7 +92,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; end else Y = MOD.zeros(ncon, batch_size) - hprod_vals = BOI.lagrangian_hprod!(bm, X, Θ, Y, V) + hprod_vals = BNK.lagrangian_hprod!(bm, X, Θ, Y, V) @test size(hprod_vals) == (nvar, batch_size) @test all(isfinite, hprod_vals) for i in 1:batch_size @@ -106,60 +106,60 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Batch Size Validation" begin X_large = MOD.randn(nvar, batch_size + 1) - @test_throws AssertionError BOI.objective!(bm, X_large) + @test_throws AssertionError BNK.objective!(bm, X_large) if ncon > 0 - @test_throws AssertionError BOI.constraints!(bm, X_large) + @test_throws AssertionError BNK.constraints!(bm, X_large) end - @test_throws AssertionError BOI.objective_gradient!(bm, X_large) + @test_throws AssertionError BNK.objective_gradient!(bm, X_large) if ncon > 0 V_jprod = MOD.randn(nvar, batch_size + 1) - @test_throws AssertionError BOI.constraints_jprod!(bm, X_large, V_jprod) + @test_throws AssertionError BNK.constraints_jprod!(bm, X_large, V_jprod) V_jtprod = MOD.randn(ncon, batch_size + 1) - @test_throws AssertionError BOI.constraints_jtprod!(bm, X_large, V_jtprod) + @test_throws AssertionError BNK.constraints_jtprod!(bm, X_large, V_jtprod) end V_hprod = MOD.randn(nvar, batch_size + 1) if ncon > 0 Y_large = MOD.randn(ncon, batch_size + 1) - @test_throws AssertionError BOI.lagrangian_hprod!(bm, X_large, Y_large, V_hprod) + @test_throws AssertionError BNK.lagrangian_hprod!(bm, X_large, Y_large, V_hprod) else Y_large = MOD.zeros(ncon, batch_size + 1) - @test_throws AssertionError BOI.lagrangian_hprod!(bm, X_large, Y_large, V_hprod) + @test_throws AssertionError BNK.lagrangian_hprod!(bm, X_large, Y_large, V_hprod) end end @testset "Dimension Validation" begin X_wrong = MOD.randn(nvar + 1, batch_size) - @test_throws DimensionMismatch BOI.objective!(bm, X_wrong) + @test_throws DimensionMismatch BNK.objective!(bm, X_wrong) if nθ > 0 Θ_wrong = MOD.randn(nθ + 1, batch_size) - @test_throws DimensionMismatch BOI.objective!(bm, X, Θ_wrong) + @test_throws DimensionMismatch BNK.objective!(bm, X, Θ_wrong) end if ncon > 0 V_jprod_wrong = MOD.randn(nvar + 1, batch_size) - @test_throws DimensionMismatch BOI.constraints_jprod!(bm, X, V_jprod_wrong) + @test_throws DimensionMismatch BNK.constraints_jprod!(bm, X, V_jprod_wrong) V_jtprod_wrong = MOD.randn(ncon + 1, batch_size) - @test_throws DimensionMismatch BOI.constraints_jtprod!(bm, X, V_jtprod_wrong) + @test_throws DimensionMismatch BNK.constraints_jtprod!(bm, X, V_jtprod_wrong) Y_wrong = MOD.randn(ncon + 1, batch_size) V_hprod = MOD.randn(nvar, batch_size) - @test_throws DimensionMismatch BOI.lagrangian_hprod!(bm, X, Y_wrong, V_hprod) + @test_throws DimensionMismatch BNK.lagrangian_hprod!(bm, X, Y_wrong, V_hprod) end V_hprod_wrong = MOD.randn(nvar + 1, batch_size) if ncon > 0 Y = MOD.randn(ncon, batch_size) - @test_throws DimensionMismatch BOI.lagrangian_hprod!(bm, X, Y, V_hprod_wrong) + @test_throws DimensionMismatch BNK.lagrangian_hprod!(bm, X, Y, V_hprod_wrong) else Y = MOD.zeros(ncon, batch_size) - @test_throws DimensionMismatch BOI.lagrangian_hprod!(bm, X, Y, V_hprod_wrong) + @test_throws DimensionMismatch BNK.lagrangian_hprod!(bm, X, Y, V_hprod_wrong) end end end diff --git a/test/config.jl b/test/config.jl index f8eb7fc..0774ad6 100644 --- a/test/config.jl +++ b/test/config.jl @@ -9,44 +9,44 @@ Θ = OpenCL.randn(nθ, batch_size) @testset "Minimal" begin - bm_minimal = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:minimal)) - @test_throws ArgumentError BOI.objective_gradient!(bm_minimal, X, Θ) - @test_throws ArgumentError BOI.constraints_jacobian!(bm_minimal, X, Θ) + bm_minimal = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:minimal)) + @test_throws ArgumentError BNK.objective_gradient!(bm_minimal, X, Θ) + @test_throws ArgumentError BNK.constraints_jacobian!(bm_minimal, X, Θ) if ncon > 0 Y = OpenCL.randn(ncon, batch_size) - @test_throws ArgumentError BOI.lagrangian_hessian!(bm_minimal, X, Θ, Y) + @test_throws ArgumentError BNK.lagrangian_hessian!(bm_minimal, X, Θ, Y) end end @testset "Mat-vec" begin model_with_prod = create_luksan_vlcek_model(5; M = 1, prod = true) - bm_partial = BOI.BatchModel(model_with_prod, batch_size, config=BOI.BatchModelConfig(obj=true, cons=true, grad=true, jac=true, hess=true, jprod=false, jtprod=false, hprod=false)) + bm_partial = BNK.BatchModel(model_with_prod, batch_size, config=BNK.BatchModelConfig(obj=true, cons=true, grad=true, jac=true, hess=true, jprod=false, jtprod=false, hprod=false)) if ncon > 0 V = OpenCL.randn(nvar, batch_size) - @test_throws ArgumentError BOI.constraints_jprod!(bm_partial, X, Θ, V) + @test_throws ArgumentError BNK.constraints_jprod!(bm_partial, X, Θ, V) V_t = OpenCL.randn(ncon, batch_size) - @test_throws ArgumentError BOI.constraints_jtprod!(bm_partial, X, Θ, V_t) + @test_throws ArgumentError BNK.constraints_jtprod!(bm_partial, X, Θ, V_t) end V_h = OpenCL.randn(nvar, batch_size) if ncon > 0 Y = OpenCL.randn(ncon, batch_size) - @test_throws ArgumentError BOI.lagrangian_hprod!(bm_partial, X, Θ, Y, V_h) + @test_throws ArgumentError BNK.lagrangian_hprod!(bm_partial, X, Θ, Y, V_h) else Y = OpenCL.zeros(ncon, batch_size) - @test_throws ArgumentError BOI.lagrangian_hprod!(bm_partial, X, Θ, Y, V_h) + @test_throws ArgumentError BNK.lagrangian_hprod!(bm_partial, X, Θ, Y, V_h) end end @testset "Output" begin - bm_no_gradout = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(obj=true, cons=true, grad=false, jac=false, hess=false, jprod=false, jtprod=false, hprod=false)) + bm_no_gradout = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(obj=true, cons=true, grad=false, jac=false, hess=false, jprod=false, jtprod=false, hprod=false)) G = OpenCL.randn(nvar, batch_size) - @test_throws ArgumentError BOI.objective_gradient!(bm_no_gradout, X, Θ, G) - @test_throws ArgumentError BOI.objective_gradient!(bm_no_gradout, X, Θ) + @test_throws ArgumentError BNK.objective_gradient!(bm_no_gradout, X, Θ, G) + @test_throws ArgumentError BNK.objective_gradient!(bm_no_gradout, X, Θ) - bm_no_cons = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(obj=true, cons=false, grad=false, jac=false, hess=false, jprod=false, jtprod=false, hprod=false)) - @test_throws ArgumentError BOI.constraints!(bm_no_cons, X, Θ) + bm_no_cons = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(obj=true, cons=false, grad=false, jac=false, hess=false, jprod=false, jtprod=false, hprod=false)) + @test_throws ArgumentError BNK.constraints!(bm_no_cons, X, Θ) end end diff --git a/test/test_diff.jl b/test/test_diff.jl index af70d39..402c17c 100644 --- a/test/test_diff.jl +++ b/test/test_diff.jl @@ -1,5 +1,5 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) - bm = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:full)) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) nvar = model.meta.nvar ncon = model.meta.ncon @@ -9,13 +9,13 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) Θ_gpu = MOD.randn(nθ, batch_size) @testset "objective!" begin - y = BOI.objective!(bm, X_gpu, Θ_gpu) + y = BNK.objective!(bm, X_gpu, Θ_gpu) @test size(y) == (batch_size,) function f_gpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.objective!(bm, X, Θ)) + return sum(BNK.objective!(bm, X, Θ)) end params = vcat(X_gpu, Θ_gpu) @@ -27,13 +27,13 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) ncon == 0 && return @testset "constraints!" begin - y = BOI.constraints!(bm, X_gpu, Θ_gpu) + y = BNK.constraints!(bm, X_gpu, Θ_gpu) @test size(y) == (ncon, batch_size) function f_gpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.constraints!(bm, X, Θ)) + return sum(BNK.constraints!(bm, X, Θ)) end params = vcat(X_gpu, Θ_gpu) @@ -44,7 +44,7 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) end function test_diff_cpu(model::ExaModel, batch_size::Int) - bm = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:full)) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) nvar = model.meta.nvar ncon = model.meta.ncon @@ -54,13 +54,13 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) Θ_cpu = randn(nθ, batch_size) @testset "objective! CPU" begin - y = BOI.objective!(bm, X_cpu, Θ_cpu) + y = BNK.objective!(bm, X_cpu, Θ_cpu) @test size(y) == (batch_size,) function f_cpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.objective!(bm, X, Θ)) + return sum(BNK.objective!(bm, X, Θ)) end params = vcat(X_cpu, Θ_cpu) @@ -77,13 +77,13 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) ncon == 0 && return @testset "constraints! CPU" begin - y = BOI.constraints!(bm, X_cpu, Θ_cpu) + y = BNK.constraints!(bm, X_cpu, Θ_cpu) @test size(y) == (ncon, batch_size) function f_cpu(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - return sum(BOI.constraints!(bm, X, Θ)) + return sum(BNK.constraints!(bm, X, Θ)) end params = vcat(X_cpu, Θ_cpu) diff --git a/test/test_viols.jl b/test/test_viols.jl index b80eca2..f7caa9d 100644 --- a/test/test_viols.jl +++ b/test/test_viols.jl @@ -1,6 +1,6 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; atol::Float64=1e-10, rtol::Float64=1e-10, MOD=OpenCL) - bm = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:violations)) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:violations)) nvar = model.meta.nvar ncon = model.meta.ncon @@ -22,7 +22,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; @testset "All Violations" begin if nθ > 0 - Vc, Vb = BOI.all_violations!(bm, X, Θ) + Vc, Vb = BNK.all_violations!(bm, X, Θ) @test size(Vc) == (ncon, batch_size) @test size(Vb) == (nvar, batch_size) @test all(>=(0), Vc) @@ -41,7 +41,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; end end - Vc, Vb = BOI.all_violations!(bm, X) + Vc, Vb = BNK.all_violations!(bm, X) @test size(Vc) == (ncon, batch_size) @test size(Vb) == (nvar, batch_size) @test all(>=(0), Vc) @@ -52,8 +52,8 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; @testset "Constraint Violations" begin if ncon > 0 - V_cons = BOI.constraints!(bm, X, Θ) - Vc = BOI.constraint_violations!(bm, V_cons) + V_cons = BNK.constraints!(bm, X, Θ) + Vc = BNK.constraint_violations!(bm, V_cons) @test size(Vc) == (ncon, batch_size) @test all(>=(0), Vc) @@ -82,7 +82,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; end @testset "Bound Violations" begin - Vb = BOI.bound_violations!(bm, X) + Vb = BNK.bound_violations!(bm, X) @test size(Vb) == (nvar, batch_size) @test all(>=(0), Vb) @test all(isfinite, Vb) @@ -114,42 +114,42 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; X_feasible[uunconstr, :] .= model.meta.lvar[uunconstr] .+ 0.1 @assert all(isfinite, X_feasible) end - Vb_feasible = BOI.bound_violations!(bm, X_feasible) + Vb_feasible = BNK.bound_violations!(bm, X_feasible) @test all(==(0), Vb_feasible) end end @testset "Dimension Validation" begin X_wrong = MOD.randn(nvar + 1, batch_size) - @test_throws DimensionMismatch BOI.all_violations!(bm, X_wrong) - @test_throws DimensionMismatch BOI.bound_violations!(bm, X_wrong) + @test_throws DimensionMismatch BNK.all_violations!(bm, X_wrong) + @test_throws DimensionMismatch BNK.bound_violations!(bm, X_wrong) if nθ > 0 Θ_wrong = MOD.randn(nθ + 1, batch_size) - @test_throws DimensionMismatch BOI.all_violations!(bm, X, Θ_wrong) + @test_throws DimensionMismatch BNK.all_violations!(bm, X, Θ_wrong) end if ncon > 0 V_wrong = MOD.randn(ncon + 1, batch_size) - @test_throws DimensionMismatch BOI.constraint_violations!(bm, V_wrong) + @test_throws DimensionMismatch BNK.constraint_violations!(bm, V_wrong) end end @testset "Batch Size Validation" begin X_large = MOD.randn(nvar, batch_size + 1) - @test_throws AssertionError BOI.all_violations!(bm, X_large) - @test_throws AssertionError BOI.bound_violations!(bm, X_large) + @test_throws AssertionError BNK.all_violations!(bm, X_large) + @test_throws AssertionError BNK.bound_violations!(bm, X_large) if ncon > 0 V_large = MOD.randn(ncon, batch_size + 1) - @test_throws AssertionError BOI.constraint_violations!(bm, V_large) + @test_throws AssertionError BNK.constraint_violations!(bm, V_large) end end end end function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) - bm = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:viol_grad)) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:viol_grad)) nvar = model.meta.nvar ncon = model.meta.ncon @@ -163,7 +163,7 @@ function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; function f_all_viols(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - Vc, Vb = BOI.all_violations!(bm, X, Θ) + Vc, Vb = BNK.all_violations!(bm, X, Θ) return sum(Vc) + sum(Vb) end @@ -178,8 +178,8 @@ function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; function f_cons_viols(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - V_cons = BOI.constraints!(bm, X, Θ) - Vc = BOI.constraint_violations!(bm, V_cons) + V_cons = BNK.constraints!(bm, X, Θ) + Vc = BNK.constraint_violations!(bm, V_cons) return sum(Vc) end @@ -192,7 +192,7 @@ function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; @testset "Bound Violations Sum" begin function f_bound_viols(X) - Vb = BOI.bound_violations!(bm, X) + Vb = BNK.bound_violations!(bm, X) return sum(Vb) end @@ -204,7 +204,7 @@ function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; end function test_violations_differentiability_cpu(model::ExaModel, batch_size::Int) - bm = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:viol_grad)) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:viol_grad)) nvar = model.meta.nvar nθ = length(model.θ) @@ -217,7 +217,7 @@ function test_violations_differentiability_cpu(model::ExaModel, batch_size::Int) function f_all_viols(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - Vc, Vb = BOI.all_violations!(bm, X, Θ) + Vc, Vb = BNK.all_violations!(bm, X, Θ) return sum(Vc) + sum(Vb) end @@ -234,7 +234,7 @@ function test_violations_differentiability_cpu(model::ExaModel, batch_size::Int) @testset "Bound Violations Sum" begin function f_bound_viols(X) - Vb = BOI.bound_violations!(bm, X) + Vb = BNK.bound_violations!(bm, X) return sum(Vb) end @@ -261,13 +261,13 @@ function test_violations_config_errors(MOD=OpenCL) Θ = MOD.randn(nθ, batch_size) @testset "Config Errors" begin - bm_no_viols = BOI.BatchModel(model, batch_size, config=BOI.BatchModelConfig(:minimal)) - @test_throws ArgumentError BOI.all_violations!(bm_no_viols, X, Θ) - @test_throws ArgumentError BOI.bound_violations!(bm_no_viols, X) + bm_no_viols = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:minimal)) + @test_throws ArgumentError BNK.all_violations!(bm_no_viols, X, Θ) + @test_throws ArgumentError BNK.bound_violations!(bm_no_viols, X) if ncon > 0 V = MOD.randn(ncon, batch_size) - @test_throws ArgumentError BOI.constraint_violations!(bm_no_viols, V) + @test_throws ArgumentError BNK.constraint_violations!(bm_no_viols, V) end end end From e9fd663e6d4b26f3def539112b7ece30d451255e Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 20:50:29 -0400 Subject: [PATCH 04/10] fix viols --- src/api/viols.jl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/api/viols.jl b/src/api/viols.jl index 18c1799..47c1853 100644 --- a/src/api/viols.jl +++ b/src/api/viols.jl @@ -4,9 +4,7 @@ Compute all constraint and variable violations for a batch of solutions. """ function all_violations!(bm::BatchModel, X::AbstractMatrix) - V = constraints!(bm, X) - - Vc = constraint_violations!(bm, V) + Vc = constraint_violations!(bm, X) Vb = bound_violations!(bm, X) return Vc, Vb @@ -18,9 +16,7 @@ end Compute all constraint and variable violations for a batch of solutions and parameters. """ function all_violations!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix) - V = constraints!(bm, X, Θ) - - Vc = constraint_violations!(bm, V) + Vc = constraint_violations!(bm, X, Θ) Vb = bound_violations!(bm, X) return Vc, Vb From 63943939e1d1ec49e23bba6448413865a8bb11b0 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 20:54:39 -0400 Subject: [PATCH 05/10] fix cons --- test/test_viols.jl | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/test/test_viols.jl b/test/test_viols.jl index f7caa9d..6b046c7 100644 --- a/test/test_viols.jl +++ b/test/test_viols.jl @@ -52,8 +52,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; @testset "Constraint Violations" begin if ncon > 0 - V_cons = BNK.constraints!(bm, X, Θ) - Vc = BNK.constraint_violations!(bm, V_cons) + Vc = BNK.constraint_violations!(bm, X, Θ) @test size(Vc) == (ncon, batch_size) @test all(>=(0), Vc) @@ -130,8 +129,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; end if ncon > 0 - V_wrong = MOD.randn(ncon + 1, batch_size) - @test_throws DimensionMismatch BNK.constraint_violations!(bm, V_wrong) + @test_throws DimensionMismatch BNK.constraint_violations!(bm, X_wrong) end end @@ -141,8 +139,7 @@ function test_violations_correctness(model::ExaModel, batch_size::Int; @test_throws AssertionError BNK.bound_violations!(bm, X_large) if ncon > 0 - V_large = MOD.randn(ncon, batch_size + 1) - @test_throws AssertionError BNK.constraint_violations!(bm, V_large) + @test_throws AssertionError BNK.constraint_violations!(bm, X_large) end end end @@ -178,8 +175,7 @@ function test_violations_differentiability_gpu(model::ExaModel, batch_size::Int; function f_cons_viols(params) X = params[1:nvar, :] Θ = params[nvar+1:end, :] - V_cons = BNK.constraints!(bm, X, Θ) - Vc = BNK.constraint_violations!(bm, V_cons) + Vc = BNK.constraint_violations!(bm, X, Θ) return sum(Vc) end @@ -267,7 +263,7 @@ function test_violations_config_errors(MOD=OpenCL) if ncon > 0 V = MOD.randn(ncon, batch_size) - @test_throws ArgumentError BNK.constraint_violations!(bm_no_viols, V) + @test_throws ArgumentError BNK._constraint_violations!(bm_no_viols, V) end end end From d8977c2fe7c58ec50914ed4f40e7484cf99d517c Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 21:00:10 -0400 Subject: [PATCH 06/10] fix cons cont --- ext/BNKChainRulesCore.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/BNKChainRulesCore.jl b/ext/BNKChainRulesCore.jl index b3fb93f..10023bc 100644 --- a/ext/BNKChainRulesCore.jl +++ b/ext/BNKChainRulesCore.jl @@ -57,8 +57,8 @@ function ChainRulesCore.rrule(::typeof(BatchNLPKernels.constraints!), bm::BatchM end -function ChainRulesCore.rrule(::typeof(BatchNLPKernels.constraint_violations!), bm::BatchModel, V) - Vc = BatchNLPKernels.constraint_violations!(bm, V) +function ChainRulesCore.rrule(::typeof(BatchNLPKernels._constraint_violations!), bm::BatchModel, V) + Vc = BatchNLPKernels._constraint_violations!(bm, V) function constraint_violations_pullback(Ȳ) Ȳ = ChainRulesCore.unthunk(Ȳ) From 23ab46aebcbd0432ddf9ee09725cd88e412c54e7 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 21:17:02 -0400 Subject: [PATCH 07/10] no param crc tests --- test/test_diff.jl | 63 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 62 insertions(+), 1 deletion(-) diff --git a/test/test_diff.jl b/test/test_diff.jl index 402c17c..f2ca46c 100644 --- a/test/test_diff.jl +++ b/test/test_diff.jl @@ -23,6 +23,19 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) @test grad isa AbstractMatrix @test size(grad) == size(params) end + + @testset "objective! no params" begin + y = BNK.objective!(bm, X_gpu) + @test size(y) == (batch_size,) + + function f_gpu(X) + return sum(BNK.objective!(bm, X)) + end + + grad = DI.gradient(f_gpu, AutoZygote(), X_gpu) + @test grad isa AbstractMatrix + @test size(grad) == size(X_gpu) + end ncon == 0 && return @@ -41,6 +54,19 @@ function test_diff_gpu(model::ExaModel, batch_size::Int; MOD=OpenCL) @test grad isa AbstractMatrix @test size(grad) == size(params) end + + @testset "constraints! no params" begin + y = BNK.constraints!(bm, X_gpu, Θ_gpu) + @test size(y) == (ncon, batch_size) + + function f_gpu(X) + return sum(BNK.constraints!(bm, X)) + end + + grad = DI.gradient(f_gpu, AutoZygote(), X_gpu) + @test grad isa AbstractMatrix + @test size(grad) == size(X_gpu) + end end function test_diff_cpu(model::ExaModel, batch_size::Int) @@ -74,6 +100,24 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) end end + @testset "objective! no params" begin + y = BNK.objective!(bm, X_cpu) + @test size(y) == (batch_size,) + + function f_cpu(X) + return sum(BNK.objective!(bm, X)) + end + + grad = DI.gradient(f_cpu, AutoZygote(), X_cpu) + @test grad isa AbstractMatrix + @test size(grad) == size(X_cpu) + + @testset "FiniteDifferences objective! no params" begin + gradfd = DI.gradient(f_cpu, AutoFiniteDifferences(fdm=FiniteDifferences.central_fdm(3,1)), X_cpu) + @test gradfd ≈ grad atol=1e-4 rtol=1e-4 + end + end + ncon == 0 && return @testset "constraints! CPU" begin @@ -96,8 +140,25 @@ function test_diff_cpu(model::ExaModel, batch_size::Int) @test gradfd[1:nvar,:] ≈ grad[1:nvar,:] atol=1e-4 rtol=1e-4 end end -end + @testset "constraints! no params" begin + y = BNK.constraints!(bm, X_cpu, Θ_cpu) + @test size(y) == (ncon, batch_size) + + function f_cpu(X) + return sum(BNK.constraints!(bm, X)) + end + + grad = DI.gradient(f_cpu, AutoZygote(), X_cpu) + @test grad isa AbstractMatrix + @test size(grad) == size(X_cpu) + + @testset "FiniteDifferences constraints! no params" begin + gradfd = DI.gradient(f_cpu, AutoFiniteDifferences(fdm=FiniteDifferences.central_fdm(3,1)), X_cpu) + @test gradfd ≈ grad atol=1e-4 rtol=1e-4 + end + end +end @testset "AD rules - Luksan" begin cpu_models, names = create_luksan_models(CPU()) From 5d9716806079bc092a83dbd72a71125a30a7cf98 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 21:34:14 -0400 Subject: [PATCH 08/10] jac/hess coord tests --- test/api.jl | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/test/api.jl b/test/api.jl index d0edb53..8478493 100644 --- a/test/api.jl +++ b/test/api.jl @@ -5,6 +5,8 @@ function test_batch_model(model::ExaModel, batch_size::Int; nvar = model.meta.nvar ncon = model.meta.ncon + nnzh = model.meta.nnzh + nnzj = model.meta.nnzj nθ = length(model.θ) X = MOD.randn(nvar, batch_size) @@ -46,6 +48,20 @@ function test_batch_model(model::ExaModel, batch_size::Int; @allowscalar @test grad_vals[:, i] ≈ grad_single atol=atol rtol=rtol end end + + @testset "Jacobian" begin + if ncon > 0 + jac_vals = BNK.constraints_jacobian!(bm, X, Θ) + @test size(jac_vals) == (nnzj, batch_size) + @test all(isfinite, jac_vals) + for i in 1:batch_size + @allowscalar nθ > 0 && (model.θ .= Θ[:, i]) + jac_single = similar(jac_vals, nnzj) + @allowscalar ExaModels.jac_coord!(model, X[:, i], jac_single) + @allowscalar @test jac_vals[:, i] ≈ jac_single atol=atol rtol=rtol + end + end + end @testset "Jacobian-Vector Product" begin if ncon > 0 @@ -76,7 +92,20 @@ function test_batch_model(model::ExaModel, batch_size::Int; end end end - + + @testset "Hessian" begin + if ncon > 0 + hess_vals = BNK.objective_hessian!(bm, X, Θ) + @test size(hess_vals) == (nnzh, batch_size) + @test all(isfinite, hess_vals) + for i in 1:batch_size + @allowscalar nθ > 0 && (model.θ .= Θ[:, i]) + hess_single = similar(hess_vals, nnzh) + @allowscalar ExaModels.hess_coord!(model, X[:, i], hess_single) + @allowscalar @test hess_vals[:, i] ≈ hess_single atol=atol rtol=rtol + end + end + end @testset "Hessian-Vector Product" begin V = MOD.randn(nvar, batch_size) if ncon > 0 From c864221ce460b65e81e543914d5acd2da54368e0 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 21:49:59 -0400 Subject: [PATCH 09/10] typo --- test/api.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/api.jl b/test/api.jl index 8478493..cafa7c6 100644 --- a/test/api.jl +++ b/test/api.jl @@ -95,7 +95,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Hessian" begin if ncon > 0 - hess_vals = BNK.objective_hessian!(bm, X, Θ) + hess_vals = BNK.lagrangian_hessian!(bm, X, Θ) @test size(hess_vals) == (nnzh, batch_size) @test all(isfinite, hess_vals) for i in 1:batch_size From 6e8cd0341e5da7138ffba36781f9b3ce4d413f0f Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 8 Jul 2025 22:12:49 -0400 Subject: [PATCH 10/10] typo cont. --- test/api.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/api.jl b/test/api.jl index cafa7c6..a1d9afe 100644 --- a/test/api.jl +++ b/test/api.jl @@ -95,13 +95,14 @@ function test_batch_model(model::ExaModel, batch_size::Int; @testset "Hessian" begin if ncon > 0 - hess_vals = BNK.lagrangian_hessian!(bm, X, Θ) + Y = MOD.randn(ncon, batch_size) + hess_vals = BNK.lagrangian_hessian!(bm, X, Θ, Y) @test size(hess_vals) == (nnzh, batch_size) @test all(isfinite, hess_vals) for i in 1:batch_size @allowscalar nθ > 0 && (model.θ .= Θ[:, i]) hess_single = similar(hess_vals, nnzh) - @allowscalar ExaModels.hess_coord!(model, X[:, i], hess_single) + @allowscalar ExaModels.hess_coord!(model, X[:, i], Y[:, i], hess_single) @allowscalar @test hess_vals[:, i] ≈ hess_single atol=atol rtol=rtol end end