A bit better public api

[ferrine]

API change

The purpose of this PR is to allow users to use the library in general setting. The current API does not provide enough flexibility in terms of such kind of operations:

• vector transport from x to y
• retraction of order k in Manifold.retr
• retraction+vector transport approximating retraction at given order

These simple features require workarounds in development if we choose the way to maintain the very early API. Therefore I propose to change methods signatures a bit (the change in possible old code is mechanical) and introduce new optional arguments.

Breaking changes

Retraction

-    def retr(self, x, u, t):
+    def retr(self, x, u, t=1., order=None):

Vector transport

-    def transp(self, x, u, t, v, *more):
+    def transp(self, x, v, *more, u=None, t=1., y=None, order=None):

Retraction + vector transport

-    def retr_transp(self, x, u, t, v, *more):
+    def retr_transp(self, x, v, *more, u, t=1.0, order=None):

order is an order of retraction approximation, by default uses the simplest that is usually a first order approximation. Possible choices depend on a concrete manifold and -1 stays for exponential map

exponential map is assumed to be a separate method

Migration

Vector transport

-u = manifold.transp(u, 1.0, v)
+u = manifold.transp(v, u=u, t=1.0)

Retraction + vector transport

-p, v = manifold.retr_transp(p, u, 1.0, v)
+p, v = manifold.retr_transp(p, v, u=u, t=1.0)

New methods

Some new methods are introduced in the base Manifold class:

def expmap(self, x, u, t=1.0):

This method returns best possible approximation (of maximum order) for retraction map. Exponential map is assumed to be ideal and has order -1.

Same holds for exponential map + vector transport

def expmap_transp(self, x, v, *more, u, t=1.0):

It will be possible to specify the default order on per manifold basis with the following method

def set_default_order(self, order):

Developer API changes

Base class

From now I propose to use a metaclass that tracks and registers retractions. Right after a class creation, it filters dir(cls) and looks for special declared methods. If a method is not implemented, then it should be geoopt.base.not_implemented. geoopt.base.not_implemented is just a placeholder for a function that raises not implemented error.

Special private functions contain the following:

• r"^_retr(\d+)?$" for retraction of a given order (if int postfix provided)
• r"^_retr(\d+)?_transp$" for retraction and transport
• r"^_expmap$" for exponential map (retraction with order -1)
• r"^_expmap_transp$" for exponential map + vector transport (retraction with order -1)
• r"^_transp_follow(\d+)?$" vector transport that uses direction + retraction rather that the final point
• r"^_transp_follow_expmap$" vector transport that uses direction + exponential map rather that the final point (retraction+transport with order -1)

After this all is registered in MethodDict (default dict with not_implemented as missing value)

retractoins = MethodDict()
retractoins_transport = MethodDict()
transports_follow = MethodDict()

With this dict it comes possible to define generic dispatch methods for different orders of approximations like this:

    def retr(self, x, u, t=1.0, order=None):
        t = self.broadcast_scalar(t)
        return self._retr_funcs[order](self, x, u, t)

As you see, we avoid weird code that makes use of if or getattr with handling exceptions.

Exponential map is dispatched in the same way

    def expmap(self, x, u, t=1.0):
        t = self.broadcast_scalar(t)
        return self._retr_funcs[-1](self, x, u, t)

To make retr_transp work in case of _transp2y is much more efficient than _transp_follow there is a class attribute _retr_transp_default_preference to indicate this. The attribute should be present in the class definition if differs from default provided in Manifold

Deprecations

methods like

def _transp_one(self, x, u, t, v):
def _transp_many(self, x, u, t, *vs):

are deprecated in favour of their unified alternatives

def _transp_follow(self, x, v, *more, u, t):
def _transp2y(self, x, v, *more, y):
# and others

transp2y method appears to be very useful sometimes and allows public interface to have optional y in transp

[ferrine]

This PR helps with #33