def unpickle(filename=None):
"""Quick function to unpickle a population.
Args:
filename (str, optional): Define the path to the pickled population,
or give the population name
Returns:
Population: Population class
"""
# Find population file
if os.path.isfile(filename):
f = open(filename, 'rb')
else:
# Find standard population files
try:
name = filename.lower()
p = paths.populations() + f'{name}.p'
f = open(p, 'rb')
except FileNotFoundError:
s = 'Pickled population file "{0}" does not exist'.format(filename)
raise FileNotFoundError(s)
# Unpickle
pop = pickle.load(f)
f.close()
return pop
def save(self, extention='p'):
"""
Write out source properties as data file.
Args:
extention (str): Type of file to save.
Choice from ['csv', 'dat', 'p']
"""
# Check if a population has been a survey name
if not self.name:
file_name = 'pop'
else:
file_name = self.name.lower()
path = paths.populations() + file_name
# Set file types
if extention == 'p':
path += '.p'
self.to_pickle(path)
if extention == 'csv':
path += '.csv'
self._to_csv(path)
elif extention == 'dat':
path += '.dat'
self._to_csv(path, sep=' ')
def save(self, path=None):
"""
Write out source properties as data file.
Args:
path (str): Path to which to save.
"""
if path is None:
# Check if a population has been a survey name
if not self.name:
file_name = 'pop'
else:
file_name = self.name.lower()
if self.uid:
file_name += f'_{self.uid}'
path = paths.populations() + f'{file_name}.p'
self.to_pickle(path)
def plot(*pops,
files=[],
frbcat=True,
show=True,
no_browser=False,
mute=True,
port=5006,
print_command=False):
"""
Plot populations with bokeh. Has to save populations before plotting.
Args:
*pops (Population, optional): Add the populations you would like to
see plotted
files (list, optional): List of population files to plot.
frbcat (bool, optional): Whether to plot frbcat parameters. Defaults to
True
show (bool, optional): Whether to display the plot or not. Mainly used
for debugging purposes. Defaults to True.
port (int): The port on which to launch Bokeh
print_command (bool): Whether to show the command do_plot is running
"""
if len(pops) > 0:
# Save populations
for pop in pops:
if type(pop) == str:
name = pop
else:
pop.name = pop.name.lower() + '_for_plotting'
name = pop.name
pop.save()
# Save location
file_name = name + '.p'
out = os.path.join(paths.populations(), file_name)
files.append(out)
# Command for starting up server
command = 'nice -n 19'.split(' ')
if show:
command.extend('bokeh serve --show'.split(' '))
else:
command.append('python3')
# Command for changing port
if port != 5006:
command.append(f'--port={port}')
# Add script path
script = 'plot.py'
out = os.path.join(os.path.dirname(__file__), script)
command.append(out)
# For the arguments
command.append('--args')
# Add frbcat
command.append('-frbcat')
if frbcat is False:
command.append('False')
if frbcat is True:
command.append('True')
elif type(frbcat) == str and len(frbcat) > 0:
command.append(f'{frbcat}')
# Add in populations
for f in files:
command.append(f'"{f}"')
# Let people know what's happening
pprint('Plotting populations')
if print_command:
pprint(' '.join(command))
pprint('Press Ctrl+C to quit')
# Add method to gracefully quit plotting
try:
with open(os.devnull, 'w') as f:
if mute:
out = f
else:
out = None
subprocess.run(command, stderr=out, stdout=out)
except KeyboardInterrupt:
print(' ')
sys.exit()
def merge(self):
"""Merge populations."""
pprint('Merging populations')
self.pops = []
for s in self.surveys:
pops = []
files = [f'{self.base_name}_{s.name}_{uid}' for uid in self.uids]
for f in files:
pops.append(unpickle(f))
# Main population
mp = pops[0]
# Merge each parameter
for attr in mp.frbs.__dict__.keys():
parm = getattr(mp.frbs, attr)
if type(parm) is np.ndarray:
parms = []
for pop in pops:
parms.append(getattr(pop.frbs, attr))
try:
merged_parm = np.concatenate(parms, axis=0)
except ValueError:
# Check maximum size values should be padded to
max_size = max([p.shape[1] for p in parms])
new_parms = []
# Ensure matrices are the same shapes by padding them
for p in parms:
if p.shape[1] != max_size:
padded_p = np.zeros((p.shape[0], max_size))
padded_p[:] = np.nan
padded_p[:, :p.shape[1]] = p
new_parms.append(padded_p)
else:
new_parms.append(p)
merged_parm = np.concatenate(new_parms, axis=0)
setattr(mp.frbs, attr, merged_parm)
# Add up detections
for pop in pops[1:]:
mp.source_rate.faint += pop.source_rate.faint
mp.source_rate.late += pop.source_rate.late
mp.source_rate.out += pop.source_rate.out
mp.source_rate.det += pop.source_rate.det
if mp.repeaters:
mp.burst_rate.faint += pop.burst_rate.faint
mp.burst_rate.late += pop.burst_rate.late
mp.burst_rate.out += pop.burst_rate.out
mp.burst_rate.det += pop.burst_rate.det
mp.burst_rate.pointing += pop.burst_rate.pointing
# Recalculate detection rates
mp.calc_rates(s)
# Save the main population as one big population
mp.uid = None
mp.save()
# Remove all of the smaller ones
for f in files:
p = paths.populations() + f'{f}.p'
os.remove(p)
self.pops.append(mp)