effectiveFisher.py

# Copyright (C) 2013  Evan Ochsner
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""
Module of routines to compute an effective Fisher matrix and related utilities,
such as finding a region of interest and laying out a grid over it
"""

from lalinference.rapid_pe import lalsimutils as lsu
import numpy as np
from scipy.optimize import leastsq, brentq
from scipy.linalg import eig, inv

__author__ = "Evan Ochsner <evano@gravity.phys.uwm.edu>"

### This is the original effective Fisher ###
# def effectiveFisher(residual_func, *flat_grids):
#     """
#     Fit a quadratic to the ambiguity function tabulated on a grid.
#     Inputs:
#         - a pointer to a function to compute residuals, e.g.
#           z(x1, ..., xN) - fit
#           for N-dimensions, this is called 'residualsNd'
#         - N+1 flat arrays of length K. N arrays for each on N parameters,
#           plus 1 array of values of the overlap
#     Returns:
#         - flat array of upper-triangular elements of the effective Fisher matrix
# 
#     Example:
#     x1s = [x1_1, ..., x1_K]
#     x2s = [x2_1, ..., x2_K]
#     ...
#     xNs = [xN_1, ..., xN_K]
# 
#     gamma = effectiveFisher(residualsNd, x1s, x2s, ..., xNs, rhos)
#     gamma
#         [g_11, g_12, ..., g_1N, g_22, ..., g_2N, g_33, ..., g_3N, ..., g_NN]
#     """
#     x0 = np.ones(len(flat_grids))
#     fitgamma = leastsq(residual_func, x0=x0, args=tuple(flat_grids))
#     return fitgamma[0]

#### CHECK ####
def effectiveFisher(residual_func, *flat_grids):
    """
    Fit a quadratic to the ambiguity function tabulated on a grid.
    Inputs:
        - a pointer to a function to compute residuals, e.g.
          z(x1, ..., xN) - fit
          for N-dimensions, this is called 'residualsNd'
        - N+1 flat arrays of length K. N arrays for each on N parameters,
          plus 1 array of values of the overlap
    Returns:
        - flat array of upper-triangular elements of the effective Fisher matrix

    Example:
    x1s = [x1_1, ..., x1_K]
    x2s = [x2_1, ..., x2_K]
    ...
    xNs = [xN_1, ..., xN_K]

    gamma = effectiveFisher(residualsNd, x1s, x2s, ..., xNs, rhos)
    gamma
        [g_11, g_12, ..., g_1N, g_22, ..., g_2N, g_33, ..., g_3N, ..., g_NN]
    """
    x0 = np.ones(len(flat_grids) + 2)
    fitgamma = leastsq(residual_func, x0=x0, args=tuple(flat_grids))
    return fitgamma[0]
#### CHECK ####

def evaluate_ip_on_grid(hfSIG, P, IP, param_names, grid):
    """
    Evaluate IP.ip(hsig, htmplt) everywhere on a multidimensional grid
    """
    Nparams = len(param_names)
    Npts = len(grid)
    assert len(grid[0])==Nparams
    return np.array([update_params_ip(hfSIG, P, IP, param_names, grid[i])
            for i in xrange(Npts)])


def update_params_ip(hfSIG, P, IP, param_names, vals):
    """
    Update the values of 1 or more member of P, recompute norm_hoff(P),
    and return IP.ip(hfSIG, norm_hoff(P))

    Inputs:
        - hfSIG: A COMPLEX16FrequencySeries of a fixed, unchanging signal
        - P: A ChooseWaveformParams object describing a varying template
        - IP: An InnerProduct object
        - param_names: An array of strings of parameters to be updated.
            e.g. [ 'm1', 'm2', 'incl' ]
        - vals: update P to have these parameter values. Must have as many
            vals as length of param_names, ordered the same way

    Outputs:
        - A COMPLEX16FrequencySeries, same as norm_hoff(P, IP)
    """
    hfTMPLT = update_params_norm_hoff(P, IP, param_names, vals)
    if IP.full_output == True:
        rho, rhoSeries, rhoIdx, rhoArg = IP.ip(hfSIG, hfTMPLT)
        return rho
    else:
        return IP.ip(hfSIG, hfTMPLT)


