import numpy as np

import numpy.random as npr

import bilby

import bilby.gw.conversion as bc

import matplotlib.pyplot as mp

import matplotlib.cm as mpcm

import matplotlib.colors as mpc

import os

import os.path as osp

import errno

import getdist as gd

import getdist.plots as gdp

import corner

import copy

import lalsimulation as lalsim

import sys

import sklearn.decomposition as skd

#import matplotlib

#matplotlib.use('TkAgg')

def comp_masses_to_chirp_q(m_1, m_2):

return m_c, q_inv

def chirp_q_to_comp_masses(m_c, q_inv):

return m_1, m_2

def prior_change_jac(m_c, q_inv):

return jac

def lalinf_adiabatic_index(gammas, x):

return np.exp(logGamma)

def lalinf_sd_gamma_check(gammas):

return True

def lalinf_sd_gamma_check_slow(gammas):

return True

def lal_inf_eos_physical_check(gammas, verbose=False):

return True

def lal_inf_sd_gammas_mass_to_lambda(gammas, mass_m_sol):

return 2.0 / 3.0 * love / comp ** 5

def lal_inf_sd_gammas_fam(gammas):

return fam

def lal_inf_sd_gammas_mass_to_lambda(fam, mass_m_sol):

return 2.0 / 3.0 * love / comp ** 5

def emcee_log_prior(gammas):

return 0.0

def emcee_log_post(gammas):

return log_prior + log_like

def emcee_log_prior_wysocki(thetas):

return 0.0

def emcee_log_post_wysocki(thetas):

return log_prior + log_like

def dd2_lambda_from_mass(m):

return 1.60491e6 - 23020.6 * m**-5 + 194720. * m**-4 - 658596. * m**-3 \

+ 1.33938e6 * m**-2 - 1.78004e6 * m**-1 - 992989. * m + 416080. * m**2 \

- 112946. * m**3 + 17928.5 * m**4 - 1263.34 * m**5

# plot settings

lw = 1.5

mp.rc('font', family = 'serif')

mp.rcParams['text.latex.preamble'] = [r'\boldmath']

mp.rcParams['axes.linewidth'] = lw

mp.rcParams['lines.linewidth'] = lw

# fix a weird LaTeX bug with exponents

mp.rcParams['text.latex.preamble'] = r'\newcommand{\mathdefault}[1][]{}'

# constants

m_sol_kg = 1.988409902147041637325262574352366540e30 # LAL_MSUN_SI = LAL_GMSUN_SI / LAL_G_SI

lal_mrsun_si = 1.476625061404649406193430731479084713e3 # LAL_MRSUN_SI = LAL_GMSUN_SI / (LAL_C_SI * LAL_C_SI)

big_g = 6.67430e-11 # m^3/kg/s^2

c = 2.998e8 # m/s

log_zero = -1.0e10

# @TODO: incorporate maximum mass per EOS?

# needs some thinking. there's a maximum mass used in the

# initial sampling (which, in turn, is complicated by the

# fact we had to use M_c and q): this is the maximum mass

# specified in the prior. 1) if the EOS permits larger

# masses than this, should we alter the prior volume and

# normalization? 2) if the EOS's maximum mass is lower

# than the prior, we need to reject prior draws higher

# than the EOS's m_max. should we instead draw from an

# EOS-specific range each time? 3) do we need to adjust

# the prior volume each time?

# @TODO: KDE version

# @TODO: skip duff IMRPhenom run

# @TODO: store log prior separately

# @TODO: parallelize w/ MPI

# EOS inference settings

n_inds = 4

prior_ranges = [[0.2, 2.0], [-1.6, 1.7], [-0.6, 0.6], [-0.02, 0.02]]

c_s_max = 1.1

ns_mass_max_kg = 1.97 * m_sol_kg

reparam = True

emcee_sample = True

if emcee_sample:

import emcee

n_walkers = 2 * n_inds

else:

n_samples = 12 # 160 # @TODO: update w/ mass samples too

n_m_samples = 100

# data sample settings

use_mpi = False

snr_thresh = 12.0

duration = 32.0 # 8.0

