cosmic-popsynth 3.4.17__cp310-cp310-macosx_14_0_arm64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
cosmic/test_sample.py ADDED
@@ -0,0 +1,559 @@
1
+ """Unit test for cosmic
2
+ """
3
+
4
+ __author__ = 'Katie Breivik <katie.breivik@gmail.com>'
5
+
6
+ import os
7
+ import unittest
8
+ import numpy as np
9
+ import pandas as pd
10
+ from cosmic.sample import InitialBinaryTable
11
+ from cosmic.sample.sampler.independent import Sample
12
+ from cosmic.sample.sampler.multidim import MultiDim
13
+ from cosmic.sample.sampler.cmc import CMCSample
14
+ from cosmic.sample.cmc import elson
15
+ from cosmic.sample.initialcmctable import InitialCMCTable
16
+ from scipy.optimize import curve_fit
17
+ from cosmic.utils import a_from_p, get_porb_norm
18
+
19
+ SAMPLECLASS = Sample()
20
+ MULTIDIMSAMPLECLASS = MultiDim()
21
+ CMCSAMPLECLASS = CMCSample()
22
+ TEST_DATA_DIR = os.path.join(os.path.split(__file__)[0], 'data')
23
+
24
+ #distribution slopes/power laws
25
+ KROUPA_93_HI = -2.7
26
+ KROUPA_93_MID = -2.2
27
+ KROUPA_93_LO = -1.3
28
+ KROUPA_01_HI = -2.3
29
+ KROUPA_01_LO = -1.3
30
+ SALPETER_55 = -2.35
31
+ SANA12_PORB_POWER_LAW = -0.55
32
+ FLAT_SLOPE = 0.0
33
+ # Since we integrate up to 0.7 in ecc, the slope is 4.0 rounded at the 0th decimal
34
+ THERMAL_SLOPE = 4.0
35
+ SANA12_ECC_POWER_LAW = -0.45
36
+ SANA12_ECC_POWER_LAW_ROUND = -0.5
37
+ MEAN_RAGHAVAN = 4.9
38
+ SIGMA_RAGHAVAN = 2.3
39
+ CLOSE_BINARY_FRAC = 0.42
40
+
41
+
42
+ N_BINARY_SELECT = 85
43
+ VANHAAFTEN_BINFRAC_MAX = 0.9989087986493874
44
+ VANHAAFTEN_BINFRAC_MIN = 0.6192803136799157
45
+ OFFNER_MASS_RANGES = [(0.075,0.15), (0.15,0.30), (0.3,0.6), (0.75,1.00), (0.85,1.25), (1.00,1.25), (1.6,2.4), (3,5), (5,8), (8,17), (17,50)]
46
+ OFFNER_DATA = [0.19, 0.23, 0.30, 0.42, 0.47, 0.50, 0.68, 0.81, 0.89, 0.93, 0.96]
47
+ OFFNER_ERRORS = [0.03, 0.02, 0.02, 0.03, 0.03, 0.04, 0.07, 0.06, 0.05, 0.04, 0.04]
48
+ MULTIDIM_BINFRAC_MAX = 0.6146916774140262
49
+ MULTIDIM_BINFRAC_MIN = 0.13786300908773025
50
+ CONST_SFR_SUM = 460028.2453521937
51
+ BURST_SFR_SUM = 946002.8245352195
52
+ KSTAR_SOLAR = 1.0
53
+ MOE_TOTAL_MASS = 20.27926225850954
54
+ METALLICITY_1000 = 0.02
55
+ METALLICITY_13000 = 0.02*0.15
56
+
57
+ QMIN_LOWMASS = 0.1
58
+ QMIN_HIGHMASS = 0.7
59
+ M2MIN_LOWMASS = 0.08
60
+ M2MIN_HIGHMASS = 12.0
61
+ BINFRAC_LOWMASS = 499139
62
+ BINFRAC_HIGHMASS = 2205
63
+
64
+ KING_TEST_DATA = np.load(os.path.join(TEST_DATA_DIR, "cmc_king_test.npz"))
65
+ ELSON_TEST_DATA = np.load(os.path.join(TEST_DATA_DIR, "cmc_elson_test.npz"))
66
+ PLUMMER_TEST_DATA = np.load(os.path.join(TEST_DATA_DIR, "cmc_plummer_test.npz"))
67
+ R_PLUMMER_TEST_ARRAY, VR_PLUMMER_TEST_ARRAY, VT_PLUMMER_TEST_ARRAY = PLUMMER_TEST_DATA["arr_0"], PLUMMER_TEST_DATA["arr_1"], PLUMMER_TEST_DATA["arr_2"]
68
+ R_ELSON_TEST_ARRAY, VR_ELSON_TEST_ARRAY, VT_ELSON_TEST_ARRAY = ELSON_TEST_DATA["arr_0"], ELSON_TEST_DATA["arr_1"], ELSON_TEST_DATA["arr_2"]
69
+ R_KING_TEST_ARRAY, VR_KING_TEST_ARRAY, VT_KING_TEST_ARRAY = KING_TEST_DATA["arr_0"], KING_TEST_DATA["arr_1"], KING_TEST_DATA["arr_2"]
70
+
71
+ REFF_TEST_ARRAY = np.array([3.94190562, 5.99895482])
72
+
73
+ SINGLES_CMC_FITS, BINARIES_CMC_FITS = InitialCMCTable.read(filename=os.path.join(TEST_DATA_DIR, "input_cmc.fits"))
74
+ SINGLES_CMC_HDF5, BINARIES_CMC_HDF5 = InitialCMCTable.read(filename=os.path.join(TEST_DATA_DIR, "input_cmc.hdf5"))
75
+
76
+ def power_law_fit(data, n_bins=100):
77
+ def line(x, a, b):
78
+ return x*a + b
79
+ def center_bins(bins):
80
+ mid_bin = []
81
+ for bin_lo, bin_hi in zip(bins[:-1], bins[1:]):
82
+ mid_bin.append(bin_lo + (bin_hi-bin_lo)/2)
83
+ return mid_bin
84
+
85
+ hist, bins = np.histogram(data, bins=n_bins)
86
+ bins = center_bins(bins)
87
+
88
+ xdata = np.log10(bins)
89
+ ydata = np.log10(hist)
90
+
91
+ popt, pcov = curve_fit(line, xdata, ydata)
92
+
93
+ slope, intercept = popt[0], popt[1]
94
+
95
+ return slope
96
+
97
+ def linear_fit(data):
98
+ def line(x, a, b):
99
+ return x*a + b
100
+ def center_bins(bins):
101
+ mid_bin = []
102
+ for bin_lo, bin_hi in zip(bins[:-1], bins[1:]):
103
+ mid_bin.append(bin_lo + (bin_hi-bin_lo)/2)
104
+ return mid_bin
105
+
106
+ hist, bins = np.histogram(data, bins=100, density=True)
107
+ bins = center_bins(bins)
108
+ popt, pcov = curve_fit(line, bins, hist)
109
+
110
+ slope, intercept = popt[0], popt[1]
111
+
112
+ return slope
113
+
114
+ class TestSample(unittest.TestCase):
115
+ """`TestCase` for the cosmic Sample class, which generates several
116
+ independent initial parameters drawn from specified distributions
117
+ """
118
+
119
+ def test_sample_primary_kroupa93(self):
120
+ np.random.seed(2)
121
+ # Check that the sample_primary function samples mass correctly
122
+ mass, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa93', size=1000000)
123
+ mass_hi = mass[mass > 1.0]
124
+ # filter out highest masses because they kill us with the histograms
125
+ mass_hi = mass_hi[mass_hi < 10.0]
126
+ mass_mid = mass[(mass <= 1.0)]
127
+ mass_mid = mass_mid[mass_mid > 0.5]
128
+ mass_lo = mass[mass <= 0.5]
129
+
130
+ #few bins for the most massive stars
131
+ power_slope_hi = power_law_fit(mass_hi, n_bins=50)
132
+ power_slope_mid = power_law_fit(mass_mid)
133
+ power_slope_lo = power_law_fit(mass_lo)
134
+
135
+ self.assertEqual(np.round(power_slope_hi, 1), KROUPA_93_HI)
136
+ self.assertEqual(np.round(power_slope_mid, 1), KROUPA_93_MID)
137
+ self.assertEqual(np.round(power_slope_lo, 1), KROUPA_93_LO)
138
+
139
+
140
+ def test_sample_primary_kroupa01(self):
141
+ np.random.seed(2)
142
+ # Check that the sample_primary function samples mass correctly
143
+ mass, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa01', size=1000000)
144
+ mass_hi = mass[mass > 0.5]
145
+ # filter out highest masses because they kill us with the histograms
146
+ mass_hi = mass_hi[mass_hi < 10.0]
147
+ mass_lo = mass[mass <= 0.5]
148
+
149
+ #few bins for the most massive stars
150
+ power_slope_hi = power_law_fit(mass_hi, n_bins=50)
151
+ power_slope_lo = power_law_fit(mass_lo)
152
+
153
+ self.assertEqual(np.round(power_slope_hi, 1), KROUPA_01_HI)
154
+ self.assertEqual(np.round(power_slope_lo, 1), KROUPA_01_LO)
155
+
156
+ def test_sample_primary_salpeter55(self):
157
+ np.random.seed(3)
158
+ # Check that the sample_primary function samples mass correctly
159
+ mass, total_mass = SAMPLECLASS.sample_primary(primary_model='salpeter55', size=10000000)
160
+ #filter on mass to get better statistics
161
+ power_slope = power_law_fit(mass[mass < 1.0], n_bins=50)
162
+ self.assertEqual(np.round(power_slope, 2), SALPETER_55)
163
+
164
+ def test_sample_secondary(self):
165
+ np.random.seed(2)
166
+ # Check that the sample_secondary function samples secondary mass correctly
167
+ mass1, total_mass = SAMPLECLASS.sample_primary(primary_model='salpeter55', size=10000000)
168
+ ind_massive, = np.where(mass1 > 15.0)
169
+ ind_not_massive, = np.where(mass1 < 15.0)
170
+
171
+ np.random.seed(2)
172
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass=mass1, qmin=0.1)
173
+ q = mass2[ind_massive] / mass1[ind_massive]
174
+ slope = linear_fit(q)
175
+ self.assertEqual(np.round(slope, 1), FLAT_SLOPE)
176
+
177
+ np.random.seed(2)
178
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass=mass1, qmin=-1)
179
+ q = mass2[ind_not_massive] / mass1[ind_not_massive]
180
+ slope = linear_fit(q)
181
+ self.assertEqual(np.round(slope, 1), FLAT_SLOPE)
182
+
183
+ np.random.seed(2)
184
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass=mass1, m2_min=0.08)
185
+ q = mass2[ind_massive] / mass1[ind_massive]
186
+ slope = linear_fit(q)
187
+ self.assertEqual(np.round(slope, 1), FLAT_SLOPE)
188
+
189
+ np.random.seed(2)
190
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass=mass1, qmin=0.1, m2_min=0.08)
191
+ q = mass2[ind_massive] / mass1[ind_massive]
192
+ slope = linear_fit(q)
193
+ self.assertEqual(np.round(slope, 1), FLAT_SLOPE)
194
+
195
+ def test_sample_q(self):
196
+ """Test you can sample different mass ratio distributions"""
197
+ np.random.seed(2)
198
+ mass1, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa01', size=10000000)
199
+ for slope in [0, 1, 2]:
200
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass=mass1, q_power_law=slope, qmin=0.0)
201
+ q = mass2 / mass1
202
+ fit_slope = power_law_fit(q)
203
+ self.assertEqual(np.round(fit_slope, 1), slope)
204
+
205
+ def test_binary_select(self):
206
+ np.random.seed(2)
207
+ # Check that the binary select function chooses binarity properly
208
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=np.arange(1,100), binfrac_model=0.0)
209
+ self.assertEqual(len(m1_b), 0)
210
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=np.arange(1,100), binfrac_model=1.0)
211
+ self.assertEqual(len(m1_b), 99)
212
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=np.arange(1,100), binfrac_model='vanHaaften')
213
+ self.assertEqual(len(m1_b), N_BINARY_SELECT)
214
+
215
+ def test_binary_fraction(self):
216
+ np.random.seed(2)
217
+ # Check that the binary fraction tracking is correct
218
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=np.arange(1,100), binfrac_model=0.5)
219
+ self.assertEqual(binfrac.max(), 0.5)
220
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=np.arange(1,100), binfrac_model='vanHaaften')
221
+ self.assertEqual(binfrac.max(), VANHAAFTEN_BINFRAC_MAX)
222
+ self.assertEqual(binfrac.min(), VANHAAFTEN_BINFRAC_MIN)
223
+
224
+ test_fracs = []
225
+ test_errs = []
226
+ primary_mass = np.array([float(x) for x in np.logspace(np.log10(0.08), np.log10(150), num=100000)])
227
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=primary_mass, binfrac_model='offner22')
228
+ for i in range(len(OFFNER_MASS_RANGES)):
229
+ low, high = OFFNER_MASS_RANGES[i][0], OFFNER_MASS_RANGES[i][1]
230
+ offner_value = OFFNER_DATA[i]
231
+ offner_error = OFFNER_ERRORS[i]
232
+ bins_count = len(m1_b[(m1_b >= low) & (m1_b <= high)])
233
+ singles_count = len(m1_s[(m1_s >= low) & (m1_s <= high)])
234
+ bin_frac = bins_count / (bins_count + singles_count)
235
+ error = abs(offner_value - bin_frac)
236
+ self.assertLess(error, offner_error)
237
+
238
+ def test_msort(self):
239
+ np.random.seed(2)
240
+ mass1, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa01', size=1000000)
241
+ # Check that qmin_msort and m2_min_msort are workings as expected
242
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass = mass1, qmin=0.1, m2_min=0.08, msort=15, qmin_msort=0.7, m2_min_msort=12)
243
+ ind_light, = np.where(mass1 < 15.0)
244
+ ind_massive, = np.where(mass1 >= 15.0)
245
+ m2_light = mass2[ind_light]
246
+ m2_massive = mass2[ind_massive]
247
+ q_light = mass2[ind_light]/mass1[ind_light]
248
+ q_massive = mass2[ind_massive]/mass1[ind_massive]
249
+ assert m2_light.min() > M2MIN_LOWMASS
250
+ assert m2_massive.min() > M2MIN_HIGHMASS
251
+ assert q_light.min() > QMIN_LOWMASS
252
+ assert q_massive.min() > QMIN_HIGHMASS
253
+ # Check that the binary fraction tracking is correct when using msort
254
+ m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=mass1, binfrac_model=0.5, msort=15, binfrac_model_msort=0.8)
255
+ assert np.sum(binfrac==0.5) == BINFRAC_LOWMASS
256
+ assert np.sum(binfrac==0.8) == BINFRAC_HIGHMASS
257
+
258
+ #
259
+ def test_sample_porb(self):
260
+ # next do Sana12
261
+ np.random.seed(4)
262
+ mass1, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa01', size=100000)
263
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass = mass1, qmin=0.1)
264
+ rad1 = SAMPLECLASS.set_reff(mass=mass1, metallicity=0.02)
265
+ rad2 = SAMPLECLASS.set_reff(mass=mass2, metallicity=0.02)
266
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
267
+ mass1, mass2, rad1, rad2, 'sana12', size=mass1.size
268
+ )
269
+ power_slope = power_law_fit(np.log10(porb))
270
+ self.assertEqual(np.round(power_slope, 2), SANA12_PORB_POWER_LAW)
271
+
272
+ # now some custom power laws
273
+ for slope in [-0.5, 0.5, 1]:
274
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
275
+ mass1, mass2, rad1, rad2, porb_model={
276
+ "min": 0.15, "max": 5, "slope": slope
277
+ }, size=mass1.size
278
+ )
279
+ power_slope = power_law_fit(np.log10(porb))
280
+ self.assertEqual(np.round(power_slope, 1), slope)
281
+
282
+ np.random.seed(5)
283
+ # next do Renzo+19
284
+ m1_high = mass1+15
285
+ rad1_high = SAMPLECLASS.set_reff(mass=m1_high, metallicity=0.02)
286
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
287
+ m1_high, mass2, rad1_high, rad2, 'renzo19', size=m1_high.size
288
+ )
289
+ porb_cut = porb[porb > 2.5]
290
+ power_slope = power_law_fit(np.log10(porb_cut))
291
+ self.assertAlmostEqual(np.round(power_slope, 2), SANA12_PORB_POWER_LAW)
292
+
293
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
294
+ mass1, mass2, rad1, rad2, 'renzo19', size=mass1.size
295
+ )
296
+ ind_log_uniform, = np.where(mass1 <= 15)
297
+ porb_log_uniform = porb[ind_log_uniform]
298
+ uniform_slope = linear_fit(np.log10(porb_log_uniform))
299
+ self.assertEqual(np.round(uniform_slope, 1), FLAT_SLOPE)
300
+
301
+ # next check the log uniform
302
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
303
+ mass1, mass2, rad1, rad2, 'log_uniform', size=mass1.size
304
+ )
305
+ power_slope = linear_fit(np.log10(porb))
306
+ sep = a_from_p(porb, mass1, mass2)
307
+ sep = sep[sep > 10]
308
+ uniform_slope = linear_fit(np.log10(sep))
309
+ self.assertEqual(np.round(uniform_slope, 1), FLAT_SLOPE)
310
+
311
+ # next check raghavan10
312
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
313
+ mass1, mass2, rad1, rad2, 'raghavan10', size=mass1.size
314
+ )
315
+ log_porb_mean = np.mean(np.log10(porb))
316
+ log_porb_sigma = np.std(np.log10(porb))
317
+ self.assertTrue(np.round(log_porb_mean, 1) >= MEAN_RAGHAVAN-0.15)
318
+ self.assertEqual(np.round(log_porb_sigma, 0), np.round(SIGMA_RAGHAVAN, 0))
319
+
320
+ # next check moe19
321
+ from cosmic.utils import get_met_dep_binfrac
322
+ from scipy.interpolate import interp1d
323
+ from scipy.stats import kstest
324
+ metallicity = 0.001
325
+ # this is a metallicity dependent population:
326
+ binfrac = get_met_dep_binfrac(metallicity)
327
+ print(binfrac)
328
+ mass1, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa01', size=100000)
329
+ (mass1_binaries, mass_single, binfrac_binaries, binary_index,) = SAMPLECLASS.binary_select(
330
+ mass1, binfrac_model=binfrac,
331
+ )
332
+ mass2_binaries = SAMPLECLASS.sample_secondary(
333
+ primary_mass=mass1_binaries, qmin=0.1
334
+ )
335
+ rad1 = SAMPLECLASS.set_reff(mass=mass1_binaries, metallicity=metallicity)
336
+ rad2 = SAMPLECLASS.set_reff(mass=mass2_binaries, metallicity=metallicity)
337
+
338
+
339
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
340
+ mass1_binaries, mass2_binaries, rad1, rad2, 'moe19', size=mass1_binaries.size, met=metallicity
341
+ )
342
+
343
+ binary_frac = len(porb) / (len(mass1))
344
+ self.assertAlmostEqual(np.round(binary_frac, 2), binfrac)
345
+
346
+
347
+
348
+ def test_sample_ecc(self):
349
+ np.random.seed(2)
350
+ # Check that the sample_ecc function samples ecc properly
351
+
352
+ # first sample orbital periods
353
+ np.random.seed(4)
354
+ mass1, total_mass = SAMPLECLASS.sample_primary(primary_model='kroupa01', size=100000)
355
+ mass2 = SAMPLECLASS.sample_secondary(primary_mass = mass1, qmin=0.1)
356
+ rad1 = SAMPLECLASS.set_reff(mass=mass1, metallicity=0.02)
357
+ rad2 = SAMPLECLASS.set_reff(mass=mass2, metallicity=0.02)
358
+ porb,aRL_over_a = SAMPLECLASS.sample_porb(
359
+ mass1, mass2, rad1, rad2, 'sana12', size=mass1.size
360
+ )
361
+
362
+ # now we feed aRL_over_a into sample_ecc
363
+ ecc = SAMPLECLASS.sample_ecc(aRL_over_a, ecc_model='thermal', size=mass1.size)
364
+ ecc_cut = ecc[ecc < 0.7]
365
+ slope = linear_fit(ecc_cut)
366
+ self.assertEqual(np.round(slope, 0), THERMAL_SLOPE)
367
+
368
+ ecc = SAMPLECLASS.sample_ecc(aRL_over_a, ecc_model='sana12', size=mass1.size)
369
+ ecc_cut = ecc[ecc < 0.91]
370
+ power_slope = power_law_fit(ecc_cut)
371
+ self.assertEqual(np.round(power_slope, 1), SANA12_ECC_POWER_LAW_ROUND)
372
+
373
+ ecc = SAMPLECLASS.sample_ecc(aRL_over_a, ecc_model='circular', size=mass1.size)
374
+ self.assertEqual(np.mean(ecc), 0.0)
375
+
376
+
377
+ def test_sample_SFH(self):
378
+ np.random.seed(2)
379
+ # Check that the sample SFH function samples SFH='const' correctly
380
+ times, met = SAMPLECLASS.sample_SFH(SF_start=10000.0,\
381
+ SF_duration=10000.0,\
382
+ met = 0.02, size=100)
383
+ self.assertEqual(times.sum(), CONST_SFR_SUM)
384
+ self.assertAlmostEqual(np.mean(met), 0.02)
385
+
386
+ np.random.seed(2)
387
+ # Check that the sample SFH function samples SFH='burst' correctly
388
+ times, met = SAMPLECLASS.sample_SFH(SF_start=10000.0,\
389
+ SF_duration=1000.0,\
390
+ met = 0.02, size=100)
391
+ self.assertEqual(times.sum(), BURST_SFR_SUM)
392
+ self.assertAlmostEqual(np.mean(met), 0.02)
393
+
394
+ # Check that the sample SFH function samples SFH='delta_burst' correctly
395
+ times, met = SAMPLECLASS.sample_SFH(SF_start=10000.0,\
396
+ SF_duration=0.0,\
397
+ met = 0.02, size=100)
398
+ self.assertEqual(times.sum(), 100*10000.0)
399
+ self.assertAlmostEqual(np.mean(met), 0.02)
400
+
401
+ def test_set_kstar(self):
402
+ # Check that the kstar is selected properly
403
+ kstar = SAMPLECLASS.set_kstar(pd.DataFrame([1.0, 1.0, 1.0, 1.0, 1.0]))
404
+ self.assertEqual(np.mean(kstar), KSTAR_SOLAR)
405
+
406
+ def test_Moe_sample(self):
407
+ # Test the multidim sampler and system-by-system binary fraction
408
+ m1, m2, porb, ecc, mass_singles, mass_binaries, n_singles, n_binaries, binfrac = MULTIDIMSAMPLECLASS.initial_sample(rand_seed = 2, size=10, nproc=1, mp_seeds=[0])
409
+ self.assertEqual(np.sum(mass_singles), MOE_TOTAL_MASS)
410
+ self.assertAlmostEqual(binfrac.max(), MULTIDIM_BINFRAC_MAX)
411
+ self.assertAlmostEqual(binfrac.min(), MULTIDIM_BINFRAC_MIN)
412
+
413
+ def test_sample_MultiDim_SFH(self):
414
+ np.random.seed(2)
415
+ # Check that the sample SFH function samples SFH='const' correctly
416
+ times, met = MULTIDIMSAMPLECLASS.sample_SFH(SF_start=10000.0,\
417
+ SF_duration=10000.0,\
418
+ met = 0.02, size=100)
419
+ self.assertEqual(times.sum(), CONST_SFR_SUM)
420
+ self.assertAlmostEqual(np.mean(met), 0.02)
421
+
422
+ np.random.seed(2)
423
+ # Check that the sample SFH function samples SFH='burst' correctly
424
+ times, met = MULTIDIMSAMPLECLASS.sample_SFH(SF_start=10000.0,\
425
+ SF_duration=1000.0,\
426
+ met = 0.02, size=100)
427
+ self.assertEqual(times.sum(), BURST_SFR_SUM)
428
+ self.assertAlmostEqual(np.mean(met), 0.02)
429
+
430
+ # Check that the sample SFH function samples SFH='delta_burst' correctly
431
+ times, met = MULTIDIMSAMPLECLASS.sample_SFH(SF_start=10000.0,\
432
+ SF_duration=0.0,\
433
+ met = 0.02, size=100)
434
+ self.assertEqual(times.sum(), 100*10000.0)
435
+ self.assertAlmostEqual(np.mean(met), 0.02)
436
+
437
+ def test_set_kstar_MultiDim(self):
438
+ # Check that the kstar is selected properly
439
+ kstar = MULTIDIMSAMPLECLASS.set_kstar(pd.DataFrame([1.0, 1.0, 1.0, 1.0, 1.0]))
440
+ self.assertEqual(np.mean(kstar), KSTAR_SOLAR)
441
+
442
+ def test_sampling_binfrac_zero(self):
443
+ # check that you can't sample based on size with a binary fraction of 0.0
444
+ it_fails = False
445
+ try:
446
+ InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
447
+ primary_model='kroupa01', ecc_model='thermal',
448
+ porb_model='sana12', binfrac_model=0.0,
449
+ SF_start=10.0, SF_duration=0.0, met=0.02,
450
+ size=1000)
451
+ except ValueError:
452
+ it_fails = True
453
+ self.assertTrue(it_fails)
454
+
455
+ def test_sampling_targets_bad_input(self):
456
+ # check that you get an error if you don't supply a size or total_mass
457
+ it_fails = False
458
+ try:
459
+ InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
460
+ primary_model='kroupa01', ecc_model='thermal',
461
+ porb_model='sana12', binfrac_model=0.5,
462
+ SF_start=10.0, SF_duration=0.0, met=0.02,
463
+ sampling_target="total_mass")
464
+ except ValueError:
465
+ it_fails = True
466
+ self.assertTrue(it_fails)
467
+
468
+ it_fails = False
469
+ try:
470
+ InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
471
+ primary_model='kroupa01', ecc_model='thermal',
472
+ porb_model='sana12', binfrac_model=0.5,
473
+ SF_start=10.0, SF_duration=0.0, met=0.02,
474
+ total_mass=None, sampling_target="total_mass")
475
+ except ValueError:
476
+ it_fails = True
477
+ self.assertTrue(it_fails)
478
+
479
+ it_fails = False
480
+ try:
481
+ InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
482
+ primary_model='kroupa01', ecc_model='thermal',
483
+ porb_model='sana12', binfrac_model=0.5,
484
+ SF_start=10.0, SF_duration=0.0, met=0.02,
485
+ size=None, total_mass=None, sampling_target="size")
486
+ except ValueError:
487
+ it_fails = True
488
+ self.assertTrue(it_fails)
489
+
490
+ def test_sampling_targets_size(self):
491
+ # check that you can sample based on size
492
+ for size in np.random.randint(100, 1000, size=100):
493
+ samples = InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
494
+ primary_model='kroupa01', ecc_model='thermal',
495
+ porb_model='sana12', binfrac_model=0.5,
496
+ SF_start=10.0, SF_duration=0.0, met=0.02,
497
+ size=size)
498
+ self.assertGreaterEqual(len(samples[0]), size)
499
+
500
+ def test_sampling_targets_mass(self):
501
+ # check that you can sample based on total mass
502
+ for mass in np.random.randint(100, 1000, size=100):
503
+ samples = InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
504
+ primary_model='kroupa01', ecc_model='thermal',
505
+ porb_model='sana12', binfrac_model=0.5,
506
+ SF_start=10.0, SF_duration=0.0, met=0.02,
507
+ size=None, total_mass=mass, sampling_target="total_mass")
508
+ self.assertGreaterEqual(samples[1] + samples[2], mass)
509
+
510
+ def test_sampling_targets_mass_trimmed(self):
511
+ # check that you can sample based on total mass and that the samples are trimmed properly
512
+ for mass in np.random.randint(100, 1000, size=100):
513
+ samples = InitialBinaryTable.sampler('independent', np.arange(16), np.arange(16),
514
+ primary_model='kroupa01', ecc_model='thermal',
515
+ porb_model='sana12', binfrac_model=0.5,
516
+ SF_start=10.0, SF_duration=0.0, met=0.02,
517
+ size=None, total_mass=mass, sampling_target="total_mass",
518
+ trim_extra_samples=True)
519
+ self.assertLessEqual(abs(samples[1] + samples[2] - mass), 300)
520
+
521
+ class TestCMCSample(unittest.TestCase):
522
+ def test_plummer_profile(self):
523
+ np.random.seed(2)
524
+ r, vr, vt = CMCSAMPLECLASS.set_r_vr_vt('plummer',N=100, r_max=300)
525
+ np.testing.assert_allclose(VR_PLUMMER_TEST_ARRAY, vr, rtol=1e-5)
526
+ np.testing.assert_allclose(VT_PLUMMER_TEST_ARRAY, vt, rtol=1e-5)
527
+ np.testing.assert_allclose(R_PLUMMER_TEST_ARRAY, r, rtol=1e-5)
528
+
529
+ def test_elson_profile(self):
530
+ np.random.seed(2)
531
+ r, vr, vt = CMCSAMPLECLASS.set_r_vr_vt('elson',N=100, r_max=300, gamma=3)
532
+ np.testing.assert_allclose(VR_ELSON_TEST_ARRAY, vr, rtol=1e-5)
533
+ np.testing.assert_allclose(VT_ELSON_TEST_ARRAY, vt, rtol=1e-5)
534
+ np.testing.assert_allclose(R_ELSON_TEST_ARRAY, r, rtol=1e-5)
535
+
536
+ def test_king_profile(self):
537
+ np.random.seed(2)
538
+ r, vr, vt = CMCSAMPLECLASS.set_r_vr_vt('king',N=100, w_0=5)
539
+ np.testing.assert_allclose(VR_KING_TEST_ARRAY, vr, rtol=1e-5)
540
+ np.testing.assert_allclose(VT_KING_TEST_ARRAY, vt, rtol=1e-5)
541
+ np.testing.assert_allclose(R_KING_TEST_ARRAY, r, rtol=1e-5)
542
+
543
+ def test_set_reff(self):
544
+ reff = CMCSAMPLECLASS.set_reff(mass=np.array([10.0, 20.0]), metallicity=0.02)
545
+ np.testing.assert_allclose(REFF_TEST_ARRAY, reff)
546
+
547
+ def test_cmc_sampler(self):
548
+ np.random.seed(2)
549
+ # Test generating CMC initial conditions and test saving the output to files
550
+ Singles, Binaries = InitialCMCTable.sampler('cmc', binfrac_model=0.2, primary_model='kroupa01', ecc_model='sana12', porb_model='sana12', cluster_profile='plummer', met=0.014, size=20, params=os.path.join(TEST_DATA_DIR,'Params.ini'), gamma=4, r_max=100, qmin=0.1)
551
+ InitialCMCTable.write(Singles, Binaries, filename="input.hdf5")
552
+ InitialCMCTable.write(Singles, Binaries, filename="input.fits")
553
+ Singles, Binaries = InitialCMCTable.read(filename="input.fits")
554
+ # read the test files and compare to the static unit tests files
555
+ pd.testing.assert_frame_equal(Singles, SINGLES_CMC_FITS)
556
+ pd.testing.assert_frame_equal(Binaries, BINARIES_CMC_FITS)
557
+ Singles, Binaries = InitialCMCTable.read(filename="input.hdf5")
558
+ pd.testing.assert_frame_equal(Singles, SINGLES_CMC_HDF5)
559
+ pd.testing.assert_frame_equal(Binaries, BINARIES_CMC_HDF5)