def update_params_norm_hoff(P, IP, param_names, vals, verbose=False):
    """
    Update the values of 1 or more member of P and recompute norm_hoff(P).

    Inputs:
        - P: A ChooseWaveformParams object
        - IP: An InnerProduct object
        - param_names: An array of strings of parameters to be updated.
            e.g. [ 'm1', 'm2', 'incl' ]
        - vals: update P to have these parameter values. Must be array-like
            with same length as param_names, ordered the same way
    Outputs:
        - A COMPLEX16FrequencySeries, same as norm_hoff(P, IP)
    """
    special_params = []
    special_vals = []
    assert len(param_names)==len(vals)
    for i, val in enumerate(vals):
        if hasattr(P, param_names[i]): # Update an attribute of P...
            setattr(P, param_names[i], val)
        else: # Either an incorrect param name, or a special case...
            special_params.append(param_names[i])
            special_vals.append(val)

    # Check allowed special cases of params not in P, e.g. Mc and eta
    if special_params==['Mc','eta']:
        m1, m2 = lsu.m1m2(special_vals[0],
                lsu.sanitize_eta(special_vals[1])) # m1,m2 = m1m2(Mc,eta)
        setattr(P, 'm1', m1)
        setattr(P, 'm2', m2)
    elif special_params==['eta','Mc']:
        m1, m2 = lsu.m1m2(lsu.sanitize_eta(special_vals[1]), special_vals[0])
        setattr(P, 'm1', m1)
        setattr(P, 'm2', m2)
    elif special_params==['Mc']:
        eta = lsu.sanitize_eta(lsu.symRatio(P.m1, P.m2))
        m1, m2 = lsu.m1m2(special_vals[0], eta)
        setattr(P, 'm1', m1)
        setattr(P, 'm2', m2)
    elif special_params==['eta']:
        Mc = lsu.mchirp(P.m1, P.m2)
        m1, m2 = lsu.m1m2(Mc, lsu.sanitize_eta(special_vals[0]))
        setattr(P, 'm1', m1)
        setattr(P, 'm2', m2)
    elif special_params != []:
        print special_params
        raise Exception

    if verbose==True: # for debugging - make sure params change properly
        P.print_params()
    return lsu.norm_hoff(P, IP)



def find_effective_Fisher_region(P, IP, target_match, param_names,param_bounds):
    """
    Example Usage:
        find_effective_Fisher_region(P, IP, 0.9, ['Mc', 'eta'], [[mchirp(P.m1,P.m2)-lal.MSUN_SI,mchirp(P.m1,P.m2)+lal.MSUN_SI], [0.05, 0.25]])
    Arguments:
        - P: a ChooseWaveformParams object describing a target signal
        - IP: An inner product class to compute overlaps.
                Should have deltaF and length consistent with P
        - target_match: find parameter variation where overlap is target_match.
                Should be a real number in [0,1]
        - param_names: array of string names for members of P to vary.
                Should have length N for N params to be varied
                e.g. ['Mc', 'eta']
        - param_bounds: array of max variations of each param in param_names
                Should be an Nx2 array for N params to be varied
    
    Returns:
        Array of boundaries of a hypercube meant to encompass region where
                match is >= target_match.
                e.g. [ [3.12,3.16] , [0.12, 0.18] ]

    N.B. Only checks variations along parameter axes. If params are correlated,
    may get better matches off-axis, and the hypercube won't fully encompass
    region where target_match is achieved. Therefore, allow a generous
    safety factor in your value of 'target_match'.
    """
    TOL = 1.e-6 # Don't need to be very precise for this...
    Nparams = len(param_names)
    assert len(param_bounds) == Nparams
    param_cube = []
    hfSIG = lsu.norm_hoff(P, IP)
    for i, param in enumerate(param_names):
        PT = P.copy()
        if param=='Mc':
            param_peak = lsu.mchirp(P.m1, P.m2)
        elif param=='eta':
            param_peak = lsu.symRatio(P.m1, P.m2)
        else:
            param_peak = getattr(P, param)
        func = lambda x: update_params_ip(hfSIG, PT, IP, [param], [x]) - target_match
        try:
            min_param = brentq(func, param_peak, param_bounds[i][0], xtol=TOL)
        except ValueError:
            print "\nWarning! Value", param_bounds[i][0], "of", param,\
                    "did not bound target match", target_match, ". Using",\
                    param_bounds[i][0], "as the lower bound of", param,\
                    "range for the effective Fisher region.\n"
            min_param = param_bounds[i][0]
        try:
            max_param = brentq(func, param_peak, param_bounds[i][1], xtol=TOL)
        except ValueError:
            print "\nWarning! Value", param_bounds[i][1], "of", param,\
                    "did not bound target match", target_match, ". Using",\
                    param_bounds[i][1], "as the upper bound of", param,\
                    "range for the effective Fisher region.\n"
            max_param = param_bounds[i][1]
        param_cube.append( [min_param, max_param] )

    return param_cube

