Jacobians (#47)

* Jacobians

* semicolons

* expand docs

* rename basis arguments, generalize allocate_jacobian

* add codecov token?

* bump Julia version on CI

.github/workflows/ci.yml

@@ -24,5 +24,8 @@ jobs:
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: ./lcov.info
+          name: codecov-umbrella
           fail_ci_if_error: false
         if: ${{ matrix.os =='ubuntu-latest' }}

.github/workflows/format.yml

@@ -13,7 +13,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
-          version: 1.6
+          version: "1.10"
       - name: Install JuliaFormatter and format
         run: |
           julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'

Changelog.md

@@ -5,6 +5,15 @@ All notable changes to this Julia package will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.3.13] unreleased
+
+### Added
+
+* `jacobian_exp_argument` and its mutating variant, `jacobian_exp_argument!`.
+* `jacobian_exp_basepoint` and its mutating variant, `jacobian_exp_basepoint!`.
+* `jacobian_log_argument` and its mutating variant, `jacobian_log_argument!`.
+* `jacobian_log_basepoint` and its mutating variant, `jacobian_log_basepoint!`.
+
 ## [0.3.12] September 5, 2024
 
 ### Added

Project.toml

@@ -1,7 +1,7 @@
 name = "ManifoldDiff"
 uuid = "af67fdf4-a580-4b9f-bbec-742ef357defd"
 authors = ["Seth Axen <[email protected]>", "Mateusz Baran <[email protected]>", "Ronny Bergmann <[email protected]>"]
-version = "0.3.12"
+version = "0.3.13"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

docs/src/index.md

@@ -14,5 +14,7 @@ For example
 * `adjoint_differential_f` for pullbacks
 * `adjoint_differential_f_variable` if `f` has multiple variables / parameters
 * `f_derivative` for ``f'``
+* `jacobian_f` for Jacobian matrix of ``f``.
+* `jacobian_f_variable` if `f` has multiple parameters.
 
 the scheme is not completely fixed but tries to follow the mathematical notation.

docs/src/library.md

@@ -44,6 +44,14 @@ Pages = ["Jacobi_fields.jl"]
 Order = [:type, :function, :constant]
 ```
 
+## Jacobians
+
+```@autodocs
+Modules = [ManifoldDiff]
+Pages = ["jacobians.jl"]
+Order = [:type, :function, :constant]
+```
+
 ## Riemannian differentials
 
 ```@autodocs

src/ManifoldDiff.jl

@@ -173,6 +173,7 @@ include("gradients.jl")
 include("Jacobi_fields.jl")
 include("proximal_maps.jl")
 include("subgradients.jl")
+include("jacobians.jl")
 
 include("riemannian_diff.jl")
 include("embedded_diff.jl")

src/differentials.jl

@@ -51,7 +51,7 @@ end
     Z = differential_exp_argument(M, p, X, Y)
     differential_exp_argument!(M, Z, p, X, Y)
 
-computes ``D_X\exp_pX[Y]`` (in place of `Z`).
+Compute ``D_X\exp_pX[Y]`` (in place of `Z`).
 Note that ``X ∈  T_X(T_p\mathcal M) = T_p\mathcal M`` is still a tangent vector.
 
 # See also

src/jacobians.jl