sampling_frequency = 2048.

minimum_frequency = 20.0 # 40.0

reference_frequency = 14.0 # 50.0

min_network = False

if min_network:

ifo_list = ['H1', 'L1', 'V1', 'K1-']

else:

ifo_list = ['H1+', 'L1+', 'V1+', 'K1+', 'A1']

use_polychord = True

use_weighted_samples = True

imp_sample = True

if use_polychord:

n_live = 1000 # 500

else:

n_live = 1000

zero_spins = False

remnants_only = True

min_remnant_mass = 0.01

tight_loc = False

fixed_ang = True

sample_z = True

redshift_rate = True

uniform_bh_masses = True

uniform_ns_masses = True

low_metals = True

broad_bh_spins = True

seobnr_waveform = False

if seobnr_waveform:

waveform_approximant = 'SEOBNRv4_ROM_NRTidalv2_NSBH'

aligned_spins = True

else:

waveform_approximant = 'IMRPhenomPv2_NRTidal'

aligned_spins = False

lam_det_test = False

old = False

outdir = 'outdir'

support_thresh = 1.0e-3

# BH mass and spin prior limits

if uniform_bh_masses:

m_min_bh = 2.5

if low_metals:

m_max_bh = 40.0

else:

m_max_bh = 12.0

else:

m_min_bh = 5.0

m_max_bh = 20.0

spin_min_bh = 0.0

if broad_bh_spins:

spin_max_bh = 0.99

else:

spin_max_bh = 0.5

# NS mass and spin prior limits

m_min_ns = 1.0

if uniform_ns_masses:

m_max_ns = 2.42

else:

m_max_ns = 2.0

spin_min_ns = 0.0

spin_max_ns = 0.05

# conversions

m_c_min, _ = comp_masses_to_chirp_q(m_min_bh, m_min_ns)

m_c_max, _ = comp_masses_to_chirp_q(m_max_bh, m_max_ns)

_, q_inv_min = comp_masses_to_chirp_q(m_max_bh, m_min_ns)

_, q_inv_max = comp_masses_to_chirp_q(m_min_bh, m_max_ns)

# getdist settings

if old:

gd_ranges={'a_1':(0.0, 0.8), 'iota':(0.0, np.pi), \

'q':(0.02, 0.4), 'lambda_2':(0.0, 4000.0), \

'lambda_tilde':(0.0, None)}

else:

gd_ranges={'a_1':(spin_min_bh, spin_max_bh), \

'iota':(0.0, np.pi), \

'q':(q_inv_min, q_inv_max), \

'lambda_2':(0.0, 4500.0), \

'lambda_tilde':(0.0, None)}

# set up identical within-chain MPI processes

if use_mpi:

import mpi4py.MPI as mpi

n_procs = mpi.COMM_WORLD.Get_size()

rank = mpi.COMM_WORLD.Get_rank()

elif len(sys.argv) > 1:

if len(sys.argv) == 3:

n_procs = int(sys.argv[1])

rank = int(sys.argv[2])

if rank > n_procs or rank < 1:

exit('ERROR: 1 <= rank <= number of processes.')

rank -= 1

else:

exit('ERROR: please call using ' + \

'"python sim_nsbh_analysis.py <n_procs> <rank>" format ' + \

'to specify number of processes and rank without MPI. ' + \

'NB: rank should be one-indexed.')

else:

n_procs = 1

rank = 0

# set rank-specific random seed

npr.seed(221216 + rank * 10)

# filename stub

if ifo_list == ['H1', 'L1', 'V1']:

ifo_str = ''

else:

ifo_str = '_'.join(ifo_list) + '_'

label_str = 'nsbh_pop_' + ifo_str + \

'd_{:04.1f}_mf_{:4.1f}_rf_{:4.1f}'

if sample_z:

label_str += '_dndz'

if redshift_rate:

label_str += '_rr'

if uniform_bh_masses:

label_str += '_ubhmp_{:.1f}_{:.1f}'.format(m_min_bh, m_max_bh)

if uniform_ns_masses:

label_str += '_unsmp_{:.1f}_{:.1f}'.format(m_min_ns, m_max_ns)