#
# Routines to make various types of grids for arbitrary dimension
#
def make_regular_1d_grids(param_ranges, pts_per_dim):
    """
    Inputs: 
        - param_ranges is an array of parameter bounds, e.g.:
        [ [p1_min, p1_max], [p2_min, p2_max], ..., [pN_min, pN_max] ]
        - pts_per_dim is either:
            a) an integer - use that many pts for every parameter
            b) an array of integers of same length as param_ranges, e.g.
                [ N1, N2, ..., NN ]
                the n-th entry is the number of pts for the n-th parameter

    Outputs:
        outputs N separate 1d arrays of evenly spaced values of that parameter,
        where N = len(param_ranges)
    """
    Nparams = len(param_ranges)
    assert len(pts_per_dim)
    grid1d = []
    for i in range(Nparams):
        MIN = param_ranges[i][0]
        MAX = param_ranges[i][1]
        STEP = (MAX-MIN)/(pts_per_dim[i]-1)
        EPS = STEP/100.
        grid1d.append( np.arange(MIN,MAX+EPS,STEP) )

    return tuple(grid1d)

def multi_dim_meshgrid(*arrs):
    """
    Version of np.meshgrid generalized to arbitrary number of dimensions.
    Taken from: http://stackoverflow.com/questions/1827489/numpy-meshgrid-in-3d
    """
    arrs = tuple(reversed(arrs))
    lens = map(len, arrs)
    dim = len(arrs)

    sz = 1
    for s in lens:
        sz*=s

    ans = []    
    for i, arr in enumerate(arrs):
        slc = [1]*dim
        slc[i] = lens[i]
        arr2 = np.asarray(arr).reshape(slc)
        for j, sz in enumerate(lens):
            if j!=i:
                arr2 = arr2.repeat(sz, axis=j) 
        ans.append(arr2)

    #return tuple(ans)
    return tuple(ans[::-1])

def multi_dim_flatgrid(*arrs):
    """
    Creates flattened versions of meshgrids.
    Returns a tuple of arrays of values of individual parameters
    at each point in a grid, returned in a flat array structure.

    e.g.
    x = [1,3,5]
    y = [2,4,6]
    X, Y = multi_dim_flatgrid(x, y)
    returns:
    X
        [1,1,1,3,3,3,5,5,5]
    Y
        [2,4,6,2,4,6,2,4,6]
    """
    outarrs = multi_dim_meshgrid(*arrs)
    return tuple([ outarrs[i].flatten() for i in xrange(len(outarrs)) ])

def multi_dim_grid(*arrs):
    """
    Creates an array of values of all pts on a multidimensional grid.

    e.g.
    x = [1,3,5]
    y = [2,4,6]
    multi_dim_grid(x, y)
    returns:
    [[1,2], [1,4], [1,6],
     [3,2], [3,4], [3,6],
     [5,2], [5,4], [5,6]]
    """
    temp = multi_dim_flatgrid(*arrs)
    return np.transpose( np.array(temp) )


#
# Routines for least-squares fit
#
def residuals2d(gamma, y, x1, x2):
    g11 = gamma[0]
    g12 = gamma[1]
    g22 = gamma[2]
    return y - (1. - g11*x1*x1/2. - g12*x1*x2 - g22*x2*x2/2.)

def evalfit2d(x1, x2, gamma):
    g11 = gamma[0]
    g12 = gamma[1]
    g22 = gamma[2]
    return 1. - g11*x1*x1/2. - g12*x1*x2 - g22*x2*x2/2.

