EM brute (#4299)

* Space before colons in eos_utils.py help messages

* Enable extrapolation in eos_utils.py when NS is out of bounds

* Same check on shape in primary_mass and secondary mass

* Function to swap masses and spins of object 1 and object2

* Allow user to request for remnant_mass calculations: 1) extrapolation, 2) swapping object 1 and object 2 if mass2 > mass1, 3) treat as a (dark) BBH and NSBH with NS mass greater than a value probided by the user

* Renamned swap_companions to ensure_obj1_is_primary and generalized it

* codeclimate

* codeclimate

* Fixed line that was too long

* Fixed docstring and using keywords explicitly in all remnant_mass_* calls

* Renamed max_ns_mass as ns_bh_mass_boundary and improved help messages of functions

* Improved docstings

* Comprehensions in pycbc_brute_bank

* Complying to naming convention for spin components in spherical coordinates

* Update conversions.py

Fixing single code climate complaint.

bin/bank/pycbc_brute_bank

@@ -92,8 +92,7 @@ if args.input_config is not None:
 
 fdict = {}
 if args.fixed_params:
-    for p, v in zip(args.fixed_params, args.fixed_values):
-        fdict[p] = v
+    fdict = {p: v for (p, v) in zip(args.fixed_params, args.fixed_values)}
 
 class Shrinker(object):
     def __init__(self, data):
@@ -321,18 +320,16 @@ if args.input_file:
 
 
 def draw(rtype):
-    params = {}
 
     if rtype == 'uniform':
         if args.input_config is None:
-            for name, pmin, pmax in zip(args.params, args.min, args.max):
-                params[name] = numpy.random.uniform(pmin, pmax, size=size)
+            params = {name: numpy.random.uniform(pmin, pmax, size=size)
+                      for name, pmin, pmax in zip(args.params, args.min, args.max)}
         else:
             # `draw_samples_from_config` has its own fixed seed, so must overwrite it.
             random_seed = numpy.random.randint(low=0, high=2**32-1)
             samples = draw_samples_from_config(args.input_config, size, random_seed)
-            for name in samples.fieldnames:
-                params[name] = samples[name]
+            params = {name: samples[name] for name in samples.fieldnames}
             # Add `static_args` back.
             if static_args is not None:
                 for k in static_args.keys():
@@ -351,8 +348,7 @@ def draw(rtype):
 
         kde = gaussian_kde(bdata)
         points = kde.resample(size=size)
-        for k, v in zip(p, points):
-            params[k] = v
+        params = {k: v for k, v in zip(p, points)}
 
         # Add `static_args` back, some transformations may need them.
         if args.input_config is not None and static_args is not None:
@@ -392,8 +388,7 @@ def draw(rtype):
     else:
         l = dists_joint.contains(params)
 
-    for k in params:
-        params[k] = params[k][l]
+    params = {k: params[k][l] for k in params}
 
     return params
 
@@ -405,22 +400,19 @@ def cdraw(rtype, ts, te):
         t = tau0_from_mass1_mass2(p['mass1'], p['mass2'],
                                   args.tau0_cutoff_frequency)
         l = (t < te) & (t > ts)
-        for k in p:
-            p[k] = p[k][l]
+        p = {k: p[k][l] for k in p}
 
     i = 0
     while len(p[list(p.keys())[0]]) < size:
 
         tp = draw(rtype)
-        for k in p:
-            p[k] = numpy.concatenate([p[k], tp[k]])
+        p = {k: numpy.concatenate([p[k], tp[k]]) for k in p}
 
         if  len(p[list(p.keys())[0]]) > 0:
             t = tau0_from_mass1_mass2(p['mass1'], p['mass2'],
                                       args.tau0_cutoff_frequency)
             l = (t < te) & (t > ts)
-            for k in p:
-                p[k] = p[k][l]
+            p = {k: p[k][l] for k in p}
 
         i += 1
         if i > 1000:

pycbc/conversions.py

@@ -106,6 +106,8 @@ def sec_to_year(sec):
 def primary_mass(mass1, mass2):
     """Returns the larger of mass1 and mass2 (p = primary)."""
     mass1, mass2, input_is_array = ensurearray(mass1, mass2)
+    if mass1.shape != mass2.shape:
+        raise ValueError("mass1 and mass2 must have same shape")
     mp = copy.copy(mass1)
     mask = mass1 < mass2
     mp[mask] = mass2[mask]
@@ -441,8 +443,60 @@ def lambda_from_mass_tov_file(mass, tov_file, distance=0.):
     return lambdav
 
 
