JuliaManifolds / Manifolds.jl

Manifolds.jl provides a library of manifolds aiming for an easy-to-use and fast implementation.
https://juliamanifolds.github.io/Manifolds.jl
MIT License
368 stars 52 forks source link

Implementing a custom metric on Euclidean space #632

Closed arielgoodwin closed 1 year ago

arielgoodwin commented 1 year ago

Hi, I'm trying to implement a custom metric on $\mathbb{R}^n$. My goal is to test out the built-in ODE geodesic solver when I fully specify the metric as a matrix function (e.g., g= local_metric(M, p) where $p \in \mathbb{R}^n$, $M = \mathbb{R}^n$ and local_metric is implemented by me and returns an $n \times n$ matrix). From the documentation it seems like the types and functionality for this exist, but I could not figure out how to actually implement it. Do you have any further information or examples I could look it?

kellertuer commented 1 year ago

Hi, this seems to fit better to Manifolds.jl (I will move it there after posting this).

Cool, that you want to do that! I think you would also need to define inner. for your manifold. Luckily, I just recently did that for an example over on ManoptExamples.jl. Here https://github.com/JuliaManifolds/ManoptExamples.jl/blob/a7f2165e2e132807953ae122756cd05308d48fe2/src/objectives/Rosenbrock.jl#L108 I define a new struct, and you can see a few more functions I defined, especially for optimization I also needed exp and log.

You can also see that you would just use M=MetricManifold(Euclidean(n), YourMetricType()) to “switch” to your metric, see the rendered example https://juliamanifolds.github.io/ManoptExamples.jl/stable/examples/Difference-of-Convex-Rosenbrock/ how it is used (especially the Riemannian Gradient Descent section).

Let me know whether you have further questions :)

arielgoodwin commented 1 year ago

This is a very nice and helpful example, thank you for sharing! It is actually not unlike what I'm working on, where I am trying to test a metric derived from the graph of a function and using it for optimization. Great timing, I will try to work from this today and close the issue thereafter unless more questions arise.

kellertuer commented 1 year ago

If you end up with a nice small example, we could also consider adding it here in Manifolds.jl as a small tutorial (“How to define your own metric”).

arielgoodwin commented 1 year ago

I think I have the MetricManifold established as in your example, defining something like M = MetricManifold($\mathbb{R}^n$,my_metric) where my_metric is defined beforehand. I've also implemented inner and local_metric. How then do I call the exponential map in such a way that it just tries to evaluate numerically?

As I understand there is a numerical solver that it can fall back to when you don't implement it explicitly. What would that function call look like, e.g.: exp(M, p, v)? When I do this, I receive DomainError with ExponentialRetraction(): You can not use the exponential map as an inner method to solve the ode for the exponential map.

kellertuer commented 1 year ago

Yes, in principle, when you have loaded NLSolve you could run the ode-solver fallback we have.

Still, said solver needs a way to “walk around” on the manifold, i.e. a retraction. You could take the exponential map from the original manifold (i.e. + ;) ). Hm, I have to think a bit. The easiest I probably to use, see https://juliamanifolds.github.io/ManifoldsBase.jl/stable/retractions.html#ManifoldsBase.NLSolveInverseRetraction (where the default is default_retraction_method which falls back to exp – yes that is not the most clever default, but we do not have a better one.

So here's a trick that should work.

You define a https://juliamanifolds.github.io/ManifoldsBase.jl/stable/retractions.html#ManifoldsBase.ApproximateRetraction

retract!(::YourMetricManifold, q, p, X, ::ApproximateRetraction) = (q .= p+X)

and set the default retraction to this one

default_retraction_method(::YourMetricManifold) = ApproximateRetration()

then the ode solver uses the + you just defined as an approximation for the ODE solver (I think, just typed not tested ;) ). Remember to import both functions before you overwrite them.

arielgoodwin commented 1 year ago

So the solver, to approximate the exponential map, needs to move between points on the manifold as part of its procedure. Seems sensible, and since the underlying manifold is really R^n we can just feed it the most natural retraction +. I import Manifolds.retract!, Manifolds.default_retraction_method

Now we need to define this + retraction so that it can be used by the solver. You suggest I define an approximate retraction, so I initialize: struct approx_ret <: AbstractInverseRetractionMethod end