def residuals3d(gamma, y, x1, x2, x3):
    g11 = gamma[0]
    g12 = gamma[1]
    g13 = gamma[2]
    g22 = gamma[3]
    g23 = gamma[4]
    g33 = gamma[5]
    return y - (1. - g11*x1*x1/2. - g12*x1*x2 - g13*x1*x3
            - g22*x2*x2/2. - g23*x2*x3 - g33*x3*x3/2.)

def evalfit3d(x1, x2, x3, gamma):
    g11 = gamma[0]
    g12 = gamma[1]
    g13 = gamma[2]
    g22 = gamma[3]
    g23 = gamma[4]
    g33 = gamma[5]
    return 1. - g11*x1*x1/2. - g12*x1*x2 - g13*x1*x3\
            - g22*x2*x2/2. - g23*x2*x3 - g33*x3*x3/2.

def residuals4d(gamma, y, x1, x2, x3, x4):
    g11 = gamma[0]
    g12 = gamma[1]
    g13 = gamma[2]
    g14 = gamma[3]
    g22 = gamma[4]
    g23 = gamma[5]
    g24 = gamma[6]
    g33 = gamma[7]
    g34 = gamma[8]
    g44 = gamma[9]
    return y - (1. - g11*x1*x1/2. - g12*x1*x2 - g13*x1*x3 - g14*x1*x4
            - g22*x2*x2/2. - g23*x2*x3 - g24*x2*x4 - g33*x3*x3/2. - g34*x3*x4
            - g44*x4*x4/2.)

def evalfit4d(x1, x2, x3, x4, gamma):
    g11 = gamma[0]
    g12 = gamma[1]
    g13 = gamma[2]
    g14 = gamma[3]
    g22 = gamma[4]
    g23 = gamma[5]
    g24 = gamma[6]
    g33 = gamma[7]
    g34 = gamma[8]
    g44 = gamma[9]
    return 1. - g11*x1*x1/2. - g12*x1*x2 - g13*x1*x3 - g14*x1*x4\
            - g22*x2*x2/2. - g23*x2*x3 - g24*x2*x4 - g33*x3*x3/2. - g34*x3*x4\
            - g44*x4*x4/2.

def residuals5d(gamma, y, x1, x2, x3, x4, x5):
    g11 = gamma[0]
    g12 = gamma[1]
    g13 = gamma[2]
    g14 = gamma[3]
    g15 = gamma[4]
    g22 = gamma[5]
    g23 = gamma[6]
    g24 = gamma[7]
    g25 = gamma[8]
    g33 = gamma[9]
    g34 = gamma[10]
    g35 = gamma[11]
    g44 = gamma[12]
    g45 = gamma[13]
    g55 = gamma[14]
    return y - (1. - g11*x1*x1/2. - g12*x1*x2 - g13*x1*x3 - g14*x1*x4
            - g15*x1*x5 - g22*x2*x2/2. - g23*x2*x3 - g24*x2*x4 - g25*x2*x5
            - g33*x3*x3/2. - g34*x3*x4 - g35*x3*x5 - g44*x4*x4/2. - g45*x4*x5
            - g55*x5*x5/2.)

def evalfit5d(x1, x2, x3, x4, x5, gamma):
    g11 = gamma[0]
    g12 = gamma[1]
    g13 = gamma[2]
    g14 = gamma[3]
    g15 = gamma[4]
    g22 = gamma[5]
    g23 = gamma[6]
    g24 = gamma[7]
    g25 = gamma[8]
    g33 = gamma[9]
    g34 = gamma[10]
    g35 = gamma[11]
    g44 = gamma[12]
    g45 = gamma[13]
    g55 = gamma[14]
    return 1. - g11*x1*x1/2. - g12*x1*x2 - g13*x1*x3 - g14*x1*x4\
            - g15*x1*x5 - g22*x2*x2/2. - g23*x2*x3 - g24*x2*x4 - g25*x2*x5\
            - g33*x3*x3/2. - g34*x3*x4 - g35*x3*x5 - g44*x4*x4/2. - g45*x4*x5\
            - g55*x5*x5/2.


