Rework FH problem: reduce allocations and improve efficiency

Project.toml

@@ -24,29 +24,33 @@ MLDatasets = "^0.7.4"
 NLPModels = "0.16, 0.17, 0.18, 0.19, 0.20, 0.21"
 NLPModelsModifiers = "0.7.3"
 Noise = "0.2"
+OrdinaryDiffEqVerner = "1.13.0"
 Requires = "1"
+ReverseDiff = "1.16.2"
 julia = "^1.6.0"
 
 [extras]
 ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a"
-DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
 ProximalOperators = "a725b495-10eb-56fe-b38b-717eba820537"
 QuadraticModels = "f468eda6-eac5-11e8-05a5-ff9e497bcd19"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 ShiftedProximalOperators = "d4fd37fa-580c-4e43-9b30-361c21aae263"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = [
   "ADNLPModels", 
-  "DifferentialEquations", 
   "LinearOperators", 
   "MLDatasets", 
   "NLPModelsModifiers", 
+  "OrdinaryDiffEqVerner", 
   "ProximalOperators", 
   "QuadraticModels", 
+  "ReverseDiff", 
   "ShiftedProximalOperators", 
   "Test"
 ]

docs/src/index.md

@@ -45,7 +45,7 @@ Calling `model = bpdn_nls_model()` returns the same problem modeled explicitly a
 
 ### Fitzhugh-Nagumo data-fitting problem
 
-If `ADNLPModels` and `DifferentialEquations` have been imported, `model = fh_model()` returns a model representing the over-determined nonlinear least-squares residual
+If `ADNLPModels`, `OrdinaryDiffEqVerner` and `ReverseDiff` have been imported, `model = fh_model()` returns a model representing the over-determined nonlinear least-squares residual
 ```math
 f(x) = \tfrac{1}{2} \|F(x)\|_2^2,
 ```

src/RegularizedProblems.jl

@@ -24,10 +24,12 @@ function __init__()
     include("testset_group_lasso.jl")
   end
   @require ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a" begin
-    @require DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" begin
-      include("fh_model.jl")
-      @require ProximalOperators = "a725b495-10eb-56fe-b38b-717eba820537" begin
-        include("testset_fh.jl")
+    @require OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2" begin
+      @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
+        include("fh_model.jl")
+        @require ProximalOperators = "a725b495-10eb-56fe-b38b-717eba820537" begin
+          include("testset_fh.jl")
+        end
       end
     end
   end

src/fh_model.jl

@@ -15,50 +15,56 @@ function FH_smooth_term(; abstol = 1e-14, reltol = 1e-14)
   tspan = (0.0, 20.0)
   savetime = 0.2
 
-  pars_FH = [0.5, 0.08, 1.0, 0.8, 0.7]
-  prob_FH = DifferentialEquations.ODEProblem(FH_ODE, u0, tspan, pars_FH)
+  x0 = [0.0, 0.2, 1.0, 0.0, 0.0]
+  p = copy(x0)
+  prob = OrdinaryDiffEqVerner.ODEProblem(FH_ODE, u0, tspan, x0)
+  alg = OrdinaryDiffEqVerner.Vern9()
+  integrator = OrdinaryDiffEqVerner.init(prob, alg, abstol = abstol, reltol = reltol, p = p, saveat = savetime)
+  reverse_diff_integrator = OrdinaryDiffEqVerner.init(prob, alg, abstol = abstol, reltol = reltol, p = ReverseDiff.TrackedReal.(p, 0.0), saveat = savetime)
+  OrdinaryDiffEqVerner.solve!(integrator)
 
-  x0 = [0, 0.2, 1.0, 0, 0]
-  prob_VDP = DifferentialEquations.ODEProblem(FH_ODE, u0, tspan, x0)
-  sol_VDP = DifferentialEquations.solve(prob_VDP, reltol = 1e-6, saveat = savetime)
 
   # add random noise to vdP solution
-  t = sol_VDP.t
-  b = vec(sol_VDP)
-  noise = 0.1 * randn(size(b))
-  data = noise .+ b
-
-  # solve FH with parameters p
-  function simulate(p, abstol = 1e-14, reltol = 1e-14)
-    temp_prob = DifferentialEquations.remake(prob_FH, p = p)
-    sol = DifferentialEquations.solve(
-      temp_prob,
-      DifferentialEquations.Vern9(),
-      abstol = abstol,
-      reltol = reltol,
-      saveat = savetime,
-    )
-    # if any((sol.retcode != :Success for s in sol))
-    #   @warn "ODE solution failed with parameters" p'
-    #   error("ODE solution failed")
-    # end
-    return vec(sol)
-  end
+  noise = 0.1 * randn(length(integrator.sol.u), 2)
+  noise = collect(eachrow(noise))
+  data = noise + integrator.sol.u
+  temp = zeros(eltype(data[1]), length(data) * 2)
+  reverse_diff_temp = ReverseDiff.TrackedReal.(temp, 0.0)
 
   # define residual vector
-  function residual(p, args...)
-    F = simulate(p, args...)
-    F .-= data
+  function residual!(F, x :: AbstractVector{T}) where{T <: Real}
+    OrdinaryDiffEqVerner.reinit!(integrator, u0)
+    integrator.p .= x
+    OrdinaryDiffEqVerner.solve!(integrator)
+    @inbounds for i in 1:length(integrator.sol.u)
+      F[i] = integrator.sol.u[i][1] - data[i][1]
+      F[i + length(integrator.sol.u)] = integrator.sol.u[i][2] - data[i][2]
+    end
+    return F
+  end
+
+  function residual!(F, x :: AbstractVector{T}) where{T <: ReverseDiff.TrackedReal}
+    OrdinaryDiffEqVerner.reinit!(reverse_diff_integrator, u0)
+    reverse_diff_integrator.p .= x
+    OrdinaryDiffEqVerner.solve!(reverse_diff_integrator)
+    @inbounds for i in 1:length(reverse_diff_integrator.sol.u)
+      F[i] = reverse_diff_integrator.sol.u[i][1] - data[i][1]
+      F[i + length(reverse_diff_integrator.sol.u)] = reverse_diff_integrator.sol.u[i][2] - data[i][2]
+    end
     return F
   end
 
-  # misfit = ‖residual‖² / 2
-  function misfit(p, args...)
-    F = residual(p, args...)
-    return dot(F, F) / 2
+  function obj(x :: AbstractVector{T}) where{T <: Real}
+    residual!(temp, x)
+    return dot(temp, temp) / 2
   end
 
-  return data, simulate, residual, misfit, x0
+  function obj(x :: AbstractVector{T}) where{T <: ReverseDiff.TrackedReal}
+    residual!(reverse_diff_temp, x)
+    return dot(reverse_diff_temp, reverse_diff_temp) / 2
+  end
+
+  return data, residual!, obj, x0
 end
 
 """
@@ -77,17 +83,29 @@ oscillator with fixed, but unknown, parameters.
 ## Keyword Arguments
 
 All keyword arguments are passed directly to the `ADNLPModel` (or `ADNLSModel`)
-constructure, e.g., to set the automatic differentiation backend.
+constructor.
 
 ## Return Value
 
 An instance of an `ADNLPModel` that represents the Fitzhugh-Nagumo problem, an instance
 of an `ADNLSModel` that represents the same problem, and the exact solution.
 """
 function fh_model(; kwargs...)
-  data, simulate, resid, misfit, x0 = FH_smooth_term()
+  data, resid!, misfit, x0 = FH_smooth_term()
   nequ = 202
-  ADNLPModels.ADNLPModel(misfit, ones(5); kwargs...),
-  ADNLPModels.ADNLSModel(resid, ones(5), nequ; kwargs...),
-  x0
-end
+  nlp = ADNLPModels.ADNLPModel(misfit, ones(5); 
+    hessian_backend = ADNLPModels.ReverseDiffADHessian,
+    gradient_backend = ADNLPModels.ReverseDiffADGradient,
+    kwargs...
+  )
+  nls = ADNLPModels.ADNLSModel!(resid!, ones(5), nequ; 
+    jacobian_residual_backend = ADNLPModels.ReverseDiffADJacobian,
+    jprod_residual_backend = ADNLPModels.ReverseDiffADJprod, 
+    jtprod_residual_backend = ADNLPModels.ReverseDiffADJtprod,
+    hessian_backend = ADNLPModels.ReverseDiffADHessian,
+    hessian_residual_backend = ADNLPModels.ReverseDiffADHessian,
+    gradient_backend = ADNLPModels.ReverseDiffADGradient,
+    kwargs...
+  )
+  return nlp, nls, x0
+end

test/runtests.jl

@@ -1,11 +1,11 @@
-using LinearAlgebra, Test
+using LinearAlgebra, ReverseDiff, Test
 using ADNLPModels,
-  DifferentialEquations,
   LinearOperators,
   ManualNLPModels,
   MLDatasets,
   NLPModels,
   NLPModelsModifiers,
+  OrdinaryDiffEqVerner,
   QuadraticModels
 using RegularizedProblems