LSR1: remove allocations, make more generic

Author: dpo

README.md

@@ -1,7 +1,7 @@
 # A [Julia](http://julialang.org) Linear Operator Package
 
 [![DOI](https://zenodo.org/badge/20136006.svg)](https://zenodo.org/badge/latestdoi/20136006)
-[![Build Status](https://travis-ci.org/JuliaSmoothOptimizers/LinearOperators.jl.svg?branch=master)](https://travis-ci.org/JuliaSmoothOptimizers/LinearOperators.jl)
+[![Build Status](https://github.com/JuliaSmoothOptimizers/NLPModels.jl/workflows/CI/badge.svg?branch=master)
 [![Build status](https://ci.appveyor.com/api/projects/status/kp1o6ejuu6kgskvp/branch/master?svg=true)](https://ci.appveyor.com/project/dpo/linearoperators-jl/branch/master)
 [![Build Status](https://api.cirrus-ci.com/github/JuliaSmoothOptimizers/LinearOperators.jl.svg)](https://cirrus-ci.com/github/JuliaSmoothOptimizers/LinearOperators.jl)
 [![Coveralls](https://coveralls.io/repos/JuliaSmoothOptimizers/LinearOperators.jl/badge.svg?branch=master&service=github)](https://coveralls.io/github/JuliaSmoothOptimizers/LinearOperators.jl?branch=master)

src/lbfgs.jl

@@ -237,7 +237,7 @@ end
 Extract the diagonal of a L-BFGS operator in forward mode.
 """
 function diag(op :: LBFGSOperator{T}) where T
-  d = zeros(T, op.nrow)
+  d = Vector{T}(undef, op.nrow)
   diag!(op, d)
 end
 

src/lsr1.jl

@@ -5,25 +5,27 @@ mutable struct LSR1Data{T}
   mem :: Int
   scaling :: Bool
   scaling_factor :: T
-  s   :: Matrix{T}
-  y   :: Matrix{T}
+  s   :: Vector{Vector{T}}
+  y   :: Vector{Vector{T}}
   ys  :: Vector{T}
-  a   :: Matrix{T}
+  a   :: Vector{Vector{T}}
   as  :: Vector{T}
   insert :: Int
+  Ax :: Vector{T}
 end
 
 function LSR1Data(T :: DataType, n :: Int; mem :: Int=5, scaling :: Bool=true, inverse :: Bool=false)
   inverse && @warn "inverse LSR1 operator not yet implemented"
   LSR1Data{T}(max(mem, 1),
               scaling,
               convert(T, 1),
-              zeros(T, n, mem),
-              zeros(T, n, mem),
+              [zeros(T, n) for _ = 1 : mem],
+              [zeros(T, n) for _ = 1 : mem],
               zeros(T, mem),
-              zeros(T, n, mem),
+              [zeros(T, n) for _ = 1 : mem],
               zeros(T, mem),
-              1)
+              1,
+              Vector{T}(undef, n))
 end
 
 LSR1Data(n :: Int; kwargs...) = LSR1Data(Float64, n; kwargs...)
@@ -60,15 +62,18 @@ function LSR1Operator(T :: DataType, n :: Int; kwargs...)
   function lsr1_multiply(data :: LSR1Data, x :: AbstractArray)
     # Multiply operator with a vector.
 
-    result_type = promote_type(T, eltype(x))
-    q = convert(Array{result_type}, copy(x))
+    q = data.Ax
+    q .= x
 
     data.scaling && (q ./= data.scaling_factor)  # q = B₀ * x
 
     for i = 1 : data.mem
       k = mod(data.insert + i - 2, data.mem) + 1
       if data.ys[k] != 0
-        @views q .+= dot(data.a[:, k], x) / data.as[k] * data.a[:, k]
+        ax = dot(data.a[k], x) / data.as[k]
+        for j ∈ eachindex(q)
+          q[j] += ax * data.a[k][j]
+        end
       end
     end
     return q
@@ -97,16 +102,20 @@ function push!(op :: LSR1Operator, s :: AbstractVector, y :: AbstractVector)
   Bs = op * s
   ymBs = y - Bs
   ys = dot(y, s)
+  sNorm = norm(s)
+  yy = dot(y, y)
 
-  well_defined = abs(dot(ymBs, s)) ≥ 1.0e-8 + 1.0e-8 * norm(ymBs)^2
+  ϵ = eps(eltype(op))
+  well_defined = abs(dot(ymBs, s)) ≥ ϵ + ϵ * norm(ymBs) * sNorm
 
   sufficient_curvature = true
   scaling_condition = true
   if data.scaling
-    sufficient_curvature = abs(ys) ≥ 1.0e-8
+    yNorm = √yy
+    sufficient_curvature = abs(ys) ≥ ϵ * yNorm * sNorm
     if sufficient_curvature
-      scaling_factor = ys / dot(y, y)
-      scaling_condition = norm(y - s / scaling_factor) >= 1.0e-8
+      scaling_factor = ys / yy
+      scaling_condition = norm(y - s / scaling_factor) >= ϵ *  yNorm * sNorm
     end
   end
 
@@ -116,12 +125,12 @@ function push!(op :: LSR1Operator, s :: AbstractVector, y :: AbstractVector)
     return op
   end
 
-  data.s[:, data.insert] .= s
-  data.y[:, data.insert] .= y
+  data.s[data.insert] .= s
+  data.y[data.insert] .= y
   data.ys[data.insert] = ys
 
   # update scaling factor
-  data.scaling && (data.scaling_factor = ys / dot(y, y))
+  data.scaling && (data.scaling_factor = ys / yy)
 
   # update next insertion position
   data.insert = mod(data.insert, data.mem) + 1
@@ -130,14 +139,15 @@ function push!(op :: LSR1Operator, s :: AbstractVector, y :: AbstractVector)
   for i = 1 : data.mem
     k = mod(data.insert + i - 2, data.mem) + 1
     if data.ys[k] != 0
-      data.a[:, k] .= data.y[:, k] - data.s[:, k] / data.scaling_factor  # = y - B₀ * s
+      data.a[k] .= data.y[k] - data.s[k] / data.scaling_factor  # = y - B₀ * s
       for j = 1 : i-1
         l = mod(data.insert + j - 2, data.mem) + 1
         if data.ys[l] != 0
-          @views data.a[:, k] .-= dot(data.a[:, l], data.s[:, k]) / data.as[l] * data.a[:, l]
+          as = dot(data.a[l], data.s[k]) / data.as[l]
+          data.a[k] .-= as * data.a[l]
         end
       end
-      @views data.as[k] = dot(data.a[:, k], data.s[:, k])
+      data.as[k] = dot(data.a[k], data.s[k])
     end
   end
 
@@ -147,21 +157,27 @@ end
 
 """
     diag(op)
+    diag!(op, d)
 
 Extract the diagonal of a L-SR1 operator in forward mode.
 """
 function diag(op :: LSR1Operator{T}) where T
+  d = Vector{T}(undef, op.nrow)
+  diag!(op, d)
+end
+
+function diag!(op :: LSR1Operator{T}, d) where T
   op.inverse && throw("only the diagonal of a forward L-SR1 approximation is available")
   data = op.data
 
-  d = ones(T, op.nrow)
+  fill!(d, 1)
   data.scaling && (d ./= data.scaling_factor)
 
   for i = 1 : data.mem
     k = mod(data.insert + i - 2, data.mem) + 1
     if data.ys[k] != 0.0
       for j = 1 : op.nrow
-        d[j] += data.a[j, k]^2 / data.as[k]
+        d[j] += data.a[k][j]^2 / data.as[k]
       end
     end
   end
@@ -174,10 +190,12 @@ end
 Reset the given LSR1 data.
 """
 function reset!(data :: LSR1Data{T}) where T
-  fill!(data.s, 0)
-  fill!(data.y, 0)
+  for i = 1 : data.mem
+    fill!(data.s[i], 0)
+    fill!(data.y[i], 0)
+    fill!(data.a[i], 0)
+  end
   fill!(data.ys, 0)
-  fill!(data.a, 0)
   fill!(data.as, 0)
   data.scaling_factor = T(1)
   data.insert = 1

test/test_lsr1.jl

@@ -77,6 +77,23 @@ function test_lsr1()
       @test eltype(B * v) == T
     end
   end
+
+  @testset "LSR1 allocations" begin
+    n = 100
+    mem = 20
+    B = LSR1Operator(n, mem=mem)
+    for _ = 1 :2:n
+      s = rand(n)
+      y = rand(n)
+      push!(B, s, y)
+    end
+    x = rand(n)
+    B * x  # warmup
+    nallocs = @allocated B * x
+    @test nallocs == 0
+    nallocs = @allocated diag!(B, x)
+    @test nallocs == 0
+  end
 end
 
 test_lsr1()