morphokinematicsdiagnostics.py

import numpy as np
from scipy import interpolate,linalg
"""
NOTE: This program was originally written in Python 2.7 by Adrien Thob, but was modified slightly to be compatible with Python 3.7.0. 
Adrien's more extensive documentation etc. for this code is included in his GitHub repository, cited in the paper.
"""

def cumsummedian(a,weights=None):
    """
    Compute the weighted median.

    Returns the median of the array elements.

    Parameters
    ----------
    a : array_like, shape (n, )
        Input array or object that can be converted to an array.
    weights : {array_like, shape (n, ), None}, optional
        Input array or object that can be converted to an array.

    Returns
    -------
    median : float

    """
    if weights is None:
        weights=np.ones(np.array(a).shape)
    A = np.array(a).astype('float')
    W = np.array(weights).astype('float')
    if not(np.product(np.isnan(A))):
        I = np.argsort(A)
        cumweight = np.hstack([0,np.cumsum(W[I])])
        X = np.hstack([0,(cumweight[:-1]+cumweight[1:])/(2*cumweight[-1]),1])
        Y = np.hstack([np.min(A),A[I],np.max(A)])
        P = interpolate.interp1d(X,Y)(0.5)
        return float(P)
    else:
        return np.nan

def center_and_unloop(XYZ,XYZ0,BoxL=np.inf):
    """
    Center and unloop the input coordinates.

    Returns the centered and unlooped coordinates.

    Parameters
    ----------
    XYZ : array_like of dtype float, shape (n, 3)
        Particles coordinates (in unit of length L) such that XYZ[:,0] = X,
        XYZ[:,1] = Y & XYZ[:,2] = Z
    XYZ0 : array_like of dtype float, shape (3, )
        Centre coordinates (in unit of length L) such that XYZ0[0] = X0,
        XYZ0[1] = Y0 & XYZ0[2] = Z0
    BoxL : float, optional
        Length of the looping cubical box. Default is infinity

    Returns
    -------
    XYZ_out : array of dtype float, shape (n, 3)
        Centered and unlooped particles coordinates (in unit of length L) such
        that XYZ[:,0] = X, XYZ[:,1] = Y & XYZ[:,2] = Z

    """
    XYZ_out = XYZ.copy()
    XYZ_out-=XYZ0
    if np.isfinite(BoxL):
        XYZ_out+=BoxL/2.
        XYZ_out%=BoxL
        XYZ_out-=BoxL/2.
    return XYZ_out

