Preallocate on LineModel

Author: abelsiqueira

src/linesearch/line_model.jl

@@ -1,7 +1,7 @@
 import NLPModels: obj, grad, grad!, objgrad!, objgrad, hess
 
 export LineModel
-export obj, grad, derivative, grad!, objgrad!, objgrad, derivative!, hess, redirect!
+export obj, grad, derivative, grad!, objgrad!, objgrad, derivative!, hess, hess!, redirect!
 
 """A type to represent the restriction of a function to a direction.
 Given f : R → Rⁿ, x ∈ Rⁿ and a nonzero direction d ∈ Rⁿ,
@@ -12,17 +12,23 @@ represents the function ϕ : R → R defined by
 
     ϕ(t) := f(x + td).
 """
-mutable struct LineModel <: AbstractNLPModel
+mutable struct LineModel{T} <: AbstractNLPModel
   meta::NLPModelMeta
   counters::Counters
   nlp::AbstractNLPModel
-  x::AbstractVector
-  d::AbstractVector
+  x::Vector{T}
+  d::Vector{T}
+  xt::Vector{T}
 end
 
-function LineModel(nlp::AbstractNLPModel, x::AbstractVector, d::AbstractVector)
-  meta = NLPModelMeta(1, x0 = zeros(eltype(x), 1), name = "LineModel to $(nlp.meta.name))")
-  return LineModel(meta, Counters(), nlp, x, d)
+function LineModel(
+  nlp::AbstractNLPModel,
+  x::AbstractVector{T},
+  d::AbstractVector{T};
+  xt::AbstractVector{T} = similar(x),
+) where {T}
+  meta = NLPModelMeta(1, x0 = zeros(T, 1), name = "LineModel to $(nlp.meta.name))")
+  return LineModel{T}(meta, Counters(), nlp, x, d, xt)
 end
 
 """`redirect!(ϕ, x, d)`
@@ -40,7 +46,8 @@ end
 """
 function obj(f::LineModel, t::AbstractFloat)
   NLPModels.increment!(f, :neval_obj)
-  return obj(f.nlp, f.x + t * f.d)
+  @. f.xt = f.x + t * f.d
+  return obj(f.nlp, f.xt)
 end
 
 """`grad(f, t)` evaluates the first derivative of the `LineModel`
@@ -53,7 +60,8 @@ i.e.,
 """
 function grad(f::LineModel, t::AbstractFloat)
   NLPModels.increment!(f, :neval_grad)
-  return dot(grad(f.nlp, f.x + t * f.d), f.d)
+  @. f.xt = f.x + t * f.d
+  return dot(grad(f.nlp, f.xt), f.d)
 end
 derivative(f::LineModel, t::AbstractFloat) = grad(f, t)
 
@@ -69,7 +77,8 @@ The gradient ∇f(x + td) is stored in `g`.
 """
 function grad!(f::LineModel, t::AbstractFloat, g::AbstractVector)
   NLPModels.increment!(f, :neval_grad)
-  return dot(grad!(f.nlp, f.x + t * f.d, g), f.d)
+  @. f.xt = f.x + t * f.d
+  return dot(grad!(f.nlp, f.xt, g), f.d)
 end
 derivative!(f::LineModel, t::AbstractFloat, g::AbstractVector) = grad!(f, t, g)
 
@@ -86,7 +95,8 @@ The gradient ∇f(x + td) is stored in `g`.
 function objgrad!(f::LineModel, t::AbstractFloat, g::AbstractVector)
   NLPModels.increment!(f, :neval_obj)
   NLPModels.increment!(f, :neval_grad)
-  fx, gx = objgrad!(f.nlp, f.x + t * f.d, g)
+  @. f.xt = f.x + t * f.d
+  fx, gx = objgrad!(f.nlp, f.xt, g)
   return fx, dot(gx, f.d)
 end
 
@@ -112,5 +122,20 @@ i.e.,
 """
 function hess(f::LineModel, t::AbstractFloat)
   NLPModels.increment!(f, :neval_hess)
-  return dot(f.d, hprod(f.nlp, f.x + t * f.d, f.d))
+  @. f.xt = f.x + t * f.d
+  return dot(f.d, hprod(f.nlp, f.xt, f.d))
+end
+
+"""Evaluate the second derivative of the `LineModel`
+
+    ϕ(t) := f(x + td),
+
+i.e.,
+
+    ϕ"(t) = dᵀ∇²f(x + td)d.
+"""
+function hess!(f::LineModel, t::AbstractFloat, Hv::AbstractVector)
+  NLPModels.increment!(f, :neval_hess)
+  @. f.xt = f.x + t * f.d
+  return dot(f.d, hprod!(f.nlp, f.xt, f.d, Hv))
 end

test/runtests.jl

@@ -8,6 +8,7 @@ using ADNLPModels, NLPModels
 using LinearAlgebra, Logging, Test
 
 include("dummy_solver.jl")
+include("simple_model.jl")
 
 include("test_auxiliary.jl")
 include("test_linesearch.jl")

test/simple_model.jl