@@ -0,0 +1,233 @@
+
+"""
+    allocate_jacobian(
+        M_domain::AbstractManifold,
+        M_codomain::AbstractManifold,
+        f,
+        p;
+        basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+        basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+    )
+
+Allocate Jacobian of function `f` with given domain and codomain at point `p`.
+`basis_domain` and `basis_codomain` denote bases of tangent spaces at, respectively, `p` and `f(p)`.
+"""
+function allocate_jacobian(
+    M_domain::AbstractManifold,
+    M_codomain::AbstractManifold,
+    f,
+    p;
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    n = number_of_coordinates(M_domain, basis_domain)
+    m = number_of_coordinates(M_codomain, basis_codomain)
+    J = zeros(float(number_eltype(p)), n, m)
+    return J
+end
+
+function jacobian_exp_argument!(
+    M::AbstractManifold,
+    J::AbstractMatrix,
+    p,
+    X,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    Bb = get_basis(M, p, basis_domain)
+    Vs = get_vectors(M, p, Bb)
+    Z = similar(X)
+    p_exp = exp(M, p, X)
+    for i in 1:length(Vs)
+        differential_exp_argument!(M, Z, p, X, Vs[i])
+        get_coordinates!(M, view(J, :, i), p_exp, Z, basis_codomain)
+    end
+    return J
+end
+
+@doc raw"""
+    jacobian_exp_argument(
+        M::AbstractManifold,
+        p,
+        X,
+        basis_domain::AbstractBasis=DefaultOrthonormalBasis(),
+        basis_codomain::AbstractBasis=DefaultOrthonormalBasis(),
+    )
+
+Compute Jacobian of the exponential map with respect to its argument (tangent vector).
+Differential of the exponential map is here considered as a function from ``T_p \mathcal{M}``
+to ``T_{\exp_p X} \mathcal{M}``. Jacobian coefficients are represented in basis `basis_domain`
+in the domain and in `basis_codomain` in the codomain.
+"""
+function jacobian_exp_argument(
+    M::AbstractManifold,
+    p,
+    X,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    J = allocate_jacobian(
+        M,
+        M,
+        jacobian_exp_argument,
+        X;
+        basis_domain = basis_domain,
+        basis_codomain = basis_codomain,
+    )
+    jacobian_exp_argument!(M, J, p, X, basis_domain, basis_codomain)
+    return J
+end
+
+function jacobian_exp_basepoint!(
+    M::AbstractManifold,
+    J::AbstractMatrix,
+    p,
+    X,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    Bb = get_basis(M, p, basis_domain)
+    Vs = get_vectors(M, p, Bb)
+    Z = similar(X)
+    p_exp = exp(M, p, X)
+    for i in 1:length(Vs)
+        differential_exp_basepoint!(M, Z, p, X, Vs[i])
+        get_coordinates!(M, view(J, :, i), p_exp, Z, basis_codomain)
+    end
+    return J
+end
+
+@doc raw"""
+    jacobian_exp_basepoint(
+        M::AbstractManifold,
+        p,
+        X,
+        basis_domain::AbstractBasis=DefaultOrthonormalBasis(),
+        basis_codomain::AbstractBasis=DefaultOrthonormalBasis(),
+    )
+
+Compute Jacobian of the exponential map with respect to the basepoint.
+Differential of the exponential map is here considered as a function from ``T_p \mathcal{M}``
+to ``T_{\exp_p X} \mathcal{M}``. Jacobian coefficients are represented in basis `basis_domain`
+in the domain and in `basis_codomain` in the codomain.
+"""
+function jacobian_exp_basepoint(
+    M::AbstractManifold,
+    p,
+    X,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    J = allocate_jacobian(
+        M,
+        M,
+        jacobian_exp_basepoint,
+        X;
+        basis_domain = basis_domain,
+        basis_codomain = basis_codomain,
+    )
+    jacobian_exp_basepoint!(M, J, p, X, basis_domain, basis_codomain)
+    return J
+end
+
+function jacobian_log_argument!(
+    M::AbstractManifold,
+    J::AbstractMatrix,
+    p,
+    q,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    Bb = get_basis(M, q, basis_domain)
+    Vs = get_vectors(M, q, Bb)
+    Z = zero_vector(M, p)
+    for i in 1:length(Vs)
+        differential_log_argument!(M, Z, p, q, Vs[i])
+        get_coordinates!(M, view(J, :, i), p, Z, basis_codomain)
+    end
+    return J
+end
+
+@doc raw"""
+    jacobian_log_argument(
+        M::AbstractManifold,
+        p,
+        q,
+        basis_domain::AbstractBasis=DefaultOrthonormalBasis(),
+        basis_codomain::AbstractBasis=DefaultOrthonormalBasis(),
+    )
+
+Compute Jacobian of the logarithmic map with respect to its argument (point `q`).
+Differential of the logarithmic map is here considered as a function from ``T_q \mathcal{M}``
+to ``T_p \mathcal{M}``. Jacobian coefficients are represented in basis `basis_domain`
+in the domain and in `basis_codomain` in the codomain.
+"""
+function jacobian_log_argument(
+    M::AbstractManifold,
+    p,
+    q,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    J = allocate_jacobian(
+        M,
+        M,
+        jacobian_log_argument,
+        q;
+        basis_domain = basis_domain,
+        basis_codomain = basis_codomain,
+    )
+    jacobian_log_argument!(M, J, p, q, basis_domain, basis_codomain)
+    return J
+end
+
+function jacobian_log_basepoint!(
+    M::AbstractManifold,
+    J::AbstractMatrix,
+    p,
+    q,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    Bb = get_basis(M, p, basis_domain)
+    Vs = get_vectors(M, p, Bb)
+    Z = zero_vector(M, p)
+    for i in 1:length(Vs)
+        differential_log_basepoint!(M, Z, p, q, Vs[i])
+        get_coordinates!(M, view(J, :, i), p, Z, basis_codomain)
+    end
+    return J
+end
+
+@doc raw"""
+    jacobian_log_basepoint(
+        M::AbstractManifold,
+        p,
+        q,
+        basis_domain::AbstractBasis=DefaultOrthonormalBasis(),
+        basis_codomain::AbstractBasis=DefaultOrthonormalBasis(),
+    )
+
+Compute Jacobian of the logarithmic map with respect to the basepoint.
+Differential of the logarithmic map is here considered as a function from ``T_q \mathcal{M}``
+to ``T_p \mathcal{M}``. Jacobian coefficients are represented in basis `basis_domain`
+in the domain and in `basis_codomain` in the codomain.
+"""
+function jacobian_log_basepoint(
+    M::AbstractManifold,
+    p,
+    q,
+    basis_domain::AbstractBasis = DefaultOrthonormalBasis(),
+    basis_codomain::AbstractBasis = DefaultOrthonormalBasis(),
+)
+    J = allocate_jacobian(
+        M,
+        M,
+        jacobian_log_basepoint,
+        q;
+        basis_domain = basis_domain,
+        basis_codomain = basis_codomain,
+    )
+    jacobian_log_basepoint!(M, J, p, q, basis_domain, basis_codomain)
+    return J
+end

test/runtests.jl

@@ -17,4 +17,5 @@ using RecursiveArrayTools
     include("test_derivatives.jl")
     include("test_proximal_maps.jl")
     include("test_subgradients.jl")
+    include("test_jacobians.jl")
 end

test/test_jacobians.jl

@@ -0,0 +1,72 @@
+using Manifolds, ManifoldsBase, ManifoldDiff, Test
+using FiniteDifferences
+
+function finitedifferences_jacobian(
+    M_arg::AbstractManifold,
+    M_val::AbstractManifold,
+    f,
+    p;
+    B_arg::AbstractBasis = DefaultOrthonormalBasis(),
+    B_val::AbstractBasis = DefaultOrthonormalBasis(),
+    m = central_fdm(5, 1),
+)
+    fv = f(p)
+    function fj(_c)
+        Y = get_vector(M_arg, p, _c, B_arg)
+        return get_coordinates(M_val, fv, log(M_val, fv, f(exp(M_arg, p, Y))), B_val)
+    end
+    return FiniteDifferences.jacobian(m, fj, zeros(manifold_dimension(M_arg)))[1]
+end
+
+@testset "Jacobians" begin
+    M = Sphere(2)
+    p = [0.0, 0.0, 1.0]
+    X = [0.0, 1.0, 0.0]
+    q = [0.0, 1.0, 0.0]
+
+    Jfnd = finitedifferences_jacobian(TangentSpace(M, p), M, Y -> exp(M, p, Y), X)
+
+    J = ManifoldDiff.jacobian_exp_argument(M, p, X)
+    @test J ≈ Jfnd atol = 1e-8
+
+    ManifoldDiff.differential_exp_argument(M, p, X, [1.0, 0.0, 0.0])
+
+    Jc = similar(J)
+    ManifoldDiff.jacobian_exp_argument!(M, Jc, p, X)
+    @test Jc ≈ Jfnd atol = 1e-8
+
+    Jfnd = finitedifferences_jacobian(
+        M,
+        M,
+        q -> exp(M, q, parallel_transport_to(M, p, X, q)),
+        p,
+    )
+    J = ManifoldDiff.jacobian_exp_basepoint(M, p, X)
+    @test J ≈ Jfnd atol = 1e-8
+
+    Jc = similar(J)
+    ManifoldDiff.jacobian_exp_basepoint!(M, Jc, p, X)
+    @test Jc ≈ Jfnd atol = 1e-8
+
+    Jfnd = finitedifferences_jacobian(M, TangentSpace(M, p), q -> log(M, p, q), q)
+
+    J = ManifoldDiff.jacobian_log_argument(M, p, q)
+    @test J ≈ Jfnd atol = 1e-8
+
+    Jc = similar(J)
+    ManifoldDiff.jacobian_log_argument!(M, Jc, p, q)
+    @test Jc ≈ Jfnd atol = 1e-8
+
+    Jref = finitedifferences_jacobian(
+        M,
+        TangentSpace(M, p),
+        r -> parallel_transport_to(M, r, log(M, r, q), p),
+        p,
+    )
+    J = ManifoldDiff.jacobian_log_basepoint(M, p, q)
+    @test J ≈ Jref atol = 1e-8
+
+    Jc = similar(J)
+    ManifoldDiff.jacobian_log_basepoint!(M, Jc, p, q)
+    @test Jc ≈ Jref atol = 1e-8
+end