more changes

    ez = E(z)

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

class LVKDistribution():

    def __init__(self, npoints=int(1e4)):

        self.mvec = xp.linspace(2, 100, npoints)

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

    def __call__(self, m, truths):

        return ligo_combined_pm(m, self.mvec, self.dm, truths)

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

        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['M_s'] = test_df['M'] / (1 + test_df['z'])

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

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

    for i in range(len(events)):

        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['M_s'] = test_df['M'] / (1 + test_df['z'])

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

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

 No newline at end of file

import matplotlib.pyplot as plt

from selection_function.model import *

from emri_comfi.model import *

from emri_comfi.sample_population import construct_cosmo_splines

import cupy as xp

import pandas as pd

import numpy as np

from nessai.flowsampler import FlowSampler

from nessai.model import Model

from nessai.utils import setup_logger

from EMRI_DET.validate import run_on_dataset

from EMRI_DET.nn.model_creation import load_mlp

from emri_comfi.nn.validate import run_on_dataset

from emri_comfi.nn.model_creation import load_mlp

import torch

from collections import namedtuple

import corner 

        except FileNotFoundError:

            break

        out['M'].append(df['M'].to_numpy().tolist())

        out['mu'].append(df['mu_act'].to_numpy().tolist())

        out['M'].append(df['M_s'].to_numpy().tolist())

        out['mu'].append(df['mu_s'].to_numpy().tolist())

        out['z'].append(df['z'].to_numpy().tolist())

        i += 1

    out['M'] = xp.asarray(out['M'])

    return out

class FirstModel(Model):

    def __init__(self, names, bounds, ppop, data, selection_function):

    def __init__(self, names, bounds, ppop, data, selection_function, pmu=None):

        self.names = names

        self.bounds = bounds

        self.ppop = ppop

        self.n_posteriors = self.in_shape[0]

        self.samples_per_posterior = self.in_shape[1]

        self.selection_function = selection_function

        self.pmu = pmu

    def log_prior(self, x):

        xd = dict()

        for n in self.names:

            xd[n] = xp.asarray(x[n].item())

        log_l = xp.sum(self.per_event(xd), axis=-1)

        selfact = self.selection_factor(xd)

        selfact = 0#self.selection_factor(xd)

        return (log_l + selfact).get()

        # rpart = self.rate_part(xd)

        return -ndet + self.n_posteriors*xp.log(ndet)

    def per_event(self, xd):

        ppop_evaluated = self.ppop(self.data['mu'], xd) / pz(self.data['z']) / zjac(self.data['z'])#**2

        # ppop_evaluated = self.ppop(self.data['mu'], xd) / pz(self.data['z']) / zjac(self.data['z'])#**2

        ppop_evaluated = self.ppop(self.data["M"], self.data["mu"], self.data["z"], self.pmu, pmu_params={"truths":xd}) / pz(self.data['z']) / zjac(self.data['z'])#**2

        return -np.log(self.samples_per_posterior) + xp.log(xp.sum(ppop_evaluated.reshape(self.in_shape),axis=-1))

def model_sf(model, ordered_names, params):

device="cuda:0"

# data = assemble_big_data('population_generation/populations/second_attempt/catalogue/{}.csv')

data = assemble_big_data('population_generation/populations/withbeta/first_fake/catalogue/{}.csv')

data = assemble_big_data('population_generation/population_datasets/test/full_gaussian_catalogue/{}.csv')

def E(z):

    omega_m = 0.3089

zjacobians = xp.asarray(1/E(zin.get())) + xp.asarray([I(zval) for zval in zin.get()])

zjac = lambda z: xp.interp(z, zin, zjacobians)

# pm = justLIGO()

# bounds=dict(alpha=[-4, 12], mmin=[2, 10], mmax=[30, 100], lam=[0,1],mpp=[20,50],sigpp=[1,10],delta_m=[0,10], rate=[100, 30000])

pm = justLIGO_wbeta()

bounds=dict(alpha=[-4, 12], beta = [-2, 7], mmin=[2, 10], mmax=[30, 100], lam=[0,1],mpp=[20,50],sigpp=[1,10],delta_m=[0,10], rate=[100, 3000])

splines = construct_cosmo_splines(4.5)

selection_function_nn = load_mlp("withbeta_fixed_selnet_log_8L_LR1e-4_5B", device=device, get_state_dict=True, outdir="selection_function/models").to(device)

selection_function_nn.eval()

selection_function_nn.xscalevals = torch.as_tensor(selection_function_nn.xscalevals, device=device).float()

selection_function_nn.yscalevals = torch.as_tensor(selection_function_nn.yscalevals, device=device).float()

#pm = LVKDistribution()

pm = OverallDistribution(1e4, 1e7, 2, 100, 4.5, dVdz_spline=splines[2])

pmu = LVKDistribution()

bounds=dict(alpha=[-4, 12], beta = [-2, 7], mmin=[2, 10], mmax=[30, 100], lam=[0,1],mpp=[20,50],sigpp=[1,10],delta_m=[0,10])#, rate=[100, 3000])

ordered_names = list(bounds.keys())[:-1]

print(ordered_names)

selnet = partial(model_sf, selection_function_nn, ordered_names)

# selection_function_nn = load_mlp("withbeta_fixed_selnet_log_8L_LR1e-4_5B", device=device, get_state_dict=True, outdir="selection_function/models").to(device)

# selection_function_nn.eval()

# selection_function_nn.xscalevals = torch.as_tensor(selection_function_nn.xscalevals, device=device).float()

# selection_function_nn.yscalevals = torch.as_tensor(selection_function_nn.yscalevals, device=device).float()

# ordered_names = list(bounds.keys())[:-1]

# print(ordered_names)

# selnet = partial(model_sf, selection_function_nn, ordered_names)

# bounds=dict(delta_m=[0, 10], rate=[100, 3000])    

model = FirstModel(list(bounds.keys()), bounds, pm, data, selnet)

# model = FirstModel(list(bounds.keys()), bounds, pm, data, lambda x: 1)

output = './ligo_only_fakepops_withbeta/pop1_wpz2'

# model = FirstModel(list(bounds.keys()), bounds, pm, data, selnet)

model = FirstModel(list(bounds.keys()), bounds, pm, data, lambda x: 1, pmu=pmu)

output = './sampler_results/full_lhood/testing/full_flipp0'

logger = setup_logger(output=output, log_level='WARNING')

fs = FlowSampler(model, output=output, resume=False, seed=1729)

        return summary

truths = {'alpha': 3.40401232, 'beta':1.07679192, 'mmin': 5.08383396, 'mmax': 86.85433456,

'lam': 0.03868563, 'mpp': 33.72549972,'sigpp': 3.55853002, 'delta_m': 4.82841, 'rate':1.3e3}

'lam': 0.03868563, 'mpp': 33.72549972,'sigpp': 3.55853002, 'delta_m': 4.82841}#, 'rate':1.3e3}

defaults_kwargs = dict(

            bins=50, smooth=0.9, label_kwargs=dict(fontsize=16),

    del samples["it"]  # only in nessai 0.5.x

except:

    pass

fig = corner.corner(samples, truths=truths, **defaults_kwargs)

axes = fig.get_axes()

for i, par in enumerate(samples.keys()):