mors.star

def Percentile(Mstar=None,Omega=None,Prot=None,percentile=None,MstarDist=None,OmegaDist=None,ProtDist=None,params=params.paramsDefault):
    """Gets rotation rate of percentile or percentile of rotation rate in the model rotation distribution.

    This function can be used for two purposes
        1. to determine the percentile in a rotation distribution of a star given its mass and rotation rate
        2. to determine the rotation rate of a star in a rotation distribution given its mass and percentile
    In the first case, the user should specify the rotation rate using either the Omega or Prot keyword arguments (given
    in OmegaSun and days respectively) and the mass. In the second case, the user should specify the percentile using the
    percentile keyword argument (in the range 0 to 100) and the mass. The mass should be specified in Msun using the Mstar
    keyword argument. These should only be given as floats. The user can also specify the rotation distribution using the
    MstarDist and the OmegaDist or ProtDist keyword arguments. If this is not done, then the code will assume the 1 Myr
    rotation distribution from the empirical model distribution used in Johnstone et al. (2020).

    Parameters
    ----------
    Mstar : float
        Stellar mass in Msun.
    Omega : float , optional
        Rotation rate of star in OmegaSun.
    Prot : float , optional
        Rotation period of star in days.
    percentile : float , optional
        percentile in distribution (between 0 and 100).
    MstarDist : numpy.ndarray , optional
        Array of masses of stars in distribution.
    OmegaDist : numpy.ndarray , optional
        Array of rotation rates of stars in distribution in OmegaSun.
    ProtDist : numpy.ndarray , optional
        Array of rotation periods of stars in distribution in days.
    params : dict , optional
        Parameter dictionary used for getting width of bin to use for distribution.

    Returns
    ----------
    result : float
        Either rotation rate for percentile or percentile for rotation rate.

    """

    # Make sure Mstar was set
    if Mstar is None:
        raise Exception( "keyword argument Mstar must be set" )

    # If percentile is not set, make sure rotation rate is set
    if percentile is None:
        # Not both
        if not Omega is None:
            if not Prot is None:
                raise Exception( "keyword arguments Omega and Prot cannot both be set" )
        # At least one
        if ( Omega is None ) and ( Prot is None ):
            raise Exception( "must set either Omega, Prot, or percentile keyword arguments" )
        # Set Omega is not set then set it
        if Omega is None:
            Omega = phys._Omega(Prot)

    # Setup the distribution
    if MstarDist is None:
        MstarDist , OmegaDist = misc.ModelCluster()
    else:

        # Check that one of OmegaDist and ProtDist are set
        if ( OmegaDist is None ) and ( ProtDist is None ):
            raise Exception( "one of keyword arguments OmegaDist and ProtDist must be set" )

        # Check that OmegaDist and ProtDist are not both set
        if ( not OmegaDist is None ) and ( not ProtDist is None ):
            raise Exception( "keyword arguments OmegaDist and ProtDist cannot both be set" )

        # Setup OmegaDist if ProtDist was set
        if OmegaDist is None:
            OmegaDist = misc._Omega(ProtDist)

        # Make sure arrays are same length
        if not ( len(MstarDist) == len(OmegaDist) ):
            raise Exception( "MstarDist and OmegaDist (or ProtDist) must be same length" )

    # Work out which one to do
    if not percentile is None:

        # percentile was set, so get corresponding rotation rates
        result = _OmegaPercentile(Mstar,percentile,MstarDist,OmegaDist,params)

    else:

        # Rotation rate was set, so get corresponding percentile
        result = _PerPercentile(Mstar,Omega,MstarDist,OmegaDist,params)


    return result

def _CheckInputMstar(Mstar):
    """Takes stellar mass, checks if value is a valid input mass."""

    # Make sure Mstar was set
    if Mstar is None:
        raise Exception("stellar mass not given")

    # Make sure it is a float
    #if not isinstance(Mstar,float):
        #raise Exception("stellar mass must be given as float")

    # Make sure it above minimum
    if ( Mstar < MstarMin ):
        raise Exception( "stellar mass cannot be less than lower limit of "+str(MstarMin)+" Msun" )

    # Make sure it above minimum
    if ( Mstar > MstarMax ):
        raise Exception( "stellar mass cannot be greater than upper limit of "+str(MstarMax)+" Msun" )

    return

