Skip to content

Commit ac472a9

Browse files
dehanndehann
authored
Merge pull request #1732 from JuliaRobotics/23Q3/enh/factorjacobian
autodiff for manifolds, new test file
2 parents 71a8f37 + 36412b4 commit ac472a9

File tree

5 files changed

+285
-8
lines changed

5 files changed

+285
-8
lines changed

‎Project.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
9797
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
9898
Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
9999
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
100+
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
100101

101102
[targets]
102-
test = ["DifferentialEquations", "Flux", "Graphs", "Manopt", "InteractiveUtils", "Interpolations", "LineSearches", "Pkg", "Rotations", "Test"]
103+
test = ["DifferentialEquations", "Flux", "Graphs", "Manopt", "InteractiveUtils", "Interpolations", "LineSearches", "Pkg", "Rotations", "Test", "Zygote"]

‎src/parametric/services/ParametricManoptDev.jl

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -116,26 +116,42 @@ function (jacF!::JacF_RLM!)(
116116
q = jacF!.q
117117

118118
fill!(X0, 0)
119+
120+
# ϵ = getPointIdentity(M)
121+
# function jaccost(res, Xc)
122+
# exp!(M, q, ϵ, get_vector!(M, X, p, Xc, basis_domain))
123+
# compose!(M, q, p, q)
124+
# jacF!.costF!(M, res, q)
125+
# end
126+
119127

120128
# TODO maybe move to struct
121129
colorvec = matrix_colors(J)
122130

123-
# ϵ = getPointIdentity(M)
124-
# function jaccost(res, Xc)
125-
# exp!(M, q, ϵ, get_vector!(M, X, p, Xc, basis_domain))
126-
# compose!(M, q, p, q)
127-
# jacF!.costF!(M, res, q)
128-
# end
129-
131+
# TBD would the non-mutating staticarrays version be better?
130132
FiniteDiff.finite_difference_jacobian!(
131133
J,
132134
(res,Xc) -> jacF!.costF!(M, res, exp!(M, q, p, get_vector!(M, X, p, Xc, basis_domain))),
133135
X0;
134136
colorvec
135137
)
138+
139+
# ManifoldDiff._jacobian!(
140+
# J,
141+
# (Xc)->jacF!.costF!(M, jacF!.res, exp!(M, q, p, get_vector!(M, X, p, Xc, basis_domain))),
142+
# X0,
143+
# ManifoldDiff.default_differential_backend()
144+
# )
145+
136146
return J
137147
end
138148

149+
struct FactorGradient{A <: AbstractMatrix}
150+
manifold::AbstractManifold
151+
JacF!::JacF_RLM!
152+
J::A
153+
end
154+
139155
function getSparsityPattern(fg)
140156
biadj = getBiadjacencyMatrix(fg)
141157

‎src/services/FactorGradients.jl

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,75 @@
11
# utilities for calculating the gradient over factors
22

3+
4+
5+
function factorJacobian(
6+
fg,
7+
faclabel::Symbol,
8+
fro_p = first.(getVal.(fg, getVariableOrder(fg, faclabel), solveKey = :parametric))
9+
)
10+
faclabels = Symbol[faclabel;]
11+
frontals = ls(fg, faclabel)
12+
separators = Symbol[] # setdiff(ls(fg), frontals)
13+
14+
# get the subgraph formed by all frontals, separators and fully connected factors
15+
varlabels = union(frontals, separators)
16+
17+
filter!(faclabels) do fl
18+
return issubset(getVariableOrder(fg, fl), varlabels)
19+
end
20+
21+
facs = getFactor.(fg, faclabels)
22+
23+
# so the subgraph consists of varlabels(frontals + separators) and faclabels
24+
25+
varIntLabel = OrderedDict(zip(varlabels, collect(1:length(varlabels))))
26+
27+
# varIntLabel_frontals = filter(p->first(p) in frontals, varIntLabel)
28+
# varIntLabel_separators = filter(p->first(p) in separators, varIntLabel)
29+
30+
calcfacs = map(f->IIF.CalcFactorManopt(f, varIntLabel), facs)
31+
32+
# get the manifold and variable types
33+
frontal_vars = getVariable.(fg, frontals)
34+
vartypes, vartypecount, vartypeslist = IIF.getVariableTypesCount(frontal_vars)
35+
36+
PMs = map(vartypes) do vartype
37+
N = vartypecount[vartype]
38+
G = getManifold(vartype)
39+
return IIF.NPowerManifold(G, N)
40+
end
41+
M = ProductManifold(PMs...)
42+
43+
#
44+
#FIXME
45+
@assert length(M.manifolds) == 1 "#FIXME, this only works with 1 manifold type component"
46+
MM = M.manifolds[1]
47+
48+
# inital values and separators from fg
49+
# fro_p = first.(getVal.(frontal_vars, solveKey = :parametric))
50+
# sep_p::Vector{eltype(fro_p)} = first.(getVal.(fg, separators, solveKey = :parametric))
51+
sep_p = Vector{eltype(fro_p)}()
52+
53+
#cost and jacobian functions
54+
# cost function f: M->ℝᵈ for Riemannian Levenberg-Marquardt
55+
costF! = IIF.CostF_RLM!(calcfacs, fro_p, sep_p)
56+
# jacobian of function for Riemannian Levenberg-Marquardt
57+
jacF! = IIF.JacF_RLM!(MM, costF!)
58+
59+
num_components = length(jacF!.res)
60+
61+
p0 = deepcopy(fro_p)
62+
63+
# initial_residual_values = zeros(num_components)
64+
J = zeros(num_components, manifold_dimension(MM))
65+
66+
jacF!(MM, J, p0)
67+
68+
J
69+
end
70+
71+
72+
373
export getCoordSizes
474
export checkGradientsToleranceMask, calcPerturbationFromVariable
575

‎test/manifolds/factordiff.jl

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
2+
# dev test script for factor gradients using ForwardDiff and ManifoldDiff
3+
4+
using ManifoldDiff
5+
using Manifolds
6+
using IncrementalInference
7+
using StaticArrays
8+
using Zygote
9+
using Test
10+
11+
##
12+
13+
14+
@testset "test factorJacobian" begin
15+
##
16+
17+
# manual EuclidDistance
18+
z = 10.0
19+
f_eucliddist(x1,x2; z) = z - norm(x2 - x1)
20+
x0, x1 = [0.0; 0.0], [10.0, 0.0]
21+
J_ = Zygote.jacobian(()->f_eucliddist(x0, x1; z), Zygote.Params([x0, x1]))
22+
23+
Jx0 = J_[x0]
24+
Jx1 = J_[x1]
25+
26+
27+
##
28+
29+
fg = LocalDFG(;
30+
solverParams = SolverParams(;
31+
graphinit=false
32+
)
33+
)
34+
35+
addVariable!.(fg, [:x0; :x1], Position2)
36+
f = addFactor!(fg, [:x0; :x1], EuclidDistance(Normal(z,1.0)))
37+
38+
p1 = [SA[x1[1];x1[2]] for _ in 1:1]
39+
40+
setVal!(fg, :x1, p1, solveKey=:parametric)
41+
42+
J = IIF.factorJacobian(fg, :x0x1f1)
43+
44+
@test isapprox( Jx0, J[1:1,1:2]; atol=1e-8)
45+
@test isapprox( Jx1, J[1:1,3:4]; atol=1e-8)
46+
47+
48+
##
49+
end
50+
##
51+
52+
53+
# ##
54+
# @testset "using RoME; FiniteDiff.jacobian of SpecialEuclidean(2) factor" begin
55+
# ##
56+
57+
# fg = LocalDFG(;
58+
# solverParams = SolverParams(;
59+
# graphinit=false
60+
# )
61+
# )
62+
63+
# addVariable!.(fg, [:x0; :x1], Pose2)
64+
# f = addFactor!(fg, [:x0; :x1], Pose2Pose2(MvNormal([10;0;pi/2],[1 0 0; 0 1 0; 0 0 1.0])))
65+
66+
# p1 = [ArrayPartition([10; 0.0], [0 1; -1 0.0]) for _ in 1:1]
67+
68+
# setVal!(fg, :x1, p1, solveKey=:parametric)
69+
70+
# J = IIF.factorJacobian(fg, :x0x1f1)
71+
72+
# @test isapprox( Jx0, J[1:1,1:2]; atol=1e-8)
73+
# @test_broken isapprox( Jx1, J[1:1,3:4]; atol=1e-8)
74+
75+
76+
# ##
77+
# end
78+
# ##
79+
80+
81+
##
82+
@testset "ManifoldDiff.jacobian of SpecialEuclidean(2) factor" begin
83+
##
84+
85+
86+
M = SpecialEuclidean(2)
87+
z = ArrayPartition(SA[10.0; 0.0], SMatrix{2,2}(0.0, -1.0, 1.0, 0.0))
88+
89+
p1 = ArrayPartition(SA[0.0; 0.0], SMatrix{2,2}(1, 0, 0, 1.))
90+
e0 = identity_element(M, p1)
91+
p2 = exp(M, e0, hat(M, e0, [10,0,pi/2]))
92+
93+
94+
function resid_SE2(X, p, q)
95+
q̂ = Manifolds.compose(M, p, exp(M, identity_element(M, p), X)) #for groups
96+
return vee(M, q, log(M, q, q̂))
97+
end
98+
99+
100+
# finitediff setup
101+
r_backend = ManifoldDiff.TangentDiffBackend(
102+
ManifoldDiff.FiniteDifferencesBackend()
103+
)
104+
105+
Me = Euclidean(3)
106+
107+
function _factorJac!(J, z, p1, p2)
108+
g1(p_) = resid_SE2(z, p_, p2)
109+
g2(p_) = resid_SE2(z, p1, p_)
110+
111+
J[1:3,1:3] = ManifoldDiff.jacobian(M, Me, g1, p1, r_backend)
112+
J[1:3,4:6] = ManifoldDiff.jacobian(M, Me, g2, p2, r_backend)
113+
114+
J
115+
end
116+
# f_SE2_z(z_) = resid_SE2(z_, e0, p2)
117+
# f_SE2_x0(p_) = resid_SE2(z, e0, p_)
118+
# f_SE2_x0(p_) = resid_SE2(z, e0, p_)
119+
120+
J = zeros(3,6)
121+
122+
J_ = _factorJac!(J, z, p1, p2)
123+
# @profview _factorJac!(J, z, p1, p2)
124+
125+
if false
126+
z_backend = ManifoldDiff.TangentDiffBackend(
127+
ManifoldDiff.ZygoteDiffBackend()
128+
)
129+
g = ManifoldDiff.jacobian(M, Euclidean(3), f_SE2_x0, p1, z_backend)
130+
else
131+
@info "ManifoldDiff.ZygoteDiffBackend usage still under development (23Q3)"
132+
end
133+
134+
135+
136+
##
137+
end
138+
139+
# cost(p_) = distance(M, e0, p_) # ManifoldDiff.gradient(M, cost, p, r_backend)
140+
# cost(p)
141+
# ManifoldDiff.gradient(M, cost, p, r_backend)
142+
143+
144+
145+
146+
147+
148+
149+
150+
151+
152+
# ##
153+
# @testset "CODE STILL UNDER DEV:::Zygote on SpecialEuclidean(2)" begin
154+
# ##
155+
156+
# # manual PosePose
157+
# M = SpecialEuclidean(2)
158+
# z = ArrayPartition(SA[10.0; 0.0], SMatrix{2,2}(0.0, -1.0, 1.0, 0.0))
159+
# e0 = identity_element(M, z)
160+
161+
# # modified from IIF/test/testSpecialEuclidean2Mani.jl
162+
# function f_SE2(X, p, q)
163+
# q̂ = Manifolds.compose(M, p, exp(M, identity_element(M, p), X)) #for groups
164+
# Xc = zeros(3)
165+
# vee!(M, Xc, q, log(M, q, q̂))
166+
# return Xc
167+
# # return (Xc'*Xc)[1]
168+
# end
169+
170+
# Xc_0 = [0.0; 0.0; 0.0] # deepcopy(e0)
171+
# Xc_1 = [10.0; 0.0; pi/2] # deepcopy(z)
172+
173+
# J_ = Zygote.jacobian(
174+
# ()->begin
175+
# f_SE2(
176+
# log(M, e0, z),
177+
# exp(M, e0, hat(M, e0, Xc_0)),
178+
# exp(M, e0, hat(M, e0, Xc_1))
179+
# )
180+
# end,
181+
# Zygote.Params([Xc_0, Xc_1])
182+
# )
183+
184+
# Jx0 = J_[x0]
185+
# Jx1 = J_[x1]
186+
187+
188+
# ##
189+
# end

‎test/runtests.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ end
1212
if TEST_GROUP in ["all", "basic_functional_group"]
1313
# more frequent stochasic failures from numerics
1414
include("manifolds/manifolddiff.jl")
15+
include("manifolds/factordiff.jl")
1516
include("testSpecialEuclidean2Mani.jl")
1617
include("testEuclidDistance.jl")
1718

0 commit comments

Comments
 (0)