def plot_dists(surv_pop, telescope):
"""Plot the fluence and DM distribution of a surveyed population.
Args:
surv_pop (Population): Population from which to plot
telescope (str): Name of the telescope with which to compare the
distribution. Necessary for Frbcat.
"""
# Use a nice font for axes
plt.rc('text', usetex=True)
# Plot dm distribution
f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
dm_frbpoppy = surv_pop.frbs.dm
pprint(f'Number of detected FRBs: {len(dm_frbpoppy)}')
ax1.step(*hist(dm_frbpoppy), where='mid', linestyle='dashed')
df = Frbcat().df
dm_frbcat = df[df.telescope == telescope].dm
ax1.step(*hist(dm_frbcat), where='mid')
# Compare distributions
ks = ks_2samp(dm_frbpoppy, dm_frbcat)
text = fr'$p={round(ks[1], 2)}$'
anchored_text = AnchoredText(text, loc=1, borderpad=1., frameon=False)
ax1.add_artist(anchored_text)
ax1.set_xlabel(r'DM ($\textrm{pc}\ \textrm{cm}^{-3}$)')
ax1.set_ylabel('Fraction')
ax1.set_ylim([0, 1.1])
ax1.set_xlim([0, 3500])
# Plot fluence distributions
fluence_frbpoppy = surv_pop.frbs.fluence
ax2.step(*hist(fluence_frbpoppy, bins='log'),
where='mid',
label='frbpoppy',
linestyle='dashed')
fluence_frbcat = df[df.telescope == telescope].fluence
ax2.step(*hist(fluence_frbcat, bins='log'), where='mid', label='frbcat')
# Compare distributions
ks = ks_2samp(fluence_frbpoppy, fluence_frbcat)
text = fr'$p={round(ks[1], 2)}$'
anchored_text = AnchoredText(text, loc=1, borderpad=1., frameon=False)
ax2.add_artist(anchored_text)
ax2.set_xlabel(r'Fluence (Jy ms)')
ax2.set_ylim([0, 1.1])
ax2.set_xlim([5e-1, 1e4])
plt.xscale('log')
plt.figlegend(loc='upper center', ncol=2, framealpha=1)
plt.tight_layout()
plt.savefig(f'plots/frbpoppy_{telescope}.pdf')
plt.clf()
def get_data():
"""Get the population data."""
# Construct population
pop = CosmicPopulation(n_srcs=SIZE, n_days=1, name='standard_candle')
pop.set_dist(model='sfr', z_max=2.5, H_0=67.74, W_m=0.3089, W_v=0.6911)
pop.set_dm_host(model='constant', value=100)
pop.set_dm_igm(model='ioka', slope=1000, std=None)
pop.set_dm_mw(model='ne2001')
pop.set_emission_range(low=10e6, high=10e9)
pop.set_lum(model='constant', value=1e36)
pop.set_w(model='constant', value=1.)
pop.set_si(model='constant', value=0)
pop.generate()
# Survey population
pops = {}
for b in BEAMPATTERNS:
pprint(f'Surveying with {b} beampattern')
n_s = 0
bp = b
if b.startswith('airy'):
bp, n_s = b.split('-')
n_s = int(n_s)
survey = Survey(name='perfect-small')
# Prevent beam from getting larger than the sky
survey.set_beam(model=bp, n_sidelobes=n_s, size=10)
surv_pop = SurveyPopulation(pop, survey)
print(surv_pop.source_rate)
pops[b] = surv_pop
return pops
def complex_rates(remake=REMAKE, alphas=ALPHAS, size=SIZE, surveys=SURVEYS):
"""Calculate expected rates for a complex populations."""
rates = defaultdict(list)
# Don't always regenerate a population
if remake is False:
for alpha in alphas:
for s in surveys:
surv_rates = unpickle(f'complex_alpha_{alpha}_{s}').source_rate
pprint(f'Alpha:{alpha:.2}, Survey: {s}, Det: {surv_rates.det}')
rate = (surv_rates.det / surv_rates.days)
rates[s].append(rate)
else:
pops = []
for alpha in alphas:
if alpha <= -1.0 and ADAPTATIVE_SCALING:
size = 1e7
if alpha <= -1.5 and ADAPTATIVE_SCALING:
size = 1e8
pop = CosmicPopulation.complex(size)
pop.set_dist(model='vol_co',
z_max=2.5,
alpha=alpha,
H_0=67.74,
W_m=0.3089,
W_v=0.6911)
pop.set_lum(model='powerlaw', low=1e40, high=1e45, power=-1)
pop.name = f'complex_alpha_{alpha}'
pops.append(pop)
# Set up surveys
ss = []
for s in surveys:
survey = Survey(name=s)
survey.set_beam(model='airy', n_sidelobes=1)
ss.append(survey)
surv_pops = LargePopulation(pop, *ss).pops
for i, s in enumerate(surveys):
surv_rates = surv_pops[i].source_rate
pprint(f'Alpha:{alpha:.2}, Survey: {s}, Det: {surv_rates.det}')
rate = (surv_rates.det / surv_rates.days)
rates[s].append(rate)
# Scale rates to first survey in list
for s in surveys:
if s != surveys[0]:
norm = []
for i, r in enumerate(rates[s]):
norm.append(r / rates[surveys[0]][i])
rates[s] = norm
rates[surveys[0]] = [r / r for r in rates[surveys[0]]]
return rates
def simple_rates(remake=REMAKE, alphas=ALPHAS, size=SIZE, surveys=SURVEYS):
"""Calculate expected rates for a simple populations."""
rates = defaultdict(list)
# Don't always regenerate a population
if remake is False:
for alpha in alphas:
for s in surveys:
surv_rates = unpickle(f'simple_alpha_{alpha}_{s}').rates()
pprint(f'Alpha:{alpha:.2}, Survey: {s}, Det: {surv_rates.det}')
rate = (surv_rates.det / surv_rates.days)
rates[s].append(rate)
else:
pops = []
for alpha in alphas:
pop = CosmicPopulation.simple(size)
pop.alpha = alpha
pop.name = f'simple_alpha_{alpha}'
pops.append(pop)
# Set up surveys
ss = []
for s in surveys:
survey = Survey(name=s,
gain_pattern='perfect',
n_sidelobes=0.5)
ss.append(survey)
surv_pops = LargePopulation(pop, *ss).pops
for i, s in enumerate(surveys):
surv_rates = surv_pops[i].rates()
pprint(f'Alpha:{alpha:.2}, Survey: {s}, Det: {surv_rates.det}')
rate = (surv_rates.det / surv_rates.days)
rates[s].append(rate)
# Scale rates to HTRU
for s in surveys:
if s != 'htru':
norm = []
for i, r in enumerate(rates[s]):
norm.append(r / rates['htru'][i])
rates[s] = norm
rates['htru'] = [r / r for r in rates['htru']]
return rates
def calc_gofs(self, run):
# For each requested run
self.so = SimulationOverview()
par_set = self.so.df[self.so.df.run == run].par_set.iloc[0]
pprint(f'Calculating goodness of fit for run {run}, par set {par_set}')
pars = self.run_pars[par_set]
values = []
# Loop through all combination of parameters
for values, group in self.so.df[self.so.df.run == run].groupby(pars):
pprint(f' - {list(zip(pars, values))}')
# Calculate goodness of fit values for each simulation
for row_ix, row in group.iterrows():
survey_name = row.survey
uuid = row.uuid
pop = unpickle(f'mc/run_{run}/{uuid}')
# Apply a DM cutoff
mask = (pop.frbs.dm <= 950)
pop.frbs.apply(mask)
pop.source_rate.det = pop.n_sources() * pop.source_rate.f_area
dm_gof = self.dm(pop, survey_name)
snr_gof = self.snr(pop, survey_name)
self.so.df.at[row_ix, 'dm_gof'] = dm_gof
self.so.df.at[row_ix, 'snr_gof'] = snr_gof
if pop.n_sources() == 0:
self.so.df.at[row_ix, 'weight'] = 0
self.so.df.at[row_ix, 'n_det'] = pop.n_sources()
pprint(f' - No sources in {survey_name}')
continue
# Find corresponding rate normalisation population uuid
norm_mask = dict(zip(pars, values))
norm_mask['survey'] = self.norm_surv
norm_mask['run'] = run
k = norm_mask.keys()
v = norm_mask.values()
norm_uuid = group.loc[group[k].isin(v).all(axis=1), :].uuid
norm_uuid = norm_uuid.values[0]
rate_diff, n_det = self.rate(pop, survey_name, norm_uuid, run)
# Get rate weighting
self.so.df.at[row_ix, 'weight'] = rate_diff
self.so.df.at[row_ix, 'n_det'] = n_det
pprint(f'Saving the results for run {run}')
# Best matching in terms of rates
max_w = np.nanmax(self.so.df.weight)
self.so.df.loc[self.so.df.weight == 1e3]['weight'] = max_w
self.so.save()
def limit_ra_dec(pop, pointings):
"""Doesn't work at boundary coordinates."""
pprint('Limiting coordinates')
def sample(n_gen):
u = np.random.uniform
ra = u(0, 360, n_gen)
dec = np.rad2deg(np.arccos(u(-1, 1, n_gen))) - 90
return ra, dec
def accept(ra, dec):
coords = np.full(len(ra), False)
r = np.sqrt(40 / np.pi)
for p_ra, p_dec in pointings:
limits = go.separation(ra, dec, p_ra, p_dec) < r
coords[limits] = True
return coords
# Limit population to smaller area
# Sample RA, dec
ra, dec = sample(pop.n_gen)
mask = accept(ra, dec)
reject, = np.where(~mask)
while reject.size > 0:
fill_ra, fill_dec = sample(reject.size)
mask = accept(fill_ra, fill_dec)
ra[reject[mask]] = fill_ra[mask]
dec[reject[mask]] = fill_dec[mask]
reject = reject[~mask]
frbs = pop.frbs
frbs.ra = ra
frbs.dec = dec
# Convert to galactic coordinates
frbs.gl, frbs.gb = go.radec_to_lb(frbs.ra, frbs.dec, frac=True)
pop.gen_gal_coords()
return pop
def get_probabilities(normalise=False):
"""Get the number of bursts per maximum time.
Args:
normalise (type): Whether to normalise at each maximum time, such that
there's a number of bursts at each time stamp having a maximum
probability of one.
Returns:
array: Number of bursts per maximum time.
"""
days = np.arange(1, N_DAYS, 1)
m_bursts = np.arange(0, M_BURSTS, 1)
xx, yy = np.meshgrid(m_bursts, days)
prob = np.full((len(days), len(m_bursts)), np.nan)
# Mask any frbs over the maximum time
time = td.clustered(r=R, k=K, n_srcs=N_SRCS, n_days=N_DAYS, z=0)
pprint('Masking days')
for i, day in reversed(list(enumerate(days))):
time[(time > day)] = np.nan
m_bursts = (~np.isnan(time)).sum(1)
unique, counts = np.unique(m_bursts, return_counts=True)
try:
prob[(i, unique)] = counts
except IndexError:
pprint('Please ensure M is large enough.')
exit()
prob = prob.T / N_SRCS
if normalise:
# Normalise at each maximum time (highest chance becomes one)
prob = prob / np.nanmax(prob, axis=0)
return prob
def bw_from_epn():
"""Scrape bandwidth from EPN website."""
pprint('Getting data from EPN')
bw_db = {'name': [], 'bw': []}
url = 'http://www.epta.eu.org/epndb/list.php'
def clean(l):
return float(l.split('>')[1].split(' ')[0])
with urllib.request.urlopen(url) as resp:
data = resp.read().decode().replace(' ', ' ')
for line in data.split('\n'):
if line.startswith('<li>'):
name = line.split()[0].split('>')[-1]
freqs = [clean(l) for l in line.split('<a ')[1:]]
bw = (max(freqs) - min(freqs))
bw_db['name'].append(name)
bw_db['bw'].append(bw)
pprint('Finished getting data from EPN')
bw_db = pd.DataFrame.from_dict(bw_db)
return bw_db
def __str__(self):
"""How to print the class."""
# Set up title
r = '{:20.19} {:>10} {:>10}\n'
t = r.format(self.name, 'Days', 'FRBs')
line = '-'*len(t.split('\n')[-2].strip()) + '\n'
t += line
# Format rates
rdays = round(self.days, 3)
t += r.format('In population', rdays, round(self.tot()))
t += r.format('Detected', rdays, round(self.det, 3))
t += r.format('Too late', rdays, round(self.late, 3))
t += r.format('Too faint', rdays, round(self.faint, 3))
t += r.format('Outside survey', rdays, round(self.out, 3))
t += r.format('/Gpc^3', 365.25, round(self.vol, 3))
t += r.format('Expected', round(self.exp, 4), 1)
t += line
return pprint(t, output=False)
def iter_run(i):
r = CosmicPopulation(N_SRCS, n_days=N_DAYS, repeaters=True)
r.set_dist(model='vol_co', z_max=1.0)
r.set_dm_host(model='gauss', mean=100, std=200)
r.set_dm_igm(model='ioka', slope=1000, std=None)
r.set_dm(mw=True, igm=True, host=True)
r.set_emission_range(low=100e6, high=10e9)
r.set_lum(model='powerlaw',
per_source='different',
low=1e40,
high=1e45,
power=0)
r.set_si(model='gauss', mean=-1.4, std=1)
r.set_w(model='lognormal', per_source='different', mean=0.1, std=1)
rate = lognormal(RATE, 2, N_SRCS)
if LARGE_POP:
rate = lognormal(RATE, 2, int(MAX_SIZE)) # Not completely kosher
r.set_time(model='poisson', rate=rate)
# Set up survey
s = Survey('chime-frb', n_days=N_DAYS)
s.set_beam(model='chime-frb')
s.gen_pointings() # To ensure each sub pop has the same pointings
# Only generate FRBs in CHIME's survey region
r.set_direction(model='uniform',
min_ra=s.ra_min,
max_ra=s.ra_max,
min_dec=s.dec_min,
max_dec=s.dec_max)
if LARGE_POP:
surv_pop = LargePopulation(r, s, max_size=MAX_SIZE).pops[0]
else:
r.generate()
surv_pop = SurveyPopulation(r, s)
surv_pop.name = f'cosmic_chime_longer_{i}'
surv_pop.save()
print(surv_pop.source_rate)
print(surv_pop.burst_rate)
pprint(f'i: {i}')
pprint(f'# one-offs: {surv_pop.n_one_offs()}')
pprint(f'# repeaters: {surv_pop.n_repeaters()}')
def gen_par_set_4(self,
parallel=True,
alpha=-1.5,
si=0,
li=-1,
lum_min=1e40,
lum_max=1e40,
w_mean=np.nan,
w_std=np.nan,
run=np.nan):
dm_igm_slopes = np.linspace(800, 1200, 11)
dm_hosts = np.linspace(0, 500, 11)
# Put all options into a dataframe
self.so.df = self.so.df[self.so.df.run != run]
opt = np.meshgrid(dm_igm_slopes, dm_hosts, self.survey_ix)
options = np.array(opt).T.reshape(-1, 3)
cols = ('dm_igm_slope', 'dm_host', 'survey')
df = pd.DataFrame(options, columns=cols)
df['run'] = run
df['par_set'] = 4
df['uuid'] = [uuid.uuid4() for _ in range(len(df.index))]
df['date'] = datetime.today()
self.so.append(df)
self.so.map_surveys(self.survey_ix, self.survey_names)
self.so.save()
# Remove previous par_set of the same number
if not self.set_up_dirs(run=run):
fs = f'{frbpoppy.paths.populations()}mc/run_{run}/*'
for f in glob(fs):
os.remove(f)
pop = CosmicPopulation.complex(self.pop_size)
if not np.isnan(alpha):
pop.set_dist(model='vol_co', z_max=1.0, alpha=alpha)
pop.set_si(model='constant', value=si)
if not np.isnan(lum_min):
pop.set_lum(model='powerlaw', low=lum_min, high=lum_max, index=li)
if not np.isnan(w_mean):
pop.set_w(model='lognormal', mean=w_mean, std=w_std)
pop.generate()
def adapt_pop(e):
dm_igm_slope, dm_host = e
t_pop = deepcopy(pop)
t_pop.set_dm_igm(model='ioka', slope=dm_igm_slope)
t_pop.gen_dm_igm()
t_pop.set_dm_host(model='constant', value=dm_host)
t_pop.gen_dm_host()
t_pop.frbs.dm = t_pop.frbs.dm_mw + t_pop.frbs.dm_igm
t_pop.frbs.dm += t_pop.frbs.dm_host
for survey in self.surveys:
surv_pop = SurveyPopulation(t_pop, survey)
# Get unique identifier
mask = (self.so.df.par_set == 4)
mask &= (self.so.df.run == run)
mask &= (self.so.df.dm_igm_slope == dm_igm_slope)
mask &= (self.so.df.dm_host == dm_host)
mask &= (self.so.df.survey == survey.name)
uuid = self.so.df[mask].uuid.iloc[0]
surv_pop.name = f'mc/run_{run}/{uuid}'
surv_pop.save()
n_cpu = min([4, os.cpu_count() - 1])
pprint(f'{os.cpu_count()} CPUs available')
mg = np.meshgrid(dm_igm_slopes, dm_hosts)
loop = np.array(mg).T.reshape(-1, 2)
if parallel:
Parallel(n_jobs=n_cpu)(delayed(adapt_pop)(e) for e in tqdm(loop))
else:
[adapt_pop(e) for e in tqdm(loop)]
r.set_emission_range(low=100e6, high=10e9)
r.set_lum(model='powerlaw',
per_source='different',
low=1e40,
high=1e45,
power=0)
r.set_si(model='gauss', mean=-1.4, std=1)
r.set_w(model='lognormal', per_source='different', mean=0.1, std=1)
rate = lognormal(ra, 1, int(n))
r.set_time(model='poisson', rate=rate)
# Set up survey
s = Survey('chime-frb', n_days=N_DAYS)
s.set_beam(model='chime-frb')
# Only generate FRBs in CHIME's survey region
r.set_direction(model='uniform',
min_ra=s.ra_min,
max_ra=s.ra_max,
min_dec=s.dec_min,
max_dec=s.dec_max)
r.generate()
surv_pop = SurveyPopulation(r, s)
surv_pop.name = 'cosmic_chime'
print(surv_pop.source_rate)
print(surv_pop.burst_rate)
pprint(f'# one-offs: {surv_pop.n_one_offs()}')
pprint(f'# repeaters: {surv_pop.n_repeaters()}')
plot_aa_style()
fig, ax1 = plt.subplots(1, 1)
# Fluence plot
ax1.set_xlabel('S/N')
ax1.set_xscale('log')
ax1.set_ylabel(r'\#(${>}\text{S/N}$)')
ax1.set_yscale('log')
# Update fluence plot
for i, surv_pop in enumerate(surv_pops):
name = surv_pop.name.split('_')[-1]
snr = surv_pop.frbs.snr
if snr.size == 0:
pprint(f'No FRBs in {name} population')
continue
bins, values = hist(snr, bin_type='log', norm=None)
# Cumulative sum
values = np.cumsum(values[::-1])[::-1]
# Normalise to area on sky
if not np.isnan(values.all()):
values = values * surv_pop.source_rate.f_area
plt.step(bins, values, where='mid', label=name)
plt.legend()
plt.tight_layout()
for ax in axes.flatten():
ax.set_aspect('auto')
# Get norm pop
y = 0
ys = []
names = []
rates = []
norm_sim_rate = surv_pops[0].source_rate.det
norm_real_rate = EXPECTED['parkes-htru'][0] / EXPECTED['parkes-htru'][1]
norm_rate = norm_sim_rate / norm_real_rate
for i, surv_pop in enumerate(surv_pops):
name = surv_pop.name.split('_')[-1]
pprint(name)
if surv_pop.n_sources() == 0:
print(surv_pop.source_rate)
print(f'{name} | no FRBs in population')
continue
names.append(name)
ys.append(y)
# Dimensions measure plot
ax = axes[0]
ax.set_xlabel(r'DM ($\textrm{pc}\ \textrm{cm}^{-3}$)')
ax.set_ylabel(r'\#')
ax.set_yscale('log')
bins, values = hist(surv_pop.frbs.dm,
"""Generate a repeater population and split into repeaters and one-offs."""
import numpy as np
from frbpoppy import CosmicPopulation, Survey, SurveyPopulation
from frbpoppy import split_pop, pprint, plot
DAYS = 1
r = CosmicPopulation.simple(int(1e4), n_days=DAYS, repeaters=True)
r.set_time(model='regular', rate=2)
r.set_lum(model='powerlaw', low=1e40, high=1e45, per_source='different')
r.generate()
survey = Survey('chime-frb', n_days=DAYS)
survey.set_beam(model='perfect')
surv_pop = SurveyPopulation(r, survey)
# Split population into seamingly one-off and repeater populations
mask = ((~np.isnan(surv_pop.frbs.time)).sum(1) > 1)
pop_ngt1, pop_nle1 = split_pop(surv_pop, mask)
pop_ngt1.name += ' (> 1 burst)'
pop_nle1.name += ' (1 burst)'
pops = [pop_nle1, pop_ngt1]
pprint(f'{surv_pop.n_sources()} sources detected')
pprint(f'{surv_pop.n_bursts()} bursts detected')
plot(surv_pop)
def gen_par_set_1(self,
parallel=True,
lum_min=np.nan,
lum_max=np.nan,
w_mean=np.nan,
w_std=np.nan,
dm_igm_slope=np.nan,
dm_host=np.nan,
run=0):
alphas = np.linspace(-2.5, -1, 11)
sis = np.linspace(-2, 2, 11)
lis = np.linspace(-2, 0, 11)
# Put all options into a dataframe
if 'run' in self.so.df:
self.so.df = self.so.df[self.so.df.run != run]
opt = np.meshgrid(alphas, sis, lis, self.survey_ix)
options = np.array(opt).T.reshape(-1, 4)
df = pd.DataFrame(options, columns=('alpha', 'si', 'li', 'survey'))
df['run'] = run
df['par_set'] = 1
df['uuid'] = [uuid.uuid4() for _ in range(len(df.index))]
df['date'] = datetime.today()
self.so.append(df)
self.so.map_surveys(self.survey_ix, self.survey_names)
self.so.save()
# Remove previous par_set of the same number
if not self.set_up_dirs(run=run):
fs = f'{frbpoppy.paths.populations()}mc/run_{run}/*'
for f in glob(fs):
os.remove(f)
def iter_alpha(i):
alpha = alphas[i]
pop = CosmicPopulation.complex(self.pop_size)
pop.set_dist(model='vol_co', z_max=1.0, alpha=alpha)
pop.set_lum(model='constant', value=1)
if not np.isnan(w_mean):
pop.set_w(model='lognormal', mean=w_mean, std=w_std)
if not np.isnan(dm_igm_slope):
pop.set_dm_igm(model='ioka', slope=dm_igm_slope)
pop.set_dm_host(model='constant', value=dm_host)
pop.generate()
for si in sis:
pop.set_si(model='constant', value=si)
pop.gen_si()
for li in lis:
pop.set_lum(model='powerlaw',
low=1e40,
high=1e45,
power=li)
if not np.isnan(lum_min):
pop.set_lum(model='powerlaw',
low=lum_min,
high=lum_max,
index=li)
pop.gen_lum()
for survey in self.surveys:
surv_pop = SurveyPopulation(pop, survey)
# Get unique identifier
mask = (self.so.df.par_set == 1)
mask &= (self.so.df.run == run)
mask &= (self.so.df.alpha == alpha)
mask &= (self.so.df.si == si)
mask &= (self.so.df.li == li)
mask &= (self.so.df.survey == survey.name)
uuid = self.so.df[mask].uuid.iloc[0]
surv_pop.name = f'mc/run_{run}/{uuid}'
surv_pop.save()
if parallel:
n_cpu = min([3, os.cpu_count() - 1])
pprint(f'{os.cpu_count()} CPUs available')
r = range(len(alphas))
Parallel(n_jobs=n_cpu)(delayed(iter_alpha)(i) for i in tqdm(r))
else:
[iter_alpha(i) for i in tqdm(range(len(alphas)))]
elif lum_func == 'flat pl.':
pop.set_lum(model='powerlaw', per_source='different', low=1e40,
high=1e45, power=0)
elif lum_func == 'gauss':
pop.set_lum(model='gauss', per_source='different', mean=1e42,
std=1e10)
elif lum_func == 'std. candle':
pop.set_lum(model='constant', value=1e42)
pop.generate()
surv_pop = SurveyPopulation(pop, survey)
# Split population into seamingly one-off and repeater populations
mask = ((~np.isnan(surv_pop.frbs.time)).sum(1) > 1)
pop_ngt1, pop_nle1 = split_pop(surv_pop, mask)
pprint(f'{pop_ngt1.n_bursts()} repeater bursts')
if pop_ngt1.n_bursts() < 10:
pprint('Insufficient FRB sources for plotting')
# Plot
if SCATTER:
def sample(n):
return np.random.randint(low=0, high=pop_ngt1.n_sources(), size=n)
def accept(ii):
return surv_pop.frbs.dm[ii] < 1000
ii = sample(25)
mask = accept(ii)
reject, = np.where(~mask)
while reject.size > 0:
def plot(self, run):
# Get data
# For each requested run
df = self.so.df
par_set = df[df.run == run].par_set.iloc[0]
# For each parameter
for main_par in self.run_pars[par_set]:
pprint(f'Plotting {main_par}')
other_pars = [e for e in self.run_pars[par_set] if e != main_par]
for compare_par in ['dm', 'snr']:
compare_col = f'{compare_par}_gof'
pprint(f' - {compare_col}')
for survey, group_surv in df[df.run == run].groupby('survey'):
pprint(f' - {survey}')
# Set up plot
plot_aa_style()
plt.rcParams["figure.figsize"] = (5.75373 * 3, 5.75373 * 3)
plt.rcParams['figure.max_open_warning'] = 125
n_x = group_surv[other_pars[0]].nunique()
if len(other_pars) > 1:
n_y = group_surv[other_pars[1]].nunique()
else:
n_y = 1
fig, ax = plt.subplots(n_x,
n_y,
sharex='col',
sharey='row')
groups = group_surv.groupby(other_pars)
x = -1
for i, (other_pars_vals, group) in enumerate(groups):
bins = group[main_par].values
values = group[compare_col].values
bins, values = self.add_edges_to_hist(bins, values)
if n_y > 1:
y = i % n_y
if y == 0:
x += 1
a = ax[y, x]
else:
y = i
a = ax[y]
a.step(bins, values, where='mid')
a.set_title = str(other_pars_vals)
diff = np.diff(bins)
if diff[1] != diff[0]:
a.set_xscale('log')
# Set axis label
if y == n_y - 1:
p = other_pars[0]
if isinstance(other_pars_vals, float):
val = other_pars_vals
else:
val = other_pars_vals[0]
p = p.replace('_', ' ')
a.set_xlabel(f'{p} = {val:.2}')
if x == 0:
p = other_pars[1]
val = other_pars_vals[1]
p = p.replace('_', ' ')
a.set_ylabel(f'{p} = {val:.2}')
# Set axis limits
subset = df[df.run == run][main_par]
y_subset = group_surv[compare_col].copy()
try:
low = np.nanmin(y_subset)
high = np.nanmax(y_subset)
except ValueError:
low = 0.0001
high = 1
log = False
if low > 0 and high > 0:
log = True
for a in ax.flatten():
a.set_xlim(subset.min(), subset.max())
if log:
a.set_yscale('log', nonposy='clip')
a.set_ylim(low, high)
p = main_par.replace('_', ' ')
fig.suptitle(f'{p} - {compare_par} - {survey}')
plt.tight_layout()
plt.subplots_adjust(top=0.95)
# Save to subdirectory
path_to_save = rel_path(f'./plots/mc/{main_par}_run{run}/')
if not os.path.isdir(path_to_save):
os.mkdir(path_to_save)
path_to_save += f'{compare_par}_{survey}.pdf'
plt.savefig(path_to_save)
plt.clf()
per_source='different')
r.set_time(model='poisson', rate=3)
r.set_dm_igm(model='ioka', slope=1000, std=0)
r.set_dm(mw=False, igm=True, host=False)
r.set_w('constant', value=1)
r.generate()
# Set up survey
survey = Survey('perfect', n_days=DAYS)
survey.set_beam(model='perfect')
survey.snr_limit = 1e6
surv_pop = SurveyPopulation(r, survey)
pprint(f'{r.n_bursts()}:{surv_pop.n_bursts()}')
pprint(f'{surv_pop.n_sources()} sources detected')
if r.n_bursts() < PLOTTING_LIMIT_N_SRCS:
pprint('Not sufficient FRB sources for plotting')
exit()
# Split population into seamingly one-off and repeater populations
mask = ((~np.isnan(surv_pop.frbs.time)).sum(1) > 1)
pop_rep, pop_one = split_pop(surv_pop, mask)
pop_rep.name += ' (> 1 burst)'
pop_one.name += ' (1 burst)'
if INTERACTIVE_PLOT:
plot(r, pop_rep, pop_one, tns=False, mute=False)
def generate(parallel=False):
from joblib import Parallel, delayed
from frbpoppy import CosmicPopulation, Survey, pprint
from frbpoppy import SurveyPopulation
from tqdm import tqdm
# Set up rate dataframe
paras = [ALPHAS, LIS, SIS]
vals = np.array(np.meshgrid(*paras)).T.reshape(-1, len(paras))
cols = ['alpha', 'li', 'si']
df = pd.DataFrame(vals, columns=cols)
# Set up surveys
surveys = []
for name in SURVEY_NAMES:
survey = Survey(name=name)
survey.set_beam(model='airy', n_sidelobes=1)
surveys.append(survey)
df[name] = np.nan
def iter_alpha(i, surveys=surveys, parallel=None):
alpha = ALPHAS[i]
pop = CosmicPopulation.complex(SIZE)
pop.set_dist(model='vol_co', z_max=1.0, alpha=alpha)
pop.set_lum(model='powerlaw', low=1e40, high=1e45, power=-1)
pop.generate()
for li in LIS:
pop.set_lum(model='powerlaw', low=1e40, high=1e45, power=li)
pop.gen_lum()
for si in SIS:
pop.set_si(model='constant', value=si)
pop.gen_si()
pop.name = f'complex_alpha_{alpha}_lum_{li}_si_{si}'
for survey in surveys:
surv_pop = SurveyPopulation(pop, survey)
print(surv_pop.name)
surv_pop.save()
sr = surv_pop.source_rate
rate = sr.det / sr.days
mask = (df.alpha == alpha) & (df.li == li) & (df.si == si)
if parallel is not None:
i = df[mask].index
j = SURVEY_NAMES.index(survey.name)
parallel[i, j] = rate
else:
df.loc[mask, survey.name] = rate
if parallel:
n_cpu = min([3, os.cpu_count() - 1])
pprint(f'{os.cpu_count()} CPUs available')
r = range(len(ALPHAS))
temp_path = ('./temp.mmap')
# Make a temp memmap to have a sharedable memory object
temp = np.memmap(temp_path,
dtype=np.float64,
shape=(len(vals), len(SURVEY_NAMES)),
mode='w+')
Parallel(n_jobs=n_cpu)(delayed(iter_alpha)(i, parallel=temp)
for i in tqdm(r))
for name in SURVEY_NAMES:
col = SURVEY_NAMES.index(name)
df[name] = temp[:, col]
else:
for i in tqdm(range(len(ALPHAS)), desc='Alphas'):
iter_alpha(i)
df.to_csv(CSV_PATH)
