Add gaussian catalogue support, remove old files, fix bugs

    xp = cp

except ImportError:

    xp = np

from pathlib import Path

def construct_z_sample_generator(z_max, nz=1000):

    z_samples = np.linspace(0,z_max,nz)

    mu_draws = xp.interp(xp.random.random(size), mucdf, mu_values)

    mass_samples = xp.logspace(xp.log10(mass_distribution.mmin), xp.log10(mass_distribution.mmax),grid_size)

    m_out = mass_distribution(mass_samples[:,None], mu_draws[None,:], xp.asarray(z_draws), lambda x: 1, normalise=False) # don't pass in a mu distribution

    m_out = mass_distribution(mass_samples[:,None], mu_draws[None,:], xp.asarray(z_draws), lambda x: 1) # don't pass in a mu distribution

    m_out *= 1/(m_out * xp.append(0,xp.diff(mass_samples))[:,None]).sum(axis=0)  # normalise each cdf

    m_cdf = xp.cumsum(m_out * xp.append(0,xp.diff(mass_samples))[:,None], axis=0)

        pass

    out = pd.DataFrame(samples,columns=['logM','logq','e0','dL','thetaS','phiS-phiK','thetaK', 'tplunge'])

    return out

def construct_gaussian_catalogues(dirname, noise_sigma, nsamples, threshold=20, full_noise_sigma=None, full_nsamples=None):

    if full_noise_sigma is None:

        full_noise_sigma = noise_sigma

    if full_nsamples is None:

        full_nsamples = nsamples

    events = pd.read_csv(dirname + '/all_params.csv')

    Path(dirname + '/gaussian_catalogue/').mkdir(parents=True, exist_ok=True)

    outdir = dirname + '/gaussian_catalogue/' + '{}.csv'

    Path(dirname + '/full_gaussian_catalogue/').mkdir(parents=True, exist_ok=True)

    full_outdir = dirname + '/full_gaussian_catalogue/' + '{}.csv'

    events['M'] = 10**events['logM']

    z_spl, dl_spl, dv_spl = construct_cosmo_splines(4.5)

    z_in = np.linspace(0,4.5, 1000)

    dlv = dl_spl(z_in)

    z_from_dl = CubicSpline(dlv, z_in)

    sub_events = events.loc[events['SNR'] > threshold]

    col_headers = list(events)

    for i in range(len(sub_events)):

        row = sub_events.iloc[i,:].to_numpy()

        row_ntimes = np.repeat(row[:,None], int(nsamples), axis=1)

        row_ntimes += np.random.normal(0, np.abs(row*noise_sigma), size=(row_ntimes.T.shape)).T

        test_df = pd.DataFrame(row_ntimes.T, columns=col_headers)

        test_df['z'] = z_from_dl(test_df['dL'].to_numpy())

        test_df['M'] = test_df['M'] / (1 + test_df['z'])

        test_df['mu_act'] = test_df['mu'] / (1+test_df['z'])

        test_df.to_csv(outdir.format(i), index=False)

    for i in range(len(events)):

        row = events.iloc[i,:].to_numpy()

        row_ntimes = np.repeat(row[:,None], int(full_nsamples), axis=1)

        row_ntimes += np.random.normal(0, np.abs(row*full_noise_sigma), size=(row_ntimes.T.shape)).T

        test_df = pd.DataFrame(row_ntimes.T, columns=col_headers)

        test_df['z'] = z_from_dl(test_df['dL'].to_numpy())

        test_df['M'] = test_df['M'] / (1 + test_df['z'])

        test_df['mu_act'] = test_df['mu'] / (1+test_df['z'])

        test_df.to_csv(full_outdir.format(i), index=False)

 No newline at end of file

from emri_comfi.sample_population import construct_gaussian_catalogues

construct_gaussian_catalogues(

    "./population_datasets/test", 

    noise_sigma=1e-3, nsamples = int(1e3), 

    full_nsamples=int(5e0),full_noise_sigma=0.,threshold=20)

 No newline at end of file

handler.setFormatter(formatter)

logger.addHandler(handler)

outdir = 'population_datasets/{}'

outdir = 'population_datasets/test/{}'

Path(outdir[:-3]).mkdir(parents=True, exist_ok=True)

gen_wform = GenerateEMRIWaveform("FastSchwarzschildEccentricFlux", sum_kwargs=dict(pad_output=True), use_gpu=use_gpu, return_list=False,)

z_spl, dL_spl, dVdz_spl = construct_cosmo_splines(zmax)

