StochMan - Stochastic Manifolds made easier

StochMan (Stochastic Manifolds) is a collection of elementary algorithms for computations on random manifolds learned from finite noisy data. Each algorithm assume that the considered manifold model implement a specific set of interfaces.

Installation

For the latest release

pip install stochman

For master version with most recent changes we recommend:

git clone https://github.com/MachineLearningLifeScience/stochman
cd stochman
python setup.py install

API overview

StochMan includes a number of modules that each defines a set of functionalities for working with manifold data.

stochman.nnj: torch.nn with Jacobians

Key to working with Riemannian geometry is the ability to compute Jacobians. The Jacobian matrix contains the first order partial derivatives. stochman.nnj provides plug-in replacements for the many used torch.nn layers such as Linear, BatchNorm1d etc. and commonly used activation functions such as ReLU, Sigmoid etc. that enables fast computations of Jacobians between the input to the layer and the output.

import torch
from stochman import nnj

model = nnj.Sequential(nnj.Linear(10, 5),
                       nnj.ReLU())
x = torch.randn(100, 10)
y, J = model(x, jacobian=True)
print(y.shape) # output from model: torch.size([100, 5])
print(J.shape) # Jacobian between input and output: torch.size([100, 5, 10])

stochman.manifold: Interface for working with Riemannian manifolds

A manifold can be constructed simply by specifying its metric. The example below shows a toy example where the metric grows with the distance to the origin.

import torch
from stochman.manifold import Manifold

class MyManifold(Manifold):
    def metric(self, c, return_deriv=False):
        N, D = c.shape  # N is number of points where we evaluate the metric; D is the manifold dimension
        sq_dist_to_origin = torch.sum(c**2, dim=1, keepdim=True)  # Nx1
        G = (1 + sq_dist_to_origin).unsqueeze(-1) * torch.eye(D).repeat(N, 1, 1)  # NxDxD
        return G
        
model = MyManifold()
p0, p1 = torch.randn(1, 2), torch.randn(1, 2)
c, _ = model.connecting_geodesic(p0, p1)  # geodesic between two random points

If you manifold is embedded (e.g. an autoencoder) then you only have to provide a function for realizing the embedding (i.e. a decoder) and StochMan takes care of the rest (you, however, have to learn the autoencoder yourself).

import torch
from stochman.manifold import EmbeddedManifold

class Autoencoder(EmbeddedManifold):
    def embed(self, c, jacobian = False):
        return self.decode(c)
        
model = Autoencoder()
p0, p1 = torch.randn(1, 2), torch.randn(1, 2)
c, _ = model.connecting_geodesic(p0, p1)  # geodesic between two random points

stochman.geodesic: computing geodesics made easy!

Geodesics are energy-minimizing curves, and StochMan computes them as such. You can use the high-level Manifold interface or the more explicit one:

import torch
from stochman.geodesic import geodesic_minimizing_energy
from stochman.curves import CubicSpline

model = MyManifold()
p0, p1 = torch.randn(1, 2), torch.randn(1, 2)
curve = CubicSpline(p0, p1)
geodesic_minimizing_energy(curve, model)

stochman.curves: Simple curve objects

We often want to manipulate curves when computing geodesics. StochMan provides an implementation of cubic splines and discrete curves, both with the end-points fixed.

import torch
from stochman.curves import CubicSpline

p0, p1 = torch.randn(1, 2), torch.randn(1, 2)
curve = CubicSpline(p0, p1)

t = torch.linspace(0, 1, 50)
ct = curve(t)  # 50x2

Licence

Please observe the Apache 2.0 license that is listed in this repository.

BibTeX

If you want to cite the framework feel free to use this (but only if you loved it 😊):

@article{software:stochman,
  title={StochMan},
  author={Nicki S. Detlefsen and Alison Pouplin and Cilie W. Feldager and Cong Geng and Dimitris Kalatzis and Helene Hauschultz and Miguel González-Duque and Frederik Warburg and Marco Miani and Søren Hauberg},
  journal={GitHub. Note: https://github.com/MachineLearningLifeScience/stochman/},
  year={2021}
}