eval_fits.py

Go to the documentation of this file.

""" Evaluates Numerical Relativity fits such as remnant BH mass, spin, etc, for
    various models.
 
Vijay Varma, 2019.
"""
import numpy as np
import warnings
 
from lal import MSUN_SI
 
from .NRSur7dq4Remnant import NRSur7dq4Remnant
from .NRSur3dq8Remnant import NRSur3dq8Remnant
from . import quaternion_utils
 
#=============================================================================
class FitAttributes(object):
    """ Saves attributes of a particular fit.
    """
    #-------------------------------------------------------------------------
    def __init__(self, fit_class, desc, refs, model_keywords):
        """ Initializes a FitAttributes object. See usage examples below.
 
            - fit_class: Class to load for a given model. type: should be a
                daughter class of nrfits.NRFits.
            - desc: A brief description of the model, type: str.
            - refs: Paper reference, type: str.
            - model_keywords: List of strings, with one or more of
                allowed_model_keywords.
        """
 
        # Types of model keywords allowed (add more if required)
        allowed_model_keywords = ['nonspinning', 'aligned_spin', 'precessing', \
            'eccentric', 'tidal']
 
        for key in model_keywords:
            if key not in allowed_model_keywords:
                raise Exception("Invalid model_keyword: %s"%key)
 
        self.fit_class = fit_class
        self.desc = desc
        self.refs = refs
        self.model_keywords = model_keywords
 
 
#=============================================================================
 
# Collection of all implemented fits
fits_collection = {}
 
fits_collection['NRSur3dq8Remnant'] = FitAttributes( \
    fit_class = NRSur3dq8Remnant(),
    desc = 'Fits for remnant mass, spin and kick velocity for nonprecessing'
        ' BBH systems, trained on mass ratios q <= 8 and dimensionless spin '
        'magnitudes <= 0.81. This model was called surfinBH3dq8 in the paper.',
    refs = 'arxiv:1809.09125',
    model_keywords = ['aligned_spin'],
    )
 
fits_collection['NRSur7dq4Remnant'] = FitAttributes( \
    fit_class = NRSur7dq4Remnant(),
    desc = 'Fits for remnant mass, spin and kick velocity for generically'
        ' precessing BBH systems, trained on mass ratios q <= 4.01 and '
        'dimensionless spin magnitudes <= 0.81.',
    refs = 'arxiv:1905.09300',
    model_keywords = ['precessing'],
    )
 
 
#=============================================================================
def truncate_output_to_physical_limits(val_dict, behavior):
    """
    Function to truncate fit values when they go beyond the following physical
    limits:
        - FinalMass: 1, FinalMass is expected as a fraction of total mass here.
        - FinalSpin: magnitude = 1, the Kerr limit for dimensionless spins.
        - RecoilKick: magnitude = 1, as this is in units of c.
 
    Inputs: \n
    - val_dict: A dictionary of fits with one or more of the above keys. See
        eval_nrfit.
 
    - behavior (str) can be one of the following:
        * "TRUNCATE":         Rescale output magnitude to maximum limit, with an
                            info message.
        * "TRUNCATESILENT":   Rescale output magnitude to maximum limit, without
                            info message
        * "KEEP":             Keep output magnitude, even if unphysical, but
                            still give the info message
        * "IGNORE":           Keep output magnitude, even if unphysical, without
                            info message
        * "ERROR":            Abort with an error if output magnitude is
                            unphysical.
    """
 
    # Don't need to do anything
    if behavior == "IGNORE":
        return val_dict
 
    allowed_behaviors = ["ERROR","IGNORE","KEEP","TRUNCATE","TRUNCATESILENT"]
    if behavior not in allowed_behaviors:
        raise ValueError("Invalid physical_limit_violation_behavior=%s"%behavior
            + ". Should be one of ["+ ", ".join(allowed_behaviors) + "].")
 
    physical_limits = {
        "FinalMass": 1,     # Remnant can't be heavier than the total mass
        "FinalSpin": 1,     # Kerr limit
        "RecoilKick": 1,    # Speed of light limit
        }
 
    for fit_type in val_dict.keys():
        limit = physical_limits[fit_type]
        magnitude = np.linalg.norm(val_dict[fit_type])
        if magnitude > limit:
            # Error/Warning message
            message = ""
            if fit_type in ["FinalSpin", "RecoilKick"]:
                message += "%s magnitude=%.7e"%(fit_type, magnitude)
            else:
                message += "Dimensionless %s=%.7e"%(fit_type, magnitude)
            message += " is over max limit=%.2e."%(limit)
 
            if behavior == "ERROR":             # raise error
                raise RuntimeError(message + " Adapt " \
                "extra_params_dict['physical_limit_violation_behavior'] to "\
                "continue without an error.")
 
            elif behavior in ["TRUNCATE", "TRUNCATESILENT"]: # rescale magnitude
                message += " Setting to max limit."
                val_dict[fit_type] *= limit/magnitude
 
            # Only these cases get a warning
            if behavior in ["KEEP", "TRUNCATE"]:
                warnings.warn(message)
 
    return val_dict
 