if broad_bh_spins:

label_str += '_bbhsp'

if seobnr_waveform:

label_str += '_seobnr'

if aligned_spins:

label_str += '_aligned'

base_label = label_str.format(duration, minimum_frequency, \

reference_frequency)

# read injections from file

if old:

targets = np.genfromtxt('data/remnant_sorted_detected.txt', delimiter=' ')

target_ids = targets[:, 0].astype(int)

target_snrs = targets[:, 2]

if remnants_only:

target_ids = target_ids[targets[:, 1] > 0.0]

target_snrs = target_snrs[targets[:, 1] > 0.0]

else:

par_file = base_label + '.txt'

raw_pars = np.genfromtxt('data/' + par_file, \

dtype=None, names=True, delimiter=',', \

encoding=None)

det = raw_pars['snr'] >= snr_thresh

if remnants_only:

det = np.logical_and(det, raw_pars['remnant_mass'] > min_remnant_mass)

raw_pars = raw_pars[det]

ids = np.array([int(i_sim.split(':')[-1]) for i_sim in \

raw_pars['simulation_id']])

snrs = raw_pars['snr']

i_sort = np.argsort(snrs)[::-1]

target_snrs = snrs[i_sort]

target_ids = ids[i_sort]

target_iotas = raw_pars['inclination'][i_sort]

# useful grids in neutron star mass and lambda

m_ns_grid = np.linspace(m_min_ns, m_max_ns, 1000)

# loop over targets to read in the merger NS mass-lambda posteriors

n_targets = len(target_ids)

samples = []

truths = []

post_at_true_m_ns = []

m_ns_like_support = []

post_at_true_m_l_ns = []

m_l_ns_posts = []

skip = np.full(n_targets, False)

for i in range(n_targets):

# read in results file, which contains tonnes of info

label = base_label + '_inj_{:d}'.format(target_ids[i])

if use_polychord:

label = 'pc_' + label

if zero_spins:

label += '_zero_spins'

if tight_loc:

label += '_tight_loc'

elif fixed_ang:

label += '_fixed_ang'

if n_live != 1000:

label += '_nlive_{:04d}'.format(n_live)

res_file = label + '_result.json'

#print(osp.join(outdir, res_file))

if not osp.exists(osp.join(outdir, res_file)):

skip[i] = True

samples.append(None)

truths.append(None)

post_at_true_m_ns.append(None)

m_ns_like_support.append(None)

post_at_true_m_l_ns.append(None)

continue

result = bilby.result.read_in_result(filename=osp.join(outdir, res_file))

# NB: result.injection_parameters contains incorrect IMRPhenom iotas

# due to a bug in bilby!

all_pars = bc.generate_all_bns_parameters(result.injection_parameters)

truths.append([all_pars['mass_1'], \

all_pars['a_1'], 0.0, \

target_iotas[i], \

all_pars['luminosity_distance'], \

all_pars['mass_ratio'], \

all_pars['lambda_2'], \

all_pars['lambda_tilde'], \

all_pars['spin_1z'], \

all_pars['mass_2'], \

all_pars['mass_1_source'], \

all_pars['mass_2_source']])

# if sampling with polychord optionally use full, variable weight

# posterior samples: bilby takes the equal-weight posterior samples

# to build its result.posterior. there are no distance samples this

# way though!

if use_polychord and use_weighted_samples:

# set up required paramnames file

if aligned_spins:

template = 'nsbh_aligned_spins.paramnames'

else:

template = 'nsbh_precess_spins.paramnames'

gd_root = osp.join(outdir, osp.join('chains', label))

gd_pars = gd_root + '.paramnames'

#print(gd_pars)

if not osp.exists(gd_pars):

# prevent race conditions: can have all processes trying

# to create a symlink simultaneously, with the slow ones

# finding it's already been created despite not existing

# before this if statement

try:

os.symlink(template, gd_pars)

except OSError as e:

if e.errno != errno.EEXIST:

raise e

# read in samples and fill in derived parameters

gd_samples = gd.loadMCSamples(gd_root)

pars = gd_samples.getParams()