def kinematics_diagnostics(XYZ,mass,Vxyz,PBE,aperture=0.03,CoMvelocity=True):
    """
    Compute the various kinematics diagnostics.

    Returns the kinematics diagnostics for the input particles.

    Parameters
    ----------
    XYZ : array_like of dtype float, shape (n, 3)
        Particles coordinates (in unit of length L) such that XYZ[:,0] = X,
        XYZ[:,1] = Y & XYZ[:,2] = Z
    mass : array_like of dtype float, shape (n, )
        Particles masses (in unit of mass M)
    Vxyz : array_like of dtype float, shape (n, 3)
        Particles coordinates (in unit of velocity V) such that Vxyz[:,0] = Vx,
        Vxyz[:,1] = Vy & Vxyz[:,2] = Vz
    PBE : array_like of dtype float, shape (n, )
        Particles specific binding energies
    aperture : float, optional
        Aperture (in unit of length L) for the computation. Default is 0.03 L
    CoMvelocity : bool, optional
        Boolean to allow the centering of velocities by the considered particles
        centre-of-mass velocity. Default to True

    Returns
    -------
    kappa : float
        The kinetic energy fraction invested in co-rotation.
    discfrac : float
        The disc-to-total mass fraction estimated from the counter-rotating
        bulge.
    orbi : float
        The median orbital circularity of the particles values.
    vrotsig : float
        The rotation-to-dispersion ratio .
    delta : float
        The dispersion anisotropy.
    zaxis : array of dtype float, shape (3, )
        The unit vector of the momentum axis (pointing along the momentum direction).
    Momentum : float
        The momentum magnitude (in unit M.L.V).

    """
    particlesall = np.vstack([XYZ.T,mass,Vxyz.T,PBE]).T
    # Compute distances
    distancesall = np.linalg.norm(particlesall[:,:3],axis=1)
    # Restrict particles
    extract = (distancesall<aperture)
    particles = particlesall[extract].copy()
    distances = distancesall[extract].copy()
    Mass = np.sum(particles[:,3])
    if CoMvelocity:
        # Compute CoM velocty & correct
        dvVmass = np.nan_to_num(np.sum(particles[:,3][:,np.newaxis]*particles[:,4:7],axis=0)/Mass)
        particlesall[:,4:7]-=dvVmass
        particles[:,4:7]-=dvVmass
    # Compute momentum
    smomentums = np.cross(particles[:,:3],particles[:,4:7])
    momentum = np.sum(particles[:,3][:,np.newaxis]*smomentums,axis=0)
    Momentum = np.linalg.norm(momentum)
    # Compute cylindrical quantities
    zaxis = (momentum/Momentum)
    zheight = np.sum(zaxis*particles[:,:3],axis=1)
    cylposition = particles[:,:3]-zheight[:,np.newaxis]*[zaxis]
    cyldistances = np.sqrt(distances**2-zheight**2)
    smomentumz = np.sum(zaxis*smomentums,axis=1)
    vrots = smomentumz/cyldistances
    vrads = np.sum(cylposition*particles[:,4:7]/cyldistances[:,np.newaxis],axis=1)
    vheis = np.sum(zaxis*particles[:,4:7],axis=1)
    # Compute co-rotational kinetic energy fraction
    Mvrot2 = np.sum((particles[:,3]*vrots**2)[vrots>0])
    kappa = Mvrot2/np.sum(particles[:,3]*(np.linalg.norm(particles[:,4:7],axis=1))**2)
    # Compute disc-to-total ratio
    discfrac = 1-2*np.sum(particles[vrots<=0,3])/Mass
    # Compute orbital circularity
    sbindingenergy = particles[:,7]; sortE = np.argsort(sbindingenergy); unsortE = np.argsort(sortE)
    jzE = np.vstack([sbindingenergy,smomentumz]).T[sortE]
    orbital = (jzE[:,1]/np.maximum.accumulate(np.abs(jzE[:,1])))[unsortE]
    orbi = np.median(orbital)
    # Compute rotation-to-dispersion and dispersion anisotropy
    Vrot = np.abs(cumsummedian(vrots,weights=particles[:,3]))
    SigmaXY = np.sqrt(np.average(np.sum(particles[:,[3]]*np.vstack([vrads,vrots]).T**2,axis=0)/Mass))#
    SigmaO = np.sqrt(SigmaXY**2-.5*Vrot**2)
    SigmaZ = np.sqrt(np.average(vheis**2,weights=particles[:,3]))
    vrotsig = Vrot/SigmaO
    delta = 1-(SigmaZ/SigmaO)**2
    # Return
    return kappa,discfrac,orbi,vrotsig,delta,zaxis,Momentum

