def extract_batches(source,
batches=None,
save_plots=True,
multithread=True,
reload=False,
load_bursts=False,
load_summary=False,
basename='xrb',
param_table=None):
"""Do burst analysis on arbitrary number of batches"""
t0 = time.time()
if param_table is not None:
print('Using models from table provided')
batches = np.unique(param_table['batch'])
else:
batches = grid_tools.ensure_np_list(batches)
for batch in batches:
print_title(f'Batch {batch}')
analysis_path = grid_strings.batch_analysis_path(batch, source)
for folder in ['input', 'output']:
path = os.path.join(analysis_path, folder)
grid_tools.try_mkdir(path, skip=True)
if param_table is not None:
subset = grid_tools.reduce_table(param_table,
params={'batch': batch})
runs = np.array(subset['run'])
else:
n_runs = grid_tools.get_nruns(batch, source)
runs = np.arange(n_runs) + 1
if multithread:
args = []
for run in runs:
args.append((run, batch, source, save_plots, reload,
load_bursts, load_summary, basename))
with mp.Pool(processes=8) as pool:
pool.starmap(extract_runs, args)
else:
extract_runs(runs,
batch,
source,
reload=reload,
save_plots=save_plots,
load_bursts=load_bursts,
load_summary=load_summary,
basename=basename)
print_title('Combining run tables')
for table_name in ('summary', 'bursts'):
burst_tools.combine_run_tables(batch,
source,
table_name=table_name)
t1 = time.time()
dt = t1 - t0
print_title(f'Time taken: {dt:.1f} s ({dt/60:.2f} min)')
def get_source(self, source):
"""Returns subset of minbar table for given source
"""
if source not in self.sources:
raise ValueError(f'source "{source}" not in minbar table. '
'Look in self.sources for full list')
return grid_tools.reduce_table(self.table, params={'name': source})
def setup_table(kgrid, batches, synth_source, mc_source, mc_version, synth_version,
params=('x', 'z', 'accrate', 'qb', 'mass'),
summ_list=('rate', 'dt', 'fluence', 'peak'),
free_params=('redshift', 'd_b', 'xi_ratio'),
observables=('rate', 'peak', 'fluence', 'fper'),
update=False):
"""Sets up table of synthetic data, including input/output values
parameters
----------
kgrid : Kgrid object
batches : list
list of batches, each corresponding to an epoch. Assumes the runs in each
batch correspond to each other.
synth_source : str
mc_source : str
mc_version : int
synth_version : int
params : sequence(str)
parameters to extract from kgrid and add to the table
summ_list : sequence(str)
burst properties to extract from kgrid and add to the table
free_params : sequence(str)
free parameters to randomly choose
observables : sequence(str)
names of observables to calculate from burst properties
update : bool
whether to simply update an existing synth table,
e.g. if the models just need updating. Keeps the previously randomised params
"""
mcv = mcmc_versions.McmcVersion(source=mc_source, version=mc_version)
sub = grid_tools.reduce_table(kgrid.params, params={'batch': batches[0]})
groups = np.array(sub['run'])
table = pd.DataFrame()
# ===== For each group of epochs, setup sub-table of inputs/outputs =====
for group in groups:
group_table = initialise_group_table(group, batches)
set_param_cols(group_table, batches=batches, kgrid=kgrid, params=params)
set_summ_cols(group_table, batches=batches, kgrid=kgrid, summ_list=summ_list)
if update:
inherit_free_params(group_table, free_params=free_params, parent_group=group,
parent_source=synth_source, parent_version=synth_version)
else:
set_rand_free_params(group_table, mcv=mcv, free_params=free_params)
set_observables(group_table, observables=observables)
table = pd.concat([table, group_table], ignore_index=True)
save_table(table, source=synth_source, version=synth_version)
return table
def linregress_qnuc(source, grid_version=0):
"""Returns table of linear fits to optimal Qnuc's (versus accretion rate)
"""
param_list = ['x', 'z', 'qb', 'accdepth', 'accmass', 'mass']
full_table = load_qnuc_table(source, grid_version)
accrates = np.unique(
full_table['accrate']) # assumes all parameter-sets have this accrate
param_table = grid_tools.reduce_table(full_table,
params={'accrate': accrates[0]})
linr_table = param_table.reset_index()[param_list]
for i in range(len(linr_table)):
params = linr_table.loc[i][param_list].to_dict()
sub_table = grid_tools.reduce_table(full_table, params=params)
linr = linregress(sub_table['accrate'], sub_table['qnuc'])
linr_table.loc[i, 'm'] = linr[0]
linr_table.loc[i, 'y0'] = linr[1]
return linr_table
def plot_qnuc(source, mass, grid_version, linear=True):
"""Plots optimised Qnuc versus accretion rate, for a given (Newtonian) mass
parameters
----------
source : str
grid source (see grid_analyser.py)
mass : flt
Newtonian mass of models (as listed in grid params table)
grid_version : int
linear : bool
plot linear regression line
"""
table = qnuc_tools.load_qnuc_table(source, grid_version)
table = grid_tools.reduce_table(table, params={'mass': mass})
acc_unique = np.unique(table['accrate'])
sub_params = grid_tools.reduce_table(table,
params={'accrate': acc_unique[0]})
fig, ax = plt.subplots()
for row in sub_params.itertuples():
models = grid_tools.reduce_table(table,
params={
'x': row.x,
'z': row.z
})
ax.plot(models['accrate'],
models['qnuc'],
marker='o',
label=f'x={row.x:.3f}, z={row.z:.4f}')
if linear:
linr_table = qnuc_tools.linregress_qnuc(source, grid_version)
row = linr_table[linr_table['mass'] == mass]
x = np.array([np.min(acc_unique), np.max(acc_unique)])
y = row.m.values[0] * x + row.y0.values[0]
ax.plot(x, y, color='black', ls='--')
ax.set_title(f'mass={mass:.1f}')
ax.legend()
plt.show(block=False)
def extract_obs_data(source, version, group,
bprops=('rate', 'fluence', 'peak', 'fper')):
"""Returns summary "observed" data as dictionary, for Burstfit
"""
table = load_table(source, version)
group_table = grid_tools.reduce_table(table, params={'group': group})
obs_data = {}
for bprop in bprops:
for var in (bprop, f'u_{bprop}'):
obs_data[var] = np.array(group_table[f'{var}_obs'])
return obs_data
def multi_save(table, source, basename='xrb'):
"""Extract models from table of arbitrary batches/runs
"""
batches = np.unique(table['batch'])
t0 = time.time()
for batch in batches:
subset = grid_tools.reduce_table(table, params={'batch': batch})
runs = np.array(subset['run'])
args = []
for run in runs:
args.append((run, batch, source, basename, True))
with mp.Pool(processes=8) as pool:
pool.starmap(load_lum, args)
t1 = time.time()
dt = t1 - t0
print(f'Time taken: {dt:.1f} s ({dt/60:.2f} min)')
def predict_qnuc(params, source, linr_table=None, grid_version=0):
"""Predict optimal Qnuc for given accrate and mass
linr_table : pd.DataFrame (optional)
provide linr_table directly (if linr_table=None, will load new table)
"""
params = params.copy()
accrate = params.pop('accrate')
if linr_table is None:
linr_table = linregress_qnuc(source, grid_version)
row = grid_tools.reduce_table(linr_table, params=params)
if len(row) == 0:
raise ValueError(f'Qnuc not available for {params}')
elif len(row) > 1:
raise ValueError(
f'Qnuc not uniquely defined for given params (underdefined)')
return accrate * row.m.values[0] + row.y0.values[0]
def load_group_table(source, version, group):
"""Returns synth table of single group of epochs
"""
table = load_table(source, version)
return grid_tools.reduce_table(table, params={'group': group})