pm = OverallDistribution(mmin, mmax, mumin, mumax, zmax, dVdz_spline=dVdz_spl)

rate = pm.get_rate(ligo_combined_pm, mu_probdist_kwargs={"parameters":truths})

mvec = xp.linspace(mumin, mumax, 10000)

dm = mvec[1] - mvec[0]

rate = pm.get_rate(ligo_combined_pm, mu_probdist_kwargs={"truths":truths, "mvec":mvec,"dm":dm})

try:

    rate = rate.get().item()

except AttributeError:

from functools import partial

import numpy as np

from scipy.integrate import quad

from scipy.constants import c

try:

    import cupy as cp

    xp = cp

    from cupyx.scipy.ndimage import map_coordinates

    from cupyx.scipy.special import erf

except ImportError:

    xp = np

    from scipy.ndimage import map_coordinates

    from scipy.special import erf

def trapz(y, x=None, dx=1.0, axis=-1):

    y = xp.asanyarray(y)

    if x is None:

        d = dx

    else:

        x = xp.asanyarray(x)

        if x.ndim == 1:

            d = xp.diff(x)

            # reshape to correct shape

            shape = [1] * y.ndim

            shape[axis] = d.shape[0]

            d = d.reshape(shape)

        else:

            d = xp.diff(x, axis=axis)

    ndim = y.ndim

    slice1 = [slice(None)] * ndim

    slice2 = [slice(None)] * ndim

    slice1[axis] = slice(1, None)

    slice2[axis] = slice(None, -1)

    product = d * (y[tuple(slice1)] + y[tuple(slice2)]) / 2.0

    try:

        ret = product.sum(axis)

    except ValueError:

        ret = xp.add.reduce(product, axis)

    return ret

def trapz_cumsum(y, x=None, dx=1.0, axis=-1, flipped=False):

    y = xp.asanyarray(y)

    if x is None:

        d = dx

    else:

        x = xp.asanyarray(x)

        if x.ndim == 1:

            d = xp.diff(x)

            # reshape to correct shape

            shape = [1] * y.ndim

            shape[axis] = d.shape[0]

            d = d.reshape(shape)

        else:

            d = xp.diff(x, axis=axis)

    ndim = y.ndim

    slice1 = [slice(None)] * ndim

    slice2 = [slice(None)] * ndim

    slice1[axis] = slice(1, None)

    slice2[axis] = slice(None, -1)

    product = d * (y[tuple(slice1)] + y[tuple(slice2)]) / 2.0

    if not flipped:

        ret = xp.cumsum(product, axis=axis)

    elif flipped:

        ret = xp.cumsum(xp.flip(product, axis=axis), axis=axis)

    return ret

def dL(z):

    '''

    Luminosity distance from redshift assuming default cosmology.

    :param z: Redshift

    :return: dL, in Gpc.

    '''

    h = 0.6774

    omega_m = 0.3089

    omega_lambda = 1 - omega_m

    dH = 1e-5 * c / h

    def E(z):

        return np.sqrt((omega_m * (1 + z) ** (3) + omega_lambda))

    def I(z):

        fact = lambda x: 1 / E(x)

        integral = quad(fact, 0, z)

        return integral[0]

    return (1+z) * dH * I(z) * 1e-3

def truncnorm(xx, mu, sigma, high, low):

    # breakpoint()

    x_op = xp.repeat(xx, mu.size).reshape((xx.size,mu.size)).T

    hi_op = xp.repeat(high, xx.size).reshape((high.size,xx.size))

    lo_op = xp.repeat(low, xx.size).reshape((low.size,xx.size))

    norm = 2**0.5 / np.pi**0.5 / sigma

    norm /= erf((high - mu) / 2**0.5 / sigma) + erf((mu - low) / 2**0.5 / sigma)  #vector of norms

    try:

        prob = xp.exp(-xp.power(xx[None,:] - mu[:,None], 2) / (2 * sigma[:,None]**2)) # array of dims len(xx) * len(mu)

        prob *= norm[:,None]  # should be fine considering dimensionality

        prob[x_op < lo_op] = 0

        prob[x_op > hi_op] = 0

    except IndexError:

        prob = xp.exp(-xp.power(xx - mu, 2) / (2 * sigma**2)) # vector of len(xx)

        prob *= norm

        prob *= (xx <= high) & (xx >= low)

    return prob