m_1, m_2 = chirp_q_to_comp_masses(pars.chirp_mass, \

pars.mass_ratio)

gd_samples.addDerived(m_1, name='mass_1', label=r'm_{\rm BH}')

gd_samples.addDerived(m_2, name='mass_2', label=r'm_{\rm NS}')

if aligned_spins:

gd_samples.addDerived(np.abs(pars.chi_1), name='a_1', label='a_1')

# optionally importance sample the input mass priors

if imp_sample:

# extract posterior samples relevant for reweighting

m_c_samples = pars.chirp_mass

q_inv_samples = pars.mass_ratio

m_1_samples, m_2_samples = \

chirp_q_to_comp_masses(m_c_samples, q_inv_samples)

# define importance weights. we want to convert from the

# prior we used to sample, which is uniform in chirp mass

# and mass to the prior we used to simulate, which is

# uniform in component masses. we technically used a different

# redshift prior too, but they're almost identical in practice.

# note that our sampling prior extends into regions where the

# component-mass prior is zero: apply tiny weights to these

# values

mass_weights = prior_change_jac(m_c_samples, q_inv_samples)

m_1_mask = np.logical_and(m_1_samples >= m_min_bh, \

m_1_samples <= m_max_bh)

m_2_mask = np.logical_and(m_2_samples >= m_min_ns, \

m_2_samples <= m_max_ns)

m12_mask = np.logical_and(m_1_mask, m_2_mask)

mass_weights *= m12_mask

mass_weights += ~m12_mask * 1.0e-10

weights = mass_weights / np.sum(mass_weights)

gd_samples.reweightAddingLogLikes(-np.log(weights))

# build up list of samples objects

samples.append(gd_samples)

else:

# test posterior is correctly sampled

distance_label = r'(d_L - d_L^{\rm true})/d_L^{\rm true}'

try:

delta_distance = \

(result.posterior.luminosity_distance - \

result.injection_parameters['luminosity_distance']) / \

result.injection_parameters['luminosity_distance']

except ValueError:

skip[i] = True

samples.append(None)

post_at_true_m_ns.append(None)

m_ns_like_support.append(None)

post_at_true_m_l_ns.append(None)

continue

# optionally importance sample the input mass priors

if imp_sample:

# extract posterior samples relevant for reweighting

m_c_samples = result.posterior.chirp_mass

q_inv_samples = result.posterior.mass_ratio

m_1_samples, m_2_samples = \

chirp_q_to_comp_masses(m_c_samples, q_inv_samples)

# define importance weights. we want to convert from the

# prior we used to sample, which is uniform in chirp mass

# and mass to the prior we used to simulate, which is

# uniform in component masses. we technically used a different

# redshift prior too, but they're almost identical in practice.

# note that our sampling prior extends into regions where the

# component-mass prior is zero: apply tiny weights to these

# values

mass_weights = prior_change_jac(m_c_samples, q_inv_samples)

m_1_mask = np.logical_and(m_1_samples >= m_min_bh, \

m_1_samples <= m_max_bh)

m_2_mask = np.logical_and(m_2_samples >= m_min_ns, \

m_2_samples <= m_max_ns)

m12_mask = np.logical_and(m_1_mask, m_2_mask)

mass_weights *= m12_mask

mass_weights += ~m12_mask * 1.0e-10

weights = mass_weights / np.sum(mass_weights)

else:

weights = np.ones(len(result.posterior.luminosity_distance))

weights /= np.sum(weights)

# convert to GetDist MCSamples object

if aligned_spins:

gd_samples = np.array([result.posterior.a_1, \

result.posterior.mass_1, \

delta_distance, \

result.posterior.iota, \

result.posterior.mass_ratio, \

result.posterior.lambda_2, \

result.posterior.lambda_tilde, \

result.posterior.chi_1, \

result.posterior.mass_2, \

result.posterior.mass_1_source, \

result.posterior.mass_2_source]).T

