best_fit_ellipsoid.py

# # ls_ellipsoid and polyToParams3D functions are taken and adapted from
# http://www.juddzone.com/ALGORITHMS/least_squares_3D_ellipsoid.html
# # FindBestFittingEllipsoid3D and GetConvexHull3D are adapted from a discussion at
# https://stackoverflow.com/questions/58501545/python-fit-3d-ellipsoid-oblate-prolate-to-3d-points


from scipy.spatial import ConvexHull
import numpy as np_
from numpy.linalg import eig, inv

from brick.general.type import site_h, array_t
from brick.component.soma import soma_t


def ls_ellipsoid(xx: array_t, yy: array_t, zz: array_t) -> array_t:  # finds the ellipsoid
    """
    least squares fit to a 3D-ellipsoid
     Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz  = 1

    Note that sometimes it is expressed as a solution to
     Ax^2 + By^2 + Cz^2 + 2Dxy + 2Exz + 2Fyz + 2Gx + 2Hy + 2Iz  = 1
    where the last six terms have a factor of 2 in them
    This is in anticipation of forming a matrix with the polynomial coefficients.
    Those terms with factors of 2 are all off diagonal elements.  These contribute
    two terms when multiplied out (symmetric) so would need to be divided by two
    """

    # change xx from vector of length N to Nx1 matrix so we can use hstack
    x = xx[:, np_.newaxis]
    y = yy[:, np_.newaxis]
    z = zz[:, np_.newaxis]

    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz = 1
    J = np_.hstack((x * x, y * y, z * z, x * y, x * z, y * z, x, y, z))
    K = np_.ones_like(x)  # column of ones

    # np.hstack performs a loop over all samples and creates
    # a row in J for each x,y,z sample:
    # J[ix,0] = x[ix]*x[ix]
    # J[ix,1] = y[ix]*y[ix]
    # etc.

    JT = J.transpose()
    JTJ = np_.dot(JT, J)
    InvJTJ = np_.linalg.inv(JTJ);
    ABC = np_.dot(InvJTJ, np_.dot(JT, K))

    # Rearrange, move the 1 to the other side
    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz - 1 = 0
    #    or
    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz + J = 0
    #  where J = -1
    ellipsoid_coef = np_.append(ABC, -1)

    return (ellipsoid_coef)


def polyToParams3D(ellipsoid_coef: array_t, spherical_coordinates: bool = False, printMe: bool = False) -> tuple:
    """
    gets 3D parameters of an ellipsoid.
    convert the polynomial form of the 3D-ellipsoid to parameters
    center, axes, and transformation matrix
    vec is the vector whose elements are the polynomial
    coefficients A..J
    returns (center, axes, rotation matrix)

    Algebraic form: X.T * Amat * X --> polynomial form
    """

    if printMe: print('\npolynomial\n', ellipsoid_coef)

    Amat = np_.array(
        [
            [ellipsoid_coef[0], ellipsoid_coef[3] / 2.0, ellipsoid_coef[4] / 2.0, ellipsoid_coef[6] / 2.0],
            [ellipsoid_coef[3] / 2.0, ellipsoid_coef[1], ellipsoid_coef[5] / 2.0, ellipsoid_coef[7] / 2.0],
            [ellipsoid_coef[4] / 2.0, ellipsoid_coef[5] / 2.0, ellipsoid_coef[2], ellipsoid_coef[8] / 2.0],
            [ellipsoid_coef[6] / 2.0, ellipsoid_coef[7] / 2.0, ellipsoid_coef[8] / 2.0, ellipsoid_coef[9]]
        ])

    if printMe: print('\nAlgebraic form of polynomial\n', Amat)

    # See B.Bartoni, Preprint SMU-HEP-10-14 Multi-dimensional Ellipsoidal Fitting
    # equation 20 for the following method for finding the center
    A3 = Amat[0:3, 0:3]
    A3inv = inv(A3)
    ofs = ellipsoid_coef[6:9] / 2.0
    center = -np_.dot(A3inv, ofs)
    if printMe: print('\nCenter at:', center)

    # Center the ellipsoid at the origin
    Tofs = np_.eye(4)
    Tofs[3, 0:3] = center
    R = np_.dot(Tofs, np_.dot(Amat, Tofs.T))
    if printMe: print('\nAlgebraic form translated to center\n', R, '\n')

    R3 = R[0:3, 0:3]
    s1 = -R[3, 3]
    R3S = R3 / s1

    (el, ec) = eig(R3S)
    print(el, ec)

    # Sorting in croissant order of the eigenvalues and eigenvectors
    sorting_vec = np_.argsort(el)
    el2 = el[sorting_vec]
    ec2 = ec[:, sorting_vec]

    # Calculating cartesian axes
    recip2 = 1.0 / np_.abs(el2)
    axes2 = np_.sqrt(recip2)
    if printMe: print('\nAxes are\n', axes2, '\n')

    # Calculating the orientation
    inve = inv(ec2)  # inverse is actually the transpose here
    if printMe: print('\nRotation matrix\n', inve)

    if spherical_coordinates:
        # Calculating spherical axes
        recip = 1.0 / np_.abs(el)
        axes = np_.sqrt(recip)
        r = np_.sqrt(sum(axe**2 for axe in axes))
        phi = np_.arctan(axes[1]/axes[0])
        theta = np_.arctan(axes[2]/r)
        spherical_axes = (r, theta, phi)

        if printMe: print('\nAxes in spherical coord are\n', spherical_axes, '\n')

        return (center, axes2, inve, spherical_axes)

    else:
        return (center, axes2, inve)


def GetConvexHull3D(soma_sites: site_h) -> tuple:
    #
    three_D = np_.stack((soma_sites[0], soma_sites[1], soma_sites[2]), axis=-1)
    hull_three_D = ConvexHull(three_D)
    volume_of_CH = hull_three_D.volume

    len_hull = len(hull_three_D.vertices)
    hull = np_.zeros((len_hull, 3))

    for i in range(len(hull_three_D.vertices)):
        hull[i] = three_D[hull_three_D.vertices[i]]

    convex_hull = np_.transpose(hull)

    return convex_hull, volume_of_CH


def FindBestFittingEllipsoid3D(soma: soma_t) -> tuple:
    """
    Find the best fitting ellipsoid for the data points in 3D based on their convex hull.
    Return Tuple[ellipsoid coefficients, ellipsoid center, ellipsoid axes, ellipsoid orientation]
    """

    # get convex hull
    convex_hull = GetConvexHull3D(soma.sites)[0]

    # fit ellipsoid on the convex hull
    # # get ellipsoid polynomial coefficients
    ellipsoid_coef = ls_ellipsoid(convex_hull[0], convex_hull[1], convex_hull[2])
    # # get ellipsoid 3D parameters
    center, axes, orientation = polyToParams3D(ellipsoid_coef, False)

    return ellipsoid_coef, center, axes, orientation