def morphological_diagnostics(XYZ,mass,Vxyz,aperture=0.03,CoMvelocity=True,reduced_structure=True):
    """
    Compute the morphological diagnostics through the (reduced or not) inertia tensor.

    Returns the morphological diagnostics for the input particles.

    Parameters
    ----------
    ----------
    XYZ : array_like of dtype float, shape (n, 3)
        Particles coordinates (in unit of length L) such that XYZ[:,0] = X,
        XYZ[:,1] = Y & XYZ[:,2] = Z
    mass : array_like of dtype float, shape (n, )
        Particles masses (in unit of mass M)
    Vxyz : array_like of dtype float, shape (n, 3)
        Particles coordinates (in unit of velocity V) such that Vxyz[:,0] = Vx,
        Vxyz[:,1] = Vy & Vxyz[:,2] = Vz
    aperture : float, optional
        Aperture (in unit of length L) for the computation. Default is 0.03 L
    CoMvelocity : bool, optional
        Boolean to allow the centering of velocities by the considered particles
        centre-of-mass velocity. Default to True
    reduced_structure : bool, optional
        Boolean to allow the computation to adopt the iterative reduced form of the
        inertia tensor. Default to True

    Returns
    -------
    ellip : float
        The ellipticity parameter 1-c/a.
    triax : float
        The triaxiality parameter (a^2-b^2)/(a^2-c^2).
    Transform : array of dtype float, shape (3, 3)
        The orthogonal matrix representing the 3 axes as unit vectors: in real-world
        coordinates, Transform[0] = major, Transform[1] = inter, Transform[2] = minor. 
    abc : array of dtype float, shape (3, )
        The corresponding (a,b,c) lengths (in unit of length L).

    """
    particlesall = np.vstack([XYZ.T,mass,Vxyz.T]).T
    # Compute distances
    distancesall = np.linalg.norm(particlesall[:,:3],axis=1)
    # Restrict particles
    extract = (distancesall<aperture)
    particles = particlesall[extract].copy()
    distances = distancesall[extract].copy()
    Mass = np.sum(particles[:,3])
    # Compute kinematic diagnostics
    if CoMvelocity:
        # Compute CoM velocty, correct
        dvVmass = np.nan_to_num(np.sum(particles[:,3][:,np.newaxis]*particles[:,4:7],axis=0)/Mass)
        particlesall[:,4:7]-=dvVmass
        particles[:,4:7]-=dvVmass
    # Compute momentum
    smomentums = np.cross(particlesall[:,:3],particlesall[:,4:7])
    momentum = np.sum(particles[:,3][:,np.newaxis]*smomentums[extract],axis=0)
    # Compute morphological diagnostics
    s = 1; q = 1; Rsphall = 1+reduced_structure*(distancesall-1); stop = False
    while not('structure' in locals()) or (reduced_structure and not(stop)):
        particles = particlesall[extract].copy()
        Rsph = Rsphall[extract]; Rsph/=np.median(Rsph)
        # Compute structure tensor
        structure = np.sum((particles[:,3]/Rsph**2)[:,np.newaxis,np.newaxis]*(np.matmul(particles[:,:3,np.newaxis],particles[:,np.newaxis,:3])),axis=0)/np.sum(particles[:,3]/Rsph**2)
        # Diagonalise structure tensor
        eigval,eigvec = linalg.eigh(structure)
        # Get structure direct oriented orthonormal base
        eigvec[:,2]*=np.round(np.sum(np.cross(eigvec[:,0],eigvec[:,1])*eigvec[:,2]))
        # Return minor axe
        structmainaxe = eigvec[:,np.argmin(eigval)].copy()
        # Permute base and align Y axis with minor axis in momentum direction
        sign = int(np.sign(np.sum(momentum*structmainaxe)+np.finfo(float).tiny))
        structmainaxe *= sign
        intro_arr = np.array([1,sign,1])
        second_val = (np.argmin(eigval)+np.array([(3+sign)/2,0,(3-sign)/2]))%3
        temp = intro_arr*(eigvec[:,second_val.astype(int)])
        eigval = eigval[second_val.astype(int)]
        # Permute base to align Z axis with major axis
        foo = (np.argmax(eigval)/2)*2
        intro_arr = np.array([(-1)**(1+foo/2),1,1])
        integer_foo_arr = np.array([2-foo, 1, foo]).astype(int)
        second_val = temp[:,integer_foo_arr]
        temp = intro_arr*second_val
        eigval = eigval[integer_foo_arr]
        # Compute change of basis matrix
        transform = linalg.inv(temp)
        stop = (np.max((1-np.sqrt(eigval[:2]/eigval[2])/np.array([q,s]))**2)<1e-4)
        if (reduced_structure and not(stop)):
            q,s = np.sqrt(eigval[:2]/eigval[2])
            Rsphall = linalg.norm(np.matmul(transform,particlesall[:,:3,np.newaxis])[:,:,0]/np.array([q,s,1]),axis=1)
            extract = (Rsphall<aperture/(q*s)**(1/3.))
    Transform = transform.copy()
    ellip = 1-np.sqrt(eigval[1]/eigval[2])
    triax = (1-eigval[0]/eigval[2])/(1-eigval[1]/eigval[2])
    Transform = Transform[...,[2,0,1],:]#so that transform[0] = major, transform[1] = inter, transform[2] = minor
    abc = np.sqrt(eigval[[2,0,1]])
    # Return
    return ellip,triax,Transform,abc