From 70dcdf7ae1d9c4c2449b6ea6a72d3c02e538ad64 Mon Sep 17 00:00:00 2001
From: Francesco Alemanno <50984334+francescoalemanno@users.noreply.github.com>
Date: Wed, 23 Mar 2022 19:20:25 +0100
Subject: [PATCH] add inplace version

---
 Project.toml      |  2 +-
 src/FixedPoint.jl | 63 ++++++++++++++++++++++++++++++++++++++++++++++-
 test/runtests.jl  | 27 ++++++++++++++++++++
 3 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 44a0bc5..dddf968 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "FixedPoint"
 uuid = "3325f569-5a18-4e7d-8356-246b69339eea"
 authors = ["Francesco Alemanno <francescoalemanno710@gmail.com>"]
-version = "1.0.0"
+version = "1.0.1"
 
 [compat]
 julia = "^1.5"
diff --git a/src/FixedPoint.jl b/src/FixedPoint.jl
index a1a4964..13c1217 100644
--- a/src/FixedPoint.jl
+++ b/src/FixedPoint.jl
@@ -54,5 +54,66 @@ function afps(
     end
     (x = x_n, error = grad_norm(f(x_n) - x_n), iters = runs)
 end
-export afps
+
+
+"""
+    afps!(f!, x; iters::Int = 5000, vel::Float64 = 0.9, ep::Float64 = 0.01, tol::Float64 = 1e-12, grad_norm=x->maximum(abs,x))
+
+solve equation `f(x) = x` according to:
+
+    `f!` : inplace version of function to find fixed point for, calling `f!(out,x)` should amount to writing `out = f(x)`
+
+    `x` : initial condition, ideally it should be close to the final solution
+
+    `vel` : amount of Nesterov acceleration in [0,1]
+
+    `ep` : learning rate, typically in ]0,1[
+
+    `tol` : absolute tolerance on |f(x)-x|
+
+    `grad_norm` : function to evaluate the norm for |f(x)-x|
+
+returns a named tuple (x, error, iters) where:
+
+    `x` : is the solution found for f(x)=x
+
+    `error` : is the norm of f(x)-x at the solution point
+
+    `iters` : total number of iterations performed
+"""
+function afps!(
+    f!::Fun,
+    x_n::Mat;
+    iters::Int = 5000,
+    vel::T = 0.9,
+    ep::T = 0.01,
+    tol::T = 1e-12,
+    grad_norm = x -> maximum(abs, x),
+) where {T <: Number, Mat<:Union{AbstractArray{T},T}, Fun <: Function}
+    v_n = zero(x_n)
+    trial = zero(x_n)
+    g = zero(x_n)
+    β = vel
+    ϵ = ep
+    runs = 0
+    for _ = 1:iters
+        #prediction using velocity
+        x_n .+= β .* v_n
+        #eval gradient
+        f!(g,x_n)
+        g .-= x_n
+        #correction
+        x_n .+= ϵ .* g
+        #update velocity 
+        v_n .= β .* v_n .+ ϵ .* g
+        runs += 1
+        if grad_norm(g) < tol
+            break
+        end
+    end
+    (x = x_n, error = grad_norm(g), iters = runs)
+end
+
+
+export afps,afps!
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index f5f7d3e..23e3ac4 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -15,3 +15,30 @@ end
     s = afps(x -> 2 - x^2 + x, 1.3)
     @test s.x ≈ √2
 end
+
+
+@testset "Inplace Vector Equation" begin
+    Ts = LinRange(0.01, 2.0, 500)
+    βs = 1 ./ Ts
+    function f!(out,x)
+        @. out = tanh(βs * x)
+    end
+    x  = zero(βs) .+ 1 
+    mag = afps!(
+        f!,
+        x,
+        grad_norm = x -> maximum(abs, x),
+        iters = 5000,
+    )
+    @test mag.error < 1e-4
+    @test maximum(abs, x .- tanh.(βs .* x)) < 1e-4
+end
+
+@testset "Inplace Equation" begin
+    x = [1.3]
+    function f!(out, x)
+        out[1] = 2 - x[1]^2 + x[1]
+    end
+    s = afps!(f!, x)
+    @test s.x[1] ≈ √2
+end