Make Values.TangentVector == VectorValues

Author: Marc Rasi

Examples/Pose2SLAMG2O/main.swift

@@ -36,7 +36,7 @@ struct G2OFactorGraph: G2OReader {
   var graph: NonlinearFactorGraph = NonlinearFactorGraph()
 
   public mutating func addInitialGuess(index: Int, pose: Pose2) {
-    initialGuess.insert(index, AnyDifferentiable(pose))
+    initialGuess.insert(index, pose)
   }
 
   public mutating func addMeasurement(frameIndex: Int, measuredIndex: Int, pose: Pose2) {
@@ -72,11 +72,7 @@ func main() {
       dx.insert(i, Vector(zeros: 3))
     }
     optimizer.optimize(gfg: gfg, initial: &dx)
-    for i in 0..<val.count {
-      var p = val[i].baseAs(Pose2.self)
-      p.move(along: Vector3(dx[i]))
-      val[i] = AnyDifferentiable(p)
-    }
+    val.move(along: dx)
     print("Current error: \(problem.graph.error(val))")
   }
   print("Final error: \(problem.graph.error(val))")

Sources/SwiftFusion/Inference/BetweenFactor.swift

@@ -67,7 +67,7 @@ public struct BetweenFactor<T: LieGroup>: NonlinearFactor where T.TangentVector:
   /// Returns the `error` of the factor.
   @differentiable(wrt: values)
   public func error(_ values: Values) -> Double {
-    let actual = values[key1].baseAs(T.self).inverse() * values[key2].baseAs(T.self)
+    let actual = values[key1, as: T.self].inverse() * values[key2, as: T.self]
     let error = difference.local(actual)
     // TODO: It would be faster to call `error.squaredNorm` because then we don't have to pay
     // the cost of a conversion to `Vector`. To do this, we need a protocol
@@ -78,19 +78,13 @@ public struct BetweenFactor<T: LieGroup>: NonlinearFactor where T.TangentVector:
   @differentiable(wrt: values)
   public func errorVector(_ values: Values) -> T.Coordinate.LocalCoordinate {
     let error = difference.local(
-      values[key1].baseAs(T.self).inverse() * values[key2].baseAs(T.self)
+      values[key1, as: T.self].inverse() * values[key2, as: T.self]
     )
 
     return error
   }
 
   public func linearize(_ values: Values) -> JacobianFactor {
-    let j = jacobian(of: self.errorVector, at: values)
-
-    let j1 = Matrix(stacking: (0..<j.count).map { i in (j[i]._values[values._indices[key1]!].base as! T.TangentVector).vector } )
-    let j2 = Matrix(stacking: (0..<j.count).map { i in (j[i]._values[values._indices[key2]!].base as! T.TangentVector).vector } )
-
-    // TODO: remove this negative sign
-    return JacobianFactor(keys, [j1, j2], errorVector(values).vector.scaled(by: -1))
+    return JacobianFactor(of: self.errorVector, at: values)
   }
 }

Sources/SwiftFusion/Inference/JacobianFactor.swift

@@ -104,3 +104,40 @@ public struct JacobianFactor: LinearFactor {
     return result
   }
 }
+
+extension JacobianFactor {
+  /// Creates a `JacobianFactor` by linearizing the error function `f` at `p`.
+  public init<R: VectorConvertible & TangentStandardBasis>(
+    of f: @differentiable (Values) -> R,
+    at p: Values
+  ) {
+    // Compute the rows of the jacobian.
+    let (value, pb) = valueWithPullback(at: p, in: f)
+    let rows = R.tangentStandardBasis.map { pb($0) }
+
+    // Construct empty matrices with the correct shape.
+    assert(rows.count > 0)
+    var matrices = Dictionary<Int, Matrix>(uniqueKeysWithValues: rows[0].keys.map { key in
+      let row = rows[0][key]
+      var matrix = Matrix([], rowCount: 0, columnCount: row.dimension)
+      matrix.reserveCapacity(rows.count * row.dimension)
+      return (key, matrix)
+    })
+
+    // Fill in the matrix entries.
+    for row in rows {
+      for key in row.keys {
+        matrices[key]!.append(row: row[key])
+      }
+    }
+
+    // Return the jacobian factor with the matrices and value.
+    let orderedKeys = Array(matrices.keys)
+    self = JacobianFactor(
+      orderedKeys,
+      orderedKeys.map { matrices[$0]! },
+      // TODO: remove this negative sign
+      value.vector.scaled(by: -1)
+    )
+  }
+}

Sources/SwiftFusion/Inference/PriorFactor.swift

