Alternative fix to extra allocations in get_gradient! from Jacobian (#…

…301) * Alternative fix to extra allocations in get_gradient! from Jacobian * docs * add new function to docs. * add some tests * bump version --------- Co-authored-by: Ronny Bergmann <[email protected]>
JuliaManifolds · Oct 8, 2023 · 445e833 · 445e833 · kellertuer · Oct 8, 2023
1 parent 48f5ab0
commit 445e833
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Manopt"
 uuid = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
 authors = ["Ronny Bergmann <[email protected]>"]
-version = "0.4.37"
+version = "0.4.38"
 
 [deps]
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"

diff --git a/docs/src/plans/objective.md b/docs/src/plans/objective.md
@@ -130,6 +130,12 @@ and internally
 get_gradient_function
 ```
 
+#### Internal Helpers
+
+```@docs
+get_gradient_from_Jacobian!
+```
+
 ### Subgradient Objective
 
 ```@docs

diff --git a/src/plans/nonlinear_least_squares_plan.jl b/src/plans/nonlinear_least_squares_plan.jl
@@ -70,6 +70,33 @@ function get_cost(
     return 1//2 * norm(residual_values)^2
 end
 
+"""
+    get_gradient_from_Jacobian!(
+        M::AbstractManifold,
+        X,
+        nlso::NonlinearLeastSquaresObjective{InplaceEvaluation},
+        p,
+        Jval=zeros(nlso.num_components, manifold_dimension(M)),
+    )
+
+Compute gradient of [`NonlinearLeastSquaresObjective`](@ref) `nlso` at point `p` in place of
+`X`, with temporary Jacobian stored in the optional argument `Jval`.
+"""
+function get_gradient_from_Jacobian!(
+    M::AbstractManifold,
+    X,
+    nlso::NonlinearLeastSquaresObjective{InplaceEvaluation},
+    p,
+    Jval=zeros(nlso.num_components, manifold_dimension(M)),
+)
+    basis_p = _maybe_get_basis(M, p, nlso.jacobian_tangent_basis)
+    nlso.jacobian!!(M, Jval, p; basis_domain=basis_p)
+    residual_values = zeros(nlso.num_components)
+    nlso.f(M, residual_values, p)
+    get_vector!(M, X, p, transpose(Jval) * residual_values, basis_p)
+    return X
+end
+
 function get_gradient(
     M::AbstractManifold, nlso::NonlinearLeastSquaresObjective{AllocatingEvaluation}, p
 )
@@ -101,12 +128,7 @@ end
 function get_gradient!(
     M::AbstractManifold, X, nlso::NonlinearLeastSquaresObjective{InplaceEvaluation}, p
 )
-    basis_p = _maybe_get_basis(M, p, nlso.jacobian_tangent_basis)
-    Jval = zeros(nlso.num_components, manifold_dimension(M))
-    nlso.jacobian!!(M, Jval, p; basis_domain=basis_p)
-    residual_values = zeros(nlso.num_components)
-    nlso.f(M, residual_values, p)
-    get_vector!(M, X, p, transpose(Jval) * residual_values, basis_p)
+    get_gradient_from_Jacobian!(M, X, nlso, p)
     return X
 end
 

diff --git a/src/solvers/LevenbergMarquardt.jl b/src/solvers/LevenbergMarquardt.jl
@@ -206,8 +206,9 @@ function initialize_solver!(
     dmp::DefaultManoptProblem{mT,<:NonlinearLeastSquaresObjective{InplaceEvaluation}},
     lms::LevenbergMarquardtState,
 ) where {mT<:AbstractManifold}
-    get_objective(dmp).f(get_manifold(dmp), lms.residual_values, lms.p)
-    lms.X = get_gradient(dmp, lms.p)
+    M = get_manifold(dmp)
+    get_objective(dmp).f(M, lms.residual_values, lms.p)
+    get_gradient_from_Jacobian!(M, lms.X, get_objective(dmp), lms.p, lms.jacF)
     return lms
 end
 

diff --git a/test/solvers/test_Levenberg_Marquardt.jl b/test/solvers/test_Levenberg_Marquardt.jl
@@ -217,6 +217,7 @@ end
     X_r2 = similar(x0)
     get_gradient!(p_r2, X_r2, x0)
     @test isapprox(X_r2, [270.3617451389837, 677.6730784956912])
+    @test isapprox(get_gradient(p_r2, x0), [270.3617451389837, 677.6730784956912])
 
     p_r2_mut = DefaultManoptProblem(
         M,
@@ -228,6 +229,7 @@ end
     X_r2 = similar(x0)
     get_gradient!(p_r2_mut, X_r2, x0)
     @test isapprox(X_r2, [270.3617451389837, 677.6730784956912])
+    @test isapprox(get_gradient(p_r2_mut, x0), [270.3617451389837, 677.6730784956912])
 
     @testset "errors" begin
         @test_throws ArgumentError LevenbergMarquardtState(