#=============================================================================
def check_extra_params_and_set_defaults(extra_params_dict):
    """ Does some sanity checks on extra_params_dict.
    If any of the default_keywords are not specified, sets them to the
    default values.
    """
 
    # NOTE: To add a new key to extra_params_dict, set a default value here
    # and update the documentation of eval_nrfit
    default_keywords = {
        'Lambda1': None,
        'Lambda2': None,
        'eccentricity': None,
        'mean_anomaly': None,
        'unlimited_extrapolation': False,
        'physical_limit_violation_behavior': "ERROR",
        }
 
    if extra_params_dict is None:
        extra_params_dict = {}
 
    # Sanity checks
    for key in extra_params_dict.keys():
        if key not in default_keywords.keys():
            raise ValueError('Invalid key %s in extra_params_dict. '%(key)
            + "Should be one of ["+ ", ".join(default_keywords.keys()) + "].")
 
    # set to default if keyword is not specified
    for key in default_keywords:
        if key not in extra_params_dict.keys():
            extra_params_dict[key] = default_keywords[key]
 
    return extra_params_dict
 
 
 
#=============================================================================
def eval_nrfit(m1, m2, chiA_vec, chiB_vec, model_name, fit_types_list, f_ref=-1,
        extra_params_dict=None):
    r"""
    Evaluates Numerical Relativity fits for a given model.
 
    - m1:
        mass of object 1 in kg. \n
        This needs to be in kg for models that allow f_ref != -1. When
        f_ref = -1, if other units are used for m1/m2, the remnant mass will be
        returned in the same units. \n
    - m2:
        mass of the object 2 in kg. \n
    - chiA_vec:
        dimensionless spin (3-vector) of object 1 at the reference epoch. \n
    - chiB_vec:
        dimensionless spin (3-vector) of object 2 at the reference epoch. \n
        This follows the same convention as the waveform interface (see
        https://dcc.ligo.org/T1800226/public), where the spin components are
        defined as:
        * \f$\chi_z = \chi \cdot \hat{L}\f$, where \f$L\f$ is the orbital
            angular momentum vector at the reference epoch.
        * \f$\chi_x = \chi \cdot \hat{n}\f$, where \f$n =\f$ body2 -> body1
            is the separation vector pointing from body2 to body1 at the
            reference epoch.
        * \f$\chi_y = \chi \cdot (\hat{L} \times \hat{n})\f$. \n
        These spin components are frame-independent as they are defined using
        vector inner products. One can also think of the above spins as being
        defined in the following reference frame: The positive z-axis is along
        the orbital angular momentum at the reference epoch. The separation
        vector from the object 2 to object 1 at the reference epoch is along
        the positive x-axis. The y-axis completes the right-handed triad. The
        returned vectors such as final spin and kick velocity are also defined
        in the same frame.
    - model_name:
        model used for fit. \n
        See lalsimulation.nrfits.eval_fits.fits_collection.keys() for all
        available fits. Details about a particular model, including its
        validity, can be found by doing the following: \n
            >> from lalsimulation.nrfits.eval_fits import fits_collection \n
            >> help(fits_collection[model_name].fit_class)
    - fit_types_list:
        List of fits to evaluate. \n
        Allowed values for elements are
        "FinalMass", "FinalSpin", "RecoilKick", and "PeakLuminosity". \n
        Example: fit_types_list = ["FinalMass", "FinalSpin"]
    - f_ref:
        reference frequency (in Hz) used to set the reference epoch. \n
        The reference epoch is set such that the orbital frequency is equal to
        f_ref/2. If f_ref = -1, the reference epoch is taken to be the
        time/frequency at which the fits were constructed. This can be
        different for different models, see the documentation for the specific
        model. Default: f_ref = -1.
    - extra_params_dict:
        Any additional parameters required for a specific model.
        Default: None. \n
        Allowed keys for extra_params_dict are:
        * Lambda1:
            Tidal deformability of object 1. Default: None.
        * Lambda2:
            Tidal deformability of object 2. Default: None.
        * eccentricity:
            Eccentricity at the reference epoch. Default: None.
        * mean_anomaly:
            Mean anomaly at the reference epoch. Default: None.
        * unlimited_extrapolation:
            Some models will raise an error if evaluated well outside
            their calibration range. If unlimited_extrapolation = True,
            then these errors are not raised and the model will still
            produce an output. USE AT YOUR OWN RISK!! Default: False.
        * physical_limit_violation_behavior:
            What to do if the fit values violate physical limits
            FinalMass > total mass, |FinalSpin| > 1, or |RecoilKick| > 1. \n
            Allowed options are (Default: "ERROR"):
            - "ERROR":
                Abort with an error message.
            - "TRUNCATE":
                Rescale output magnitude to maximum limit, with a warning.
            - "TRUNCATESILENT":
                Rescale output magnitude to maximum limit, without
                warning.
            - "KEEP":
                Keep output magnitude, but with a warning.
            - "IGNORE":
                Keep output magnitude, without warning.
 
 
    - Returns: return_dict. \n
        Dictionary of values corresponding to the keys in fit_types_list. \n
        Example: mf = return_dict["FinalMass"]
    """
 
    ### Sanity checks
    if model_name not in fits_collection.keys():
        raise ValueError("Invalid model_name=%s. "%model_name \
            + "Should be one of ["+ ", ".join(fits_collection.keys()) + "].")
 
    if m1 < 0.09 * MSUN_SI and f_ref != -1:
        warnings.warn("Small value of m1 = %e (kg) = %e (Msun) requested. "
            "When f_ref != -1, component masses must be in kgs, perhaps you "
            "are using different units?"%(m1, m1/MSUN_SI))
 
    if m2 < 0.09 * MSUN_SI and f_ref != -1:
        warnings.warn("Small value of m2 = %e (kg) = %e (Msun) requested. "
            "When f_ref != -1, component masses must be in kgs, perhaps you "
            "are using different units?"%(m2, m2/MSUN_SI))
 
    if m1 <= 0 or m2 <= 0:
        raise ValueError("Got nonpositive mass: m1=%.3e, m2=%.3e"%(m1,m2))
 
    chiA_vec = np.atleast_1d(chiA_vec)
    chiB_vec = np.atleast_1d(chiB_vec)
    if len(chiA_vec) != 3 or len(chiB_vec) != 3:
        raise TypeError("Expected input spins to be 3-vectors.")
 
    if np.linalg.norm(chiA_vec) > 1:
        raise ValueError("Invalid spin magnitude |chiA_vec|=%.3f."%( \
            np.linalg.norm(chiA_vec)))
 
    if np.linalg.norm(chiB_vec) > 1:
        raise ValueError("Invalid spin magnitude |chiB_vec|=%.3f."%( \
            np.linalg.norm(chiB_vec)))
 
    if not type(fit_types_list) == list:
        raise TypeError("fit_types_list should be a list.")
 
 
    # do sanity checks on extra_params_dict and set default values if
    # required.
    extra_params_dict = check_extra_params_and_set_defaults(extra_params_dict)
 
    # Some further sanity checks to makes sure extra_params_dict is
    # compatible with the given model
    if (extra_params_dict['Lambda1'] is not None) \
            or (extra_params_dict['Lambda2'] is not None):
        if 'tidal' not in fits_collection[model_name].model_keywords:
            raise ValueError("This model does not allow Lambda1/Lambda2.")
 
    if (extra_params_dict['eccentricity'] is not None) \
            or (extra_params_dict['mean_anomaly'] is not None):
        if 'eccentric' not in fits_collection[model_name].model_keywords:
            raise ValueError("This model does not allow eccentricity or "
                    "mean_anomaly.")
 
    if 'aligned_spin' in fits_collection[model_name].model_keywords:
        if np.linalg.norm(chiA_vec[:2]) > 0 or np.linalg.norm(chiB_vec[:2]) > 0:
            raise ValueError("This model only allows nonprecessing spins")
 
    if 'nonspinning' in fits_collection[model_name].model_keywords:
        if np.linalg.norm(chiA_vec) > 0 or np.linalg.norm(chiB_vec) > 0:
            raise ValueError("This model only allows zero spins")
 
 
    swapped_labels = False
    # If m1 < m2, we switch the labels of the two objects, and then rotate the
    # spins by pi about the z-direction. This amounts to a rigid rotation of
    # the full system about the z-axis by pi. After computing the remnant
    # properties, the final spin and kick vectors are rotated by pi to undo
    # this switch.
    if m1 < m2:
        temp = m1
        m1 = m2
        m2 = temp
        temp = chiA_vec
        chiA_vec = chiB_vec
        chiB_vec = temp
        chiA_vec = quaternion_utils.rotate_in_plane(chiA_vec, np.pi)
        chiB_vec = quaternion_utils.rotate_in_plane(chiB_vec, np.pi)
        swapped_labels = True
 
    # Compute NR fit quantities
    val_dict = fits_collection[model_name].fit_class(m1, m2, chiA_vec, chiB_vec,
            f_ref, fit_types_list, extra_params_dict)
 
    # If the output breaks physical limits, truncate it if requested.
    val_dict = truncate_output_to_physical_limits(val_dict, \
        extra_params_dict['physical_limit_violation_behavior'])
 
    # So far FinalMass was dimless, now rescale to same units as total mass
    if 'FinalMass' in val_dict.keys():
        val_dict['FinalMass'] *= (m1 + m2)
 
    # Rotate remnant vectors by pi if the component lables were swapped
    if swapped_labels:
        if 'FinalSpin' in val_dict.keys():
            val_dict['FinalSpin'] = quaternion_utils.rotate_in_plane( \
                val_dict['FinalSpin'], np.pi)
        if 'RecoilKick' in val_dict.keys():
            val_dict['RecoilKick'] = quaternion_utils.rotate_in_plane( \
                val_dict['RecoilKick'], np.pi)
 
    return val_dict