# Convenience function to return eigenvalues and eigenvectors of a matrix
def eigensystem(matrix):
    """
    Given an array-like 'matrix', returns:
        - An array of eigenvalues
        - An array of eigenvectors
        - A rotation matrix that rotates the eigenbasis
            into the original basis

    Example:
        mat = [[1,2,3],[2,4,5],[3,5,6]]
        evals, evecs, rot = eigensystem(mat)
        evals
            array([ 11.34481428+0.j,  -0.51572947+0.j,   0.17091519+0.j]
        evecs
            array([[-0.32798528, -0.59100905, -0.73697623],
                   [-0.73697623, -0.32798528,  0.59100905],
                   [ 0.59100905, -0.73697623,  0.32798528]])
        rot
            array([[-0.32798528, -0.73697623,  0.59100905],
                   [-0.59100905, -0.32798528, -0.73697623],
                   [-0.73697623,  0.59100905,  0.32798528]]))

    This allows you to translate between original and eigenbases:

        If [v1, v2, v3] are the components of a vector in eigenbasis e1, e2, e3
        Then:
            rot.dot([v1,v2,v3]) = [vx,vy,vz]
        Will give the components in the original basis x, y, z

        If [wx, wy, wz] are the components of a vector in original basis z, y, z
        Then:
            inv(rot).dot([wx,wy,wz]) = [w1,w2,w3]
        Will give the components in the eigenbasis e1, e2, e3

        inv(rot).dot(mat).dot(rot)
            array([[evals[0], 0,        0]
                   [0,        evals[1], 0]
                   [0,        0,        evals[2]]])

    Note: For symmetric input 'matrix', inv(rot) == evecs
    """
    evals, emat = eig(matrix)
    return evals, np.transpose(emat), emat

def array_to_symmetric_matrix(gamma):
    """
    Given a flat array of length N*(N+1)/2 consisting of
    the upper right triangle of a symmetric matrix,
    return an NxN numpy array of the symmetric matrix

    Example:
        gamma = [1, 2, 3, 4, 5, 6]
        array_to_symmetric_matrix(gamma)
            array([[1,2,3],
                   [2,4,5],
                   [3,5,6]])
    """
    length = len(gamma)
    if length==3: # 2x2 matrix
        g11 = gamma[0]
        g12 = gamma[1]
        g22 = gamma[2]
        return np.array([[g11,g12],[g12,g22]])
    if length==6: # 3x3 matrix
        g11 = gamma[0]
        g12 = gamma[1]
        g13 = gamma[2]
        g22 = gamma[3]
        g23 = gamma[4]
        g33 = gamma[5]
        return np.array([[g11,g12,g13],[g12,g22,g23],[g13,g23,g33]])
    if length==10: # 4x4 matrix
        g11 = gamma[0]
        g12 = gamma[1]
        g13 = gamma[2]
        g14 = gamma[3]
        g22 = gamma[4]
        g23 = gamma[5]
        g24 = gamma[6]
        g33 = gamma[7]
        g34 = gamma[8]
        g44 = gamma[9]
        return np.array([[g11,g12,g13,g14],[g12,g22,g23,g24],
            [g13,g23,g33,g34],[g14,g24,g34,g44]])
    if length==15: # 5x5 matrix
        g11 = gamma[0]
        g12 = gamma[1]
        g13 = gamma[2]
        g14 = gamma[3]
        g15 = gamma[4]
        g22 = gamma[5]
        g23 = gamma[6]
        g24 = gamma[7]
        g25 = gamma[8]
        g33 = gamma[9]
        g34 = gamma[10]
        g35 = gamma[11]
        g44 = gamma[12]
        g45 = gamma[13]
        g55 = gamma[14]
        return np.array([[g11,g12,g13,g14,g15],[g12,g22,g23,g24,g25],
            [g13,g23,g33,g34,g35],[g14,g24,g34,g44,g45],[g15,g25,g5,g45,g55]])



