From 0f1658b1a92e792dbbe74e5debf03c82994a6771 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Tue, 21 Apr 2020 21:54:34 -0400
Subject: [PATCH 01/13] Add Nonlinear FG interface

---
 .../Inference/AnyNonlinearFactor.swift        | 152 ++++++++++++++++++
 .../SwiftFusion/Inference/BetweenFactor.swift |  82 ++++++++++
 Sources/SwiftFusion/Inference/Factor.swift    |  26 ++-
 .../Inference/JacobianFactor.swift            |   5 +-
 .../Inference/NonlinearFactorGraph.swift      |  50 ++++++
 Sources/SwiftFusion/Inference/Values.swift    |  71 ++++++++
 .../Inference/NonlinearFactorGraphTests.swift |  31 ++++
 7 files changed, 413 insertions(+), 4 deletions(-)
 create mode 100644 Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
 create mode 100644 Sources/SwiftFusion/Inference/BetweenFactor.swift
 create mode 100644 Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
 create mode 100644 Sources/SwiftFusion/Inference/Values.swift
 create mode 100644 Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift

diff --git a/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift b/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
new file mode 100644
index 00000000..731946e0
--- /dev/null
+++ b/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
@@ -0,0 +1,152 @@
+//===----------------------------------------------------------------------===//
+// `AnyNonlinearFactor`
+//===----------------------------------------------------------------------===//
+
+internal protocol _AnyNonlinearFactorBox {
+  typealias ScalarType = Double
+  // `Differentiable` requirements.
+  mutating func _move(along direction: AnyDerivative)
+
+  /// The underlying base value, type-erased to `Any`.
+  var _typeErasedBase: Any { get }
+
+  /// Returns the underlying value unboxed to the given type, if possible.
+  func _unboxed<U: NonlinearFactor>(to type: U.Type) -> U?
+  
+  var _keys: Array<Int> { get }
+  
+  @differentiable(wrt: values)
+  func _error(_ values: Values) -> ScalarType
+  
+  func _linearize(_ values: Values) -> JacobianFactor
+}
+
+internal struct _ConcreteNonlinearFactorBox<T: NonlinearFactor>: _AnyNonlinearFactorBox
+{
+  var _keys: Array<Int> {
+    get {
+      _base.keys
+    }
+  }
+  
+  /// The underlying base value.
+  var _base: T
+
+  init(_ base: T) {
+    self._base = base
+  }
+
+  /// The underlying base value, type-erased to `Any`.
+  var _typeErasedBase: Any {
+    return _base
+  }
+
+  func _unboxed<U: NonlinearFactor>(to type: U.Type) -> U? {
+    return (self as? _ConcreteNonlinearFactorBox<U>)?._base
+  }
+
+  mutating func _move(along direction: AnyDerivative) {
+    guard
+      let directionBase =
+        direction.base as? T.TangentVector
+    else {
+      _derivativeTypeMismatch(T.self, type(of: direction.base))
+    }
+    _base.move(along: directionBase)
+  }
+  
+  @differentiable(wrt: values)
+  func _error(_ values: Values) -> ScalarType {
+    _base.error(values)
+  }
+  
+  func _linearize(_ values: Values) -> JacobianFactor {
+    _base.linearize(values)
+  }
+}
+
+public struct AnyNonlinearFactor: NonlinearFactor {
+  internal var _box: _AnyNonlinearFactorBox
+
+  internal init(_box: _AnyNonlinearFactorBox) {
+    self._box = _box
+  }
+
+  /// The underlying base value.
+  public var base: Any {
+    return _box._typeErasedBase
+  }
+
+  /// Creates a type-erased derivative from the given derivative.
+  @differentiable
+  public init<T: NonlinearFactor>(_ base: T) {
+    self._box = _ConcreteNonlinearFactorBox<T>(base)
+  }
+
+  @inlinable
+  @derivative(of: init)
+  internal static func _vjpInit<T: NonlinearFactor>(
+    _ base: T
+  ) -> (value: AnyNonlinearFactor, pullback: (AnyDerivative) -> T.TangentVector)
+  {
+    return (AnyNonlinearFactor(base), { v in v.base as! T.TangentVector })
+  }
+
+  @inlinable
+  @derivative(of: init)
+  internal static func _jvpInit<T: NonlinearFactor>(
+    _ base: T
+  ) -> (
+    value: AnyNonlinearFactor, differential: (T.TangentVector) -> AnyDerivative
+  ) {
+    return (AnyNonlinearFactor(base), { dbase in AnyDerivative(dbase) })
+  }
+
+  public typealias TangentVector = AnyDerivative
+
+  public mutating func move(along direction: TangentVector) {
+    _box._move(along: direction)
+  }
+}
+
+extension AnyNonlinearFactor {
+  @differentiable
+  public func baseAs<T: NonlinearFactor>(_ t: T.Type) -> T {
+    base as! T
+  }
+
+  @derivative(of: baseAs)
+  @usableFromInline
+  func jvpBaseAs<T: NonlinearFactor>(_ t: T.Type) -> (
+    value: T,
+    differential: (AnyDerivative) -> T.TangentVector
+  ) {
+    (baseAs(t), { $0.base as! T.TangentVector })
+  }
+
+  @derivative(of: baseAs)
+  @usableFromInline
+  func vjpBaseAs<T: NonlinearFactor>(_ t: T.Type) -> (
+    value: T,
+    pullback: (T.TangentVector) -> AnyDerivative
+  ) {
+    (baseAs(t), { AnyDerivative($0) })
+  }
+}
+
+extension AnyNonlinearFactor {
+  @differentiable(wrt: values)
+  public func error(_ values: Values) -> ScalarType {
+    _box._error(values)
+  }
+  
+  public var keys: Array<Int> {
+    get {
+      _box._keys
+    }
+  }
+  
+  public func linearize(_ values: Values) -> JacobianFactor {
+    _box._linearize(values)
+  }
+}
diff --git a/Sources/SwiftFusion/Inference/BetweenFactor.swift b/Sources/SwiftFusion/Inference/BetweenFactor.swift
new file mode 100644
index 00000000..dba8baa4
--- /dev/null
+++ b/Sources/SwiftFusion/Inference/BetweenFactor.swift
@@ -0,0 +1,82 @@
+// Copyright 2019 The SwiftFusion Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import TensorFlow
+
+/// A `NonlinearFactor` that calculates the difference of two Values of the same type
+///
+/// Input is a dictionary of `Key` to `Value` pairs, and the output is the scalar
+/// error value
+///
+/// Interpretation
+/// ================
+/// `Input`: the input values as key-value pairs
+///
+public struct BetweenFactor: NonlinearFactor {
+  @noDerivative
+  public var keys: Array<Int> = []
+  public var difference: Pose2
+  public typealias Output = Error
+  
+  public init (_ key1: Int, _ key2: Int, _ difference: Pose2) {
+    keys = [key1, key2]
+    self.difference = difference
+  }
+  typealias ScalarType = Double
+  
+  /// TODO: `Dictionary` still does not conform to `Differentiable`
+  /// Tracking issue: https://bugs.swift.org/browse/TF-899
+//  typealias Input = Dictionary<UInt, Tensor<ScalarType>>
+
+// I want to build a general differentiable dot product
+//  @differentiable(wrt: (a, b))
+//  static func dot<T: Differentiable & KeyPathIterable>(_ a: T, _ b: T) -> Double {
+//    let squared = a.recursivelyAllKeyPaths(to: Double.self).map { a[keyPath: $0] * b[keyPath: $0] }
+//
+//    return squared.differentiableReduce(0.0, {$0 + $1})
+//  }
+//
+//  @derivative(of: dot)
+//  static func _vjpDot<T: Differentiable & KeyPathIterable>(_ a: T, _ b: T) -> (
+//    value: Double,
+//    pullback: (Double) -> (T.TangentVector, T.TangentVector)
+//  ) {
+//    return (value: dot(a, b), pullback: { v in
+//      ((at.scaled(by: v), bt.scaled(by: v)))
+//    })
+//  }
+  
+  /// Returns the `error` of the factor.
+  @differentiable(wrt: values)
+  public func error(_ values: Values) -> Double {
+    let error = (values[keys[1]].baseAs(Pose2.self) * values[keys[0]].baseAs(Pose2.self).inverse()) * difference.inverse()
+    
+    return error.t.norm + error.rot.theta * error.rot.theta
+  }
+  
+  @differentiable(wrt: values)
+  public func errorVector(_ values: Values) -> Vector3 {
+    let error = (values[keys[1]].baseAs(Pose2.self) * values[keys[0]].baseAs(Pose2.self).inverse()) * difference.inverse()
+    
+    return Vector3(error.rot.theta, error.t.x, error.t.y)
+  }
+  
+  public func linearize(_ values: Values) -> JacobianFactor {
+    let j = jacobian(of: self.errorVector, at: values)
+    
+    let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[0].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    let j2 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[1].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    
+    return JacobianFactor(keys, [j1, j2], errorVector(values).tensor.reshaped(to: [3, 1]))
+  }
+}
diff --git a/Sources/SwiftFusion/Inference/Factor.swift b/Sources/SwiftFusion/Inference/Factor.swift
index 435f30d7..8d572ba4 100644
--- a/Sources/SwiftFusion/Inference/Factor.swift
+++ b/Sources/SwiftFusion/Inference/Factor.swift
@@ -15,7 +15,7 @@ import TensorFlow
 
 /// The most general factor protocol.
 public protocol Factor {
-  var keys: Array<Int> { get set }
+  var keys: Array<Int> { get }
 }
 
 /// A `LinearFactor` corresponds to the `GaussianFactor` in GTSAM.
