From 923cc421d90699b3921fab138ccc5bf372f4b755 Mon Sep 17 00:00:00 2001
From: Magalame <12373104+Magalame@users.noreply.github.com>
Date: Sat, 2 Mar 2019 05:05:51 -0400
Subject: [PATCH] Add support for inplace jacobian and model for unidimensional
 input (#102)

* add support for inplace jacobian

* make range calls compatible with nightly

* Added entry in readme about inplacejac. Removed superfluous option in one lmfit

* correction to f!_from_f

* update

* added support for inplacef

* update readme

-introduce the `inplacef` option
-generalize use of `@.` for jacobians
-introduce use of `@views` for slices

* correct bug

* update inplace tests

- removed callable type tests (obsolete)
- added more tests to cover all inplace cases

* unified inplacef and inplacejac, adapted tests

* Removed LOAD_PATH modification

* Update curve_fit.jl

* Remove unused f!_from_f

* Remove reference to univariate only in README.md.

* Remove Function annotations.

* Apply AbstractArray annotations

* Update curve_fit.jl
---
 README.md                 |  19 ++++-
 src/curve_fit.jl          |  86 ++++++++++++++++-------
 test/curve_fit_inplace.jl | 144 ++++++++++++++++++++++++++++++++++++++
 test/runtests.jl          |   3 +-
 4 files changed, 223 insertions(+), 29 deletions(-)
 create mode 100644 test/curve_fit_inplace.jl

diff --git a/README.md b/README.md
index 7ad35e1..3d8ee6d 100644
--- a/README.md
+++ b/README.md
@@ -60,9 +60,9 @@ confidence_inter = confidence_interval(fit, 0.05)
 # The finite difference method is used above to approximate the Jacobian.
 # Alternatively, a function which calculates it exactly can be supplied instead.
 function jacobian_model(x,p)
-    J = Array{Float64}(undef, length(x),length(p))
-    J[:,1] = exp.(-x.*p[2])    #dmodel/dp[1]
-    J[:,2] = -x.*p[1].*J[:,1]  #dmodel/dp[2]
+    J = Array{Float64}(undef, length(x), length(p))
+    @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
+    @. @views J[:,2] = -x*p[1]*J[:,1] #dmodel/dp[2], thanks to @views we don't allocate memory for the J[:,1] slice
     J
 end
 fit = curve_fit(model, jacobian_model, xdata, ydata, p0)
@@ -81,6 +81,19 @@ The default is to calculate the Jacobian using a central finite differences sche
 fit = curve_fit(model, xdata, ydata, p0; autodiff=:forwarddiff)
 ```
 
+Inplace model and jacobian 
+-------------------------
+It is possible to either use an inplace model, or an inplace model *and* an inplace jacobian. It might be pertinent to use this feature when `curve_fit` is slow, or consumes a lot of memory
+```
+model_inplace(F, x, p) = (@. F = p[1] * exp(-x * p[2]))
+
+function jacobian_inplace(J::Array{Float64,2},x,p)
+        @. J[:,1] = exp(-x*p[2])     
+        @. @views J[:,2] = -x*p[1]*J[:,1] 
+    end
+fit = curve_fit(model_inplace, jacobian_inplace, xdata, ydata, p0; inplace = true)
+```
+
 Existing Functionality
 ----------------------
 
diff --git a/src/curve_fit.jl b/src/curve_fit.jl
index 4cc8348..50ca3f7 100755
--- a/src/curve_fit.jl
+++ b/src/curve_fit.jl
@@ -15,15 +15,27 @@ StatsBase.weights(lfr::LsqFitResult) = lfr.wt
 StatsBase.residuals(lfr::LsqFitResult) = lfr.resid
 mse(lfr::LsqFitResult) = rss(lfr)/dof(lfr)
 
-# provide a method for those who have their own Jacobian function
-function lmfit(f, g, p0, wt; autodiff = :finite, kwargs...)
+# provide a method for those who have their own (non inplace) Jacobian function
+function lmfit(f, g, p0::AbstractArray, wt::AbstractArray; kwargs...)
     r = f(p0)
-    autodiff = autodiff == :forwarddiff ? :forward : autodiff
     R = OnceDifferentiable(f, g, p0, similar(r); inplace = false)
     lmfit(R, p0, wt; kwargs...)
 end
 
-function lmfit(f, p0, wt; autodiff = :finite, kwargs...)
+#for inplace f and inplace g
+function lmfit(f!, g!, p0::AbstractArray, wt::AbstractArray, r::AbstractArray; kwargs...)
+    R = OnceDifferentiable(f!, g!, p0, similar(r); inplace = true)
+    lmfit(R, p0, wt; kwargs...)
+end
+
+#for inplace f only
+function lmfit(f!, p0::AbstractArray, wt::AbstractArray, r::AbstractArray; autodiff = :finite, kwargs...)
+    autodiff = autodiff == :forwarddiff ? :forward : autodiff
+    R = OnceDifferentiable(f!, p0, similar(r); inplace = true, autodiff = autodiff)
+    lmfit(R, p0, wt; kwargs...)
+end
+
+function lmfit(f, p0::AbstractArray, wt::AbstractArray; autodiff = :finite, kwargs...)
     # this is a convenience function for the curve_fit() methods
     # which assume f(p) is the cost functionj i.e. the residual of a
     # model where
@@ -44,7 +56,7 @@ function lmfit(f, p0, wt; autodiff = :finite, kwargs...)
     lmfit(R, p0, wt; kwargs...)
 end
 
-function lmfit(R::OnceDifferentiable, p0, wt; autodiff = :finite, kwargs...)
+function lmfit(R::OnceDifferentiable, p0::AbstractArray, wt::AbstractArray; autodiff = :finite, kwargs...)
     results = levenberg_marquardt(R, p0; kwargs...)
     p = minimizer(results)
     return LsqFitResult(p, value!(R, p), jacobian!(R, p), converged(results), wt)
@@ -82,40 +94,66 @@ fit = curve_fit(model, xdata, ydata, p0)
 """
 function curve_fit end
 
-function curve_fit(model::Function, xpts::AbstractArray, ydata::AbstractArray, p0; kwargs...)
+function curve_fit(model, xpts::AbstractArray, ydata::AbstractArray, p0::AbstractArray; inplace = false, kwargs...)
     # construct the cost function
-    f(p) = model(xpts, p) - ydata
     T = eltype(ydata)
-    lmfit(f,p0,T[]; kwargs...)
+
+    if inplace
+        f! = (F,p)  -> (model(F,xpts,p); @. F = F - ydata)
+        lmfit(f!, p0, T[], ydata; kwargs...)
+    else
+        f = (p) -> model(xpts, p) - ydata
+        lmfit(f,p0,T[]; kwargs...)
+    end
 end
 
-function curve_fit(model::Function, jacobian_model::Function,
-            xpts::AbstractArray, ydata::AbstractArray, p0; kwargs...)
-    f(p) = model(xpts, p) - ydata
-    g(p) = jacobian_model(xpts, p)
+function curve_fit(model, jacobian_model,
+            xpts::AbstractArray, ydata::AbstractArray, p0::AbstractArray; inplace = false, kwargs...)
+
     T = eltype(ydata)
-    lmfit(f, g, p0, T[]; kwargs...)
+
+    if inplace
+        f! = (F,p) -> (model(F,xpts,p); @. F = F - ydata)
+        g! = (G,p)  -> jacobian_model(G, xpts, p)
+        lmfit(f!, g!, p0, T[], similar(ydata); kwargs...)
+    else 
+        f = (p) -> model(xpts, p) - ydata
+        g = (p) -> jacobian_model(xpts, p)
+        lmfit(f, g, p0, T[]; kwargs...)
+    end
 end
 
-function curve_fit(model::Function, xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T}, p0; kwargs...) where T
+function curve_fit(model, xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T}, p0::AbstractArray; inplace = false, kwargs...) where T
     # construct a weighted cost function, with a vector weight for each ydata
     # for example, this might be wt = 1/sigma where sigma is some error term
     u = sqrt.(wt) # to be consistant with the matrix form
-
-    f(p) = u .* ( model(xpts, p) - ydata )
-    lmfit(f,p0,wt; kwargs...)
+    
+    if inplace
+        f! = (F,p) -> (model(F,xpts,p); @. F = u*(F - ydata))
+        lmfit(f!, p0, wt, ydata; kwargs...)
+    else
+        f = (p)  -> u .* ( model(xpts, p) - ydata )
+        lmfit(f,p0,wt; kwargs...)
+    end
 end
 
-function curve_fit(model::Function, jacobian_model::Function,
-            xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T}, p0; kwargs...) where T
+function curve_fit(model, jacobian_model,
+            xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T}, p0; inplace = false, kwargs...) where T
 
     u = sqrt.(wt) # to be consistant with the matrix form
-    f(p) = u .* ( model(xpts, p) - ydata )
-    g(p) = u .* ( jacobian_model(xpts, p) )
-    lmfit(f, g, p0, wt; kwargs...)
+
+    if inplace
+        f! = (F,p) -> (model(F,xpts,p); @. F = u*(F - ydata))
+        g! = (G,p) -> (jacobian_model(G, xpts, p); @. G = u*G )
+        lmfit(f!, g!, p0, wt, ydata; kwargs...)
+    else 
+        f = (p) -> u .* ( model(xpts, p) - ydata )
+        g = (p) -> u .* ( jacobian_model(xpts, p) )
+        lmfit(f, g, p0, wt; kwargs...)
+    end
 end
 
-function curve_fit(model::Function, xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T,2}, p0; kwargs...) where T
+function curve_fit(model, xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T,2}, p0; kwargs...) where T
     # as before, construct a weighted cost function with where this
     # method uses a matrix weight.
     # for example: an inverse_covariance matrix
@@ -129,7 +167,7 @@ function curve_fit(model::Function, xpts::AbstractArray, ydata::AbstractArray, w
     lmfit(f,p0,wt; kwargs...)
 end
 
-function curve_fit(model::Function, jacobian_model::Function,
+function curve_fit(model, jacobian_model,
             xpts::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T,2}, p0; kwargs...) where T
     u = cholesky(wt).U
 
diff --git a/test/curve_fit_inplace.jl b/test/curve_fit_inplace.jl
new file mode 100644
index 0000000..6331815
--- /dev/null
+++ b/test/curve_fit_inplace.jl
@@ -0,0 +1,144 @@
+let
+    # fitting noisy data to an exponential model
+    # TODO: Change to `.-x` when 0.5 support is dropped
+    @. model(x, p) = p[1] * exp(-x * p[2])
+    model_inplace(F, x, p) = (@. F = p[1] * exp(-x * p[2]))
+
+    # some example data
+    Random.seed!(12345)
+    xdata = range(0, stop=10, length=500000)
+    ydata = model(xdata, [1.0, 2.0]) + 0.01*randn(length(xdata))
+    p0 = [0.5, 0.5]
+
+    # if your model is differentiable, it can be faster and/or more accurate
+    # to supply your own jacobian instead of using the finite difference
+    function jacobian_model(x,p)
+        J = Array{Float64}(undef, length(x), length(p))
+        @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
+        @. @views J[:,2] = -x*p[1]*J[:,1] 
+        J
+    end
+
+    function jacobian_model_inplace(J::Array{Float64,2},x,p)
+        @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
+        @. @views J[:,2] = -x*p[1]*J[:,1] 
+    end
+
+
+    f(p) = model(xdata, p) - ydata
+    g(p) = jacobian_model(xdata, p)
+    df = OnceDifferentiable(f, g, p0, similar(ydata); inplace = false);
+    evalf(x) = NLSolversBase.value!!(df, x)
+    evalg(x) = NLSolversBase.jacobian!!(df, x)
+    r = evalf(p0);
+    j = evalg(p0);
+
+    f_inplace = (F,p) -> (model_inplace(F, xdata, p); @. F = F - ydata)
+    g_inplace = (G,p) -> jacobian_model_inplace(G, xdata, p)
+    df_inplace = OnceDifferentiable(f_inplace, g_inplace, p0, similar(ydata); inplace = true);
+    evalf_inplace(x) = NLSolversBase.value!!(df_inplace,x)
+    evalg_inplace(x) = NLSolversBase.jacobian!!(df_inplace,x)
+    r_inplace = evalf_inplace(p0);
+    j_inplace = evalg_inplace(p0);
+
+
+    @test r == r_inplace
+    @test j == j_inplace
+
+    println("--------------\nPerformance of non-inplace")
+    println("\t Evaluation function")
+
+    stop = 8 #8 because the tests afterwards will call the eval function 8 or 9 times, so it makes it easy to compare
+    step = 1
+
+    @time for i=range(0,stop=stop,step=step)
+        evalf(p0);
+    end
+
+    println("\t Jacobian function")
+    @time for i=range(0,stop=stop,step=step)
+        evalg(p0);
+    end
+
+    println("--------------\nPerformance of inplace")
+    println("\t Evaluation function")
+    @time for i=range(0,stop=stop,step=step)
+        evalf_inplace(p0);
+    end
+
+    println("\t Jacobian function")
+    @time for i=range(0,stop=stop,step=step)
+        evalg_inplace(p0);
+    end
+
+    
+    curve_fit(model, xdata, ydata, p0; maxIter=100); #warmup
+    curve_fit(model_inplace, xdata, ydata, p0; inplace = true, maxIter=100);
+
+    #explicit jac
+    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=100);
+    curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, p0; inplace = true, maxIter=100);
+
+
+
+    println("--------------\nPerformance of curve_fit")
+
+    println("\t Non-inplace")
+    fit = @time curve_fit(model, xdata, ydata, p0; maxIter=100)
+    @test fit.converged
+
+    println("\t Inplace")
+    fit_inplace = @time curve_fit(model_inplace, xdata, ydata, p0; inplace = true, maxIter=100)
+    @test fit_inplace.converged
+
+    @test fit_inplace.param == fit.param
+
+
+    println("\t Non-inplace with jacobian")
+    fit_jac = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=100)
+    @test fit_jac.converged
+
+    println("\t Inplace with jacobian")
+    fit_inplace_jac = @time curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, p0; inplace = true, maxIter=100)
+    @test fit_inplace_jac.converged
+
+    @test fit_jac.param == fit_inplace_jac.param
+
+
+
+
+    # some example data
+    yvars = 1e-6*rand(length(xdata))
+    ydata = model(xdata, [1.0, 2.0]) + sqrt.(yvars) .* randn(length(xdata))
+
+    println("--------------\nPerformance of curve_fit with weights")
+
+    curve_fit(model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]);
+    curve_fit(model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100);
+
+    curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]);
+    curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100);
+
+
+    println("\t Non-inplace with weights")
+    fit_wt = @time curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; maxIter=100)
+    @test fit_wt.converged
+
+    println("\t Inplace with weights")
+    fit_inplace_wt = @time curve_fit(model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100)
+    @test fit_inplace_wt.converged
+
+    @test maximum(abs.(fit_wt.param - fit_inplace_wt.param)) < 1e-15
+
+
+    println("\t Non-inplace with jacobian with weights")
+    fit_wt_jac = @time curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; maxIter=100)
+    @test fit_wt_jac.converged
+
+    println("\t Inplace with jacobian with weights")
+    fit_inplace_wt_jac = @time curve_fit(model_inplace, jacobian_model_inplace, xdata, ydata, 1 ./ yvars, [0.5, 0.5]; inplace = true, maxIter=100)
+    @test fit_inplace_wt_jac.converged
+
+    @test maximum(abs.(fit_wt_jac.param - fit_inplace_wt_jac.param)) < 1e-15
+
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 3d03f45..0f41cab 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,12 +1,11 @@
 #
 # Correctness Tests
 #
-
 using LsqFit, Test, LinearAlgebra, Random
 using OptimBase, Calculus
 import NLSolversBase: OnceDifferentiable
 
-my_tests = [ "curve_fit.jl", "levenberg_marquardt.jl"]
+my_tests = ["curve_fit.jl", "levenberg_marquardt.jl", "curve_fit_inplace.jl"]
 
 println("Running tests:")