@@ -0,0 +1,35 @@
+mutable struct SimpleModel <: AbstractNLPModel
+  meta :: AbstractNLPModelMeta
+  counters :: Counters
+end
+
+SimpleModel(n :: Int) = SimpleModel(NLPModelMeta(n, x0=ones(n)), Counters())
+
+function NLPModels.obj(nlp::SimpleModel, x::AbstractVector)
+  increment!(nlp, :neval_obj)
+  sum(xi ^ 4 for xi in x) / 12
+end
+
+function NLPModels.grad!(nlp::SimpleModel, x::AbstractVector, g::AbstractVector)
+  increment!(nlp, :neval_grad)
+  @. g = x ^ 3 / 3
+  g
+end
+
+function NLPModels.objgrad!(nlp::SimpleModel, x::AbstractVector, g::AbstractVector)
+  increment!(nlp, :neval_obj)
+  increment!(nlp, :neval_grad)
+  @. g = x ^ 3 / 3
+  return sum(xi ^4 for xi in x) / 12, g
+end
+
+function NLPModels.hprod!(nlp::SimpleModel, x::AbstractVector{T}, v::AbstractVector, Hv::AbstractVector; obj_weight::T = one(T)) where T
+  increment!(nlp, :neval_hprod)
+  @. Hv = obj_weight * x ^ 2 * v
+  Hv
+end
+
+function NLPModels.hess(nlp::SimpleModel, x::AbstractVector{T}; obj_weight::T = one(T)) where T
+  increment!(nlp, :neval_hprod)
+  return obj_weight .* diagm(0 => x .^ 2)
+end

test/test_linesearch.jl

@@ -1,10 +1,11 @@
 @testset "Linesearch" begin
   @testset "LineModel" begin
-    nlp = ADNLPModel(x -> x[1]^2 + 4 * x[2]^2, ones(2))
+    n = 200
+    nlp = SimpleModel(n)
     x = nlp.meta.x0
-    d = -ones(2)
+    d = -ones(n)
     lm = LineModel(nlp, x, d)
-    g = zeros(2)
+    g = zeros(n)
 
     @test obj(lm, 0.0) == obj(nlp, x)
     @test grad(lm, 0.0) == dot(grad(nlp, x), d)
@@ -17,19 +18,19 @@
     @test g == grad(nlp, x + d)
     @test objgrad(lm, 0.0) == (obj(nlp, x), dot(grad(nlp, x), d))
     @test hess(lm, 0.0) == dot(d, Symmetric(hess(nlp, x), :L) * d)
+    @test hess!(lm, 0.0, g) == dot(d, hprod!(nlp, x, d, g))
 
     @test obj(lm, 1.0) == 0.0
     @test grad(lm, 1.0) == 0.0
-    @test hess(lm, 1.0) == 2d[1]^2 + 8d[2]^2
+    @test hess(lm, 1.0) == 0.0
 
-    redirect!(lm, zeros(2), ones(2))
-    @test obj(lm, 0.0) == 0.0
-    @test grad(lm, 0.0) == 0.0
-    @test hess(lm, 0.0) == 10.0
+    @test obj(lm, 0.0) ≈ n / 12
+    @test grad(lm, 0.0) ≈ -n / 3
+    @test hess(lm, 0.0) == n
 
     @test neval_obj(lm) == 5
     @test neval_grad(lm) == 8
-    @test neval_hess(lm) == 3
+    @test neval_hess(lm) == 4
   end
 
   @testset "Armijo-Wolfe" begin
@@ -63,4 +64,43 @@
     @test nbk > 0
     @test nbW > 0
   end
+
+  @testset "Don't allocate" begin
+    n = 200
+    nlp = SimpleModel(n)
+    x = nlp.meta.x0
+    g = zeros(n)
+    d = -40 * ones(n)
+    lm = LineModel(nlp, x, d)
+
+    al1 = @allocated obj(nlp, x)
+    al2 = @allocated obj(lm, 1.0)
+    @test al1 == al2
+
+    al1 = @allocated grad!(nlp, x, g)
+    al2 = @allocated grad!(lm, 1.0, g)
+    @test al1 + 16 == al2
+
+    al1 = @allocated objgrad!(nlp, x, g)
+    al2 = @allocated objgrad!(lm, 1.0, g)
+    al1 = @allocated objgrad!(nlp, x, g)
+    al2 = @allocated objgrad!(lm, 1.0, g)
+    @test al1 + 64 == al2
+
+    al1 = @allocated hprod!(nlp, x, d, g)
+    al2 = @allocated hess!(lm, 1.0, g)
+    @test al1 + 16 == al2
+
+    h₀ = obj(lm, 0.0)
+    slope = grad(lm, 0.0)
+    (t, gg, ht, nbk, nbW) = armijo_wolfe(lm, h₀, slope, g)
+    al = @allocated armijo_wolfe(lm, h₀, slope, g)
+    @test al == (8 + nbk) * 32 + 16
+
+    for bk_max = 0:8
+      (t, gg, ht, nbk, nbW) = armijo_wolfe(lm, h₀, slope, g, bk_max=bk_max)
+      al = @allocated armijo_wolfe(lm, h₀, slope, g, bk_max=bk_max)
+      @test al == min(12.5, 7 + 1.5nbk) * 32 + 208
+    end
+  end
 end