@@ -30,11 +30,33 @@ public protocol Factor {
 public protocol LinearFactor: Factor {
   typealias ScalarType = Double
   
+  /// TODO: `Dictionary` still does not conform to `Differentiable`
+  /// Tracking issue: https://bugs.swift.org/browse/TF-899
+//  typealias Input = Dictionary<UInt, Tensor<ScalarType>>
+  
+  /// Returns the `error` of the factor.
+  func error(_ values: VectorValues) -> ScalarType
+}
+
+/// A `NonlinearFactor` corresponds to the `NonlinearFactor` in GTSAM.
+///
+/// Input is a dictionary of `Key` to `Value` pairs, and the output is the scalar
+/// error value
+///
+/// Interpretation
+/// ================
+/// `Input`: the input values as key-value pairs
+///
+public protocol NonlinearFactor: Differentiable & Factor {
+  typealias ScalarType = Double
+  
   /// TODO: `Dictionary` still does not conform to `Differentiable`
   /// Tracking issue: https://bugs.swift.org/browse/TF-899
 //  typealias Input = Dictionary<UInt, Tensor<ScalarType>>
   
   /// Returns the `error` of the factor.
   @differentiable(wrt: values)
-  func error(_ indices: [Int], values: Tensor<ScalarType>) -> ScalarType
+  func error(_ values: Values) -> ScalarType
+  
+  func linearize(_ values: Values) -> JacobianFactor
 }
diff --git a/Sources/SwiftFusion/Inference/JacobianFactor.swift b/Sources/SwiftFusion/Inference/JacobianFactor.swift
index 14c4d2ab..d5bca8c3 100644
--- a/Sources/SwiftFusion/Inference/JacobianFactor.swift
+++ b/Sources/SwiftFusion/Inference/JacobianFactor.swift
@@ -39,10 +39,11 @@ import TensorFlow
 /// and `HessianFactor` conform to this protocol instead.
 public struct JacobianFactor: LinearFactor {
   
-  @differentiable(wrt: values)
-  public func error(_ indices: [Int], values: Tensor<ScalarType>) -> ScalarType {
+  // TODO(fan): correct this and add a unit test
+  public func error(_ values: VectorValues) -> ScalarType {
     ScalarType.zero
   }
+  
   public var dimension: Int {
     get {
       jacobians[0].shape.dimensions[0]
diff --git a/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift b/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
new file mode 100644
index 00000000..8532acc7
--- /dev/null
+++ b/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
@@ -0,0 +1,50 @@
+// Copyright 2019 The SwiftFusion Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import TensorFlow
+
+/// A factor graph for linear problems
+/// Factors are the Jacobians between the corresponding variables and measurements
+/// TODO(fan): Add noise model
+public struct NonlinearFactorGraph: FactorGraph {
+  public typealias KeysType = Array<Int>
+  
+  public typealias FactorsType = Array<AnyNonlinearFactor>
+  
+  public var keys: KeysType = []
+  public var factors: FactorsType = []
+  
+  /// Default initializer
+  public init() { }
+  
+  /// Convenience operator for adding factor
+  public static func += (lhs: inout Self, rhs: AnyNonlinearFactor) {
+    lhs.factors.append(rhs)
+  }
+  
+  /// linearize the nonlinear factor graph to a linear factor graph
+  public func linearize(_ values: Values) -> GaussianFactorGraph {
+    var gfg = GaussianFactorGraph()
+    
+    for i in factors {
+      let linearized = i.linearize(values)
+      
+      // Assertion for the shape of Jacobian
+      assert(linearized.jacobians.map { $0.shape.count == 2 }.reduce(true, { $0 && $1 }))
+      
+      gfg += linearized
+    }
+    
+    return gfg
+  }
+}
diff --git a/Sources/SwiftFusion/Inference/Values.swift b/Sources/SwiftFusion/Inference/Values.swift
new file mode 100644
index 00000000..a2bf33fc
--- /dev/null
+++ b/Sources/SwiftFusion/Inference/Values.swift
@@ -0,0 +1,71 @@
+// Copyright 2019 The SwiftFusion Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import TensorFlow
+
+/// The class that holds Key-Vector pairs.
+public struct Values: Differentiable & KeyPathIterable {
+  public typealias ScalarType = Double
+  var _values: [AnyDifferentiable] = []
+  
+  /// Dictionary from Key to index
+  @noDerivative
+  var _indices: Dictionary<Int, Int> = [:]
+  
+  public var keys: Dictionary<Int, Int>.Keys {
+    get {
+      _indices.keys
+    }
+  }
+  /// Default initializer
+  public init() { }
+  
+  /// The subscript operator, with some indirection
+  /// Should be replaced after Dictionary is in
+  @differentiable
+  public subscript(key: Int) -> AnyDifferentiable {
+    _values[_indices[key]!]
+  }
+  
+  /// Insert a key value pair
+  public mutating func insert(_ key: Int, _ val: AnyDifferentiable) {
+    assert(_indices[key] == nil)
+    
+    self._indices[key] = self._values.count
+    self._values.append(val)
+  }
+  
+}
+//
+//extension Values: CustomStringConvertible {
+//  public var description: String {
+//    "Values(\n\(_indices.map { "Key: \($0), J: \(_values[$1])\n"}.reduce("", { $0 + $1 }) )"
+//  }
+//}
+//
+//extension Values: Equatable {
+//  /// Order-aware comparison
+//  public static func == (lhs: VectorValues, rhs: VectorValues) -> Bool {
+//    if lhs._indices.keys != rhs._indices.keys {
+//      return false
+//    }
+//
+//    for k in lhs._indices.keys {
+//      if lhs._values[lhs._indices[k]!] != rhs._values[rhs._indices[k]!] {
+//        return false
+//      }
+//    }
+//
+//    return true
+//  }
+//}
diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
new file mode 100644
index 00000000..075014c2
--- /dev/null
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -0,0 +1,31 @@
+import SwiftFusion
+import TensorFlow
+import XCTest
+
+final class NonlinearFactorGraphTests: XCTestCase {
+  /// test ATr
+  func testBasicOps() {
+    var fg = NonlinearFactorGraph()
+    
+    let bf1 = BetweenFactor(0, 1, Pose2(0.0,0.0, 0.0))
+    
+    fg += AnyNonlinearFactor(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)))
+    
+    let gfg = fg.linearize(val)
+    print("fg = \(fg)")
+    print("fg_l = \(gfg)")
+    print("bf1 = \(bf1)")
+    print("bf1_l = \(bf1.linearize(val))")
+    
+    var vv = VectorValues()
+    
+    vv.insert(0, Tensor<Double>(shape:[3, 1], scalars: [1.0, 1.0, 0.0]))
+    vv.insert(1, Tensor<Double>(shape:[3, 1], scalars: [1.0, 1.0, 3.14]))
+    print("gfg(x) = \(gfg*vv)")
+    print(bf1.linearize(val).jacobians[0].shape)
+  }
+}

From a90be095454a4669964eff08cd0b036c5ed6b845 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Tue, 21 Apr 2020 21:58:41 -0400
Subject: [PATCH 02/13] Typo

---
 Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift b/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
index 8532acc7..73010dcb 100644
--- a/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
+++ b/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
@@ -13,8 +13,7 @@
 // limitations under the License.
 import TensorFlow
 
-/// A factor graph for linear problems
-/// Factors are the Jacobians between the corresponding variables and measurements
+/// A factor graph for nonlinear problems
 /// TODO(fan): Add noise model
 public struct NonlinearFactorGraph: FactorGraph {
   public typealias KeysType = Array<Int>

From 16ef0d20b693b46430c7958804938cd0a573d949 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Tue, 21 Apr 2020 22:00:33 -0400
Subject: [PATCH 03/13] Add custom output for Values

---
 Sources/SwiftFusion/Inference/Values.swift | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/Sources/SwiftFusion/Inference/Values.swift b/Sources/SwiftFusion/Inference/Values.swift
index a2bf33fc..a4f4ef07 100644
--- a/Sources/SwiftFusion/Inference/Values.swift
+++ b/Sources/SwiftFusion/Inference/Values.swift
@@ -13,7 +13,7 @@
 // limitations under the License.
 import TensorFlow
 
-/// The class that holds Key-Vector pairs.
+/// The class that holds Key-Value pairs.
 public struct Values: Differentiable & KeyPathIterable {
   public typealias ScalarType = Double
   var _values: [AnyDifferentiable] = []
@@ -46,16 +46,16 @@ public struct Values: Differentiable & KeyPathIterable {
   }
   
 }
-//
-//extension Values: CustomStringConvertible {
-//  public var description: String {
-//    "Values(\n\(_indices.map { "Key: \($0), J: \(_values[$1])\n"}.reduce("", { $0 + $1 }) )"
-//  }
-//}
-//
+
+extension Values: CustomStringConvertible {
+  public var description: String {
+    "Values(\n\(_indices.map { "Key: \($0), J: \(_values[$1])\n"}.reduce("", { $0 + $1 }) )"
+  }
+}
+
 //extension Values: Equatable {
 //  /// Order-aware comparison
-//  public static func == (lhs: VectorValues, rhs: VectorValues) -> Bool {
+//  public static func == (lhs: Values, rhs: Values) -> Bool {
 //    if lhs._indices.keys != rhs._indices.keys {
 //      return false
 //    }

From d6d49605aba00112ec8ac726c4cd2f4e16456646 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Tue, 21 Apr 2020 22:03:57 -0400
Subject: [PATCH 04/13] Fix test case

---
 .../Inference/NonlinearFactorGraphTests.swift          | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
index 075014c2..cc8d0298 100644
--- a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -16,16 +16,14 @@ final class NonlinearFactorGraphTests: XCTestCase {
     val.insert(1, AnyDifferentiable(Pose2(1.0, 1.0, .pi)))
     
     let gfg = fg.linearize(val)
-    print("fg = \(fg)")
-    print("fg_l = \(gfg)")
-    print("bf1 = \(bf1)")
-    print("bf1_l = \(bf1.linearize(val))")
     
     var vv = VectorValues()
     
     vv.insert(0, Tensor<Double>(shape:[3, 1], scalars: [1.0, 1.0, 0.0]))
     vv.insert(1, Tensor<Double>(shape:[3, 1], scalars: [1.0, 1.0, 3.14]))
-    print("gfg(x) = \(gfg*vv)")
-    print(bf1.linearize(val).jacobians[0].shape)
+    
+    let expected = Tensor<Double>(shape:[3, 1], scalars: [0.0, 0.0, -3.14])
+    
+    assertEqual((gfg * vv)[0], expected, accuracy: 1e-9)
   }
 }

From 0b3a196b47110a17fdf627208358ffa2a66c989e Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Tue, 21 Apr 2020 23:06:32 -0400
Subject: [PATCH 05/13] Linearization should be at 0,0 tangent

---
 .../SwiftFusion/Inference/BetweenFactor.swift | 10 ++++-
 .../Inference/GaussianFactorGraph.swift       |  6 ++-
 .../Geometry/Pose2Tests.swift                 | 39 +++++++++++++++++++
 .../Inference/NonlinearFactorGraphTests.swift | 10 +++--
 4 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/Sources/SwiftFusion/Inference/BetweenFactor.swift b/Sources/SwiftFusion/Inference/BetweenFactor.swift
index dba8baa4..b1fc76eb 100644
--- a/Sources/SwiftFusion/Inference/BetweenFactor.swift
+++ b/Sources/SwiftFusion/Inference/BetweenFactor.swift
@@ -59,14 +59,20 @@ public struct BetweenFactor: NonlinearFactor {
   /// Returns the `error` of the factor.
   @differentiable(wrt: values)
   public func error(_ values: Values) -> Double {
-    let error = (values[keys[1]].baseAs(Pose2.self) * values[keys[0]].baseAs(Pose2.self).inverse()) * difference.inverse()
+    let error = between(
+      between(values[keys[1]].baseAs(Pose2.self), values[keys[0]].baseAs(Pose2.self)),
+        difference
+    )
     
     return error.t.norm + error.rot.theta * error.rot.theta
   }
   
   @differentiable(wrt: values)
   public func errorVector(_ values: Values) -> Vector3 {
-    let error = (values[keys[1]].baseAs(Pose2.self) * values[keys[0]].baseAs(Pose2.self).inverse()) * difference.inverse()
+    let error = between(
+      between(values[keys[1]].baseAs(Pose2.self), values[keys[0]].baseAs(Pose2.self)),
+      difference
+    )
     
     return Vector3(error.rot.theta, error.t.x, error.t.y)
   }
diff --git a/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift b/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
index 8bfb11e0..04ae34e4 100644
--- a/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
+++ b/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
@@ -34,11 +34,15 @@ public struct GaussianFactorGraph: FactorGraph {
   public init() { }
   
   /// This calculates `A*x`, where x is the collection of key-values
-  /// Note A is a
   public static func * (lhs: GaussianFactorGraph, rhs: VectorValues) -> Errors {
     Array(lhs.factors.map { $0 * rhs })
   }
   
+  /// This calculates `b - A*x`, where x is the collection of key-values
+  public func residual (_ val: VectorValues) -> Errors {
+    Array(self.factors.map { $0.b - $0 * val })
+  }
+  
   /// Convenience operator for adding factor
   public static func += (lhs: inout Self, rhs: JacobianFactor) {
     lhs.factors.append(rhs)
diff --git a/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift b/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift
index 9acc3a30..1b44a5bd 100644
--- a/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift
+++ b/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift
@@ -30,6 +30,45 @@ final class Pose2Tests: XCTestCase {
     XCTAssertEqual(actual, expected)
   }
 
+  /// test the simplest compose (multiplication)
+  func testCompose() {
+    let pose1 = Pose2(Rot2(.pi/4.0), Vector2(sqrt(0.5), sqrt(0.5)))
+    let pose2 = Pose2(Rot2(.pi/2.0), Vector2(0.0, 2.0))
+
+    let actual = pose1 * pose2
+
+    let expected = Pose2(Rot2(3.0 * .pi/4.0), Vector2(-sqrt(0.5), 3.0*sqrt(0.5)))
+    
+    let _ = expected.recursivelyAllKeyPaths(to: Double.self).map {
+      XCTAssertEqual(expected[keyPath: $0], actual[keyPath: $0], accuracy: 1e-9)
+    }
+  }
+  
+  /// test the simplest compose (multiplication)
+  func testInverse() {
+    let gTl = Pose2(Rot2(.pi/2.0), Vector2(1.0, 2.0))
+
+    let actual = gTl.inverse()
+
+    let expected = Pose2(Rot2(-.pi/2.0), Vector2(-2, 1.0))
+    
+    let _ = expected.recursivelyAllKeyPaths(to: Double.self).map {
+      XCTAssertEqual(expected[keyPath: $0], actual[keyPath: $0], accuracy: 1e-9)
+    }
+  }
+  
+  /// test the between function against GTSAM
+  func testBetween() {
+    let p1 = Pose2(1.23, 2.30, 0.2)
+    let odo = Pose2(0.53, 0.39, 0.15)
+    let p2 = p1 * odo
+    
+    let expected = between(p1, p2)
+    let _ = expected.recursivelyAllKeyPaths(to: Double.self).map {
+      XCTAssertEqual(expected[keyPath: $0], odo[keyPath: $0], accuracy: 1e-9)
+    }
+  }
+  
   /// test the simplest gradient descent on Pose2
   func testBetweenDerivatives() {
     var pT1 = Pose2(Rot2(0), Vector2(1, 0)), pT2 = Pose2(Rot2(1), Vector2(1, 1))
diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
index cc8d0298..f596f14e 100644
--- a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -19,11 +19,13 @@ final class NonlinearFactorGraphTests: XCTestCase {
     
     var vv = VectorValues()
     
-    vv.insert(0, Tensor<Double>(shape:[3, 1], scalars: [1.0, 1.0, 0.0]))
-    vv.insert(1, Tensor<Double>(shape:[3, 1], scalars: [1.0, 1.0, 3.14]))
+    vv.insert(0, Tensor<Double>(shape:[3, 1], scalars: [0.0, 0.0, 0.0]))
+    vv.insert(1, Tensor<Double>(shape:[3, 1], scalars: [0.0, 0.0, 0.0]))
     
-    let expected = Tensor<Double>(shape:[3, 1], scalars: [0.0, 0.0, -3.14])
+    let expected = Tensor<Double>(shape:[3, 1], scalars: [.pi, 0.0, 0.0])
     
-    assertEqual((gfg * vv)[0], expected, accuracy: 1e-9)
+    print("gfg = \(gfg)")
+    print("error = \((gfg*vv).norm)")
+    assertEqual((gfg.residual(vv))[0], expected, accuracy: 1e-6)
   }
 }

From 05f6824d883b9f0572b0e099204b01fa5280214d Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 00:39:11 -0400
Subject: [PATCH 06/13] Working CGLS optimization

---
 .../SwiftFusion/Inference/BetweenFactor.swift |  6 +-
 .../Inference/GaussianFactorGraph.swift       |  5 +-
 .../SwiftFusion/Inference/PriorFactor.swift   | 69 +++++++++++++++++
 Sources/SwiftFusion/Inference/Values.swift    |  7 +-
 .../Geometry/Pose2Tests.swift                 | 14 ++--
 .../Inference/NonlinearFactorGraphTests.swift | 75 ++++++++++++++++++-
 6 files changed, 160 insertions(+), 16 deletions(-)
 create mode 100644 Sources/SwiftFusion/Inference/PriorFactor.swift

diff --git a/Sources/SwiftFusion/Inference/BetweenFactor.swift b/Sources/SwiftFusion/Inference/BetweenFactor.swift
index b1fc76eb..805edef4 100644
--- a/Sources/SwiftFusion/Inference/BetweenFactor.swift
+++ b/Sources/SwiftFusion/Inference/BetweenFactor.swift
@@ -80,9 +80,9 @@ public struct BetweenFactor: NonlinearFactor {
   public func linearize(_ values: Values) -> JacobianFactor {
     let j = jacobian(of: self.errorVector, at: values)
     
-    let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[0].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
-    let j2 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[1].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[keys[0]]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    let j2 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[keys[1]]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
     
-    return JacobianFactor(keys, [j1, j2], errorVector(values).tensor.reshaped(to: [3, 1]))
+    return JacobianFactor(keys, [j1, j2], -errorVector(values).tensor.reshaped(to: [3, 1]))
   }
 }
diff --git a/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift b/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
index 04ae34e4..2be190fa 100644
--- a/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
+++ b/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
@@ -38,9 +38,9 @@ public struct GaussianFactorGraph: FactorGraph {
     Array(lhs.factors.map { $0 * rhs })
   }
   
-  /// This calculates `b - A*x`, where x is the collection of key-values
+  /// This calculates `A*x - b`, where x is the collection of key-values
   public func residual (_ val: VectorValues) -> Errors {
-    Array(self.factors.map { $0.b - $0 * val })
+    Array(self.factors.map { $0 * val - $0.b })
   }
   
   /// Convenience operator for adding factor
@@ -54,7 +54,6 @@ public struct GaussianFactorGraph: FactorGraph {
     for i in r.indices {
       let JTr = factors[i].atr(r[i])
       
-      print("JTr = \(JTr)")
       vv = vv + JTr
     }
     
diff --git a/Sources/SwiftFusion/Inference/PriorFactor.swift b/Sources/SwiftFusion/Inference/PriorFactor.swift
new file mode 100644
index 00000000..53f80497
--- /dev/null
+++ b/Sources/SwiftFusion/Inference/PriorFactor.swift
@@ -0,0 +1,69 @@
+// Copyright 2019 The SwiftFusion Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import TensorFlow
+
+/// A `NonlinearFactor` that returns the difference of value and desired value
+///
+/// Input is a dictionary of `Key` to `Value` pairs, and the output is the scalar
+/// error value
+///
+/// Interpretation
+/// ================
+/// `Input`: the input values as key-value pairs
+///
+public struct PriorFactor: NonlinearFactor {
+  @noDerivative
+  public var keys: Array<Int> = []
+  public var difference: Pose2
+  public typealias Output = Error
+  
+  public init (_ key: Int, _ difference: Pose2) {
+    keys = [key]
+    self.difference = difference
+  }
+  typealias ScalarType = Double
+  
+  /// TODO: `Dictionary` still does not conform to `Differentiable`
+  /// Tracking issue: https://bugs.swift.org/browse/TF-899
+//  typealias Input = Dictionary<UInt, Tensor<ScalarType>>
+
+  /// Returns the `error` of the factor.
+  @differentiable(wrt: values)
+  public func error(_ values: Values) -> Double {
+    let error = between(
+      values[keys[0]].baseAs(Pose2.self),
+      difference
+    )
+    
+    return error.t.norm + error.rot.theta * error.rot.theta
+  }
+  
+  @differentiable(wrt: values)
+  public func errorVector(_ values: Values) -> Vector3 {
+    let error = between(
+      values[keys[0]].baseAs(Pose2.self),
+      difference
+    )
+    
+    return Vector3(error.rot.theta, error.t.x, error.t.y)
+  }
+  
+  public func linearize(_ values: Values) -> JacobianFactor {
+    let j = jacobian(of: self.errorVector, at: values)
+    
+    let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[0].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    
+    return JacobianFactor(keys, [j1], errorVector(values).tensor.reshaped(to: [3, 1]))
+  }
+}
diff --git a/Sources/SwiftFusion/Inference/Values.swift b/Sources/SwiftFusion/Inference/Values.swift
index a4f4ef07..cfd8e01b 100644
--- a/Sources/SwiftFusion/Inference/Values.swift
+++ b/Sources/SwiftFusion/Inference/Values.swift
@@ -34,7 +34,12 @@ public struct Values: Differentiable & KeyPathIterable {
   /// Should be replaced after Dictionary is in
   @differentiable
   public subscript(key: Int) -> AnyDifferentiable {
-    _values[_indices[key]!]
+    get {
+      _values[_indices[key]!]
+    }
+    set(newVal) {
+      _values[_indices[key]!] = newVal
+    }
   }
   
   /// Insert a key value pair
diff --git a/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift b/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift
index 1b44a5bd..0a0540e5 100644
--- a/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift
+++ b/Tests/SwiftFusionTests/Geometry/Pose2Tests.swift
@@ -261,8 +261,6 @@ final class Pose2Tests: XCTestCase {
 
   /// test convergence for a simple Pose2SLAM
   func testPose2SLAMWithSGD() {
-    let pi = 3.1415926
-
     let dumpjson = { (p: Pose2) -> String in
       "[ \(p.t.x), \(p.t.y), \(p.rot.theta)]"
     }
@@ -270,9 +268,9 @@ final class Pose2Tests: XCTestCase {
     // Initial estimate for poses
     let p1T0 = Pose2(Rot2(0.2), Vector2(0.5, 0.0))
     let p2T0 = Pose2(Rot2(-0.2), Vector2(2.3, 0.1))
-    let p3T0 = Pose2(Rot2(pi / 2), Vector2(4.1, 0.1))
-    let p4T0 = Pose2(Rot2(pi), Vector2(4.0, 2.0))
-    let p5T0 = Pose2(Rot2(-pi / 2), Vector2(2.1, 2.1))
+    let p3T0 = Pose2(Rot2(.pi / 2), Vector2(4.1, 0.1))
+    let p4T0 = Pose2(Rot2(.pi), Vector2(4.0, 2.0))
+    let p5T0 = Pose2(Rot2(-.pi / 2), Vector2(2.1, 2.1))
 
     var map = [p1T0, p2T0, p3T0, p4T0, p5T0]
 
@@ -285,9 +283,9 @@ final class Pose2Tests: XCTestCase {
 
         // Odometry measurements
         let p2T1 = between(between(map[1], map[0]), Pose2(2.0, 0.0, 0.0))
-        let p3T2 = between(between(map[2], map[1]), Pose2(2.0, 0.0, pi / 2))
-        let p4T3 = between(between(map[3], map[2]), Pose2(2.0, 0.0, pi / 2))
-        let p5T4 = between(between(map[4], map[3]), Pose2(2.0, 0.0, pi / 2))
+        let p3T2 = between(between(map[2], map[1]), Pose2(2.0, 0.0, .pi / 2))
+        let p4T3 = between(between(map[3], map[2]), Pose2(2.0, 0.0, .pi / 2))
+        let p5T4 = between(between(map[4], map[3]), Pose2(2.0, 0.0, .pi / 2))
 
         // Sum through the errors
         let error = self.e_pose2(p2T1) + self.e_pose2(p3T2) + self.e_pose2(p4T3) + self.e_pose2(p5T4)
diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
index f596f14e..6dc62948 100644
--- a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -25,7 +25,80 @@ final class NonlinearFactorGraphTests: XCTestCase {
     let expected = Tensor<Double>(shape:[3, 1], scalars: [.pi, 0.0, 0.0])
     
     print("gfg = \(gfg)")
-    print("error = \((gfg*vv).norm)")
+    print("error = \(gfg.residual(vv).norm)")
     assertEqual((gfg.residual(vv))[0], expected, accuracy: 1e-6)
   }
+  
+  /// test CGLS iterative solver
+  func testCGLSPose2SLAM() {
+    // Initial estimate for poses
+    let p1T0 = Pose2(Rot2(0.2), Vector2(0.5, 0.0))
+    let p2T0 = Pose2(Rot2(-0.2), Vector2(2.3, 0.1))
+    let p3T0 = Pose2(Rot2(.pi / 2), Vector2(4.1, 0.1))
+    let p4T0 = Pose2(Rot2(.pi), Vector2(4.0, 2.0))
+    let p5T0 = Pose2(Rot2(-.pi / 2), Vector2(2.1, 2.1))
+
+    let map = [p1T0, p2T0, p3T0, p4T0, p5T0]
+
+    var fg = NonlinearFactorGraph()
+    
+    fg += AnyNonlinearFactor(BetweenFactor(1, 0, Pose2(2.0, 0.0, .pi / 2)))
+    fg += AnyNonlinearFactor(BetweenFactor(2, 1, Pose2(2.0, 0.0, .pi / 2)))
+    fg += AnyNonlinearFactor(BetweenFactor(3, 2, Pose2(2.0, 0.0, .pi / 2)))
+    fg += AnyNonlinearFactor(BetweenFactor(4, 3, Pose2(2.0, 0.0, .pi / 2)))
+//    fg += AnyNonlinearFactor(BetweenFactor(4, 0, Pose2(0.0, 0.0, 0)))
+    
+    // fg += AnyNonlinearFactor(PriorFactor(0, Pose2(0.0, 0.0, 0.0)))
+    
+    var val = Values()
+    
+    for i in 0..<5 {
+      val.insert(i, AnyDifferentiable(map[i]))
+    }
+    
+    for _ in 0..<50 {
+      let gfg = fg.linearize(val)
+      
+      let optimizer = CGLS(precision: 1e-6, max_iteration: 500)
+      
+      var dx = VectorValues()
+      
+      for i in 0..<5 {
+        dx.insert(i, Tensor<Double>(shape: [3, 1], scalars: [0, 0, 0]))
+      }
+      
+      optimizer.optimize(gfg: gfg, initial: &dx)
+      
+//      print("dx = \(dx)")
+      for i in 0..<5 {
+        var p = val[i].baseAs(Pose2.self)
+//        print("p = \(p), dp = \(dx[i])")
+        p.move(along: Vector3(dx[i].reshaped(toShape: [3])))
+        val[i] = AnyDifferentiable(p)
+        print("p'[\(i)] = \(val[i])")
+      }
+    }
+    
+    let dumpjson = { (p: Pose2) -> String in
+      "[ \(p.rot.theta), \(p.t.x), \(p.t.y)]"
+    }
+    
+    print("map_init = [")
+    for v in map.indices {
+      print("\(dumpjson(map[v]))\({ () -> String in if v == map.indices.endIndex - 1 { return "" } else { return "," } }())")
+    }
+    print("]")
+    
+    let map_final = (0..<5).map { val[$0].baseAs(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))
+
+    // Test condition: P_5 should be identical to P_1 (close loop)
+    XCTAssertEqual(p5T1.t.norm, 0.0, accuracy: 1e-2)
+  }
 }

From 81e5ccf95c4d418fd8a9f45e5733fc6b19ec2650 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 01:30:23 -0400
Subject: [PATCH 07/13] PriorFactor is working

---
 Sources/SwiftFusion/Inference/PriorFactor.swift             | 3 ++-
 .../Inference/NonlinearFactorGraphTests.swift               | 6 ++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/Sources/SwiftFusion/Inference/PriorFactor.swift b/Sources/SwiftFusion/Inference/PriorFactor.swift
index 53f80497..7dbd6fc8 100644
--- a/Sources/SwiftFusion/Inference/PriorFactor.swift
+++ b/Sources/SwiftFusion/Inference/PriorFactor.swift
@@ -64,6 +64,7 @@ public struct PriorFactor: NonlinearFactor {
     
     let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[0].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
     
-    return JacobianFactor(keys, [j1], errorVector(values).tensor.reshaped(to: [3, 1]))
+    print(j1)
+    return JacobianFactor(keys, [j1], -errorVector(values).tensor.reshaped(to: [3, 1]))
   }
 }
diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
index 6dc62948..0b394d38 100644
--- a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -46,9 +46,7 @@ final class NonlinearFactorGraphTests: XCTestCase {
     fg += AnyNonlinearFactor(BetweenFactor(2, 1, Pose2(2.0, 0.0, .pi / 2)))
     fg += AnyNonlinearFactor(BetweenFactor(3, 2, Pose2(2.0, 0.0, .pi / 2)))
     fg += AnyNonlinearFactor(BetweenFactor(4, 3, Pose2(2.0, 0.0, .pi / 2)))
-//    fg += AnyNonlinearFactor(BetweenFactor(4, 0, Pose2(0.0, 0.0, 0)))
-    
-    // fg += AnyNonlinearFactor(PriorFactor(0, Pose2(0.0, 0.0, 0.0)))
+    fg += AnyNonlinearFactor(PriorFactor(0, Pose2(0.0, 0.0, 0.0)))
     
     var val = Values()
     
@@ -56,7 +54,7 @@ final class NonlinearFactorGraphTests: XCTestCase {
       val.insert(i, AnyDifferentiable(map[i]))
     }
     
-    for _ in 0..<50 {
+    for _ in 0..<5 {
       let gfg = fg.linearize(val)
       
       let optimizer = CGLS(precision: 1e-6, max_iteration: 500)

From 6ad74675be1e6c60c99481c211dddff6c7a6a1a7 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 01:42:36 -0400
Subject: [PATCH 08/13] Add TODO

---
 Sources/SwiftFusion/Inference/BetweenFactor.swift | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Sources/SwiftFusion/Inference/BetweenFactor.swift b/Sources/SwiftFusion/Inference/BetweenFactor.swift
index 805edef4..674e39ff 100644
--- a/Sources/SwiftFusion/Inference/BetweenFactor.swift
+++ b/Sources/SwiftFusion/Inference/BetweenFactor.swift
@@ -83,6 +83,7 @@ public struct BetweenFactor: NonlinearFactor {
     let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[keys[0]]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
     let j2 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[keys[1]]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
     
+    // TODO: remove this negative sign
     return JacobianFactor(keys, [j1, j2], -errorVector(values).tensor.reshaped(to: [3, 1]))
   }
 }

From 1ab2800a36bf90f62b260a870327efcb7bef5e9b Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 01:48:18 -0400
Subject: [PATCH 09/13] Remove prints

---
 Sources/SwiftFusion/Inference/PriorFactor.swift | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Sources/SwiftFusion/Inference/PriorFactor.swift b/Sources/SwiftFusion/Inference/PriorFactor.swift
index 7dbd6fc8..a54903ac 100644
--- a/Sources/SwiftFusion/Inference/PriorFactor.swift
+++ b/Sources/SwiftFusion/Inference/PriorFactor.swift
@@ -64,7 +64,6 @@ public struct PriorFactor: NonlinearFactor {
     
     let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[0].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
     
-    print(j1)
     return JacobianFactor(keys, [j1], -errorVector(values).tensor.reshaped(to: [3, 1]))
   }
 }

From 534a48a79028b8481d06ba1fd5357d11a8ae4a3d Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 11:05:37 -0400
Subject: [PATCH 10/13] Minor changes in comments

---
 .../Inference/NonlinearFactorGraphTests.swift                | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
index 0b394d38..a3894fed 100644
--- a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -54,7 +54,7 @@ final class NonlinearFactorGraphTests: XCTestCase {
       val.insert(i, AnyDifferentiable(map[i]))
     }
     
-    for _ in 0..<5 {
+    for _ in 0..<2 {
       let gfg = fg.linearize(val)
       
       let optimizer = CGLS(precision: 1e-6, max_iteration: 500)
@@ -67,13 +67,10 @@ final class NonlinearFactorGraphTests: XCTestCase {
       
       optimizer.optimize(gfg: gfg, initial: &dx)
       
-//      print("dx = \(dx)")
       for i in 0..<5 {
         var p = val[i].baseAs(Pose2.self)
-//        print("p = \(p), dp = \(dx[i])")
         p.move(along: Vector3(dx[i].reshaped(toShape: [3])))
         val[i] = AnyDifferentiable(p)
-        print("p'[\(i)] = \(val[i])")
       }
     }
     

From 887c2349f2123457bd138abdc24e969e6666814e Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 16:35:59 -0400
Subject: [PATCH 11/13] Remove Differentiable for NonlinearFactor

---
 .../Inference/AnyNonlinearFactor.swift        | 57 -------------------
 Sources/SwiftFusion/Inference/Factor.swift    |  2 +-
 .../Applications/ManipulationTests.swift      |  8 +++
 3 files changed, 9 insertions(+), 58 deletions(-)
 create mode 100644 Tests/SwiftFusionTests/Applications/ManipulationTests.swift

diff --git a/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift b/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
index 731946e0..87c7e270 100644
--- a/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
+++ b/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
@@ -4,8 +4,6 @@
 
 internal protocol _AnyNonlinearFactorBox {
   typealias ScalarType = Double
-  // `Differentiable` requirements.
-  mutating func _move(along direction: AnyDerivative)
 
   /// The underlying base value, type-erased to `Any`.
   var _typeErasedBase: Any { get }
@@ -44,16 +42,6 @@ internal struct _ConcreteNonlinearFactorBox<T: NonlinearFactor>: _AnyNonlinearFa
   func _unboxed<U: NonlinearFactor>(to type: U.Type) -> U? {
     return (self as? _ConcreteNonlinearFactorBox<U>)?._base
   }
-
-  mutating func _move(along direction: AnyDerivative) {
-    guard
-      let directionBase =
-        direction.base as? T.TangentVector
-    else {
-      _derivativeTypeMismatch(T.self, type(of: direction.base))
-    }
-    _base.move(along: directionBase)
-  }
   
   @differentiable(wrt: values)
   func _error(_ values: Values) -> ScalarType {
@@ -78,60 +66,15 @@ public struct AnyNonlinearFactor: NonlinearFactor {
   }
 
   /// Creates a type-erased derivative from the given derivative.
-  @differentiable
   public init<T: NonlinearFactor>(_ base: T) {
     self._box = _ConcreteNonlinearFactorBox<T>(base)
   }
-
-  @inlinable
-  @derivative(of: init)
-  internal static func _vjpInit<T: NonlinearFactor>(
-    _ base: T
-  ) -> (value: AnyNonlinearFactor, pullback: (AnyDerivative) -> T.TangentVector)
-  {
-    return (AnyNonlinearFactor(base), { v in v.base as! T.TangentVector })
-  }
-
-  @inlinable
-  @derivative(of: init)
-  internal static func _jvpInit<T: NonlinearFactor>(
-    _ base: T
-  ) -> (
-    value: AnyNonlinearFactor, differential: (T.TangentVector) -> AnyDerivative
-  ) {
-    return (AnyNonlinearFactor(base), { dbase in AnyDerivative(dbase) })
-  }
-
-  public typealias TangentVector = AnyDerivative
-
-  public mutating func move(along direction: TangentVector) {
-    _box._move(along: direction)
-  }
 }
 
 extension AnyNonlinearFactor {
-  @differentiable
   public func baseAs<T: NonlinearFactor>(_ t: T.Type) -> T {
     base as! T
   }
-
-  @derivative(of: baseAs)
-  @usableFromInline
-  func jvpBaseAs<T: NonlinearFactor>(_ t: T.Type) -> (
-    value: T,
-    differential: (AnyDerivative) -> T.TangentVector
-  ) {
-    (baseAs(t), { $0.base as! T.TangentVector })
-  }
-
-  @derivative(of: baseAs)
-  @usableFromInline
-  func vjpBaseAs<T: NonlinearFactor>(_ t: T.Type) -> (
-    value: T,
-    pullback: (T.TangentVector) -> AnyDerivative
-  ) {
-    (baseAs(t), { AnyDerivative($0) })
-  }
 }
 
 extension AnyNonlinearFactor {
diff --git a/Sources/SwiftFusion/Inference/Factor.swift b/Sources/SwiftFusion/Inference/Factor.swift
index 8d572ba4..9c14c9e3 100644
--- a/Sources/SwiftFusion/Inference/Factor.swift
+++ b/Sources/SwiftFusion/Inference/Factor.swift
@@ -47,7 +47,7 @@ public protocol LinearFactor: Factor {
 /// ================
 /// `Input`: the input values as key-value pairs
 ///
-public protocol NonlinearFactor: Differentiable & Factor {
+public protocol NonlinearFactor: Factor {
   typealias ScalarType = Double
   
   /// TODO: `Dictionary` still does not conform to `Differentiable`
diff --git a/Tests/SwiftFusionTests/Applications/ManipulationTests.swift b/Tests/SwiftFusionTests/Applications/ManipulationTests.swift
new file mode 100644
index 00000000..d73b1501
--- /dev/null
+++ b/Tests/SwiftFusionTests/Applications/ManipulationTests.swift
@@ -0,0 +1,8 @@
+//
+//  File.swift
+//  
+//
+//  Created by Fan Jiang on 2020/4/22.
+//
+
+import Foundation

From 45b027666e0cb0a212e1abe582df5e94e2848678 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 22 Apr 2020 21:20:25 -0400
Subject: [PATCH 12/13] Add pseudo inverse

---
 Sources/SwiftFusion/Core/MathUtil.swift | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 Sources/SwiftFusion/Core/MathUtil.swift

diff --git a/Sources/SwiftFusion/Core/MathUtil.swift b/Sources/SwiftFusion/Core/MathUtil.swift
new file mode 100644
index 00000000..43a2d5df
--- /dev/null
+++ b/Sources/SwiftFusion/Core/MathUtil.swift
@@ -0,0 +1,22 @@
+//===----------------------------------------------------------------------===//
+// Pseudo inverse
+//===----------------------------------------------------------------------===//
+
+import TensorFlow
+
+public func pinv(_ m: Tensor<Double>) -> Tensor<Double> {
+  let (J_s, J_u, J_v) = m.svd(computeUV: true, fullMatrices: true)
+  
+  let m = J_v!.shape[1]
+  let n = J_u!.shape[0]
+  if (m > n) {
+    let J_ss = J_s.reciprocal.diagonal().concatenated(with: Tensor<Double>(repeating: 0, shape: [m-n, n]), alongAxis: 0)
+    return matmul(matmul(J_v!, J_ss), J_u!.transposed())
+  } else if (m < n) {
+    let J_ss = J_s.reciprocal.diagonal().concatenated(with: Tensor<Double>(repeating: 0, shape: [m, n-m]), alongAxis: 1)
+    return matmul(matmul(J_v!, J_ss), J_u!.transposed())
+  } else {
+    let J_ss = J_s.reciprocal.diagonal()
+    return matmul(matmul(J_v!, J_ss), J_u!.transposed())
+  }
+}

From cc969dedfef05595383aab9f2b55d33afbf9b244 Mon Sep 17 00:00:00 2001
From: Fan Jiang <prof.fan@foxmail.com>
Date: Wed, 29 Apr 2020 19:36:21 -0400
Subject: [PATCH 13/13] Remove AnyNonlinearFactor

---
 Sources/SwiftFusion/Geometry/File.swift       |  8 ++
 .../Inference/AnyNonlinearFactor.swift        | 95 -------------------
 .../SwiftFusion/Inference/BetweenFactor.swift | 20 ++--
 .../SwiftFusion/Inference/FactorGraph.swift   | 33 -------
 .../Inference/GaussianFactorGraph.swift       |  2 +-
 .../Inference/NonlinearFactorGraph.swift      |  6 +-
 .../SwiftFusionTests/Geometry/Rot3Tests.swift |  8 ++
 .../Inference/NonlinearFactorGraphTests.swift | 12 +--
 8 files changed, 40 insertions(+), 144 deletions(-)
 create mode 100644 Sources/SwiftFusion/Geometry/File.swift
 delete mode 100644 Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
 delete mode 100644 Sources/SwiftFusion/Inference/FactorGraph.swift
 create mode 100644 Tests/SwiftFusionTests/Geometry/Rot3Tests.swift

diff --git a/Sources/SwiftFusion/Geometry/File.swift b/Sources/SwiftFusion/Geometry/File.swift
new file mode 100644
index 00000000..e53d4b0b
--- /dev/null
+++ b/Sources/SwiftFusion/Geometry/File.swift
@@ -0,0 +1,8 @@
+//
+//  File.swift
+//  
+//
+//  Created by Fan Jiang on 2020/4/24.
+//
+
+import Foundation
diff --git a/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift b/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
deleted file mode 100644
index 87c7e270..00000000
--- a/Sources/SwiftFusion/Inference/AnyNonlinearFactor.swift
+++ /dev/null
@@ -1,95 +0,0 @@
-//===----------------------------------------------------------------------===//
-// `AnyNonlinearFactor`
-//===----------------------------------------------------------------------===//
-
-internal protocol _AnyNonlinearFactorBox {
-  typealias ScalarType = Double
-
-  /// The underlying base value, type-erased to `Any`.
-  var _typeErasedBase: Any { get }
-
-  /// Returns the underlying value unboxed to the given type, if possible.
-  func _unboxed<U: NonlinearFactor>(to type: U.Type) -> U?
-  
-  var _keys: Array<Int> { get }
-  
-  @differentiable(wrt: values)
-  func _error(_ values: Values) -> ScalarType
-  
-  func _linearize(_ values: Values) -> JacobianFactor
-}
-
-internal struct _ConcreteNonlinearFactorBox<T: NonlinearFactor>: _AnyNonlinearFactorBox
-{
-  var _keys: Array<Int> {
-    get {
-      _base.keys
-    }
-  }
-  
-  /// The underlying base value.
-  var _base: T
-
-  init(_ base: T) {
-    self._base = base
-  }
-
-  /// The underlying base value, type-erased to `Any`.
-  var _typeErasedBase: Any {
-    return _base
-  }
-
-  func _unboxed<U: NonlinearFactor>(to type: U.Type) -> U? {
-    return (self as? _ConcreteNonlinearFactorBox<U>)?._base
-  }
-  
-  @differentiable(wrt: values)
-  func _error(_ values: Values) -> ScalarType {
-    _base.error(values)
-  }
-  
-  func _linearize(_ values: Values) -> JacobianFactor {
-    _base.linearize(values)
-  }
-}
-
-public struct AnyNonlinearFactor: NonlinearFactor {
-  internal var _box: _AnyNonlinearFactorBox
-
-  internal init(_box: _AnyNonlinearFactorBox) {
-    self._box = _box
-  }
-
-  /// The underlying base value.
-  public var base: Any {
-    return _box._typeErasedBase
-  }
-
-  /// Creates a type-erased derivative from the given derivative.
-  public init<T: NonlinearFactor>(_ base: T) {
-    self._box = _ConcreteNonlinearFactorBox<T>(base)
-  }
-}
-
-extension AnyNonlinearFactor {
-  public func baseAs<T: NonlinearFactor>(_ t: T.Type) -> T {
-    base as! T
-  }
-}
-
-extension AnyNonlinearFactor {
-  @differentiable(wrt: values)
-  public func error(_ values: Values) -> ScalarType {
-    _box._error(values)
-  }
-  
-  public var keys: Array<Int> {
-    get {
-      _box._keys
-    }
-  }
-  
-  public func linearize(_ values: Values) -> JacobianFactor {
-    _box._linearize(values)
-  }
-}
diff --git a/Sources/SwiftFusion/Inference/BetweenFactor.swift b/Sources/SwiftFusion/Inference/BetweenFactor.swift
index 674e39ff..6f9cd9de 100644
--- a/Sources/SwiftFusion/Inference/BetweenFactor.swift
+++ b/Sources/SwiftFusion/Inference/BetweenFactor.swift
@@ -23,13 +23,21 @@ import TensorFlow
 /// `Input`: the input values as key-value pairs
 ///
 public struct BetweenFactor: NonlinearFactor {
+  
+  var key1: Int
+  var key2: Int
   @noDerivative
-  public var keys: Array<Int> = []
+  public var keys: Array<Int> {
+    get {
+      [key1, key2]
+    }
+  }
   public var difference: Pose2
   public typealias Output = Error
   
   public init (_ key1: Int, _ key2: Int, _ difference: Pose2) {
-    keys = [key1, key2]
+    self.key1 = key1
+    self.key2 = key2
     self.difference = difference
   }
   typealias ScalarType = Double
@@ -60,7 +68,7 @@ public struct BetweenFactor: NonlinearFactor {
   @differentiable(wrt: values)
   public func error(_ values: Values) -> Double {
     let error = between(
-      between(values[keys[1]].baseAs(Pose2.self), values[keys[0]].baseAs(Pose2.self)),
+      between(values[key2].baseAs(Pose2.self), values[key1].baseAs(Pose2.self)),
         difference
     )
     
@@ -70,7 +78,7 @@ public struct BetweenFactor: NonlinearFactor {
   @differentiable(wrt: values)
   public func errorVector(_ values: Values) -> Vector3 {
     let error = between(
-      between(values[keys[1]].baseAs(Pose2.self), values[keys[0]].baseAs(Pose2.self)),
+      between(values[key2].baseAs(Pose2.self), values[key1].baseAs(Pose2.self)),
       difference
     )
     
@@ -80,8 +88,8 @@ public struct BetweenFactor: NonlinearFactor {
   public func linearize(_ values: Values) -> JacobianFactor {
     let j = jacobian(of: self.errorVector, at: values)
     
-    let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[keys[0]]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
-    let j2 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[keys[1]]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    let j1 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[key1]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
+    let j2 = Tensor<Double>(stacking: (0..<3).map { i in (j[i]._values[values._indices[key2]!].base as! Pose2.TangentVector).tensor.reshaped(to: TensorShape([3])) })
     
     // TODO: remove this negative sign
     return JacobianFactor(keys, [j1, j2], -errorVector(values).tensor.reshaped(to: [3, 1]))
diff --git a/Sources/SwiftFusion/Inference/FactorGraph.swift b/Sources/SwiftFusion/Inference/FactorGraph.swift
deleted file mode 100644
index 6f6ffcaa..00000000
--- a/Sources/SwiftFusion/Inference/FactorGraph.swift
+++ /dev/null
@@ -1,33 +0,0 @@
-// Copyright 2019 The SwiftFusion Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-import TensorFlow
-
-/// A Factor Graph in its very general form.
-///
-/// Explanation
-/// =============
-/// A factor graph is a biparite graph that connects between *Factors* and *Values*, and
-/// the way they are stored does not matter here.
-///
-public protocol FactorGraph {
-  // TODO: find a better protocol
-  associatedtype KeysType : Collection where KeysType.Element : SignedInteger
-  associatedtype FactorsType : Collection where FactorsType.Element: Factor
-  /// TODO(fan): Is this right?
-  /// Or, do we need this at all? I think this would help to register keys to descriptions and
-  /// help debugging and serialization, but I am not sure.
-  var keys: KeysType { get }
-  
-  var factors: FactorsType { get }
-}
diff --git a/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift b/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
index 2be190fa..f9121c1a 100644
--- a/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
+++ b/Sources/SwiftFusion/Inference/GaussianFactorGraph.swift
@@ -16,7 +16,7 @@ import TensorFlow
 /// A factor graph for linear problems
 /// Factors are the Jacobians between the corresponding variables and measurements
 /// TODO(fan): Add noise model
-public struct GaussianFactorGraph: FactorGraph {
+public struct GaussianFactorGraph {
   public typealias KeysType = Array<Int>
   
   public typealias FactorsType = Array<JacobianFactor>
diff --git a/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift b/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
index 73010dcb..fdcd2512 100644
--- a/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
+++ b/Sources/SwiftFusion/Inference/NonlinearFactorGraph.swift
@@ -15,10 +15,10 @@ import TensorFlow
 
 /// A factor graph for nonlinear problems
 /// TODO(fan): Add noise model
-public struct NonlinearFactorGraph: FactorGraph {
+public struct NonlinearFactorGraph {
   public typealias KeysType = Array<Int>
   
-  public typealias FactorsType = Array<AnyNonlinearFactor>
+  public typealias FactorsType = Array<NonlinearFactor>
   
   public var keys: KeysType = []
   public var factors: FactorsType = []
@@ -27,7 +27,7 @@ public struct NonlinearFactorGraph: FactorGraph {
   public init() { }
   
   /// Convenience operator for adding factor
-  public static func += (lhs: inout Self, rhs: AnyNonlinearFactor) {
+  public static func += (lhs: inout Self, rhs: NonlinearFactor) {
     lhs.factors.append(rhs)
   }
   
diff --git a/Tests/SwiftFusionTests/Geometry/Rot3Tests.swift b/Tests/SwiftFusionTests/Geometry/Rot3Tests.swift
new file mode 100644
index 00000000..e53d4b0b
--- /dev/null
+++ b/Tests/SwiftFusionTests/Geometry/Rot3Tests.swift
@@ -0,0 +1,8 @@
+//
+//  File.swift
+//  
+//
+//  Created by Fan Jiang on 2020/4/24.
+//
+
+import Foundation
diff --git a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
index a3894fed..b9be43d8 100644
--- a/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
+++ b/Tests/SwiftFusionTests/Inference/NonlinearFactorGraphTests.swift
@@ -9,7 +9,7 @@ final class NonlinearFactorGraphTests: XCTestCase {
     
     let bf1 = BetweenFactor(0, 1, Pose2(0.0,0.0, 0.0))
     
-    fg += AnyNonlinearFactor(bf1)
+    fg += bf1
     
     var val = Values()
     val.insert(0, AnyDifferentiable(Pose2(1.0, 1.0, 0.0)))
@@ -42,11 +42,11 @@ final class NonlinearFactorGraphTests: XCTestCase {
 
     var fg = NonlinearFactorGraph()
     
-    fg += AnyNonlinearFactor(BetweenFactor(1, 0, Pose2(2.0, 0.0, .pi / 2)))
-    fg += AnyNonlinearFactor(BetweenFactor(2, 1, Pose2(2.0, 0.0, .pi / 2)))
-    fg += AnyNonlinearFactor(BetweenFactor(3, 2, Pose2(2.0, 0.0, .pi / 2)))
-    fg += AnyNonlinearFactor(BetweenFactor(4, 3, Pose2(2.0, 0.0, .pi / 2)))
-    fg += AnyNonlinearFactor(PriorFactor(0, Pose2(0.0, 0.0, 0.0)))
+    fg += BetweenFactor(1, 0, Pose2(2.0, 0.0, .pi / 2))
+    fg += BetweenFactor(2, 1, Pose2(2.0, 0.0, .pi / 2))
+    fg += BetweenFactor(3, 2, Pose2(2.0, 0.0, .pi / 2))
+    fg += BetweenFactor(4, 3, Pose2(2.0, 0.0, .pi / 2))
+    fg += PriorFactor(0, Pose2(0.0, 0.0, 0.0))
     
     var val = Values()