Recipe for defining custom manifolds

Sources/SwiftFusion/Core/EuclideanVectorSpace.swift

@@ -0,0 +1,6 @@
+/// A Euclidean vector space.
+public protocol EuclideanVectorSpace: Differentiable, VectorProtocol
+  where Self.TangentVector == Self
+{
+  // Note: This is a work in progress. We intend to add more requirements here as we need them.
+}

Sources/SwiftFusion/Core/Manifold.swift

@@ -0,0 +1,153 @@
+/// A point on a differentiable manifold with a `retract` map centered around `self`.
+///
+/// This protocol helps you define manifolds with custom tangent vectors. Instructions:
+/// 1. Define a type `C: ManifoldCoordinate`, without specifying a `TangentVector`. (Swift
+///    generates a `TangentVector` automatically, usually not the `TangentVector` that you want for
+///    your manifold).
+/// 2. Define `C.LocalCoordinate` to be the `TangentVector` type that you want for your manifold,
+///    and define `C.retract` and `C.localCoordinate` to be the retraction and inverse retration
+///    for this `TangentVector`.
+/// 3. Define a type `M: Manifold` that wraps `C`. The `Manifold` protocol automatically gives `M`
+///    the desired `TangentVector`.
+/// See "SwiftFusion/doc/DifferentiableManifoldRecipe.md" for more detailed instructions.
+public protocol ManifoldCoordinate: Differentiable {
+  /// The local coordinate type of the manifold.
+  ///
+  /// This is the `TangentVector` of the `Manifold` wrapper type.
+  ///
+  /// Note that this is not the same type as `Self.TangentVector`.
+  associatedtype LocalCoordinate: EuclideanVectorSpace
+
+  /// Diffeomorphism between a neigborhood of `Localcoordinate.zero` and `Self`.
+  ///
+  /// Satisfies the following properties:
+  /// - `retract(LocalCoordinate.zero) == self`
+  /// - There exists an open set `B` around `LocalCoordinate.zero` such that
+  ///   `localCoordinate(retract(b)) == b` for all `b \in B`.
+  @differentiable(wrt: local)
+  func retract(_ local: LocalCoordinate) -> Self
+
+  /// Inverse of `retract`.
+  ///
+  /// Satisfies the following properties:
+  /// - `localCoordinate(self) == LocalCoordinate.zero`
+  /// - There exists an open set `B` around `self` such that `localCoordinate(retract(b)) == b` for all
+  ///   `b \in B`.
+  @differentiable(wrt: global)
+  func localCoordinate(_ global: Self) -> LocalCoordinate
+}
+
+/// A point on a differentiable manifold.
+public protocol Manifold: Differentiable {
+  /// The manifold's global coordinate system.
+  associatedtype Coordinate: ManifoldCoordinate
+
+  /// The coordinate of `self`.
+  ///
+  /// Note: The distinction between `coordinateStorage` and `coordinate` is a workaround until we
+  /// can define default derivatives for protocol requirements (TF-982). Until then, implementers
+  /// of this protocol must define `coordinateStorage`, and clients of this protocol must access
+  /// coordinate`. This allows us to define default derivatives for `coordinate` that translate
+  /// between the `ManifoldCoordinate` tangent space and the `Manifold` tangent space.
+  var coordinateStorage: Coordinate { get set }
+
+  /// Creates a manifold point with coordinate `coordinateStorage`.
+  ///
+  /// Note: The distinction between `init(coordinateStorage:)` and `init(coordinate:)` is a workaround until we
+  /// can define default derivatives for protocol requirements (TF-982). Until then, implementers
+  /// of this protocol must define `init(coordinateStorage:)`, and clients of this protocol must access
+  /// init(coordinate:)`. This allows us to define default derivatives for `init(coordinate:)` that translate
+  /// between the `ManifoldCoordinate` tangent space and the `Manifold` tangent space.
+  init(coordinateStorage: Coordinate)
+}
+
+/// Methods for converting between manifolds and their coordinates.
+///
+/// To enable these, you must explicitly write
+//    public typealias TangentVector = <local coordinate type>
+/// in your manifold type.
+extension Manifold where Self.TangentVector == Coordinate.LocalCoordinate {
+  /// The coordinate of `self`.
+  @differentiable
+  public var coordinate: Coordinate {
+    return coordinateStorage
+  }
+
+  /// A custom derivative of `coordinate` that converts from the global coordinate system's
+  /// tangent vector to the local coordinate system's tangent vector, so that all functions on this
+  /// manifold using `coordinate` have derivatives involving local coordinates.
+  @derivative(of: coordinate)
+  @usableFromInline
+  func vjpCoordinate()
+    -> (value: Coordinate, pullback: (Coordinate.TangentVector) -> TangentVector)
+  {
+
+    // Explanation of this pullback:
+    //
+    // Let `f: Manifold -> Coordinate` be `f(x) = x.coordinateStorage`.
+    //
+    // `differential(at: x, in: f)` is a linear approximation of how changes in tangent vectors
+    // around `x` lead to changes in global coordinates around `x.coordinateStorage`.
+    //
+    // `x.coordinateStorage.retract: TangentVector -> Coordinate` defines _exactly_ how local
+    // coordinates around zero map to global coordinates around `x`.
+    //
+    // Therefore, `differential(at: x, in: f) = differential(at: zero, in: x.coordinateStorage.retract)`.
+    //
+    // The pullback is the dual map of the differential, so taking duals of both sides gives:
+    //   `pullback(at: x, in: f) = pullback(at: zero, in: x.coordinateStorage.retract)`.
+
+    return (
+      value: coordinateStorage,
+      pullback(at: Coordinate.LocalCoordinate.zero) { self.coordinateStorage.retract($0) }
+    )
+  }
+
+  /// Creates a manifold point with coordinate `coordinate`.
+  @differentiable
+  public init(coordinate: Coordinate) { self.init(coordinateStorage: coordinate) }
+
+  /// A custom derivative of `init(coordinate:)` that converts from the local coordinate system's
+  /// tangent vector to the global coordinate system's tangent vector, so that all functions
+  /// producing instances of this manifold using `init(coordinates:)` have derivatives involving
+  /// local coordinates.
+  @derivative(of: init(coordinate:))
+  @usableFromInline
+  static func vjpInit(coordinate: Coordinate)
+    -> (value: Self, pullback: (TangentVector) -> Coordinate.TangentVector)
+  {
+
+    // Explanation of this pullback:
+    //
+    // Let `g: Coordinate -> Manifold` be `g(x) = Self(coordinateStorage: x)`.
+    //
+    // `D_x(g)` (the derivative of `g` at `x`) is a linear approximation of how changes in global
+    // coordinates around `x` lead to changes in tangent vectors around
+    // `Self(coordinateStorage: x)`.
+    //
+    // `x.coordinateStorage.localCoordinate: Coordinate -> TangentVector` defines _exactly_ how global
+    // coordinates around `x` map to tangent vectors.
+    //
+    // Therefore, `D_x(g)` is the derivative of `x.coordinateStorage.localCoordinate`.
+
+    // Explanation of this pullback:
+    //
+    // Let `g: Coordinate -> Manifold` be `g(x) = Self(coordinateStorage: x)`.
+    //
+    // `differential(at: x, in: g)` is a linear approximation of how changes in global
+    // coordinates around `x` lead to changes in local coordinates around `Self(coordinateStorage: x)`.
+    //
+    // `x.coordinateStorage.localCoordinate: Coordinate -> LocalCoordinate` defines _exactly_ how global
+    // coordinates around `x` map to local coordinates.
+    //
+    // Therefore, `differential(at: x, in: g) = differential(at: zero, in: x.coordinateStorage.localCoordinate)`.
+    //
+    // The pullback is the dual map of the differential, so taking duals of both sides gives:
+    //   `pullback(at: x, in: g) = pullback(at: zero, in: x.coordinateStorage.localCoordinate)`.
+
+    return (
+      value: Self(coordinateStorage: coordinate),
+      pullback: pullback(at: coordinate) { coordinate.localCoordinate($0) }
+    )
+  }
+}