#
# Fill an ellipsoid uniformly in volume by points along radial spokes
#
def uniform_spoked_ellipsoid(Nrad, Nspokes, start_angles, *radii):
    """
    Return an array of pts distributed uniformly inside a
    D-dimensional ellipsoid. D is determined by the number of radii args given.

    Output:
        - cart_pts: array of pts in Cartesian coordinates
        - sph_pts: array of points in spherical coordinates
                   (r, zenith_1, .., zenith_N-2, azimuth)
    """
    D = len(radii)
    assert len(start_angles)==D-1
    if D==2:
        return uniform_spoked_ellipsoid2d(Nrad, Nspokes, start_angles[0], *radii)
    else:
        raise ValueError('Not implemented for that many dimensions')

def uniform_spoked_ellipsoid2d(Nrad, Nspokes, th0, r1, r2):
    """
    2D case of function uniform_spoked_ellipsoid
    """
    dr = 1./Nrad
    rs = np.arange(dr, 1.+dr, dr)
    dth = 2.*np.pi/Nspokes
    ths = np.arange(th0, 2.*np.pi + th0, dth)
    rrt = np.sqrt(rs)
    cart_pts = [[0.,0.]] # Put 1 pt. at origin - e.g. true parameters
    sph_pts = [[0.,0.]] # Put 1 pt. at origin - e.g. true parameters
    for r in rrt:
        for th in ths:
            x1 = r1 * r * np.cos(th)
            x2 = r2 * r * np.sin(th)
            cart_pts.append([x1, x2])
            sph_pts.append([r, th])
    return np.array(cart_pts), np.array(sph_pts)

#
# Fill an ellipsoid uniformly in volume by points along radial spokes
#
def linear_spoked_ellipsoid(Nrad, Nspokes, start_angles, *radii):
    """
    Return an array of pts distributed uniformly inside a
    D-dimensional ellipsoid. D is determined by the number of radii args given.

    Output:
        - cart_pts: array of pts in Cartesian coordinates
        - sph_pts: array of points in spherical coordinates
                   (r, zenith_1, .., zenith_N-2, azimuth)
    """
    D = len(radii)
    assert len(start_angles)==D-1
    if D==2:
        return linear_spoked_ellipsoid2d(Nrad, Nspokes, start_angles[0], *radii)
    else:
        raise ValueError('Not implemented for that many dimensions')

def linear_spoked_ellipsoid2d(Nrad, Nspokes, th0, r1, r2):
    """
    2D case of function linear_spoked_ellipsoid
    """
    dr = 1./Nrad
    rs = np.arange(dr, 1.+dr, dr)
    dth = 2.*np.pi/Nspokes
    ths = np.arange(th0, 2.*np.pi + th0, dth)
    cart_pts = [[0.,0.]] # Put 1 pt. at origin - e.g. true parameters
    sph_pts = [[0.,0.]] # Put 1 pt. at origin - e.g. true parameters
    for r in rs:
        for th in ths:
            x1 = r1 * r * np.cos(th)
            x2 = r2 * r * np.sin(th)
            cart_pts.append([x1, x2])
            sph_pts.append([r, th])
    return np.array(cart_pts), np.array(sph_pts)



#
# Functions to return points distributed randomly, uniformly inside
# an ellipsoid of arbitrary dimension
#
def uniform_random_ellipsoid(Npts, *radii):
    """
    Return an array of pts distributed randomly and uniformly inside an
    D-dimensional ellipsoid. D is determined by the number of radii args given.

    Output:
        - cart_pts: array of pts in Cartesian coordinates
        - sph_pts: array of points in spherical coordinates
                   (r, zenith_1, .., zenith_N-2, azimuth)
    """
    D = len(radii)
    if D==2:
        return uniform_random_ellipsoid2d(Npts, *radii)
    elif D==3:
        return uniform_random_ellipsoid3d(Npts, *radii)
    elif D==4:
        return uniform_random_ellipsoid4d(Npts, *radii)
    elif D==5:
        return uniform_random_ellipsoid5d(Npts, *radii)
    else:
        raise ValueError('Not implemented for that many dimensions')

def uniform_random_ellipsoid2d(Npts, r1, r2):
    """
    2D case of uniform_random_ellipsoid
    """
    r = np.random.rand(Npts)
    th = np.random.rand(Npts) * 2.*np.pi
    rrt = np.sqrt(r)
    x1 = r1 * rrt * np.cos(th)
    x2 = r2 * rrt * np.sin(th)
    cart_pts = np.transpose((x1,x2))
    sph_pts = np.transpose((rrt,th))
    origin = np.array([[0.,0.]]) # Always put a pt at ellipse center
    cart_pts = np.append(origin, cart_pts, axis=0)
    sph_pts = np.append(origin, sph_pts, axis=0)
    return cart_pts, sph_pts