samples.append(gd.MCSamples(samples=gd_samples, \

names=['a_1', 'mass_1', \

'distance', 'iota', \

'q', 'lambda_2', \

'lambda_tilde', 'chi_1', \

'mass_2', 'mass_1_source', \

'mass_2_source'], \

labels=['a_1', r'm_{\rm BH}', \

distance_label, r'\iota', \

r'm_{\rm NS}/m_{\rm BH}', \

r'\Lambda_{\rm NS}', \

r'\tilde{\Lambda}', \

r'\chi_1', r'm_{\rm NS}', \

r'm_{\rm BH}^{\rm source}', \

r'm_{\rm NS}^{\rm source}'], \

ranges=gd_ranges, weights=weights))

else:

gd_samples = np.array([result.posterior.a_1, \

result.posterior.mass_1, \

delta_distance, \

result.posterior.iota, \

result.posterior.mass_ratio, \

result.posterior.lambda_2, \

result.posterior.lambda_tilde, \

result.posterior.spin_1z, \

result.posterior.mass_2, \

result.posterior.mass_1_source, \

result.posterior.mass_2_source]).T

samples.append(gd.MCSamples(samples=gd_samples, \

names=['a_1', 'mass_1', \

'distance', 'iota', \

'q', 'lambda_2', \

'lambda_tilde', 'chi_1', \

'mass_2', 'mass_1_source', \

'mass_2_source'], \

labels=['a_1', r'm_{\rm BH}', \

distance_label, r'\iota', \

r'm_{\rm NS}/m_{\rm BH}', r'\Lambda_{\rm NS}', \

r'\tilde{\Lambda}', \

r'\chi_1', r'm_{\rm NS}', \

r'm_{\rm BH}^{\rm source}', \

r'm_{\rm NS}^{\rm source}'], \

ranges=gd_ranges, weights=weights))

# extract m_ns and m_ns-Lambda_ns GetDist posteriors. use the

# former to determine the range of m_ns over which the

# likelihood has support, which we'll use for sampling EOSs.

# we'll use the 2D posteriors in the EOS inference itself.

m_ns_post = samples[-1].get1DDensity('mass_2')

m_ns_post_grid = m_ns_post.Prob(m_ns_grid)

i_support = np.where(m_ns_post_grid > support_thresh)

if i_support[0][0] == 0:

i_support_min = 0

else:

i_support_min = i_support[0][0] - 1

if i_support[0][-1] == len(m_ns_grid) - 1:

i_support_max = len(m_ns_grid) - 1

else:

i_support_max = i_support[0][-1] + 1

m_ns_like_support.append(np.array([m_ns_grid[i_support_min], \

m_ns_grid[i_support_max]]))

post_at_true_m_ns.append(m_ns_post.Prob(truths[-1][9])[0])

m_l_ns_post = samples[-1].get2DDensity('mass_2', 'lambda_2', \

normalized=False)

post_at_true_m_l_ns.append(m_l_ns_post(truths[-1][9], \

truths[-1][6])[0, 0])

m_l_ns_posts.append(m_l_ns_post)

# choose inference algorithm: emcee or MC integration

if emcee_sample:

# use Wysocki et al (2001.01747) reparametrization

if reparam:

# read from file

gammas = np.genfromtxt('data/ns_eos_sd_gammas_wysocki.txt', \

delimiter=',')

# rescale to ~unit normal, then PCA

gammas_mean = np.mean(gammas, axis=0)

gammas_std = np.std(gammas, axis=0)

gammas_rs = (gammas - gammas_mean) / gammas_std

pca = skd.PCA(n_components=n_inds)

pca.fit(gammas_rs)

# project gammas onto PCA basis to obtain bounds

gammas_rs_tf = pca.transform(gammas_rs)

gammas_rs_tf_min = np.min(gammas_rs_tf, axis=0)

gammas_rs_tf_max = np.max(gammas_rs_tf, axis=0)

# set up walker initial conditions. the official guidance is

# to initialize a tight ball of walkers near the ML solution.

# i've not implemented that here. the first set of commented

# lines are a small ball near the prior mean. the second spreads

# the walkers across the prior. when sampling the prior, the

# former gives hardly any non-physical EOSs but then takes time

# to spread across the range. the latter is well spread out