@@ -41,7 +41,7 @@ public struct PriorFactor<T: LieGroup>: NonlinearFactor where T.TangentVector: V
   /// Returns the `error` of the factor.
   @differentiable(wrt: values)
   public func error(_ values: Values) -> Double {
-    let error = difference.local(values[keys[0]].baseAs(T.self))
+    let error = difference.local(values[keys[0], as: T.self])
     // TODO: It would be faster to call `error.squaredNorm` because then we don't have to pay
     // the cost of a conversion to `Vector`. To do this, we need a protocol
     // with a `squaredNorm` requirement.
@@ -50,14 +50,12 @@ public struct PriorFactor<T: LieGroup>: NonlinearFactor where T.TangentVector: V
 
   @differentiable(wrt: values)
   public func errorVector(_ values: Values) -> T.Coordinate.LocalCoordinate {
-    let val = values[keys[0]].baseAs(T.self)
+    let val = values[keys[0], as: T.self]
     let error = difference.local(val)
     return error
   }
-  
+
   public func linearize(_ values: Values) -> JacobianFactor {
-    let j = jacobian(of: self.errorVector, at: values)
-    let j1 = Matrix(stacking: (0..<j.count).map { i in (j[i]._values[values._indices[keys[0]]!].base as! T.TangentVector).vector } )
-    return JacobianFactor(keys, [j1], errorVector(values).vector.scaled(by: -1))
+    return JacobianFactor(of: self.errorVector, at: values)
   }
 }

Sources/SwiftFusion/Inference/Values.swift