def _InputRotation(Mstar,Age,Omega,OmegaEnv,OmegaCore,Prot,percentile,params):
    """Takes input rotation values, checks sets up values correctly."""

    # Make sure percentile and rotation are not both set
    if not percentile is None:
        if not Omega is None :
            raise Exception( "cannot set both percentile and Omega as arguments of Star" )
        if not OmegaEnv is None :
            raise Exception( "cannot set both percentile and OmegaEnv as arguments of Star" )
        if not OmegaCore is None :
            raise Exception( "cannot set both percentile and OmegaCore as arguments of Star" )
        if not Prot is None :
            raise Exception( "cannot set both percentile and Prot as arguments of Star" )

    # Make sure percentile and OmegaEnv are not both set
    if ( not percentile is None ) and ( not OmegaEnv is None ):
        raise Exception( "cannot set both percentile and OmegaEnv as arguments of Star" )

    # If percentile was set to a string, get float version
    if isinstance(percentile,str):
        if ( percentile == 'slow' ):
            percentile = 5.0
        elif ( percentile == 'medium' ):
            percentile = 50.0
        elif ( percentile == 'fast' ):
            percentile = 95.0
        else:
            raise Exception( "invalid percentile string (options are 'slow', 'medium', or 'fast')" )

    # If percentile was set, get Omega
    if not percentile is None:
        Omega = Percentile(Mstar=Mstar,percentile=percentile,params=params)

    # Make sure if Age is set that OmegaCore is not set
    if ( not Age is None ) and ( not OmegaCore is None ):
        raise Exception( "cannot set both Age and OmegaCore as arguments of Star" )

    # Make sure not both Omega and Prot were set
    if ( not Omega is None ) and ( not Prot is None ):
        raise Exception( "cannot set both Omega and Prot as arguments of Star" )

    # If Prot is set, get Omega
    if not Prot is None:
        Omega = phys._Omega(Prot)

    # Make sure at least one rotation rate is set
    if ( Omega is None ):

        # If Age is set then make sure OmegaEnv is set, otherwise make sure both OmegaEnv and OmegaCore are set
        if ( not Age is None ):
            if ( OmegaEnv is None ):
                raise Exception( "if Age is set then must set either Omega or OmegaEnv as argument of Star" )
        elif ( OmegaEnv is None ) or ( OmegaCore is None ):
            raise Exception( "must set either Omega or both OmegaEnv and OmegaCore as argument of Star" )

    else:

        # Since Omega is set, make sure both OmegaEnv and OmegaCore are not set
        if not ( ( OmegaEnv is None ) and ( OmegaCore is None ) ):
            raise Exception( "cannot set OmegaEnv and OmegaCore when Omega is set as arugment of Star" )

        # Set both OmegaEnv and OmegaCore to Omega
        OmegaEnv = Omega
        OmegaCore = Omega

    return Omega , OmegaEnv , OmegaCore

def _OmegaPercentile(Mstar,percentile,MstarDist,OmegaDist,params):
    """Gets rotation rate that corresponds to a given percentile of the rotation distribution."""

    ## If percentile was a string, get actual value

    # Get part of distribution to include
    includeStars = np.where( np.abs(MstarDist-Mstar) <= params['dMstarPer'] )

    # Need at least two stars
    if ( len(includeStars[0]) < 2 ):
        raise Exception( "need at least two stars in mass bin to get percentile" )

    # Get percentile Omega
    Omega = np.percentile( OmegaDist[includeStars] , percentile )

    return Omega

def _PerPercentile(Mstar,Omega,MstarDist,OmegaDist,params):
    """Gets percentile of the rotation distribution for a given mass and rotation rate."""

    # Get part of distribution to include
    includeStars = np.where( np.abs(MstarDist-Mstar) <= params['dMstarPer'] )

    # Need at least two stars
    if ( len(includeStars[0]) < 2 ):
        raise Exception( "need at least two stars in mass bin to get percentile" )

    # If Omega is less than minimum then make this as the 0th percentile
    if ( Omega <= np.min(OmegaDist[includeStars]) ):
        return 0.0

    # If Omega is above maximum then make this as the 100th percentile
    if ( Omega >= np.max(OmegaDist[includeStars]) ):
        return 100.0

    # Use bisector to invert numpy.percentile function
    # Initial limits
    perMin, perMax = 0.0, 100.0
    tol = 1.0e-5
    nIterMax = 10000

    # Start iterating
    for iIter in range(nIterMax):
    # Get mid percentile and corresponding Omega
        perMid = 0.5 * ( perMin + perMax )
        OmegaMid = np.percentile( OmegaDist[includeStars] , perMid )

        # See if OmegaMid is equal to or close enough to answer
        if ( abs(OmegaMid/Omega-1.0) < tol ):
            return perMid

        # See which limit to change
        if OmegaMid > Omega:
            perMax = perMid
        else:
            perMin = perMid

    # If we get here, not converged
    raise Exception(
        f"did not find percentile for Omega={Omega} (Mstar={Mstar}) after {nIterMax} iterations; last bracket=({perMin},{perMax}), last mid={perMid}")