diff --git a/README.md b/README.md index 33f0c96..84514af 100644 --- a/README.md +++ b/README.md @@ -5,15 +5,17 @@ `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)` +- `all_violations!(::BatchModel, X, Θ)` +- `constraint_violations!(::BatchModel, X, Θ)` +- `bound_violations!(::BatchModel, X)` To use these functions, first wrap your `ExaModel` in a `BatchModel`: @@ -30,7 +32,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..10023bc 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 @@ -57,8 +57,8 @@ function ChainRulesCore.rrule(::typeof(BatchNLPKernels.cons_nln_batch!), bm::Bat 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(Ȳ) diff --git a/src/BatchNLPKernels.jl b/src/BatchNLPKernels.jl index d09c930..1d8091e 100644 --- a/src/BatchNLPKernels.jl +++ b/src/BatchNLPKernels.jl @@ -9,10 +9,11 @@ const KAExtension = ExaKA.KAExtension include("interval.jl") 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! +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! 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..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 = cons_nln_batch!(bm, X) - - Vc = constraint_violations!(bm, V) + Vc = constraint_violations!(bm, X) Vb = bound_violations!(bm, X) return Vc, Vb @@ -18,20 +16,32 @@ 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, Θ) - - Vc = constraint_violations!(bm, V) + Vc = constraint_violations!(bm, X, Θ) Vb = bound_violations!(bm, X) return Vc, Vb 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..a1d9afe 100644 --- a/test/api.jl +++ b/test/api.jl @@ -1,10 +1,12 @@ 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 + nnzh = model.meta.nnzh + nnzj = model.meta.nnzj nθ = length(model.θ) X = MOD.randn(nvar, batch_size) @@ -12,7 +14,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 = BNK.objective!(bm, X, Θ) @test length(obj_vals) == batch_size @test all(isfinite, obj_vals) for i in 1:batch_size @@ -23,7 +25,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 = BNK.constraints!(bm, X, Θ) @test size(cons_vals) == (ncon, batch_size) @test all(isfinite, cons_vals) for i in 1:batch_size @@ -36,7 +38,7 @@ function test_batch_model(model::ExaModel, batch_size::Int; end @testset "Gradient" begin - grad_vals = BOI.grad_batch!(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 @@ -46,11 +48,25 @@ 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 V = MOD.randn(nvar, batch_size) - jprod_vals = BOI.jprod_nln_batch!(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 +81,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 = BNK.constraints_jtprod!(bm, X, Θ, V) @test size(jtprod_vals) == (nvar, batch_size) @test all(isfinite, jtprod_vals) for i in 1:batch_size @@ -76,12 +92,26 @@ function test_batch_model(model::ExaModel, batch_size::Int; end end end - + + @testset "Hessian" begin + if ncon > 0 + 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], Y[:, 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 Y = MOD.randn(ncon, batch_size) - hprod_vals = BOI.hprod_batch!(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 +122,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 = 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 +136,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 BNK.objective!(bm, X_large) if ncon > 0 - @test_throws AssertionError BOI.cons_nln_batch!(bm, X_large) + @test_throws AssertionError BNK.constraints!(bm, X_large) end - @test_throws AssertionError BOI.grad_batch!(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.jprod_nln_batch!(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.jtprod_nln_batch!(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.hprod_batch!(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.hprod_batch!(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.obj_batch!(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.obj_batch!(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.jprod_nln_batch!(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.jtprod_nln_batch!(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.hprod_batch!(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.hprod_batch!(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.hprod_batch!(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 163c0bc..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.grad_batch!(bm_minimal, X, Θ) - @test_throws ArgumentError BOI.jac_coord_batch!(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.hess_coord_batch!(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.jprod_nln_batch!(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.jtprod_nln_batch!(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.hprod_batch!(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.hprod_batch!(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.grad_batch!(bm_no_gradout, X, Θ, G) - @test_throws ArgumentError BOI.grad_batch!(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.cons_nln_batch!(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 6298fb3..f2ca46c 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 @@ -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 = 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.obj_batch!(bm, X, Θ)) + return sum(BNK.objective!(bm, X, Θ)) end params = vcat(X_gpu, Θ_gpu) @@ -23,17 +23,30 @@ 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 - @testset "cons_nln_batch!" begin - y = BOI.cons_nln_batch!(bm, X_gpu, Θ_gpu) + @testset "constraints!" begin + 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.cons_nln_batch!(bm, X, Θ)) + return sum(BNK.constraints!(bm, X, Θ)) end params = vcat(X_gpu, Θ_gpu) @@ -41,10 +54,23 @@ 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) - 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 @@ -53,14 +79,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 = 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.obj_batch!(bm, X, Θ)) + return sum(BNK.objective!(bm, X, Θ)) end params = vcat(X_cpu, Θ_cpu) @@ -68,22 +94,40 @@ 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 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 "cons_nln_batch! CPU" begin - y = BOI.cons_nln_batch!(bm, X_cpu, Θ_cpu) + @testset "constraints! CPU" begin + 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.cons_nln_batch!(bm, X, Θ)) + return sum(BNK.constraints!(bm, X, Θ)) end params = vcat(X_cpu, Θ_cpu) @@ -91,13 +135,30 @@ 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 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()) diff --git a/test/test_viols.jl b/test/test_viols.jl index e678db8..6b046c7 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,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, Θ) - Vc = BOI.constraint_violations!(bm, V_cons) + Vc = BNK.constraint_violations!(bm, X, Θ) @test size(Vc) == (ncon, batch_size) @test all(>=(0), Vc) @@ -82,7 +81,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 +113,40 @@ 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, X_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, X_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 +160,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 +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 = BOI.cons_nln_batch!(bm, X, Θ) - Vc = BOI.constraint_violations!(bm, V_cons) + Vc = BNK.constraint_violations!(bm, X, Θ) return sum(Vc) end @@ -192,7 +188,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 +200,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 +213,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 +230,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 +257,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