@@ -34,25 +34,64 @@ public struct Values: Differentiable & KeyPathIterable {
   public var count: Int {
     return _values.count
   }
+
+  /// MARK: - Differentiable conformance and related properties and helpers.
+
+  /// The product space of the tangent spaces of all the values.
+  public typealias TangentVector = VectorValues
+
+  /// `makeTangentVector[i]` produces a type-erased tangent vector for `values[i]`.
+  private var makeTangentVector: [(Vector) -> AnyDerivative] = []
+
+  public mutating func move(along direction: VectorValues) {
+    for key in direction.keys {
+      let index = self._indices[key]!
+      self._values[index].move(along: makeTangentVector[index](direction[key]))
+    }
+  }
 
-  /// The subscript operator, with some indirection
-  /// Should be replaced after Dictionary is in
+  /// MARK: - Value manipulation methods.
+
+  /// Access the value at `key`, with type `type`.
+  ///
+  /// Precondition: The value actually has type `type`.
   @differentiable
-  public subscript(key: Int) -> AnyDifferentiable {
+  public subscript<T: Differentiable>(key: Int, as type: T.Type) -> T
+    where T.TangentVector: VectorConvertible
+  {
     get {
-      _values[_indices[key]!]
+      return _values[_indices[key]!].baseAs(type)
     }
-    set(newVal) {
-      _values[_indices[key]!] = newVal
+    set(newValue) {
+      _values[_indices[key]!] = AnyDifferentiable(newValue)
     }
   }
-  
+
+  @derivative(of: subscript)
+  @usableFromInline
+  func vjpSubscript<T: Differentiable>(key: Int, as type: T.Type)
+    -> (value: T, pullback: (T.TangentVector) -> VectorValues)
+    where T.TangentVector: VectorConvertible
+  {
+    return (
+      self._values[self._indices[key]!].baseAs(type),
+      { (t: T.TangentVector) in
+        var vectorValues = VectorValues()
+        vectorValues.insert(key, t.vector)
+        return vectorValues
+      }
+    )
+  }
+
   /// Insert a key value pair
-  public mutating func insert(_ key: Int, _ val: AnyDifferentiable) {
+  public mutating func insert<T: Differentiable>(_ key: Int, _ val: T)
+    where T.TangentVector: VectorConvertible
+  {
     assert(_indices[key] == nil)
 
     self._indices[key] = self._values.count
-    self._values.append(val)
+    self._values.append(AnyDifferentiable(val))
+    self.makeTangentVector.append({ AnyDerivative(T.TangentVector($0)) })
   }
 
 }

Tests/SwiftFusionTests/Geometry/Pose3Tests.swift

@@ -45,7 +45,7 @@ final class Pose3Tests: XCTestCase {
     let prior_factor = PriorFactor(0, t1)
 
     var vals = Values()
-    vals.insert(0, AnyDifferentiable(t1)) // should be identity matrix
+    vals.insert(0, t1) // should be identity matrix
     // Change this to t2, still zero in upper left block
 
     let actual = prior_factor.linearize(vals).jacobians[0]
@@ -74,7 +74,7 @@ final class Pose3Tests: XCTestCase {
       let gti = Vector3(radius * cos(theta), radius * sin(theta), 0)
       let oRi = Rot3.fromTangent(Vector3(0, 0, -theta))  // negative yaw goes counterclockwise, with Z down !
       let gTi = Pose3(gRo * oRi, gti)
-      values.insert(key, AnyDifferentiable(gTi))
+      values.insert(key, gTi)
       theta = theta + dtheta
     }
     return values
@@ -83,8 +83,8 @@ final class Pose3Tests: XCTestCase {
   func testGtsamPose3SLAMExample() {
     // Create a hexagon of poses
     let hexagon = circlePose3(numPoses: 6, radius: 1.0)
-    let p0 = hexagon[0].baseAs(Pose3.self)
-    let p1 = hexagon[1].baseAs(Pose3.self)
+    let p0 = hexagon[0, as: Pose3.self]
+    let p1 = hexagon[1, as: Pose3.self]
 
     // create a Pose graph with one equality constraint and one measurement
     var fg = NonlinearFactorGraph()
@@ -101,12 +101,12 @@ final class Pose3Tests: XCTestCase {
     // Create initial config
     var val = Values()
     let s = 0.10
-    val.insert(0, AnyDifferentiable(p0))
-    val.insert(1, AnyDifferentiable(hexagon[1].baseAs(Pose3.self).global(Vector6(s * Tensor<Double>(randomNormal: [6])))))
-    val.insert(2, AnyDifferentiable(hexagon[2].baseAs(Pose3.self).global(Vector6(s * Tensor<Double>(randomNormal: [6])))))
-    val.insert(3, AnyDifferentiable(hexagon[3].baseAs(Pose3.self).global(Vector6(s * Tensor<Double>(randomNormal: [6])))))
-    val.insert(4, AnyDifferentiable(hexagon[4].baseAs(Pose3.self).global(Vector6(s * Tensor<Double>(randomNormal: [6])))))
-    val.insert(5, AnyDifferentiable(hexagon[5].baseAs(Pose3.self).global(Vector6(s * Tensor<Double>(randomNormal: [6])))))
+    val.insert(0, p0)
+    val.insert(1, hexagon[1, as: Pose3.self].global(Vector6(s * Tensor<Double>(randomNormal: [6]))))
+    val.insert(2, hexagon[2, as: Pose3.self].global(Vector6(s * Tensor<Double>(randomNormal: [6]))))
+    val.insert(3, hexagon[3, as: Pose3.self].global(Vector6(s * Tensor<Double>(randomNormal: [6]))))
+    val.insert(4, hexagon[4, as: Pose3.self].global(Vector6(s * Tensor<Double>(randomNormal: [6]))))
+    val.insert(5, hexagon[5, as: Pose3.self].global(Vector6(s * Tensor<Double>(randomNormal: [6]))))
 
     // optimize
     for _ in 0..<16 {
@@ -122,15 +122,10 @@ final class Pose3Tests: XCTestCase {
 
       optimizer.optimize(gfg: gfg, initial: &dx)
 
-      
-      for i in 0..<6 {
-        var p = val[i].baseAs(Pose3.self)
-        p.move(along: Vector6(dx[i]))
-        val[i] = AnyDifferentiable(p)
-      }
+      val.move(along: dx)
     }
 
-    let pose_1 = val[1].baseAs(Pose3.self)
+    let pose_1 = val[1, as: Pose3.self]
     assertAllKeyPathEqual(pose_1, p1, accuracy: 1e-2)
   }
 }

Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift

@@ -12,8 +12,8 @@ final class NonlinearFactorGraphTests: XCTestCase {
     fg += bf1
 
     var val = Values()
-    val.insert(0, AnyDifferentiable(Pose2(1.0, 1.0, 0.0)))
-    val.insert(1, AnyDifferentiable(Pose2(1.0, 1.0, .pi)))
+    val.insert(0, Pose2(1.0, 1.0, 0.0))
+    val.insert(1, Pose2(1.0, 1.0, .pi))
 
     let gfg = fg.linearize(val)
 
@@ -51,7 +51,7 @@ final class NonlinearFactorGraphTests: XCTestCase {
     var val = Values()
 
     for i in 0..<5 {
-      val.insert(i, AnyDifferentiable(map[i]))
+      val.insert(i, map[i])
     }
 
     for _ in 0..<3 {
@@ -67,11 +67,7 @@ final class NonlinearFactorGraphTests: XCTestCase {
 
       optimizer.optimize(gfg: gfg, initial: &dx)
 
-      for i in 0..<5 {
-        var p = val[i].baseAs(Pose2.self)
-        p.move(along: Vector3(dx[i]))
-        val[i] = AnyDifferentiable(p)
-      }
+      val.move(along: dx)
     }
 
     let dumpjson = { (p: Pose2) -> String in
@@ -84,14 +80,14 @@ final class NonlinearFactorGraphTests: XCTestCase {
     }
     print("]")
 
-    let map_final = (0..<5).map { val[$0].baseAs(Pose2.self) }
+    let map_final = (0..<5).map { val[$0, as: Pose2.self] }
     print("map = [")
     for v in map_final.indices {
       print("\(dumpjson(map_final[v]))\({ () -> String in if v == map_final.indices.endIndex - 1 { return "" } else { return "," } }())")
     }
     print("]")
 
-    let p5T1 = between(val[4].baseAs(Pose2.self), val[0].baseAs(Pose2.self))
+    let p5T1 = between(val[4, as: Pose2.self], val[0, as: Pose2.self])
 
     // Test condition: P_5 should be identical to P_1 (close loop)
     XCTAssertEqual(p5T1.t.norm, 0.0, accuracy: 1e-2)