Generate an ellipse through the MVEE method

get_ellipse.py

import torch
import torch.linalg as la


# http://stackoverflow.com/questions/14016898/port-matlab-bounding-ellipsoid-code-to-python
# http://stackoverflow.com/questions/1768197/bounding-ellipse/1768440#1768440
# https://minillinim.github.io/GroopM/dev_docs/groopm.ellipsoid-pysrc.html


def mvee(points, tol=0.01):
    """
    Finds the ellipse equation in "center form"
    (x-c).T * A * (x-c) = 1
    """
    device = torch.device('cuda')
    dot = lambda foo, bar : torch.tensordot(foo, bar, dims=1)
    N, d = points.shape
    Q = torch.column_stack((points, torch.ones(N, device=device))).T # Q.shape = (d+1, N)
    err = tol+1.0
    u = torch.ones(N, device = device)/N # u.shape = (N,)
    while err > tol:
        # assert u.sum() == 1 # invariant
        # dot(Q, torch.diag(u)) = Q * u[None, :]
        X = dot(Q * u[None, :], Q.T) # shapes: (((d+1, N), (N, N)) , (N, d+1)) = (d+1, d+1)
        # diag(dot(Y, Q)) = (Y*Q.T).sum(dim=1), where Y = dot(Q.T, la.inv(X))
        M = torch.sum(dot(Q.T, la.inv(X))*Q.T, dim = 1) # (((N, d+1), (d+1, d+1)), (d+1, N)) = (N, N); M.shape = (N,)
        jdx = torch.argmax(M)
        step_size = (M[jdx]-d-1.0)/((d+1)*(M[jdx]-1.0))
        new_u = (1-step_size)*u
        new_u[jdx] += step_size
        err = la.norm(new_u-u)
        u = new_u
    c = dot(u, points)
    # dot(points.T, torch.diag(u)) = points.T * u[None, :]
    A = la.inv(dot(points.T * u[None, :], points)
               - torch.outer(c, c))/d
    return A, c