+def ensure_obj1_is_primary(mass1, mass2, *params):
+    """
+    Enforce that the object labelled as 1 is the primary.
+
+    Parameters
+    ----------
+    mass1 : float, numpy.array
+        Mass values labelled as 1.
+    mass2 : float, numpy.array
+        Mass values labelled as 2.
+    *params :
+        The binary parameters to be swapped around when mass1 < mass2.
+        The list must have length 2N and it must be organized so that
+        params[i] and params[i+1] are the same kind of quantity, but
+        for object 1 and object 2, respsectively.
+        E.g., spin1z, spin2z, lambda1, lambda2.
+
+    Returns
+    -------
+    list :
+        A list with mass1, mass2, params as arrays, with elements, each
+        with elements re-arranged so that object 1 is the primary.
+    """
+    # Check params are 2N
+    if len(params) % 2 != 0:
+        raise ValueError("params must be 2N floats or arrays")
+    input_properties, input_is_array = ensurearray((mass1, mass2)+params)
+    # Check inputs are all the same length
+    shapes = [par.shape for par in input_properties]
+    if len(set(shapes)) != 1:
+        raise ValueError("Individual masses and params must have same shape")
+    # What needs to be swapped
+    mask = mass1 < mass2
+    # Output containter
+    output_properties = []
+    for i in numpy.arange(0, len(shapes), 2):
+        # primary (p)
+        p = copy.copy(input_properties[i])
+        # secondary (s)
+        s = copy.copy(input_properties[i+1])
+        # Swap
+        p[mask] = input_properties[i+1][mask]
+        s[mask] = input_properties[i][mask]
+        # Format and include in output object
+        output_properties.append(formatreturn(p, input_is_array))
+        output_properties.append(formatreturn(s, input_is_array))
+    # Release output
+    return output_properties
+
+
 def remnant_mass_from_mass1_mass2_spherical_spin_eos(
-        mass1, mass2, spin1a=0.0, spin1pol=0.0, eos='2H'):
+        mass1, mass2, spin1_a=0.0, spin1_polar=0.0, eos='2H',
+        spin2_a=0.0, spin2_polar=0.0, swap_companions=False,
+        ns_bh_mass_boundary=None, extrapolate=False):
     """
     Function that determines the remnant disk mass of an NS-BH system
     using the fit to numerical-relativity results discussed in
@@ -452,44 +506,80 @@ def remnant_mass_from_mass1_mass2_spherical_spin_eos(
     Stone, Loeb & Berger, PRD 87, 084053 (2013), which was originally
     devised for a previous NS-BH remnant mass fit of
     Foucart, PRD 86, 124007 (2012).
-    Note: NS spin is assumed to be 0!
+    Note: The NS spin does not play any role in this fit!
 
     Parameters
     -----------
     mass1 : float
         The mass of the black hole, in solar masses.
     mass2 : float
         The mass of the neutron star, in solar masses.
-    spin1a : float, optional
+    spin1_a : float, optional
         The dimensionless magnitude of the spin of mass1. Default = 0.
-    spin1pol : float, optional
+    spin1_polar : float, optional
         The tilt angle of the spin of mass1. Default = 0 (aligned w L).
     eos : str, optional
         Name of the equation of state being adopted. Default is '2H'.
+    spin2_a : float, optional
+        The dimensionless magnitude of the spin of mass2. Default = 0.
+    spin2_polar : float, optional
+        The tilt angle of the spin of mass2. Default = 0 (aligned w L).
+    swap_companions : boolean, optional
+        If mass2 > mass1, swap mass and spin of object 1 and 2 prior
+        to applying the fitting formula (otherwise fail). Default is False.
+    ns_bh_mass_boundary : float, optional
+        If mass2 is greater than this value, the neutron star is effectively
+        treated as a black hole and the returned value is 0. For consistency
+        with the eos, set this to the maximum mass allowed by the eos; set
+        a lower value for a more stringent cut. Default is None.
+    extrapolate : boolean, optional
+        Invoke extrapolation of NS baryonic mass and NS compactness in
+        scipy.interpolate.interp1d at low masses. If ns_bh_mass_boundary is
+        provided, it is applied at high masses, otherwise the equation of
+        state prescribes the maximum possible mass2. Default is False.
 
     Returns
     ----------
     remnant_mass: float
         The remnant mass in solar masses
     """
-    mass1, mass2, spin1a, spin1pol, input_is_array = ensurearray(
-        mass1, mass2, spin1a, spin1pol)
-    # mass1 must be greater than mass2
-    try:
-        if any(mass2 > mass1) and input_is_array:
-            raise ValueError(f'Require mass1 >= mass2')
-    except TypeError:
-        if mass2 > mass1 and not input_is_array:
-            raise ValueError(f'Require mass1 >= mass2. {mass1} < {mass2}')
-    ns_compactness, ns_b_mass = ns.initialize_eos(mass2, eos)
+    mass1, mass2, spin1_a, spin1_polar, spin2_a, spin2_polar, \
+        input_is_array = \
+        ensurearray(mass1, mass2, spin1_a, spin1_polar, spin2_a, spin2_polar)
+    # mass1 must be greater than mass2: swap the properties of 1 and 2 or fail
+    if swap_companions:
+        mass1, mass2, spin1_a, spin2_a, spin1_polar, spin2_polar = \
+            ensure_obj1_is_primary(mass1, mass2, spin1_a, spin2_a,
+                                   spin1_polar, spin2_polar)
+    else:
+        try:
+            if any(mass2 > mass1) and input_is_array:
+                raise ValueError(f'Require mass1 >= mass2')
+        except TypeError:
+            if mass2 > mass1 and not input_is_array:
+                raise ValueError(f'Require mass1 >= mass2. {mass1} < {mass2}')
     eta = eta_from_mass1_mass2(mass1, mass2)