def powerlaw(xx, lam, xmin, xmax):

    x_op = xp.repeat(xx, lam.size).reshape((xx.size,lam.size)).T

    hi_op = xp.repeat(xmax, xx.size).reshape((xmax.size,xx.size))

    lo_op = xp.repeat(xmin, xx.size).reshape((xmin.size,xx.size))

    norm = (1+lam)/(xmax**(1+lam) - xmin**(1+lam)) # vector of norms

    try:

        out =  xx[None,:]**lam[:,None] * norm[:,None] # array of dims len(xx) * len(lam)

        out[x_op < lo_op] = 0

        out[x_op > hi_op] = 0

    except IndexError:

        out =  xx**lam * norm # array of dims len(xx) * len(lam)

        out *= (xx <= xmax) & (xx >= xmin)

    return out

def smooth_exp(m, smooth_scale):

    return xp.exp((smooth_scale/m)+(smooth_scale/(m-smooth_scale)))

def smoothing(masses, mmin, delta_m):

    big_masses = xp.repeat(masses, delta_m.size).reshape((masses.size, delta_m.size)).T

    lows = xp.repeat(mmin, masses.size).reshape((mmin.size,masses.size))

    deltas = xp.repeat(delta_m, masses.size).reshape((delta_m.size,masses.size))

    try:

        ans = (1 + smooth_exp(masses[None,:] - mmin[:,None], delta_m[:,None]))**-1

        ans[big_masses < lows] = 0

        ans[big_masses > lows + deltas] = 1

    except IndexError:

        ans = (1 + smooth_exp(masses - mmin, delta_m))**-1

        ans[masses < mmin] = 0

        ans[masses > mmin + delta_m] = 1

    return ans

def ligo_ppop(m, parameters):

    tcs = two_component_single(m, parameters['alpha'],parameters['mmin'],parameters['mmax'],parameters['lam'],parameters['mpp'],parameters['sigpp'])

    smth = smoothing(m, parameters['mmin'], parameters['delta_m'])

    return tcs * smth  # NOT normalised!

def two_component_single(

    mass, alpha, mmin, mmax, lam, mpp, sigpp, gaussian_mass_maximum=100

):

    p_pow = powerlaw(mass, -alpha, mmin, mmax)  # 2d array, dims N_samp * N_hp

    p_norm = truncnorm(mass, mu=mpp, sigma=sigpp, high=xp.asarray(gaussian_mass_maximum), low=mmin)  # same as above

    try:

        prob = (1 - lam[:,None]) * p_pow + lam[:,None] * p_norm

    except IndexError:

        prob = (1 - lam) * p_pow + lam * p_norm

    return prob

def p1(m, p1norm, truths):

    p = ligo_ppop(m, truths)

    return p/p1norm

def p2_no_m1(m, beta, mmin, delta_m):

    return m**beta * smoothing(m, mmin, delta_m)

def p2_total(m, Kfunc, p1norm, beta, mmin, delta_m):

    return p2_no_m1(m, beta, mmin, delta_m)/p1norm * Kfunc(m) # integral is a function of m

def construct_I(mvec, dm, beta, mmin, delta_m):

    prob = p2_no_m1(mvec, beta, mmin, delta_m)

    cumulative_integral = xp.append(0,trapz_cumsum(prob, dx=dm, axis=-1))

    return lambda m: xp.interp(m, mvec, cumulative_integral)

def construct_K(mvec, dm, p1probs, beta, mmin, delta_m):

    Ifunc = construct_I(mvec, dm, beta, mmin, delta_m)

    prob = xp.nan_to_num(p1probs / Ifunc(mvec))

    cumulative_integral = xp.flip(xp.append(0,trapz_cumsum(prob, dx=dm, axis=-1, flipped=True)),axis=-1) # flip back

    return lambda m: xp.interp(m, mvec, cumulative_integral)

def ligo_combined_pm(m, mvec,dm, truths, return_pone_ptwo=False):

    pm1 =ligo_ppop(mvec, parameters=truths)  # cache + get p1 normalisation

    pm1_norm = trapz(pm1, dx=dm)

    pone = p1(m, pm1_norm, truths)

    Kfunc = construct_K(mvec, dm, pm1, truths["beta"], truths["mmin"], truths["delta_m"])

    ptwo = p2_total(m, Kfunc, pm1_norm, truths["beta"], truths["mmin"], truths["delta_m"])

    if return_pone_ptwo:

        return 0.5*(pone + ptwo), pone, ptwo

    else:

        return 0.5*(pone + ptwo)

def dVdz(z):

    h = 0.6774

    omega_m = 0.3089

    omega_lambda = 1 - omega_m

    def E(z):

        return np.sqrt((omega_m * (1 + z) ** (3) + omega_lambda))

    dh = 3000/h # Mpc

    dm = dh * quad(lambda x: 1/E(x),0,z)[0]

    ez = E(z)

    return 4*np.pi * dh * dm**2 / ez / (1+z) # 1+z maps time properly to redshift

