From 051c3d2f4a57cc26a8623fcc10fae91100283574 Mon Sep 17 00:00:00 2001
From: fpacaud <francoispacaud8@gmail.com>
Date: Thu, 25 Jul 2024 14:48:01 +0200
Subject: [PATCH] add ScaledModel

---
 src/NLPModelsModifiers.jl |   1 +
 src/scaled-model.jl       | 236 ++++++++++++++++++++++++++++++++++++++
 test/nlp/scaled-model.jl  |  38 ++++++
 test/runtests.jl          |   1 +
 4 files changed, 276 insertions(+)
 create mode 100644 src/scaled-model.jl
 create mode 100644 test/nlp/scaled-model.jl

diff --git a/src/NLPModelsModifiers.jl b/src/NLPModelsModifiers.jl
index 2a81e16..a713f7c 100644
--- a/src/NLPModelsModifiers.jl
+++ b/src/NLPModelsModifiers.jl
@@ -18,6 +18,7 @@ include("feasibility-form-nls.jl")
 include("feasibility-residual.jl")
 include("quasi-newton.jl")
 include("slack-model.jl")
+include("scaled-model.jl")
 include("model-interaction.jl")
 
 end # module
diff --git a/src/scaled-model.jl b/src/scaled-model.jl
new file mode 100644
index 0000000..e2aafda
--- /dev/null
+++ b/src/scaled-model.jl
@@ -0,0 +1,236 @@
+export ScaledModel
+
+struct IpoptScaling{T}
+  max_gradient::T
+end
+
+function _set_constraints_scaling!(cons, Ji, Jj, Jx, max_gradient)
+  # Return a vector storing at index i norm(∇cᵢ, Inf)
+  for (i, j, x) in zip(Ji, Jj, Jx)
+    cons[i] = max(cons[i], abs(x))
+  end
+  # Compute scaling as min(1, max_gradient / norm(∇cᵢ, Inf) )
+  for i in eachindex(cons)
+    cons[i] = min(1.0, max_gradient / cons[i])
+  end
+end
+
+function _set_jacobian_scaling!(Jx, Ji, Jj, cons)
+  k = 0
+  for (i, j) in zip(Ji, Jj)
+    Jx[k += 1] = cons[i]
+  end
+end
+
+function scale_model!(scaling::IpoptScaling{T}, nlp) where T
+  n, m = NLPModels.get_nvar(nlp), NLPModels.get_ncon(nlp)
+  nnzj = NLPModels.get_nnzj(nlp)
+  x0 = NLPModels.get_x0(nlp)
+  g = NLPModels.grad(nlp, x0)
+  scaling_obj = min(one(T), scaling.max_gradient / norm(g, Inf))
+  scaling_cons = similar(x0, m)
+  scaling_jac  = similar(x0, nnzj)
+  fill!(scaling_cons, zero(T))
+  Ji, Jj = NLPModels.jac_structure(nlp)
+  NLPModels.jac_coord!(nlp, x0, scaling_jac)
+  _set_constraints_scaling!(scaling_cons, Ji, Jj, scaling_jac, scaling.max_gradient)
+  _set_jacobian_scaling!(scaling_jac, Ji, Jj, scaling_cons)
+  return (scaling_obj, scaling_cons, scaling_jac)
+end
+
+@doc raw"""
+    ScaledModel
+
+Scale the nonlinear program
+```math
+\begin{aligned}
+       min_x  \quad & f(x)\\
+\mathrm{s.t.} \quad &  c♭ ≤ c(x) ≤ c♯ \\
+                    & x ≥ 0
+\end{aligned}
+```
+as
+```math
+\begin{aligned}
+       min_x  \quad & σf . f(x)\\
+\mathrm{s.t.} \quad &  σc . c♭ ≤ σc . c(x) ≤ σc . c♯ \\
+                    & x ≥ 0
+\end{aligned}
+```
+with ``σf`` a scalar defined as
+```
+σf = min(1, max_gradient / norm(g0, Inf))
+
+```
+and ``σc`` a vector whose size is equal to the number of constraints in the model.
+For ``i=1, ..., m``,
+```
+σc[i] = min(1, max_gradient / norm(J0[i, :], Inf))
+
+```
+
+The vector ``g0 = ∇f(x0)`` and the matrix ``J0 = ∇c(x0)`` are resp.
+the gradient and the Jacobian evaluated at the initial point ``x0``.
+
+"""
+struct ScaledModel{T, S, M} <: NLPModels.AbstractNLPModel{T, S}
+  nlp::M
+  meta::NLPModels.NLPModelMeta{T, S}
+  counters::NLPModels.Counters
+  scaling_obj::T
+  scaling_cons::S # [size m]
+  scaling_jac::S  # [size nnzj]
+  buffer_cons::S  # [size m]
+end
+
+NLPModels.show_header(io::IO, nlp::ScaledModel) =
+  println(io, "ScaledModel - Model with scaled objective and constraints")
+
+
+function ScaledModel(
+    nlp::NLPModels.AbstractNLPModel{T, S};
+    scaling=IpoptScaling(T(100)),
+) where {T, S}
+  n, m = NLPModels.get_nvar(nlp), NLPModels.get_ncon(nlp)
+  x0 = NLPModels.get_x0(nlp)
+  buffer_cons  = S(undef, m)
+  scaling_obj, scaling_cons, scaling_jac = scale_model!(scaling, nlp)
+  meta = NLPModels.NLPModelMeta(
+    n;
+    lvar=NLPModels.get_lvar(nlp),
+    uvar=NLPModels.get_uvar(nlp),
+    x0=NLPModels.get_x0(nlp),
+    y0 = NLPModels.get_y0(nlp) .* scaling_cons,
+    nnzj=NLPModels.get_nnzj(nlp),
+    nnzh=NLPModels.get_nnzh(nlp),
+    ncon=m,
+    lcon=NLPModels.get_lcon(nlp) .* scaling_cons,
+    ucon=NLPModels.get_ucon(nlp) .* scaling_cons,
+    minimize=true,
+  )
+
+  return ScaledModel(
+    nlp,
+    meta,
+    NLPModels.Counters(),
+    scaling_obj,
+    scaling_cons,
+    scaling_jac,
+    buffer_cons,
+  )
+end
+
+function NLPModels.obj(nlp::ScaledModel{T, S}, x::AbstractVector) where {T, S <: AbstractVector{T}}
+  @lencheck nlp.meta.nvar x
+  return nlp.scaling_obj * NLPModels.obj(nlp.nlp, x)
+end
+
+function NLPModels.cons!(nlp::ScaledModel, x::AbstractVector, c::AbstractVector)
+  @lencheck nlp.meta.nvar x
+  @lencheck nlp.meta.ncon c
+  NLPModels.cons!(nlp.nlp, x, c)
+  c .*= nlp.scaling_cons
+  return c
+end
+
+function NLPModels.grad!(nlp::ScaledModel, x::AbstractVector, g::AbstractVector)
+  @lencheck nlp.meta.nvar x g
+  NLPModels.grad!(nlp.nlp, x, g)
+  g .*= nlp.scaling_obj
+  return g
+end
+
+function NLPModels.jprod!(nlp::ScaledModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector)
+  @lencheck nlp.meta.nvar x v
+  @lencheck nlp.meta.ncon Jv
+  NLPModels.jprod!(nlp.nlp, x, v, Jv)
+  Jv .*= nlp.scaling_cons
+  return Jv
+end
+
+function NLPModels.jtprod!(nlp::ScaledModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector)
+  @lencheck nlp.meta.nvar x Jtv
+  @lencheck nlp.meta.ncon v
+  v_scaled = nlp.buffer_cons
+  v_scaled .= v .* nlp.scaling_cons
+  NLPModels.jtprod!(nlp.nlp, x, v_scaled, Jtv)
+  return Jtv
+end
+
+function NLPModels.jac_structure!(nlp::ScaledModel, jrows::AbstractVector, jcols::AbstractVector)
+  NLPModels.jac_structure!(nlp.nlp, jrows, jcols)
+  return jrows, jcols
+end
+
+function NLPModels.jac_coord!(nlp::ScaledModel, x::AbstractVector, jac::AbstractVector)
+  NLPModels.jac_coord!(nlp.nlp, x, jac)
+  jac .*= nlp.scaling_jac
+  return jac
+end
+
+function NLPModels.hess_structure!(nlp::ScaledModel, hrows::AbstractVector, hcols::AbstractVector)
+  @lencheck nlp.meta.nnzh hrows hcols
+  NLPModels.hess_structure!(nlp.nlp, hrows, hcols)
+  return hrows, hcols
+end
+
+function NLPModels.hess_coord!(
+    nlp::ScaledModel,
+    x::AbstractVector,
+    vals::AbstractVector;
+    obj_weight::Real=one(eltype(x)),
+)
+  @lencheck nlp.meta.nvar x
+  @lencheck nlp.meta.nnzh vals
+  σ = obj_weight * nlp.scaling_obj
+  NLPModels.hess_coord!(nlp.nlp, x, vals; obj_weight=σ)
+  return vals
+end
+
+function NLPModels.hess_coord!(
+    nlp::ScaledModel,
+    x::AbstractVector,
+    y::AbstractVector,
+    vals::AbstractVector;
+    obj_weight::Real=one(eltype(x)),
+)
+  @lencheck nlp.meta.nvar x
+  @lencheck nlp.meta.ncon y
+  @lencheck nlp.meta.nnzh vals
+  y_scaled = nlp.buffer_cons
+  y_scaled .= y .* nlp.scaling_cons
+  σ = obj_weight * nlp.scaling_obj
+  NLPModels.hess_coord!(nlp.nlp, x, y_scaled, vals; obj_weight=σ)
+  return vals
+end
+
+function NLPModels.hprod!(
+  nlp::ScaledModel,
+  x::AbstractVector,
+  v::AbstractVector,
+  hv::AbstractVector;
+  obj_weight::Real = one(eltype(x)),
+)
+  @lencheck nlp.meta.nvar x v hv
+  σ = obj_weight * nlp.scaling_obj
+  NLPModels.hprod!(nlp.nlp, x, v, hv; obj_weight = σ)
+  return hv
+end
+
+function NLPModels.hprod!(
+  nlp::ScaledModel,
+  x::AbstractVector,
+  y::AbstractVector,
+  v::AbstractVector,
+  hv::AbstractVector;
+  obj_weight::Real = one(eltype(x)),
+)
+  @lencheck nlp.meta.nvar x v hv
+  @lencheck nlp.meta.ncon y
+  y_scaled = nlp.buffer_cons
+  y_scaled .= y .* nlp.scaling_cons
+  σ = obj_weight * nlp.scaling_obj
+  NLPModels.hprod!(nlp.nlp, x, y, v, hv; obj_weight = σ)
+  return hv
+end
+
diff --git a/test/nlp/scaled-model.jl b/test/nlp/scaled-model.jl
new file mode 100644
index 0000000..604fc7e
--- /dev/null
+++ b/test/nlp/scaled-model.jl
@@ -0,0 +1,38 @@
+@testset "ScaledModel NLP tests" begin
+  @testset "API" for T in [Float64, Float32], M in [NLPModelMeta, SimpleNLPMeta]
+    nlp = ScaledModel(SimpleNLPModel(T, M))
+    σ_obj, σ_cons = nlp.scaling_obj, nlp.scaling_cons
+
+    f(x) = σ_obj * (x[1] - 2)^2 + (x[2] - 1)^2
+    ∇f(x) = [σ_obj * 2 * (x[1] - 2); σ_obj * 2 * (x[2] - 1)]
+    H(x) = T[(σ_obj * 2.0) 0; 0 (σ_obj * 2.0)]
+    c(x) = [σ_cons[1] * (x[1] - 2x[2] + 1); σ_cons[2] * (-x[1]^2 / 4 - x[2]^2 + 1)]
+    J(x) = [σ_cons[1] -2.0*σ_cons[1]; (-0.5 *σ_cons[1] * x[1]) (-2.0*σ_cons[2] * x[2])]
+    H(x, y) = H(x) + σ_cons[2] * y[2] * T[-0.5 0; 0 -2.0]
+
+    n = nlp.meta.nvar
+    m = nlp.meta.ncon
+    @test nlp.meta.x0 == T[2; 2]
+
+    x = randn(T, n)
+    y = randn(T, m)
+    v = randn(T, n)
+    w = randn(T, m)
+    Jv = zeros(T, m)
+    Jtw = zeros(T, n)
+    Hv = zeros(T, n)
+    Hvals = zeros(T, nlp.meta.nnzh)
+
+    # Basic methods
+    @test obj(nlp, x) ≈ f(x)
+    @test grad(nlp, x) ≈ ∇f(x)
+    @test hess(nlp, x) ≈ H(x)
+    @test hprod(nlp, x, v) ≈ H(x) * v
+    @test cons(nlp, x) ≈ c(x)
+    @test jac(nlp, x) ≈ J(x)
+    @test jprod(nlp, x, v) ≈ J(x) * v
+    @test jtprod(nlp, x, w) ≈ J(x)' * w
+    @test hess(nlp, x, y) ≈ H(x, y)
+    @test hprod(nlp, x, y, v) ≈ H(x, y) * v
+  end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 281269e..19cf939 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -4,6 +4,7 @@ using LinearOperators, NLPModels, NLPModelsModifiers
 include("nlp/simple-model.jl")
 include("nlp/quasi-newton.jl")
 include("nlp/slack-model.jl")
+include("nlp/scaled-model.jl")
 
 include("nls/simple-model.jl")
 include("nls/feasibility-form-nls.jl")