# from the start, but produces more non-physical EOSs.

# @TODO: could fit the true gammas, transform to the PCA space

# and initialize around there?

#mu = (gammas_rs_tf_min + gammas_rs_tf_max) / 2.0

#sigma = np.diag(((gammas_rs_tf_min - gammas_rs_tf_max) / 10.0) ** 2)

#gammas_0 = npr.multivariate_normal(mu, sigma, 10)

init_pars = np.zeros((n_walkers, n_inds))

for j in range(n_inds):

init_pars[:, j] = npr.uniform(gammas_rs_tf_min[j] * 1.1, \

gammas_rs_tf_max[j] * 1.1, \

n_walkers)

else:

# initialize walkers near prior mean. see comment above.

mu = np.array([1.0, 0.0, 0.0, 0.0])

sigma = np.diag([0.01, 0.01, 0.01, 0.001]) ** 2

init_pars = npr.multivariate_normal(mu, sigma, n_walkers)

# set up sampler

if reparam:

print('PCA sampling')

max_n_samples = 20000

sampler = emcee.EnsembleSampler(n_walkers, n_inds, \

emcee_log_post_wysocki)

else:

print('direct gamma sampling. this is probably a bad idea.')

max_n_samples = 100000

sampler = emcee.EnsembleSampler(n_walkers, n_inds, emcee_log_prior)

# sample, following emcee's monitoring and convergence example

# (https://emcee.readthedocs.io/en/stable/tutorials/monitor/).

# sample for up to max_n_samples iterations, checking convergence

# every max_n_samples / 1000 samples.

i_sample = 0

n_check = max(1, int(max_n_samples / 1000))

autocorr = np.empty(max_n_samples)

old_tau = np.inf

for sample in sampler.sample(init_pars, iterations=max_n_samples, \

progress=True):

# check convergence

if sampler.iteration % n_check:

continue

# Compute the autocorrelation time so far

# Using tol=0 means that we'll always get an estimate even

# if it isn't trustworthy

tau = sampler.get_autocorr_time(tol=0)

autocorr[i_sample] = np.mean(tau)

i_sample += 1

# Check convergence

converged = np.all(tau * 100 < sampler.iteration)

converged &= np.all(np.abs(old_tau - tau) / tau < 0.01)

if converged:

break

old_tau = tau

# store samples, reprojecting if necessary

samples = sampler.get_chain(flat=True)

if reparam:

gamma_samples = gammas_std * pca.inverse_transform(samples) + \

gammas_mean

else:

gamma_samples = samples

log_probs = sampler.get_log_prob(flat=True)

all_samples = np.concatenate((gamma_samples, log_probs[:, None]), axis=1)

print(all_samples.shape)

filename = osp.join(outdir, base_label + '_eos_emcee_samples.txt')

np.savetxt(filename, all_samples)

# plot convergence

n = n_check * np.arange(1, i_sample + 1)

y = autocorr[:i_sample]

fig, ax = mp.subplots(1, 1)

mp.plot(n, n / 100.0, "--k")

mp.plot(n, y)

#mp.xlim(0, n.max())

#mp.ylim(0, y.max() + 0.1 * (y.max() - y.min()))

mp.xlabel("number of steps")

mp.ylabel(r"mean $\hat{\tau}$")

filename = osp.join(outdir, base_label + '_eos_emcee_convergence.pdf')

fig.savefig(filename)

print(autocorr[:i_sample])

# corner plot

print(np.sum(np.isfinite(log_probs)))

fig = corner.corner(gamma_samples, plot_datapoints=True, \

plot_density=False, plot_contours=False)

filename = osp.join(outdir, base_label + '_eos_emcee_post.pdf')

fig.savefig(filename)

# @TODO: find out / fit for best gammas for DD2 for overlaying & init

# @TODO: multi-processing

exit()

else:

# draw potential SD EOSs

all_gammas = []

gammas = np.zeros(n_inds)

masses = np.zeros(n_targets)

n_samples_tot = n_samples * n_m_samples

pop_samples = np.zeros((n_inds + n_targets, n_samples_tot))

log_weights = np.zeros(n_samples_tot)