class OverallDistribution(): # bringing together all of the painful spaghetti

    def __init__(self, mmin, mmax, mumin, mumax, zmax, dVdz_spline=None, Np=10):

        self.mmin = mmin

        self.mmax = mmax

        self.mumin = mumin  # these should agree with the priors on the mu distribution to provide sufficient coverage.

        self.mumax = mumax

        self.zmax = zmax 

        self.Np = Np  # adjustable parameter for the rate, if one is so inclined.

        self._dVdz = dVdz_spline

        self._prepare_p0()

        self.cache_preliminaries()

    def __call__(self, M, mu, z, pmu, pmu_params={}): # compute probability!

        norm = self.get_rate(pmu, pmu_params)

        dn_part = self.dVdz_wrap(z) * self._mbh_distribution(M) * pmu(mu, **pmu_params)

        gam = self._gamma(mu, M)

        other_part = self._p0(z, M) * gam * self._kappa(mu, M, gam) * self._R0(M)

        # if other_part.shape[0] == 1:

        #     other_part = other_part[0]

        try:

            return dn_part * other_part / norm

        except ValueError:

            return dn_part * other_part / norm[:,None]

    def cache_preliminaries(self):

        #cache some preliminaries for the normalisation

        self.mvec = xp.linspace(self.mmin, self.mmax, 1000)   

        self.muvec = xp.linspace(self.mumin, self.mumax, 1100) 

        self.zvec = xp.linspace(0., self.zmax, 900)

        self.dm = self.mvec[1]-self.mvec[0]

        self.dmu = self.muvec[1]-self.muvec[0]

        self.dz = self.zvec[1] - self.zvec[0]

        self.get_dVdz_norm()

        self.get_Mnorm()

        self.N = self.dVdznorm * self.Mnorm

        p0star = self._integrate_over_z() # integrate over z

        thr_gamma = self._gamma(self.muvec[:,None], self.mvec)

        thr_kappa = self._kappa(self.muvec[:,None], self.mvec, thr_gamma)

        self.R = self._R0(self.mvec)*thr_gamma*thr_kappa

        all_but_mu = self._mbh_distribution(self.mvec)*p0star*self.R  # 2d

        self.integrated_over_M = trapz(all_but_mu,dx=self.dm,axis=1)  # 1d, ready to be normalised by the rate

    def get_rate(self, mu_probdist, mu_probdist_kwargs={}):  # the normalisation for the distribution

            muprobs = mu_probdist(self.muvec, **mu_probdist_kwargs)  # assume this is normalised

            combined = self.integrated_over_M * muprobs

            return self.N * trapz(combined, dx=self.dmu, axis=-1)# * 1600/1520

    def dVdz_wrap(self, z):

        try:

            z = z.get()

        except:

            pass

        vals = self._dVdz(z)

        return xp.asarray(vals)/self.dVdznorm

    def get_dVdz_norm(self):

        self.dVdznorm = 1.

        out = self.dVdz_wrap(self.zvec)

        self.dVdznorm = trapz(out, dx=self.dz) 

    def get_Mnorm(self):

        self.Mnorm = 1.

        out = self._mbh_distribution(self.mvec)

        self.Mnorm = trapz(out, dx=self.dm)

    def _mbh_distribution(self, M):

        return 0.0055 * (M/3e6)**-0.3 / M / self.Mnorm  # from babak et al. 2017, corrected from dn/dlogm to dn/dm

    def _tEMRI(self,M):  # sigmoid fitted to data

        return 10.65 / (1 + xp.exp(3.85*(xp.log10(M)-5.82))) + 2.74  # Gyr

    def _prepare_p0(self):

        p0_logmasses = xp.linspace(4.5, 7.5, 7)

        p0_outputs = xp.asarray([[0.999,0.999,0.999,0.998,0.998,0.998,0.997],[0.997,0.991, 0.98,0.96,0.95,0.94,0.9],[0.23,0.14, 0.18, 0.295, 0.4, 0.53, 0.62],[0.08,0.09,0.15, 0.29, 0.47, 0.63, 0.67],[0.018, 0.023,0.07, 0.18, 0.3, 0.43, 0.5],np.array([0.08,0.09,0.15, 0.29, 0.47, 0.63, 0.67])**3,np.array([0.08,0.09,0.15, 0.29, 0.47, 0.63, 0.67])**4])

        p0_redshifts = xp.linspace(0, 3, 7)

        self._p0_spline = partial(map_coordinates, xp.log(p0_outputs).T)

        self._p0_zmin = p0_redshifts[0].item()

        self._p0_zmax = p0_redshifts[-1].item()

        self._p0_Mmin = p0_logmasses[0].item()

        self._p0_Mmax = p0_logmasses[-1]

    def _p0(self,z,M):

        # rescale to limits

        z = xp.asarray((z - self._p0_zmin) / (self._p0_zmax - self._p0_zmin)) * 6

        M = xp.asarray((xp.log10(M) - self._p0_Mmin) / (self._p0_Mmax - self._p0_Mmin)) * 6

        M[M<0] = 0

        M[M>6] = 6  # keep M within the grid - no extrap.

        reshape = False

        if M.ndim == 2:

            reshape = True

            nz = z.size

            nm = M.size

            reshap = (nm, nz)

            M = M.repeat(nz, axis=1).flatten()

            z = z[:,None].repeat(nm,axis=1).T.flatten()

        zM = xp.vstack((z, M))

        outs = xp.exp(self._p0_spline(zM, cval=0.))

        if reshape:

            outs = outs.reshape(reshap)

        return outs

    def _p0_z_integrand(self, z, M):

        return self.dVdz_wrap(z)*self._p0(z, M)

    def _integrate_over_z(self):

        zints = xp.asarray([self._p0_z_integrand(xp.repeat(zv,self.mvec.size), self.mvec) for zv in self.zvec])

        return xp.asarray(trapz(zints, dx=self.dz, axis=0))

    def _R0(self,M):

        return 300*(M/1e6)**-0.19  # units of Gyr^-1

    def _gamma(self, mu, M):

        out = 1.2/(1+self.Np) * 10 * (M/1e6)**0.06 / mu#[None,:]

        out[out>1] = 1

        return out

    def _kappa(self, mu, M, gam):

        out = xp.exp(-1) * M / self._R0(M) / self._tEMRI(M) / (1+self.Np) / gam / mu#[None,:]

        out[out>1]=1

        return out

