Dev

nmma/core/base.py

@@ -1,6 +1,5 @@
 import inspect
-import io
-import contextlib
+import os
 import h5py
 from ast import literal_eval
 import numpy as np
@@ -15,6 +14,7 @@
 from .utils import input_obj_to_str, read_bestfit_from_posterior
 from .constants import  set_cosmology
 from .conversion import cosmology_to_distance
+from .parsing import single_messenger_analysis_parsing, nmma_base_parsing
 
 def initialisation_args_from_signature_and_namespace(_callable, namespace, prefixes = []):
     prefixes.append('')
@@ -98,12 +98,15 @@ def final_diagnostics(self, bestfit_params, args, result=None):
             The figure object containing the plot
 
         """
-        pass
+        try:
+            return self.sub_model.final_diagnostics(bestfit_params, args, result)
+        except AttributeError:
+            pass
 
     def post_process_bestfit(self, args, result=None):
         bestfit_params = read_bestfit_from_posterior(args)
         bestfit_params = self.parameter_conversion(bestfit_params)
-        self.final_diagnostics(bestfit_params, args, result)
+        return self.final_diagnostics(bestfit_params, args, result)
 
     def check_parameter_equivalencies(self, parameter_names):
         """Check for equivalent parameters and terminate if found"""
@@ -272,28 +275,29 @@ def check_priors_and_likelihood_for_nmma(priors, likelihood):
     constraints = {k: priors.pop(k) for k in priors.copy().keys()
                     if isinstance(priors[k], Constraint)}
     likelihood.constraints.update(constraints)
+
     test_draw = priors.sample(1)
     test_conversion = priors.conversion_function(test_draw)
     if len(set(test_conversion.keys()) ) != len(test_conversion.keys()):
         priors.conversion_function = priors.default_conversion_function
         likelihood.conv_functions.append(likelihood.priors.conversion_function)
+
     # add final conversions
     likelihood.setup_parameter_conversion()
     return priors, likelihood
 
 def bilby_sampling(likelihood, priors, args, injection_parameters=None, rank=0):
-    priors, likelihood = check_priors_and_likelihood_for_nmma(priors, likelihood)
-
+    if isinstance(args, dict):
+        def_args = nmma_base_parsing(single_messenger_analysis_parsing)
+        def_args.__dict__.update(args)
+        args = def_args
     # fetch the additional sampler kwargs
-    if isinstance(args.sampler_kwargs, str):
-        sampler_kwargs = literal_eval(args.sampler_kwargs) 
-    else:
-        sampler_kwargs = args.sampler_kwargs
+    sampler_kwargs = getattr(args, "sampler_kwargs", {})
     print("Running with the following additional sampler_kwargs:")
     print(sampler_kwargs)
 
     # check if it is running with reactive sampler
-    nlive = None if args.reactive_sampling else args.nlive
+    nlive = None if getattr(args, 'reactive_sampling', False) else args.nlive
     if nlive is None and args.sampler != "ultranest":
         raise ValueError("reactive sampling is only available for ultranest, "
                          "please set nlive or use ultranest sampler")
@@ -307,7 +311,7 @@ def bilby_sampling(likelihood, priors, args, injection_parameters=None, rank=0):
             sampler_kwargs["niter"] = 1
         elif args.sampler == "dynesty":
             sampler_kwargs["maxiter"] = 1
-
+    
     result = run_sampler(
         likelihood,
         priors,
@@ -327,6 +331,7 @@ def bilby_sampling(likelihood, priors, args, injection_parameters=None, rank=0):
         return
 
     try:
+        result.posterior = likelihood.posterior_conversion(result.posterior)
         result.save_to_file()
         result.save_posterior_samples()
     except FileMovedError:
@@ -359,49 +364,50 @@ def bilby_sampling(likelihood, priors, args, injection_parameters=None, rank=0):
 
     if args.bestfit or args.plot:
         likelihood.post_process_bestfit(args, result)
+    return result
 
 def multi_analysis_loop(args, analysis_setup):
-
+    USE_MPI = False
     # check if it is running under mpi
     try:
         from mpi4py import MPI
         rank = MPI.COMM_WORLD.Get_rank()
+        USE_MPI = True
     except ImportError:
         rank = 0
 
-    if rank != 0:
-        # Create buffers
-        stdout_buffer = io.StringIO()
-        stderr_buffer = io.StringIO()
-
-        # Redirect python output into buffers
-        redirect_out = contextlib.redirect_stdout(stdout_buffer)
-        redirect_err = contextlib.redirect_stderr(stderr_buffer)
-    else:
-        redirect_out = contextlib.nullcontext()
-        redirect_err = contextlib.nullcontext()
+    if rank != 0 and not getattr(args, 'verbose', False):
+        devnull = os.open(os.devnull, os.O_WRONLY)
+        os.dup2(devnull, 1)
+        os.dup2(devnull, 2)
 
-    with redirect_out, redirect_err:
-        if getattr(args, 'multi', None):
-            sub_runs = []
-            if len(args.multi) == 1:
-                arg, vals = list(args.multi.items())[0]
-                for i, val in enumerate(vals):
-                    run_args = deepcopy(args)
-                    setattr(run_args, arg, val)
-                    setattr(run_args, 'label', f"{args.label}_{i}")
-                    sub_runs.append(run_args)
-            else:
-                for run_name, changes in args.multi.items():
-                    run_args = deepcopy(args)
-                    setattr(run_args, 'label', f"{args.label}_{run_name}")
-                    for key, value in changes.items():
-                        if key not in args:
-                            raise KeyError(f"{key} not a known argument... please remove")
-                        setattr(run_args, key, value)
-                    sub_runs.append(run_args)
+    if getattr(args, 'multi', None):
+        sub_runs = []
+        if len(args.multi) == 1:
+            arg, vals = list(args.multi.items())[0]
+            for i, val in enumerate(vals):
+                run_args = deepcopy(args)
+                setattr(run_args, arg, val)
+                setattr(run_args, 'label', f"{args.label}_{i}")
+                sub_runs.append(run_args)
+        else:
+            for run_name, changes in args.multi.items():
+                run_args = deepcopy(args)
+                setattr(run_args, 'label', f"{args.label}_{run_name}")
+                for key, value in changes.items():
+                    if key not in args:
+                        raise KeyError(f"{key} not a known argument... please remove")
+                    setattr(run_args, key, value)
+                sub_runs.append(run_args)
+    else:
+        sub_runs = [args]
+    for run_args in sub_runs:
+        priors, likelihood, injection_parameters = analysis_setup(run_args)
+        priors, likelihood = check_priors_and_likelihood_for_nmma(priors, likelihood)
+        if USE_MPI and run_args.sampler =='dynesty':
+            from .mpi_setup import pbilby_sampling
+            run_function = pbilby_sampling
         else:
-            sub_runs = [args]
-        for run_args in sub_runs:
-            priors, likelihood, injection_parameters = analysis_setup(run_args)
-            bilby_sampling(likelihood, priors, run_args, injection_parameters, rank)
+            run_function = bilby_sampling
+        out = run_function(likelihood, priors, run_args, injection_parameters, rank)
+    return out

nmma/core/constants.py

@@ -8,6 +8,7 @@
 # helpers
 mc2 = const.M_sun*const.c**2
 G_per_c2 = const.G/const.c**2
+seconds_a_day = 24*3600
 
 
 
@@ -34,6 +35,11 @@
 msun_to_ergs = mc2.cgs.value
 MeV_per_fm3_to_Msun_per_km3 = 1e54/((mc2).to(u.MeV).value) # 1 MeV/fm**3 is 8.9653E-7 Msun/km**3
 
+## pulsar timing constants
+msun_s = (const.M_sun * const.G / const.c**3).value # geometrised Msun in seconds
+msun_mus = msun_s * 1e6 # Msun in microseconds, used for pulsar timing
+einstein_factor = (const.G*const.M_sun/const.c**3).value**(2/3)
+
 default_cosmology = cosmology.Planck18
 def set_cosmology(cosmology_input=None):
     """Set the cosmology for the NMMA package.