-    remnant_mass = ns.foucart18(
-        eta, ns_compactness, ns_b_mass, spin1a, spin1pol)
+    # If a maximum NS mass is not provided, accept all values and
+    # let the EOS handle this (in ns.initialize_eos)
+    if ns_bh_mass_boundary is None:
+        mask = numpy.ones(ensurearray(mass2).size[0], dtype=bool)
+    # Otherwise perform the calculation only for small enough NS masses...
+    else:
+        mask = mass2 <= ns_bh_mass_boundary
+    # ...and return 0's otherwise
+    remnant_mass = numpy.zeros(ensurearray(mass2)[0].size)
+    ns_compactness, ns_b_mass = ns.initialize_eos(mass2[mask], eos,
+                                                  extrapolate=extrapolate)
+    remnant_mass[mask] = ns.foucart18(
+            eta[mask], ns_compactness, ns_b_mass,
+            spin1_a[mask], spin1_polar[mask])
     return formatreturn(remnant_mass, input_is_array)
 
 
 def remnant_mass_from_mass1_mass2_cartesian_spin_eos(
-        mass1, mass2, spin1x=0.0, spin1y=0.0, spin1z=0.0, eos='2H'):
+        mass1, mass2, spin1x=0.0, spin1y=0.0, spin1z=0.0, eos='2H',
+        spin2x=0.0, spin2y=0.0, spin2z=0.0, swap_companions=False,
+        ns_bh_mass_boundary=None, extrapolate=False):
     """
     Function that determines the remnant disk mass of an NS-BH system
     using the fit to numerical-relativity results discussed in
@@ -515,15 +605,44 @@ def remnant_mass_from_mass1_mass2_cartesian_spin_eos(
         The dimensionless z-component of the spin of mass1. Default = 0.
     eos: str, optional
         Name of the equation of state being adopted. Default is '2H'.
+    spin2x : float, optional
+        The dimensionless x-component of the spin of mass2. Default = 0.
+    spin2y : float, optional
+        The dimensionless y-component of the spin of mass2. Default = 0.
+    spin2z : float, optional
+        The dimensionless z-component of the spin of mass2. Default = 0.
+    swap_companions : boolean, optional
+        If mass2 > mass1, swap mass and spin of object 1 and 2 prior
+        to applying the fitting formula (otherwise fail). Default is False.
+    ns_bh_mass_boundary : float, optional
+        If mass2 is greater than this value, the neutron star is effectively
+        treated as a black hole and the returned value is 0. For consistency
+        with the eos, set this to the maximum mass allowed by the eos; set
+        a lower value for a more stringent cut. Default is None.
+    extrapolate : boolean, optional
+        Invoke extrapolation of NS baryonic mass and NS compactness in
+        scipy.interpolate.interp1d at low masses. If ns_bh_mass_boundary is
+        provided, it is applied at high masses, otherwise the equation of
+        state prescribes the maximum possible mass2. Default is False.
 
     Returns
     ----------
     remnant_mass: float
         The remnant mass in solar masses
     """
-    spin1a, _, spin1pol = _cartesian_to_spherical(spin1x, spin1y, spin1z)
+    spin1_a, _, spin1_polar = _cartesian_to_spherical(spin1x, spin1y, spin1z)
+    if swap_companions:
+        spin2_a, _, spin2_polar = _cartesian_to_spherical(spin2x,
+                                                          spin2y, spin2z)
+    else:
+        size = ensurearray(spin1_a)[0].size
+        spin2_a = numpy.zeros(size)
+        spin2_polar = numpy.zeros(size)
     return remnant_mass_from_mass1_mass2_spherical_spin_eos(
-        mass1, mass2, spin1a, spin1pol, eos=eos)
+        mass1, mass2, spin1_a=spin1_a, spin1_polar=spin1_polar, eos=eos,
+        spin2_a=spin2_a, spin2_polar=spin2_polar,
+        swap_companions=swap_companions,
+        ns_bh_mass_boundary=ns_bh_mass_boundary, extrapolate=extrapolate)
 
 
 #

pycbc/io/record.py

@@ -1917,8 +1917,10 @@ def spin2_polar(self):
     def remnant_mass(self):
         """Returns the remnant mass for an NS-BH binary."""
         return conversions.remnant_mass_from_mass1_mass2_cartesian_spin_eos(
-                                     self.mass1, self.mass2, self.spin1x,
-                                     self.spin1y, self.spin1z)
+                                     self.mass1, self.mass2,
+                                     spin1x=self.spin1x,
+                                     spin1y=self.spin1y,
+                                     spin1z=self.spin1z)
 
 
 __all__ = ['FieldArray', 'WaveformArray']