From d98c8c92e22a95bc9ea8bf3b8231ed8c6c4fd081 Mon Sep 17 00:00:00 2001
From: Geoffroy Leconte <gleconte50@gmail.com>
Date: Tue, 14 Feb 2023 18:53:47 -0500
Subject: [PATCH] LM with bounds

---
 examples/demo-bpdn-constr.jl |  6 ++++++
 src/LM_alg.jl                | 17 +++++++++++++----
 test/test_bounds.jl          |  2 +-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/examples/demo-bpdn-constr.jl b/examples/demo-bpdn-constr.jl
index d4d04dd3..c4ccb9d6 100644
--- a/examples/demo-bpdn-constr.jl
+++ b/examples/demo-bpdn-constr.jl
@@ -28,6 +28,12 @@ function demo_solver(f, nls, sol, h, χ, suffix = "l0-linf")
   LMTR_out = LMTR(nls, h, χ, options, x0 = f.meta.x0)
   @info "LMTR relative error" norm(LMTR_out.solution - sol) / norm(sol)
   plot_bpdn(LMTR_out, sol, "constr-lmtr-r2-$(suffix)")
+
+  @info " using LM to solve with" h
+  reset!(nls)
+  LM_out = LM(nls, h, options, x0 = f.meta.x0)
+  @info "LM relative error" norm(LM_out.solution - sol) / norm(sol)
+  plot_bpdn(LM_out, sol, "constr-lm-r2-$(suffix)")
 end
 
 function demo_bpdn_constr(compound = 1)
diff --git a/src/LM_alg.jl b/src/LM_alg.jl
index e9e6b5ed..75bfd847 100644
--- a/src/LM_alg.jl
+++ b/src/LM_alg.jl
@@ -29,6 +29,7 @@ and σ > 0 is a regularization parameter.
 * `subsolver_logger::AbstractLogger`: a logger to pass to the subproblem solver
 * `subsolver`: the procedure used to compute a step (`PG` or `R2`)
 * `subsolver_options::ROSolverOptions`: default options to pass to the subsolver.
+* `selected::AbstractVector{<:Integer}`: (default `1:f.meta.nvar`).
 
 ### Return values
 
@@ -45,6 +46,7 @@ function LM(
   subsolver_logger::Logging.AbstractLogger = Logging.NullLogger(),
   subsolver = R2,
   subsolver_options = ROSolverOptions(),
+  selected::AbstractVector{<:Integer} = 1:(nls.meta.nvar),
 ) where {H}
   start_time = time()
   elapsed_time = 0.0
@@ -60,6 +62,12 @@ function LM(
   θ = options.θ
   σmin = options.σmin
 
+  local l_bound, u_bound
+  if has_bounds(nls)
+    l_bound = nls.meta.lvar
+    u_bound = nls.meta.uvar
+  end
+
   if verbose == 0
     ptf = Inf
   elseif verbose == 1
@@ -73,16 +81,16 @@ function LM(
   # initialize parameters
   σk = max(1 / options.ν, σmin)
   xk = copy(x0)
-  hk = h(xk)
+  hk = h(xk[selected])
   if hk == Inf
     verbose > 0 && @info "LM: finding initial guess where nonsmooth term is finite"
     prox!(xk, h, x0, one(eltype(x0)))
-    hk = h(xk)
+    hk = h(xk[selected])
     hk < Inf || error("prox computation must be erroneous")
     verbose > 0 && @debug "LM: found point where h has value" hk
   end
   hk == -Inf && error("nonsmooth term is not proper")
-  ψ = shifted(h, xk)
+  ψ = has_bounds(nls) ? shifted(h, xk, l_bound - xk, u_bound - xk, selected) : shifted(h, xk)
 
   xkn = similar(xk)
 
@@ -179,7 +187,7 @@ function LM(
     xkn .= xk .+ s
     residual!(nls, xkn, Fkn)
     fkn = dot(Fkn, Fkn) / 2
-    hkn = h(xkn)
+    hkn = h(xkn[selected])
     hkn == -Inf && error("nonsmooth term is not proper")
     mks = mk(s)
     ξ = fk + hk - mks + max(1, abs(hk)) * 10 * eps()
@@ -205,6 +213,7 @@ function LM(
 
     if η1 ≤ ρk < Inf
       xk .= xkn
+      has_bounds(nls) && set_bounds!(ψ, l_bound - xk, u_bound - xk)
 
       # update functions
       Fk .= Fkn
diff --git a/test/test_bounds.jl b/test/test_bounds.jl
index 5cc3ca4e..1998bd8a 100644
--- a/test/test_bounds.jl
+++ b/test/test_bounds.jl
@@ -27,7 +27,7 @@ for (mod, mod_name) ∈ ((x -> x, "exact"), (LSR1Model, "lsr1"), (LBFGSModel, "l
 end
 
 for (h, h_name) ∈ ((NormL0(λ), "l0"), (NormL1(λ), "l1"))
-  for solver_sym ∈ (:LMTR,)
+  for solver_sym ∈ (:LMTR, :LM)
     solver_name = string(solver_sym)
     solver = eval(solver_sym)
     @testset "bpdn-with-bounds-ls-$(solver_name)-$(h_name)" begin