lalsuite 7.26.2.dev20251210__cp312-cp312-macosx_12_0_arm64.whl

lalsimulation/gwsignal/core/utils.py

@@ -0,0 +1,329 @@
+from operator import xor
+import lal
+import lalsimulation as lalsim
+import numpy as np
+from astropy import units as u
+
+from . import parameter_conventions as pc
+import warnings
+
+def from_lal_value(val):
+    """
+    Read and return a value from LALDict
+
+    Parameters
+    ----------
+    val : `lal.DictValues` value
+        Output from a lal.DictValues
+
+    Returns
+    -------
+    lal.Value'(val) :
+    """
+    ltype = lal.ValueGetType(val)
+    if ltype == lal.CHAR_TYPE_CODE:
+        size = lal.ValueGetSize(val)
+        if size == 1:
+            return lal.ValueGetCHAR(val)
+        else:
+            return lal.ValueGetString(val)
+    if ltype == lal.I2_TYPE_CODE:
+        return lal.ValueGetINT2(val)
+    if ltype == lal.I4_TYPE_CODE:
+        return lal.ValueGetINT4(val)
+    if ltype == lal.I8_TYPE_CODE:
+        return lal.ValueGetINT8(val)
+    if ltype == lal.U2_TYPE_CODE:
+        return lal.ValueGetUINT2(val)
+    if ltype == lal.U4_TYPE_CODE:
+        return lal.ValueGetUINT4(val)
+    if ltype == lal.U8_TYPE_CODE:
+        return lal.ValueGetUINT8(val)
+    if ltype == lal.S_TYPE_CODE:
+        return lal.ValueGetREAL4(val)
+    if ltype == lal.D_TYPE_CODE:
+        return lal.ValueGetREAL8(val)
+    if ltype == lal.C_TYPE_CODE:
+        return lal.ValueGetCOMPLEX8(val)
+    if ltype == lal.Z_TYPE_CODE:
+        return lal.ValueGetCOMPLEX16(val)
+    else:
+        raise TypeError('invalid lal typecode %d' % ltype)
+
+def to_lal_dict(d):
+    """
+    Convert Python dictionary to LALDict object readble by LAL
+
+    Parameters
+    ----------
+    d : `dict`
+        Python dictionary of input parameters
+
+    Returns
+    -------
+    ldict : `LALDict`
+        LALDict object readable by LAL
+    """
+    if d is None:
+        return None
+    ldict = lal.CreateDict()
+    for k, v in d.items():
+        if k == 'ModeArray':
+            lalsim.SimInspiralWaveformParamsInsertModeArrayFromModeString(ldict, v)
+        elif k =='ModeArrayJframe':
+            lalsim.SimInspiralWaveformParamsInsertModeArrayJframeFromModeString(ldict, v)
+        else:
+            if isinstance(v, np.generic):
+                v = v.item()
+            elif type(v) is u.quantity.Quantity:
+                v_val = v.si.value
+                lal.DictInsertREAL8Value(ldict, k, v_val)
+                continue
+            if type(v) is float or lalsim.SimInspiralCheckKnownREAL8Key(k):
+                lal.DictInsertREAL8Value(ldict, k, v)
+            elif type(v) is int:
+                lal.DictInsertINT4Value(ldict, k, v)
+            elif type(v) is str:
+                lal.DictInsertStringValue(ldict, k, v)
+            else:
+                #TODO: handle other types?
+                raise TypeError
+    return ldict
+
+def from_lal_dict(ldict):
+    """
+    Convert LALDict object to a Python dictionary
+
+    Parameters
+    ----------
+    d : `LALDict`
+        LALDict object
+
+    Returns
+    -------
+    d : `dict`
+        Python dictionary
+    """
+    if ldict is None:
+        return None
+    d = dict()
+    keys = lal.DictKeys(ldict)
+    vals = lal.DictValues(ldict)
+    size = lal.ListSize(keys)
+    for i in range(size):
+        key = lal.ListItemGetStringValue(lal.ListPop(keys))
+        lal_item = lal.ListPop(vals)
+        # The ModeArray has LAL_CHAR_TYPE_CODE, but it is not a literal string
+        # from_lal_value would get confused and print weird characters, so we do the distinction below
+        if 'ModeArray' in key:
+            val = lalsim.SimInspiralModeArrayToModeString(lal.ListItemGetValue(lal_item))
+        else:
+            val = from_lal_value(lal.ListItemGetValue(lal_item))
+        d[key] = val
+    return d
+
+
+# Functions to check the parameters and/or add units to them
+
+def check_dict_parameters(waveform_dict,generic_param_dict=None):
+    """Checks the parameters used in the waveform generation routine.
+
+    Parameters
+    ----------
+        waveform_dict (dict): The dictionary of parameters used to generate the template.
+
+        generic_param_dict (dict,optional): Dictionary of extra parameter names to be accepted in order to generate
+                            non-standard waveforms. It should not include standard waveform parameters as they will
+                            be ignored. The form should be parameter_name:parameter_units and the values should be just
+                            added in the waveform_dict.
+    Raises
+    ------
+        AssertionError: If a parameter has the wrong units.
+        TypeError: If a parameter is not available to use or a dimensional parameter
+                   is passed as dimensionless.
+
+    """
+    #Check that masses are correctly input
+    CheckDeterminationOfMasses(waveform_dict)
+    #Check that spins are correctly input
+    CheckDeterminationOfSpins(waveform_dict)
+    #Check the units of the different quantities
+    default_unit_sys = 'S.I.' #International System by default
+
+
+    #Check and extend parameter list if applicable
+    if generic_param_dict is not None: full_parameter_list = np.concatenate([pc.full_parameter_list,list(generic_param_dict.keys())])
+    else : full_parameter_list = pc.full_parameter_list
+
+    # Check if python dictionary contains any key not included in this list & units of selected parameters
+    for k in waveform_dict.keys():
+    #Check all parameters are in joint list otherwise gives error
+        if k not in full_parameter_list:
+            raise(TypeError( ("Parameter %s not in accepted list of parameters"%(k))))
+    #Check the units of the parameteres are correct
+        elif k in pc.units_dict[default_unit_sys].keys():
+            try : waveform_dict[k].unit #Check if it has units at all. Otherwise will give no clue about the parameter giving error
+            except: raise(TypeError( ("Parameter {} does not have units, please pass a parameter with astropy units equivalent to u.[{}]".format(k,pc.units_dict[default_unit_sys][k]))))
+            assert waveform_dict[k].unit.is_equivalent(pc.units_dict[default_unit_sys][k]), "Parameter {} does not have proper units, units should be equivalent to u.[{}]".format(k,pc.units_dict[default_unit_sys][k])
+        elif k=='condition':
+            if int(waveform_dict[k])==0 or int(waveform_dict[k])==1:
+                continue
+            else:
+                raise(TypeError("Condition should only be 0 or 1"))
+        elif k=='lmax':
+            if waveform_dict[k]<0:
+                raise(ValueError("lmax must be >=0"))
+            continue
+        elif generic_param_dict is not None:
+            try : waveform_dict[k].unit #Check if it has units at all. Otherwise will give no clue about the parameter giving error
+            except: raise(TypeError( ("Parameter {} does not have units, please pass a parameter with astropy units equivalent to u.[{}]".format(k,generic_param_dict[k]))))
+            assert waveform_dict[k].unit.is_equivalent(generic_param_dict[k]), "Parameter {} does not have proper units, units should be equivalent to u.[{}]".format(k,generic_param_dict[k])
+
+
+
+
+def add_params_units(waveform_dict,units_sys='S.I.',generic_param_dict=None):
+    """Add units or convert to desired units system to the waveform parameters dictionary
+
+    Parameters
+    ----------
+        waveform_dict (dict): The dictionary of parameters used to generate the template.
+        units_sys (:obj:`str`, optional): System of units chosen for the given parameters.
+                   Defaults to None. If a unit system is given, try to convert and provide
+                   units to the parameters in `waveform_dict`. If default checks the parameters in `waveform_dict`
+                   have the appropriate units.
+        generic_param_dict (dict,optional): Dictionary of extra parameter names to be accepted in order to generate
+                            non-standard waveforms. It should not include standard waveform parameters as they will
+                            be ignored. The form should be parameter_name:parameter_units and the values should be just
+                            added in the waveform_dict.
+    Returns
+    -------
+        A dict the corresponding conversions to the
+        specified `units_sys` system.
+
+    Raises
+    ------
+        AssertionError: If a parameter has wrong units
+
+    """
+    # Main purpose is to check parameters names/units and convert them to astropy SI units
+    # we keep all parameter names and conventions in a separate file
+
+    if units_sys in pc.units_dict.keys():#If units system is in units_dict we give units/convert to given units
+        dict_tmp = {key:u.Quantity(value,pc.units_dict[units_sys][key]) for (key,value) in waveform_dict.items() if key in pc.units_dict[units_sys].keys()}
+    else:
+        raise(TypeError('The units system specified is not available. Available systems are {}'.format([key for key in pc.units_dict.keys()])))
+
+    # Adding units also to the non-standard parameters
+
+    if generic_param_dict is not None:
+        dict_tmp = {**dict_tmp, **{key:u.Quantity(value,generic_param_dict[key]) for (key,value) in waveform_dict.items() if key in generic_param_dict.keys()}}
+
+    #Merge dictionaries keeping those values of dict_tmp as they have been given units
+    dict_tmp = {**waveform_dict,**dict_tmp}
+
+    #Some sanity checks
+    for par in pc.mass_params_[0:4]:
+        if par in dict_tmp.keys():
+            if dict_tmp[par]>=2*10**30*u.solMass:
+                warn_string = "Are you sure you want to have a {} of {} Solar Masses?".format(par,u.Quantity(dict_tmp[par],u.solMass).value)
+                warnings.warn(warn_string)
+        if 'mass_ratio' in dict_tmp.keys():
+            if (dict_tmp['mass_ratio']<0.001) or (dict_tmp['mass_ratio']>1000.0):
+                warn_string = "Are you sure you want to have a q of {}?".format(dict_tmp['mass_ratio'])
+                warnings.warn(warn_string)
+    # Extremely high mass 2*10^30 Msol
+    # Extreme mass ratio
+
+    return dict_tmp
+
+
+
+def CheckDeterminationOfMasses(waveform_dict):
+
+    """Check mass parameters are consistent and enough to fully characterize the binary masses
+
+    Parameters
+    ----------
+        waveform_dict (dict): The dictionary of parameters used to generate the template.
+
+    Raises
+    ------
+        TypeError: whenever mass parameters are over or underspecified
+
+    """
+
+    dim_number =  0 # dimensionful-mass counter
+    nodim_number = 0 # dimensionless-mass counter
+    sym_number = 0 # symmetric masses counter
+
+    #Dictionaries
+    dimensionless_masses = ["mass_ratio", "sym_mass_ratio"]
+    dimensionful_masses = ["mass1", "mass2", "total_mass","chirp_mass", "mass_difference", "reduced_mass"]
+    symetric_masses = ["mass1", "mass2", "total_mass", "chirp_mass", "sym_mass_ratio", "reduced_mass"]
+
+
+    for param in dimensionful_masses:
+        if param in waveform_dict.keys(): dim_number += 1
+    for param in dimensionless_masses:
+        if param in waveform_dict.keys(): nodim_number += 1
+    for param in symetric_masses:
+        if param in waveform_dict.keys(): sym_number += 1
+    if ("mass1" in waveform_dict.keys()) & ("mass2" in waveform_dict.keys()): sym_number = 0
+
+    if ((dim_number == 2 and nodim_number == 0) or (dim_number == 1 and nodim_number == 1)):
+        if(sym_number == 2):
+            warn_string = "The larger object cannot be determined, assuming m1 >= m2."
+            warnings.warn(warn_string)
+    elif ((dim_number == 1 and nodim_number == 0) or dim_number == 0):
+        raise(TypeError( "Mass parameters are underspecified. Please include" \
+        " one dimensionless and one dimensionful mass parameters, or two dimensionful masses."))
+    else:
+        raise(TypeError( "Mass parameters are overspecified. Please include" \
+        " one dimensionless and one dimensionful mass parameters, or two dimensionful masses."))
+
+def CheckDeterminationOfSpins(waveform_dict):
+    """Check spin parameters are consistent and enough to fully characterize the binary spins
+
+    Parameters
+    ----------
+        waveform_dict (dict): The dictionary of parameters used to generate the template.
+
+    Raises
+    ------
+        TypeError: whenever spin parameters are over or underspecified or system of coordinates is mixed
+
+    """
+    #Counters
+    spin1_number = 0
+    spin2_number = 0
+    #Logical variables
+    cartesian_1 = False
+    cartesian_2 = False
+    spherical_1 = False
+    spherical_2 = False
+
+
+    #Suffixes for the spin parameters in the 2 different coordinates systems
+    cartesian = ['x','y','z']
+    spherical = ['_norm','_tilt','_phi']
+
+    for sfx in cartesian:
+        if 'spin1'+sfx in waveform_dict.keys(): spin1_number += 1
+        if 'spin2'+sfx in waveform_dict.keys(): spin2_number += 1
+
+    if spin1_number >0: cartesian_1=True
+    if spin2_number >0: cartesian_2=True
+
+    spin1_number = 0
+    spin2_number = 0
+
+    for sfx in spherical:
+        if 'spin1'+sfx in waveform_dict.keys(): spin1_number += 1
+        if 'spin2'+sfx in waveform_dict.keys(): spin2_number += 1
+
+    if spin1_number >0: spherical_1=True
+    if spin2_number >0: spherical_2=True
+
+    if not(xor(cartesian_1,spherical_1)) or not(xor(cartesian_2,spherical_2)):
+        raise(TypeError( "Please specify the 3 spin parameters in either spherical or cartesian coordinates."))