def uniform_random_ellipsoid3d(Npts, r1, r2, r3):
    """
    3D case of uniform_random_ellipsoid
    """
    r = np.random.rand(Npts)
    ph = np.random.rand(Npts) * 2.*np.pi
    costh = np.random.rand(Npts)*2.-1.
    sinth = np.sqrt(1.-costh*costh)
    th = np.arccos(costh)
    rrt = r**(1./3.)
    x1 = r1 * rrt * sinth * np.cos(ph)
    x2 = r2 * rrt * sinth * np.sin(ph)
    x3 = r3 * rrt * costh
    cart_pts = np.transpose((x1,x2,x3))
    cart_pts = np.transpose((rrt,th,ph))
    #### CHECK ####
    sph_pts = np.transpose((rrt,th,ph))
    #### CHECK ####
    origin = np.array([[0.,0.,0.]]) # Always put a pt at ellipse center
    cart_pts = np.append(origin, cart_pts, axis=0)
    sph_pts = np.append(origin, sph_pts, axis=0)
    return cart_pts, sph_pts

def uniform_random_ellipsoid4d(Npts, r1, r2, r3, r4):
    """
    4D case of uniform_random_ellipsoid
    """
    r = np.random.rand(Npts)
    ph = np.random.rand(Npts) * 2.*np.pi
    costh1 = np.random.rand(Npts)*2.-1.
    costh2 = np.random.rand(Npts)*2.-1.
    sinth1 = np.sqrt(1.-costh1*costh1)
    sinth2 = np.sqrt(1.-costh2*costh2)
    th1 = np.arccos(costh1)
    th2 = np.arccos(costh2)
    rrt = r**(1./4.)
    x1 = r1 * rrt * sinth1 * sinth2 * np.cos(ph)
    x2 = r2 * rrt * sinth1 * sinth2 * np.sin(ph)
    x3 = r3 * rrt * sinth1 * costh2
    x4 = r4 * rrt * costh1
    cart_pts = np.transpose((x1,x2,x3,x4))
    sph_pts = np.transpose((rrt,th1,th2,ph))
    origin = np.array([[0.,0.,0.,0.]]) # Always put a pt at ellipse center
    cart_pts = np.append(origin, cart_pts, axis=0)
    sph_pts = np.append(origin, sph_pts, axis=0)
    return cart_pts, sph_pts

def uniform_random_ellipsoid5d(Npts, r1, r2, r3, r4, r5):
    """
    5D case of uniform_random_ellipsoid
    """
    r = np.random.rand(Npts)
    ph = np.random.rand(Npts) * 2.*np.pi
    costh1 = np.random.rand(Npts)*2.-1.
    costh2 = np.random.rand(Npts)*2.-1.
    costh3 = np.random.rand(Npts)*2.-1.
    sinth1 = np.sqrt(1.-costh1*costh1)
    sinth2 = np.sqrt(1.-costh2*costh2)
    sinth3 = np.sqrt(1.-costh3*costh3)
    th1 = np.arccos(costh1)
    th2 = np.arccos(costh2)
    th3 = np.arccos(costh3)
    rrt = r**(1./5.)
    x1 = r1 * rrt * sinth1 * sinth2 * sinth3 * np.cos(ph)
    x2 = r2 * rrt * sinth1 * sinth2 * sinth3 * np.sin(ph)
    x3 = r3 * rrt * sinth1 * sinth2 * costh3
    x4 = r4 * rrt * sinth1 * costh2
    x5 = r5 * rrt * costh1
    cart_pts = np.transpose((x1,x2,x3,x4,x5))
    sph_pts = np.transpose((rrt,th1,th2,th3,ph))
    origin = np.array([[0.,0.,0.,0.,0.]]) # Always put a pt at ellipse center
    cart_pts = np.append(origin, cart_pts, axis=0)
    sph_pts = np.append(origin, sph_pts, axis=0)
    return cart_pts, sph_pts