Non-allocating rank and nuclear norm prox

Author: MohamedLaghdafHABIBOULLAH

src/shiftedNuclearnorm.jl

@@ -65,11 +65,11 @@ function prox!(
   V1 <: AbstractVector{R},
   V2 <: AbstractVector{R},
 }
+  λ = ψ.h.lambda
   ψ.sol .= q .+ ψ.xk .+ ψ.sj
   ψ.h.A .= reshape_array(ψ.sol, size(ψ.h.A))
   psvd_dd!(ψ.h.F, ψ.h.A, full = false)
-  c = sqrt(2 * ψ.λ * σ)
-  ψ.h.F.S .= max.(0, ψ.h.F.S .- ψ.λ * σ)
+  ψ.h.F.S .= max.(0, ψ.h.F.S .- λ * σ)
   for i ∈ eachindex(ψ.h.F.S)
     for j = 1:size(ψ.h.A, 1)
       ψ.h.F.U[j, i] = ψ.h.F.U[j, i] .* ψ.h.F.S[i]

src/shiftedRank.jl

@@ -65,10 +65,11 @@ function prox!(
   V1 <: AbstractVector{R},
   V2 <: AbstractVector{R},
 }
+  λ = ψ.h.lambda
   ψ.sol .= q .+ ψ.xk .+ ψ.sj
   ψ.h.A .= reshape_array(ψ.sol, size(ψ.h.A))
   psvd_dd!(ψ.h.F, ψ.h.A, full = false)
-  c = sqrt(2 * ψ.λ * σ)
+  c = sqrt(2 * λ * σ)
   for i ∈ eachindex(ψ.h.F.S)
     if ψ.h.F.S[i] <= c
       ψ.h.F.U[:, i] .= 0

test/test_allocs.jl

@@ -83,4 +83,22 @@ end
     @test @wrappedallocs(prox!(y, ψ, y, 1.0)) == 0
     @test @wrappedallocs(iprox!(y, ψ, y, d)) == 0
   end
+
+  for (op, shifted_op) ∈ zip((:Rank, :Nuclearnorm), (:ShiftedRank, :ShiftedNuclearnorm))
+    ShiftedOp = eval(shifted_op)
+    Op = eval(op)
+    m = 10
+    n = 11
+    λ = 1.0
+    γ = 5.0
+    x = vec(reshape(rand(m, n), m * n, 1))
+    q = vec(reshape(rand(m, n), m * n, 1))
+    s = vec(reshape(rand(m, n), m * n, 1))
+    F = psvd_workspace_dd(zeros(m, n), full = false)
+    h = Op(λ, ones(m, n), F)
+    f = ShiftedOp(h, x, s, true)
+    y = zeros(m * n)
+    @test @wrappedallocs(prox!(y, h, x, γ)) == 0
+    @test @wrappedallocs(prox!(y, f, q, γ)) == 0
+  end
 end