nmma/core/conversion.py

@@ -4,7 +4,7 @@
 from scipy.integrate import simpson
 from astropy import units
 from astropy import cosmology as cosmo
-from .constants import geom_msun_km, msun_to_ergs, get_cosmology, set_cosmology
+from .constants import geom_msun_km, msun_to_ergs, msun_s, get_cosmology, set_cosmology
 
 from bilby.gw.conversion import (
     component_masses_to_chirp_mass,
@@ -191,32 +191,50 @@ def convert_mtot_mni(params):
     params["mrp_c"] = (params["xmix"]*(params["mtot"]-params["mni"])-params["mrp"])
     return params
 
+############################## pulsar timing conversions ####################################
+def binary_mass_function(m_obs, m_comp, sin_i):
+    return (m_comp * sin_i)**3 / (m_obs + m_comp)**2
+
+def shapiro_delay(m_comp, sin_i):
+    "see https://arxiv.org/pdf/1007.0933.pdf"
+    shapiro_range = msun_s*1.e6 * m_comp # in microseconds
+    orthometric_ratio = sin_i/(1+np.sqrt(1-sin_i**2))
+    return shapiro_range * orthometric_ratio**3
+
+def einstein_delay_orbital_factor(orbital_period, eccentricity):
+    "see, e.g., 10.1007/978-3-662-62110-3_1, p.12 "
+    return msun_s**(2/3) * eccentricity * (orbital_period /2/np.pi)**(1/3)
+def simplified_einstein_delay(m_psr, m_comp, einstein_factor):
+    "see, e.g., 10.1007/978-3-662-62110-3_1, p.12 "
+    return einstein_factor *m_comp * (m_psr + 2*m_comp) / (m_psr + m_comp)**(4/3)
+
+def einstein_delay(m_psr, m_comp, orbital_period, eccentricity):
+    "see, e.g., 10.1007/978-3-662-62110-3_1, p.12 "
+    einstein_delay_factor = einstein_delay_orbital_factor(orbital_period, eccentricity)
+    return simplified_einstein_delay(m_psr, m_comp, einstein_delay_factor)
+
+def mass_parameters_to_sini(total_mass, mass_function, m_comp):
+    "Invert the binary mass function to get sin(i) for a given total mass and mass function"
+    return np.cbrt(mass_function * total_mass**2)/m_comp
+
 ############################## EOS-related conversions ####################################
 
-def EOS2Parameters(radii, masses, lambdas, m1_source, m2_source):
-    ### FIXME: Under what circumstance would these not simply be mass_val[-1], radius_val[-1]?
+def EOS_to_ns_parameters(radii, masses, lambdas):
     TOV_mass = masses.max(axis=-1)
     TOV_radius = radii[np.argmax(masses)]
+    R_14, R_16 = np.interp(x=[1.4, 1.6], xp=masses, fp=radii, left=0, right=0)
+
+    return TOV_mass, TOV_radius, R_14, R_16
 
+def EOS_to_system_parameters(radii, masses, lambdas, m1_source, m2_source):
     (log_lambda_1, log_lambda_2) = np.interp(x=[m1_source, m2_source],
             xp= masses, fp=np.log(lambdas), left=-np.inf, right=-np.inf)
     lambda_1 = np.exp(log_lambda_1)
     lambda_2 = np.exp(log_lambda_2)
-    try:
-        (radius_1, radius_2, R_14, R_16) = np.interp(
-                x=[m1_source, m2_source, 1.4, 1.6],
-                xp=masses, fp= radii, left =0, right=0)
+    (radius_1, radius_2) = np.interp( x=[m1_source, m2_source],
+            xp=masses, fp= radii, left =0, right=0)
 
-        return TOV_mass, TOV_radius, lambda_1, lambda_2, radius_1, radius_2, R_14, R_16
-    ## radius interpolation will raise an error if dealing with multiple sources at once
-    # In that case we return all values as corresponding arrays
-    except ValueError:
-        ref = np.ones_like(lambda_1)
-        (radius_1, radius_2, R_14, R_16) = np.interp(
-                x=[m1_source, m2_source, 1.4*ref, 1.6*ref],
-                xp=masses, fp= radii, left =0, right=0)
-
-        return ref*TOV_mass, ref*TOV_radius, lambda_1, lambda_2, radius_1, radius_2, R_14, R_16
+    return lambda_1, lambda_2, radius_1, radius_2
 
 def radii_from_qur(parameters):
     mass_1_source = parameters["mass_1_source"]
@@ -638,7 +656,7 @@ def chiBH_fitting(
 
         return chi_BH
 
-    def bns_parameter_conversion(self, converted_parameters):
+    def bns_ejecta_conversion(self, converted_parameters):
 
         # prevent the output message flooded by these warning messages
         old = np.seterr()
@@ -674,26 +692,48 @@ def bns_parameter_conversion(self, converted_parameters):
         # total eject mass
         total_ejeta_mass = 10**log10_mej_dyn + 10**log10_mej_wind
 
+        np.seterr(**old)
+        return log10_mej_dyn, log10_mej_wind, np.log10(total_ejeta_mass), log10_mdisk_fit
+
+    def grb_energy_conversion(self, converted_parameters, log10_mdisk_fit):
+
         # GRB afterglow energy
-        log10_Ejet = converted_parameters.get("log10_E0", (
-              np.log10(converted_parameters.get("ratio_epsilon", 2e-4))
-            + np.log10(1.0 - converted_parameters["ratio_zeta"])
-            + log10_mdisk_fit + np.log10(msun_to_ergs) )
-        )
+        log10_Ejet = np.log10(converted_parameters.get("ratio_epsilon", 2e-4))
+        log10_Ejet += np.log10(1.0 - converted_parameters["ratio_zeta"])
+        log10_Ejet += log10_mdisk_fit + np.log10(msun_to_ergs)
 
-        thetaCore = converted_parameters.get("thetaCore", 0.105)
+        thetaCore = converted_parameters.get("thetaCore", 0.105) ## default about 6 degree, see arxiv:2210.05695
+
+        if not any(key in converted_parameters for key in ["thetaWing", "alphaWing", "b"]):
+            return log10_Ejet - np.log10(np.sin(thetaCore/2)**2)
+
+        if "alphaWing" in converted_parameters:
+            alphaWing = converted_parameters['alphaWing'] 
+        else:
+            alphaWing = converted_parameters["thetaWing"] / converted_parameters["thetaCore"]
 
+
         if "b" in converted_parameters: # power law jet
-            alphaWing = converted_parameters.get("alphaWing", converted_parameters["thetaWing"] / converted_parameters["thetaCore"])
-            log10_E0 = np.log10(powerlaw_jet_energy_to_central_isotropic_energy_equivalent(10**log10_Ejet, thetaCore, alphaWing, converted_parameters["b"]))
-        elif "b" not in converted_parameters and ("thetaWing" in converted_parameters or "alphaWing" in converted_parameters): # gaussian jet
-            alphaWing = converted_parameters.get("alphaWing", converted_parameters["thetaWing"] / converted_parameters["thetaCore"])
-            log10_E0 = np.log10(gaussian_jet_energy_to_central_isotropic_energy_equivalent(10**log10_Ejet, thetaCore, alphaWing))
+            jet_func = powerlaw_jet_energy_to_central_isotropic_energy_equivalent
+            data = np.column_stack((10**log10_Ejet, thetaCore, alphaWing, converted_parameters["b"]))
+
+        else:
+            jet_func = gaussian_jet_energy_to_central_isotropic_energy_equivalent
+            data = np.column_stack((10**log10_Ejet, thetaCore, alphaWing))
+
+        return np.log10([jet_func(*row) for row in data])
+
+
+
+    def bns_parameter_conversion(self, parameters):
+        log10_mej_dyn, log10_mej_wind, log10_mej_total, log10_mdisk_fit = self.bns_ejecta_conversion(parameters)
+
+        if "log10_E0" in parameters:
+            log10_E0 = parameters["log10_E0"]
         else:
-            log10_E0 = log10_Ejet - np.log10(np.sin(thetaCore/2)**2)
+            log10_E0 = self.grb_energy_conversion(parameters, log10_mdisk_fit)
 
-        np.seterr(**old)
-        converted_ejecta = np.stack((log10_mej_dyn, log10_mej_wind, np.log10(total_ejeta_mass), log10_E0 ))
+        converted_ejecta = (log10_mej_dyn, log10_mej_wind, log10_mej_total, log10_E0)
 
         return  np.where(np.isfinite(converted_ejecta), converted_ejecta, -np.inf)
 
@@ -753,6 +793,9 @@ def from_dict(cls, instruction_dict):
 
         if 'em' in instruction_dict:
             conversions.append(instruction_dict['em'])
+
+        if 'custom' in instruction_dict:
+            conversions.append(instruction_dict['custom'])
 
         return cls(*conversions)
 
@@ -807,8 +850,14 @@ def identity_conversion(self, parameters):
     'log10_E0'              : r'$\log_{10}(E_0{\rm [erg]})$',
     'ratio_zeta'            : r'$\zeta$',
     'alpha'                 : r'$\alpha$',
-    'KNtheta'               : r'$\theta_{KN}$',
-    'KNphi'                 : r'$\phi_{KN}$',
+    'KNtheta'               : r'$\theta_{KN} [^\circ]$',
+    'KNphi'                 : r'$\phi_{KN} [^\circ]$',
+    # Bu parameters ##
+    'vej_dyn'               : r'$v_{\rm{dyn}}{\rm [c]}$',
+    'vej_wind'              : r'$v_{\rm{wind}}{\rm [c]}$',
+    'Ye_dyn'                : r'$Y_{e,{\rm{dyn}}}$',
+    'kappa_Ye'              : r'$\kappa_{\rm{Y_e}}$',
+    'kappa_v'               : r'$\kappa_{v}$', 
     ## GRB parameters ##
     'log10_E0'              : r'$\log_{10}(E_0{\rm [erg]})$',
     'ratio_epsilon'         : r'$\epsilon$',
@@ -826,11 +875,11 @@ def identity_conversion(self, parameters):
     'mni_c'                 : r'$M_{\rm{Ni}}/M_{\rm{tot}}$',
     'mrp_c'                 : r'$M_{\rm{rp,c}}{\rm [M_{\odot}]}$',
     ### EOS parameters ###
-    'L_sym'                 : r'$L_{\rm{sym}}{\rm [MeV]}$',
-    'K_sym'                 : r'$K_{\rm{sym}}{\rm [MeV]}$',
-    'K_sat'                 : r'$K_{\rm{sat}}{\rm [MeV]}$',
-    '3n_sat'                : r'$c_{3n_{\rm{sat}}{\rm [c]}$',
-    '5n_sat'                : r'$c_{5n_{\rm{sat}}{\rm [c]}$',
+    'L_sym'                 : r'$L_{\rm sym}{\rm [MeV]}$',
+    'K_sym'                 : r'$K_{\rm sym}{\rm [MeV]}$',
+    'K_sat'                 : r'$K_{\rm sat}{\rm [MeV]}$',
+    '3n_sat'                : r'$c^2_{3n_{\rm sat}}{\rm [c^2]}$',
+    '5n_sat'                : r'$c^2_{5n_{\rm sat}}{\rm [c^2]}$',
     'TOV_mass'              : r'$M_{\rm{TOV}}{\rm [M_{\odot}]}$',
     'R_14'                  : r'$R_{1.4}{\rm[km]}$',
     'lambda_tilde'          : r'$\tilde{\Lambda}$',