xping the code, fixing thetajn sampling

popstock/PopulationOmegaGW.py

@@ -6,6 +6,7 @@
 from .constants import z_to_dL_interpolant
 from .util import wave_energy, omega_gw
 from scipy.interpolate import interp1d
+
 from bilby.core.utils import infer_args_from_function_except_n_args
 from bilby.core.prior import Interped
 from gwpopulation.utils import xp
@@ -34,6 +35,7 @@ def __init__(self, models, mass_coordinates = ['mass_1', 'mass_ratio'], frequenc
             self.frequency_array=frequency_array
         else:
             self.frequency_array=np.arange(10, 2048)
+        self.frequency_array_xp=xp.asarray(self.frequency_array)
 
         self.models = {key.split("_model")[0]: models[key] for key in models}
 
@@ -95,10 +97,11 @@ def calculate_probabilities(self, samples, population_parameters):
         return p_masses * p_z * p_spins
 
     def set_pdraws_source(self):
+        #why this popping?
         self.pdraws = self.proposal_samples.pop("pdraw")
 
-    def calculate_pdraws(self):
-        self.proposal_samples['pdraw'] = xp.array(self.calculate_probabilities(self.proposal_samples, self.fiducial_parameters))
+    def calculate_pdraws(self, proposal_samples):
+        proposal_samples['pdraw'] = xp.array(self.calculate_probabilities(proposal_samples, self.fiducial_parameters))
 
     def calculate_p_masses(self, samples, mass_parameters):
         if hasattr(self.models['mass'], 'n_below'):
@@ -134,9 +137,9 @@ def draw_and_set_proposal_samples(self, fiducial_parameters, N_proposal_samples
         if not fiducial_parameters:
             raise ValueError("No valid parameters passed to set proposal samples.")
         self.fiducial_parameters = fiducial_parameters.copy()
-        proposal_samples = self.draw_source_proposal_samples(self.fiducial_parameters, self.N_proposal_samples)
+        proposal_samples = self.draw_source_proposal_samples(self.fiducial_parameters, N_proposal_samples)
 
-        self.set_proposal_samples(proprosal_samples=proposal_samples)
+        self.set_proposal_samples(proposal_samples=proposal_samples)
 
     def draw_mass_proposal_samples(self, population_params, N, grids = MASS_GRIDS):
 
@@ -152,7 +155,7 @@ def draw_mass_proposal_samples(self, population_params, N, grids = MASS_GRIDS):
             if hasattr(self.models['mass'], 'n_below'):
                 del self.models['mass'].n_below
                 del self.models['mass'].n_above
-            probs = self.models[self.other_mass_coord]({'mass_1': np.array([m1_sample]), self.other_mass_coord: grids[self.other_mass_coord]}, **{key: population_params[key] for key in self.model_args[self.other_mass_coord]})
+            probs = self.models[self.other_mass_coord]({'mass_1': xp.array([m1_sample]), self.other_mass_coord: grids[self.other_mass_coord]}, **{key: population_params[key] for key in self.model_args[self.other_mass_coord]})
             mass_interped = Interped(xx = grids[self.other_mass_coord], yy = probs)
             other_mass_samples.append(mass_interped.sample())
         return {'mass_1_source': mass_1_source_samples, self.other_mass_coord: xp.array(other_mass_samples)}
@@ -161,6 +164,7 @@ def draw_redshift_proposal_samples(self, z_grid, population_params, N):
 
         print("Drawing redshift samples")
         z_prob = self.models['redshift']({'redshift': z_grid}, **{key: population_params[key] for key in self.model_args['redshift']})
+        print({key: population_params[key] for key in self.model_args['redshift']})
         pz_interped = Interped(xx = z_grid, yy=z_prob)
         z_samples = pz_interped.sample(N)
 
@@ -197,13 +201,20 @@ def draw_source_proposal_samples(self, fiducial_parameters, N):
         proposal_samples.update(mass_samples)
         proposal_samples.update(z_samples)
         proposal_samples.update(spin_samples)
-        proposal_samples.update(self.calculate_pdraws())
+        self.calculate_pdraws(proposal_samples)
 
         return proposal_samples
 
     def set_proposal_samples(self, proposal_samples=None):
 
+        keys = proposal_samples.keys()
+        for key in proposal_samples:
+            if type(proposal_samples[key]) is list:
+                proposal_samples[key] = xp.array(proposal_samples[key])
+
         proposal_samples['mass_1_detector'] = proposal_samples['mass_1_source'] * (1 + proposal_samples['redshift'])
+        proposal_samples['mass_2_source'] = proposal_samples['mass_1_source'] * proposal_samples['mass_ratio']
+        proposal_samples['mass_2_detector'] = proposal_samples['mass_2_source'] * (1 + proposal_samples['redshift'])
         self.proposal_samples = proposal_samples.copy()
         self.N_proposal_samples = len(proposal_samples['pdraw'])
 
@@ -248,8 +259,8 @@ def calculate_wave_energies(self, waveform_duration=10, sampling_frequency=4096,
             if not 'phase' in inj_sample:
                 inj_sample['phase']=2*np.pi*np.random.rand()
             if not 'theta_jn' in inj_sample:
-                inj_sample['theta_jn']=np.pi*np.random.rand()
-            """
+                cos_inc = np.random.rand()*2.-1.
+                inj_sample['theta_jn']=np.arccos(cos_inc)
             if not 'a_1' in inj_sample:
                 inj_sample['a_1']=0
             if not 'a_2' in inj_sample:
@@ -258,14 +269,20 @@ def calculate_wave_energies(self, waveform_duration=10, sampling_frequency=4096,
                 inj_sample['tilt_1']=0
             if not 'tilt_2' in inj_sample:
                 inj_sample['tilt_2']=0
-            """
             use_approxed_waveform=False
             if waveform_approximant=='PC_waveform':
                 use_approxed_waveform=True
-
-            wave_energies.append(interp1d(waveform_frequencies, wave_energy(waveform_generator, inj_sample, use_approxed_waveform=use_approxed_waveform), fill_value=0, bounds_error=False, kind='cubic')(self.frequency_array))
-
-        self.wave_energies = np.array(wave_energies)
+            '''
+            waveform_frequencies = xp.asarray(waveform_frequencies)
+            wave_en = xp.asarray(wave_energy(waveform_generator, inj_sample, use_approxed_waveform=use_approxed_waveform))
+            wave_energies.append(xp.interp(self.frequency_array, waveform_frequencies, wave_en) )
+            '''
+            wave_en = wave_energy(waveform_generator, inj_sample, use_approxed_waveform=use_approxed_waveform)
+            wave_energies.append(np.interp(self.frequency_array, waveform_frequencies, wave_en) )
+            #could also do cubic interp but takes a bit longer
+            #wave_energies.append(interp1d(waveform_frequencies, wave_en, fill_value=0, bounds_error=False, kind='cubic')(frequency_array))
+
+        self.wave_energies = xp.array(wave_energies)
         self.wave_energies_calculated = True
 
     def calculate_omega_gw(self, Lambda=None, Rate_norm=None, **kwargs):
@@ -282,8 +299,10 @@ def calculate_omega_gw(self, Lambda=None, Rate_norm=None, **kwargs):
             if Lambda is None:
                 Lambda = self.fiducial_parameters
 
-            redshift_model_norm_in_Gpc3 = self.redshift_model.normalisation(Lambda)/1.e9
+            redshift_model_norm_in_Gpc3 = self.models['redshift'].normalisation(Lambda)/1.e9
             Rate_norm_in_Gpc3_per_seconds = Lambda['rate']/(60*60*24*365)
             Rate_norm = redshift_model_norm_in_Gpc3 * Rate_norm_in_Gpc3_per_seconds
 
-        self.omega_gw = omega_gw(self.frequency_array, self.wave_energies, self.weights, Rate_norm=Rate_norm)
+        frequencies = self.frequency_array_xp
+        self.omega_gw = omega_gw(frequencies, self.wave_energies, self.weights, Rate_norm=Rate_norm)
+

popstock/util.py

@@ -1,4 +1,5 @@
 import numpy as np
+from gwpopulation.utils import xp
 
 from .constants import H0
 
@@ -21,13 +22,14 @@ def wave_energy(waveform_generator, injection_parameters, use_approxed_waveform=
 
     if use_approxed_waveform:
         orient_fact = np.cos(injection_parameters['theta_jn'])**2 + ((1+np.cos(injection_parameters['theta_jn'])**2)/2)**2
-        return orient_fact*np.abs(waveform_approx_amplitude(injection_parameters, frequencies=waveform_generator.frequency_array))**2
+        return (orient_fact*np.abs(waveform_approx_amplitude(injection_parameters, frequencies=waveform_generator.frequency_array))**2)
 
     try:
         polarizations = waveform_generator.frequency_domain_strain(injection_parameters)
         # Could make this into a FrequencySeries...
-        return np.abs(polarizations['plus'])**2 + np.abs(polarizations['cross'])**2 
+        return (np.abs(polarizations['plus'])**2 + np.abs(polarizations['cross'])**2)
     except:
+        print('ouch!')
         return np.zeros_like(waveform_generator.frequency_array)
 
 
@@ -36,20 +38,20 @@ def waveform_approx_amplitude(injection_parameters, frequencies):
 
     # I-M-R frequencies
     from .constants import m_sun, mass_to_seconds_conv, G, light_speed
-    total_mass_in_kg = (injection_parameters['mass_1_source']+injection_parameters['mass_2_source']) * m_sun
-    total_mass_scaled = total_mass_in_kg * mass_to_seconds_conv
-    sym_mass_ratio = (injection_parameters['mass_1_source']*injection_parameters['mass_2_source'])/(injection_parameters['mass_1_source']+injection_parameters['mass_2_source'])**2
+    total_mass_in_kg = float((injection_parameters['mass_1_source']+injection_parameters['mass_2_source']) * m_sun)
+    total_mass_scaled = float(total_mass_in_kg * mass_to_seconds_conv)
+    sym_mass_ratio = float((injection_parameters['mass_1_source']*injection_parameters['mass_2_source'])/(injection_parameters['mass_1_source']+injection_parameters['mass_2_source'])**2)
 
     f_merg = (0.29740 * sym_mass_ratio ** 2.0 + 0.044810 * sym_mass_ratio + 0.095560) / (np.pi * total_mass_scaled)
     f_ring = (0.59411 * sym_mass_ratio ** 2.0 + 0.089794 * sym_mass_ratio + 0.19111) / (np.pi * total_mass_scaled)
     f_cut = (0.84845 * sym_mass_ratio ** 2.0 + 0.12828 * sym_mass_ratio + 0.27299) / (np.pi * total_mass_scaled)
     sigma = (0.50801 * sym_mass_ratio ** 2.0 + 0.077515 * sym_mass_ratio + 0.022369) / (np.pi * total_mass_scaled)
 
     # detector frame
-    f_merg /= (1+injection_parameters['redshift'])
-    f_ring /= (1+injection_parameters['redshift'])
-    f_cut /= (1+injection_parameters['redshift'])
-    sigma /= (1+injection_parameters['redshift'])
+    f_merg /= float(1+injection_parameters['redshift'])
+    f_ring /= float(1+injection_parameters['redshift'])
+    f_cut /= float(1+injection_parameters['redshift'])
+    sigma /= float(1+injection_parameters['redshift'])
 
     mask_insp = frequencies < f_merg
     mask_merg = (frequencies >= f_merg) & (frequencies < f_ring)
@@ -64,8 +66,8 @@ def waveform_approx_amplitude(injection_parameters, frequencies):
     #set zero frequency to zero just in case
     wave_amplitude[0] = 0
     from astropy.constants import kpc
-    dL_in_m = injection_parameters['luminosity_distance']*kpc.value*1.e3
-    const = (G*total_mass_in_kg*(1+injection_parameters['redshift']))**(5/6) * f_merg**(-7/6)/(dL_in_m)/np.pi**(2/3) * (5*sym_mass_ratio/24)**(1/2) / light_speed**(3/2)
+    dL_in_m = float(injection_parameters['luminosity_distance']*kpc.value*1.e3)
+    const = (G*total_mass_in_kg*float(1+injection_parameters['redshift']))**(5/6) * f_merg**(-7/6)/(dL_in_m)/np.pi**(2/3) * (5*sym_mass_ratio/24)**(1/2) / light_speed**(3/2)
     return const * wave_amplitude
 
 
@@ -89,6 +91,6 @@ def omega_gw(frequencies, wave_energies, weights,  Rate_norm):
     The wave energy spectrum in a np.array.
     """
     conv = frequencies**3 * 4. * np.pi**2 / (3 * H0.si.value**2)
-    weighted_energy = np.sum(weights[:, None] * wave_energies, axis=0) / len(weights)
+    weighted_energy = xp.sum(weights[:, None] * wave_energies, axis=0) / len(weights)
 
     return Rate_norm * conv * weighted_energy

population_O3.py

@@ -5,6 +5,7 @@
 from popstock.PopulationOmegaGW import PopulationOmegaGW
 from gwpopulation.models.mass import SinglePeakSmoothedMassDistribution
 from gwpopulation.models.redshift import MadauDickinsonRedshift
+from gwpopulation.utils import xp
 
 import argparse 
 import numpy as np
@@ -23,34 +24,60 @@
 """
 
 parser = argparse.ArgumentParser()
-parser.add_argument('-ns', '--number_samples',help="number of samples.",action="store", type=int, default=50000)
+parser.add_argument('-ns', '--number_samples',help="number of samples.",action="store", type=int, default=None)
 parser.add_argument('-nt', '--number_trials',help="number of population trials.",action="store", type=int, default=10000)
-parser.add_argument('-wf', '--waveform_approximant',help="Wavefrom approximant.",action="store", type=str, default='IMRPhenomD')
+parser.add_argument('-wf', '--waveform_approximant',help="Wavefrom approximant. Default is IMRPhenomD.",action="store", type=str, default='IMRPhenomD')
 parser.add_argument('-rd', '--run_directory',help="Run directory.",action="store", type=str, default='./')
+parser.add_argument('-sm', '--samples',help="Samples to use.",action="store", type=str, default=None)
 args = parser.parse_args()
 
 N_proposal_samples=args.number_samples
 N_trials=args.number_trials
 wave_approx=args.waveform_approximant 
-tag=f'{wave_approx}_{N_proposal_samples}_samples_{N_trials}_trials'
+tag=f'{wave_approx}_{N_proposal_samples}_samples_{N_trials}_trials_thetajnfix'
 rundir=Path(args.run_directory)
 
 """
 ***
 """
 
-Lambda_0 =  {'alpha': 2.5, 'beta': 1, 'delta_m': 3, 'lam': 0.04, 'mmax': 100, 'mmin': 4, 'mpp': 33, 'sigpp':5, 'gamma': 2.7, 'kappa': 5, 'z_peak': 1.9, 'rate': 15}
-
 mass_obj = SinglePeakSmoothedMassDistribution()
 redshift_obj = MadauDickinsonRedshift(z_max=10)
 
+models = {
+        'mass_model' : mass_obj,
+        'redshift_model' : redshift_obj,
+        }
 freqs = np.arange(10, 2000, 5)
-newpop = PopulationOmegaGW(mass_model=mass_obj, redshift_model=redshift_obj, frequency_array=freqs)
 
-newpop.draw_and_set_proposal_samples(Lambda_0, N_proposal_samples=N_proposal_samples)
-newpop.calculate_omega_gw(waveform_approximant=wave_approx)
+newpop = PopulationOmegaGW(models=models, frequency_array=freqs)
+
+if args.samples is not None:
+    with open(args.samples) as samples_file:
+        samples_dict = json.load(samples_file)
+    Lambda_0 = samples_dict['Lambda_0']
+    samples_dict.pop('Lambda_0')
+    if args.number_samples is None:
+        args.number_samples = len(samples_dict['redshift'])
+    else:
+        if args.number_samples is None:
+            args.number_samples = 50000
+        for key in samples_dict.keys():
+            samples_dict[key] = samples_dict[key][:args.number_samples]
+    newpop.set_proposal_samples(proposal_samples = samples_dict)
+    print(f'Using {args.number_samples} samples...')
+
+else:
+    Lambda_0 =  {'alpha': 2.5, 'beta': 1, 'delta_m': 3, 'lam': 0.04, 'mmax': 100, 'mmin': 4, 'mpp': 33, 'sigpp':5, 'gamma': 2.7, 'kappa': 3, 'z_peak': 1.9, 'rate': 15}
+    newpop.draw_and_set_proposal_samples(Lambda_0, N_proposal_samples=N_proposal_samples)
+
+newpop.calculate_omega_gw(waveform_approximant=wave_approx, Lambda=Lambda_0)
+
+if args.samples is not None:
+    np.savez(f"{os.path.join(rundir, f'omegagw_0_{tag}.npz')}", omega_gw=newpop.omega_gw, freqs=newpop.frequency_array, Lambda_0=Lambda_0)
 
-np.savez(f"{os.path.join(rundir, f'omegagw_0_{tag}.npz')}", omega_gw=newpop.omega_gw, freqs=newpop.frequency_array, fiducial_samples=newpop.proposal_samples, Lambda_0=Lambda_0, draw_dict=newpop.pdraws)
+else:
+    np.savez(f"{os.path.join(rundir, f'omegagw_0_{tag}.npz')}", omega_gw=newpop.omega_gw, freqs=newpop.frequency_array, fiducial_samples=newpop.proposal_samples, Lambda_0=Lambda_0, draw_dict=newpop.pdraws)
 
 new_omegas={}
 
@@ -73,14 +100,14 @@
             'mpp': lambda_samples['mpp'][idx],
             'sigpp': lambda_samples['sigpp'][idx],
             'rate': lambda_samples['rate'][idx],
-            'gamma': lambda_samples['lamb'][idx],
-            'kappa': 5,
-            'z_peak': 0.3*(0.5-np.random.rand())+1.9,
+            'gamma': 3.61, # lambda_samples['lamb'][idx],
+            'kappa': 3.83,
+            'z_peak': 1.04, #2.0, #0.3*(0.5-np.random.rand())+1.9,
             }
     new_omegas['Lambdas'].append(Lambda_new)
 
     newpop.calculate_omega_gw(sampling_frequency=2048, Lambda=Lambda_new)
-    new_omegas['Neff'].append( (np.sum(newpop.weights)**2)/(np.sum(newpop.weights**2)) )
+    new_omegas['Neff'].append(float( (xp.sum(newpop.weights)**2)/(xp.sum(newpop.weights**2)) ))
     new_omegas['omega_gw'].append(newpop.omega_gw.tolist())
 
 new_omegas['freqs']=newpop.frequency_array.tolist()