import torch

import torch.nn as nn

from torch.nn.init import xavier_uniform_

from emri_comfi.nn.utilities import get_script_path

import dill as pickle

from pathlib import Path

import numpy as np

class LinearModel(nn.Module):

    def __init__(self, in_features, out_features, neurons, n_layers, activation, name, out_activation=None, initialisation=xavier_uniform_, use_dropout=False,drop_p=0.25, use_bn=False):

        super().__init__()

        self.initial = initialisation

        self.name = name

        layers = [nn.Linear(in_features, neurons[0]), activation()]

        for i in range(n_layers - 1):

            layers.append(nn.Linear(neurons[i], neurons[i + 1]))

            if use_dropout:

                layers.append(nn.Dropout(drop_p))  

            layers.append(activation())

            if use_bn:

                layers.append(nn.BatchNorm1d(num_features=neurons[i+1]))

        layers.append(nn.Linear(neurons[-1], out_features))

        if out_activation is not None:

            layers.append(out_activation())

        self.layers = nn.Sequential(*layers)

        self.layers.apply(self.init_weights)

    def forward(self, x):

        return self.layers(x)

    def init_weights(self, m):

        if isinstance(m, nn.Linear):

            self.initial(m.weight)

def create_mlp(input_features, output_features, neurons, layers, activation, model_name, out_activation=None, init=xavier_uniform_, device=None, norm_type='z-score', use_dropout=False,drop_p=0.25, use_bn=False, outdir='../models'):

    if isinstance(neurons, list):

        if len(neurons) != layers:

            raise RuntimeError('Length of neuron vector does not equal number of hidden layers.')

    else:

        neurons = [neurons, ]

    model = LinearModel(input_features, output_features, neurons, layers, activation, model_name, initialisation=init, use_dropout=use_dropout,drop_p=drop_p,use_bn=use_bn, out_activation=out_activation)

    model.norm_type=norm_type

    Path(get_script_path()+f'/{outdir}/{model_name}/').mkdir(parents=True, exist_ok=True)

    pickle.dump(model, open(get_script_path()+f'/{outdir}/{model_name}/function.pickle', "wb"), pickle.HIGHEST_PROTOCOL)  # save blank model

    if device is not None:

        model = model.to(device)

    return model

def load_mlp(model_name, device, get_state_dict=False, outdir='../models'):

    model = pickle.load(open(get_script_path()+f'/{outdir}/{model_name}/function.pickle', "rb"))  # load blank model

    if get_state_dict: