Remove allocs l2 (#137)

remove allocs from ψ evaluation and prox! in shifted-L2, add tests.

Author: nathanemac

src/groupNormL2.jl

@@ -33,7 +33,7 @@ GroupNormL2(lambda::AbstractVector{R} = [one(R)], idx::I = [:]) where {R <: Real
 function (f::GroupNormL2)(x::AbstractArray{R}) where {R <: Real}
   sum_c = R(0)
   for (idx, λ) ∈ zip(f.idx, f.lambda)
-    sum_c += λ * norm(x[idx])
+    @views sum_c += λ * norm(x[idx])
   end
   return sum_c
 end

src/shiftedGroupNormL2.jl

@@ -62,15 +62,18 @@ function prox!(
   V1 <: AbstractVector{R},
   V2 <: AbstractVector{R},
 }
-  ψ.sol .= q + ψ.xk + ψ.sj
+  @. ψ.sol = q + ψ.xk + ψ.sj
   for (idx, λ) ∈ zip(ψ.h.idx, ψ.h.lambda)
-    snorm = norm(ψ.sol[idx])
+    sol_idx = view(ψ.sol, idx)
+    yv = view(y, idx)
+    snorm = norm(sol_idx)
     if snorm == 0
       y[idx] .= 0
     else
-      y[idx] .= max(1 - σ * λ / snorm, 0) .* ψ.sol[idx]
+      α = max(1 - σ * λ / snorm, 0)
+      @. yv = α * sol_idx
     end
   end
-  y .-= (ψ.xk + ψ.sj)
+  @. y -= (ψ.xk + ψ.sj)
   return y
 end

test/test_allocs.jl

@@ -40,7 +40,7 @@ end
     CompositeOp = eval(composite_op)
 
     function c!(z, x)
-      z[1] = 2*x[1] - x[4]
+      z[1] = 2 * x[1] - x[4]
       z[2] = x[2] + x[3]
     end
     function J!(z, x)
@@ -113,6 +113,22 @@ end
     @test @wrappedallocs(iprox!(y, ψ, y, d)) == 0
   end
 
+  for op ∈ (:NormL2,)
+    h = eval(op)(1.0)
+    n = 1000
+    xk = rand(n)
+    y = rand(n)
+    d = rand(n)
+
+    @test @wrappedallocs(prox!(y, h, y, 1.0)) == 0
+
+    ψ = shifted(h, xk)
+
+    @test @wrappedallocs(ψ(y)) == 0
+
+    @test @wrappedallocs(prox!(y, ψ, y, 1.0)) == 0
+  end
+
   for (op, shifted_op) ∈ zip((:Rank, :Nuclearnorm), (:ShiftedRank, :ShiftedNuclearnorm))
     ShiftedOp = eval(shifted_op)
     Op = eval(op)