Then I redefine this function per your suggestion: retract!(M, q, p, X, approx_ret) = (q .= p + X)
Here M is MetricManifold($\mathbb{R}^n$,my_metric).

Similarly I write: default_retraction_method(M) = approx_ret()

How then should I call the numerical exponential map? Sorry for all the questions, I think we're nearing the resolution. If it's helpful, I've attached the essential elements of my code to give you an idea of how I've implemented what we've talked about.

# imports
using LinearAlgebra, Manifolds, NLsolve
import Manifolds.retract!, Manifolds.default_retraction_method

# dimension parameter
n = 2

# helper function to compute inner/local_metric    
function differential(x,y)
    n = size(x)[1]
    return [Diagonal(cos.(x)) -Diagonal(y)*Diagonal(sin.(x));
    Diagonal(sin.(x)) Diagonal(y)*Diagonal(cos.(x))]  
end

# Manifolds.jl technical details
struct my_metric <: AbstractMetric end

struct approx_ret <: AbstractInverseRetractionMethod end

function inner(::MetricManifold{ℝ,Euclidean{Tuple{n},ℝ},my_metric}, p, X, Y)
    return transpose(X) * local_metric(M, p) * Y
end

function local_metric(::MetricManifold{ℝ,Euclidean{Tuple{n},ℝ},my_metric}, p)
    n = div(size(p)[1],2)
    x = p[1:n]
    y = p[n+1:2n]
    D = differential(x,y)
    return I + D'*D
end

retract!(M, q, p, X, approx_ret) = (q .= p+X) 
default_retraction_method(M) = approx_ret()

E = ℝ^n
M = MetricManifold(E,my_metric())

local_metric(M,[3;-10])

q = zeros(n)
retract!(M, q, [3;-10], [1;1], approx_ret)
kellertuer commented 1 year ago

HM, until now I just always assumed to know/anticipated what you wanted to do – now I am not sure?

Well, first things first, this line

retract!(M, q, p, X, approx_ret) = (q .= p+X) 

does not do what you think it does. None of the parameters it typed, so this introduces a retraction on any manifold (actually works for _any type Mcan have (a string, a dateable, anything) and the last parameter is also of type any (not your retraction, you could pass a string there). But now that I know more about your case, this

retract!(::MetricManifold{ℝ,Euclidean{Tuple{n},ℝ}, q, p, X, ::approx_ret) = (q .= p+X)

would be the right one – in naming, I would use capital letters for types (and maybe not abbreviations), but that is just style

but when you call it you have to instantiate the type so that is

retract!(M, q, [3;-10], [1;1], approx_ret())

but you also get the allocating one (it creates the memory q itself) for free, so

retract(M, [3;-10], [1;1], approx_ret())

also just works :)

Now – why is the retraction method you declare an inverse retraction method? The correct one would be

struct approx_ret <: AbstractRetractionMethod end

and I noticed my default retraction idea was a bit vague we actually have to define

default_retraction_method(::MetricManifold{ℝ,Euclidean{Tuple{n},ℝ},my_metric}, ::Type) = approx_ret()

and then – well we still need OrdinaryDiffEq to actually solve the ODE ;)

using OrdinaryDiffEq
exp(M, [3.0; 10.0], [1.0, 1.0])

and note that I switched to floats, I think you do not want to compute that in integer arithmetics?

...oh In theory – in practice it does not yet “find” your local metric implementation – you forgot to import that I think? Check also its signature again (sorry I am running out of time, already sitting here 30 minutes).

mateuszbaran commented 1 year ago

My goal is to test out the built-in ODE geodesic solver when I fully specify the metric as a matrix function

Manifolds.jl currently has two different ODE solvers for exponential map. Both actually work for any manifold with an affine connection, not necessarily a Riemannian one. The older one, less tested, uses christoffel_symbols_second. The newer one, which even supports switching charts, uses affine_connection. Coming from inner to either christoffel_symbol_second or affine_connection isn't automated yet IIRC, and it would be ideal if you could derive that by hand and implement one of these functions. For a custom metric on Euclidean space either approach should work but I generally prefer the new one. You can take a look at this: https://github.com/JuliaManifolds/Manifolds.jl/blob/master/tutorials/working-in-charts.jl (just assume you have only one chart, everything should work with it) and https://github.com/JuliaManifolds/Manifolds.jl/blob/master/src/manifolds/EmbeddedTorus.jl for an example of affine_connection.