def test_threshold_gain_selector():
selector = ThresholdGainSelector()
print(selector)
assert 'NectarCam' in selector.thresholds
threshold = selector.thresholds['NectarCam']
good_hg_value = 35
good_lg_value = 50
dummy_waveforms = np.ones((2, 1000, 30)) * good_lg_value
dummy_waveforms[1:] = good_hg_value #
dummy_waveforms[0, 500:, 13:15] = threshold + 10
new_waveforms, gain_mask = selector.select_gains("NectarCam",
dummy_waveforms)
assert gain_mask.shape == (1000,)
assert new_waveforms.shape == (1000, 30)
assert (new_waveforms[500:] == good_hg_value).all()
assert (new_waveforms[:500] == good_lg_value).all()
selector.select_by_sample = True
new_waveforms, gain_mask = selector.select_gains("NectarCam",
dummy_waveforms)
assert new_waveforms.shape == (1000, 30)
assert (new_waveforms[500:, 13:15] == good_hg_value).all()
assert (new_waveforms[500:, :13] == good_lg_value).all()
assert (new_waveforms[500:, 15:] == good_lg_value).all()
assert gain_mask.shape == new_waveforms.shape
# test some failures:
# Camera that doesn't have a threshold:
with pytest.raises(KeyError):
selector.select_gains("NonExistantCamera", dummy_waveforms)
# 3-gain channel input:
with pytest.raises(ValueError):
selector.select_gains("NectarCam", np.ones((3, 1000, 30)))
# 1-gain channel input:
wf0 = np.ones((1, 1000, 1))
wf1, gm = selector.select_gains("ASTRICam", wf0)
assert wf1.shape == (1000,)
assert gm.shape == (1000,)
def gain_selection(waveform, signals, peakpos, cam_id, threshold):
"""
Custom lst calibration.
Update event.dl1.tel[telescope_id] with calibrated image and peakpos
Parameters
----------
waveform: array of waveforms of the events
signals: array of calibrated pixel charges
peakpos: array of pixel peak positions
cam_id: str
threshold: int threshold to change form high gain to low gain
"""
'''
combined = signals[0].copy()
peaks = peakpos[0].copy()
for pixel in range(0, combined.size):
if np.any(waveform[0][pixel] >= threshold):
combined[pixel] = signals[1][pixel]
peaks[pixel] = peakpos[1][pixel]
'''
###Gain Selection using ctapipe GainSelector###
gainsel = ThresholdGainSelector(select_by_sample=True)
gainsel.thresholds[cam_id] = threshold
waveform, gainmask = gainsel.select_gains(cam_id, waveform)
signalmask = np.zeros(waveform.shape[0], dtype=bool)
for i in range(signalmask.size):
signalmask[i] = gainmask[i].any() == True
combined = signals[0].copy()
combined[signalmask] = signals[1][signalmask]
peaks = peakpos[0].copy()
peaks[signalmask] = peakpos[1][signalmask]
return combined, peaks
class DL1DataWriter:
"""Writes data using event sources and DL1DataDumpers.
Provides some options for controlling the output file sizes.
"""
def __init__(self,
event_source_class=None,
event_source_settings=None,
data_dumper_class=CTAMLDataDumper,
data_dumper_settings=None,
calibration_settings=None,
preselection_cut_function=None,
write_mode='parallel',
output_file_size=10737418240,
events_per_file=None,
gain_thresholds=None,
save_mc_events=False):
"""Initialize a DL1DataWriter instance.
Provides some options for controlling the output file sizes.
Parameters
----------
event_source_class : subclass of ctapipe.io.eventsource.EventSource
A subclass of EventSource which will be used to load and yield
events as DataContainers.
event_source_settings : dict
A dictionary of kwargs which will be passed into the constructor
for the EventSource.
data_dumper_class : subclass of dl1_data_writer.DL1DataDumper
A subclass of DL1DataDumper which will be used to write events from
the EventSource to output files.
data_dumper_settings : dict
A dictionary of kwargs which will be passed into the constructor
for the DL1DataDumper.
calibration_settings : dict
A dictionary of kwargs which will be passed into the constructor
for ctapipe.calib.camera.CameraCalibrator.
preselection_cut_function : function
A cut function used to determine which events in the input files
to write to the output files. Takes a
ctapipe.io.containers.DataContainer describing a single event and
returns a boolean indicating if it passes the cut. If None, no cut
will be applied.
write_mode : str
Whether to process the data with parallel threads (one per run)
or in serial. Valid options are 'serial' and 'parallel'.
output_file_size : int
Maximum size of each output file. If the total amount of input data
requested for a given output file exceeds this size, the output
will be split across multiple files.
events_per_file : int
Maximum number of events to write per output file. If the total
number of input events requested for a given output file exceeds
this number, the output will be split across multiple files.
save_mc_events : bool
Whether to save event data for all monte carlo showers, even for
events which did not trigger the array (no images were saved).
"""
self.event_source_class = event_source_class
self.event_source_settings = (event_source_settings
if event_source_settings else {})
self.data_dumper_class = data_dumper_class
self.data_dumper_settings = (data_dumper_settings
if data_dumper_settings else {})
self.data_dumper_settings['save_mc_events'] = save_mc_events
self.preselection_cut_function = preselection_cut_function
if write_mode in ['serial', 'parallel']:
self.write_mode = write_mode
self.output_file_size = output_file_size
self.events_per_file = events_per_file
self.save_mc_events = save_mc_events
if self.output_file_size:
logger.info("Max output file size set at {} bytes. Note that "
"this may increase the number of output "
"files.".format(self.output_file_size))
if self.events_per_file:
logger.info("Max number of output events per file set at {}. Note "
"that this may increase the number of output "
"files.".format(self.events_per_file))
if calibration_settings is None:
self.calibration_settings = {
'r1_product': 'HESSIOR1Calibrator',
'extractor_product': 'NeighbourPeakIntegrator'
}
else:
self.calibration_settings = calibration_settings
self.calibrator = calib.camera.calibrator.CameraCalibrator(
None, None, **self.calibration_settings)
self.gain_selector = ThresholdGainSelector(select_by_sample=True)
if gain_thresholds is None:
self.gain_thresholds = {
'LSTCam': 4094,
'NectarCam': 4094,
'ASTRICam': 4094,
}
else:
self.gain_thresholds = gain_thresholds
def process_data(self, run_list):
"""Process data from a list of runs.
If the selected write mode is parallel, creates one process for
each requested run and executes them all in parallel.
If the selected write mode is sequential, executes each run sequentially,
writing each target one by one.
Parameters
----------
run_list : list of dicts
A list of dictionaries, each containing two keys, 'inputs' and
'target'. 'inputs' points to a list of input filenames (str) which
are to be loaded. 'target' points to an output filename (str)
to which the data from the input files should be written.
"""
if self.write_mode == 'parallel':
num_processes = len(run_list)
logger.info(
"{} parallel processes requested.".format(num_processes))
logger.info("Creating processes...")
jobs = []
for i in range(0, num_processes):
process = multiprocessing.Process(target=self._process_data,
args=(run_list[i]['inputs'],
run_list[i]['target']))
jobs.append(process)
logger.info("Starting processes...")
try:
# Start all parallel processes
for j in jobs:
j.start()
# Wait for all processes to complete
for j in jobs:
j.join()
except KeyboardInterrupt:
logger.error(
"Caught keyboard interrupt, killing all processes...")
for j in jobs:
j.terminate()
elif self.write_mode == 'serial':
logger.info("Serial processing requested.")
for run in run_list:
logger.info("Starting run for target: {}...".format(
run['target']))
self._process_data(run['inputs'], run['target'])
logger.info("Done!")
@staticmethod
def _get_next_filename(output_filename, output_file_count):
"""Get the next filename in the sequence.
Parameters
----------
output_filename : str
The filename of the previous output file generated.
output_file_count : int
Number to attach to the current output file.
Returns
-------
str
Next filename in the sequence
"""
# Append a trailing digit to get next filename in sequence
dirname = os.path.dirname(output_filename)
output_filename, *extensions = os.path.basename(output_filename).split(
'.')
if re.search(r'_[0-9]+$', output_filename):
output_filename = re.sub(r'_[0-9]+$', "_" + str(output_file_count),
output_filename)
else:
output_filename = (output_filename + "_" + str(output_file_count))
for ext in extensions:
output_filename = output_filename + '.' + ext
output_filename = os.path.join(dirname, output_filename)
return output_filename
def _process_data(self, file_list, output_filename):
"""Write a single output file given a list of input files.
Parameters
----------
file_list : list
A list of input filenames (str) to read data from.
output_filename : str
Filename of the output file to write data to.
"""
output_file_count = 1
data_dumper = self.data_dumper_class(output_filename,
**self.data_dumper_settings)
for filename in file_list:
if self.event_source_class:
event_source = self.event_source_class(
filename, **self.event_source_settings)
else:
event_source = io.event_source(filename)
# Write all file-level data if not present
# Or compare to existing data if already in file
example_event = next(event_source._generator())
subarray = example_event.inst.subarray
mcheader = example_event.mcheader
data_dumper.prepare_file(filename, subarray, mcheader)
# Write all events sequentially
for event in event_source:
self.calibrator.calibrate(event)
self.combine_channels(event)
if (self.preselection_cut_function is not None
and not self.preselection_cut_function(event)):
continue
try:
data_dumper.dump_event(event)
except IOError:
logger.error("Failed to write event from file "
"{}, skipping...".format(filename))
break
max_events_reached = (
(self.events_per_file is not None)
and (data_dumper.event_index - 1 >= self.events_per_file))
max_size_reached = (
(self.output_file_size is not None) and (os.path.getsize(
data_dumper.output_filename) > self.output_file_size))
if max_events_reached or max_size_reached:
# Reset event count and increment file count
output_file_count += 1
output_filename = self._get_next_filename(
output_filename, output_file_count)
# Create a new Data Dumper pointing at a new file
# and write file-level data
# Will flush + finalize + close file owned by
# previous data dumper
data_dumper = self.data_dumper_class(
output_filename, **self.data_dumper_settings)
# Write all file-level data if not present
# Or compare to existing data if already in file
example_event = next(event_source._generator())
subarray = example_event.inst.subarray
mcheader = example_event.mcheader
data_dumper.prepare_file(filename, subarray, mcheader)
if self.save_mc_events:
for mc_event in event_source.file_.iter_mc_events():
try:
data_dumper.dump_mc_event(
mc_event, event_source.file_.header['run'])
except IOError:
logger.error("Failed to write event from file "
"{}, skipping...".format(filename))
break
# Check whether to create another file
max_events_reached = (
(self.events_per_file is not None) and
(data_dumper.event_index - 1 >= self.events_per_file))
max_size_reached = (
(self.output_file_size is not None)
and (os.path.getsize(data_dumper.output_filename) >
self.output_file_size))
if max_events_reached or max_size_reached:
# Reset event count and increment file count
output_file_count += 1
output_filename = self._get_next_filename(
output_filename, output_file_count)
# Create a new Data Dumper pointing at a new file
# and write file-level data
data_dumper = self.data_dumper_class(
output_filename, **self.data_dumper_settings)
# Write all file-level data if not present
# Or compare to existing data if already in file
example_event = next(event_source._generator())
subarray = example_event.inst.subarray
temp = io.DataContainer()
event_source.fill_mc_information(temp, mc_event)
mcheader = temp.mcheader
data_dumper.prepare_file(filename, subarray, mcheader)
def gain_selection(self, waveform, image, peakpos, cam_id, threshold):
"""
Based on the waveform and threshold, select the proper gain for each pixel.
By default, the channel 0 is kept.
If a pixel is saturated (value > threshold in the waveform), the channel 1 is used.
Parameters
----------
waveform: array of waveforms of the events
image: array of calibrated pixel charges
peakpos: array of pixel peak positions
cam_id: str
threshold: int threshold to change form high gain to low gain
Returns
-------
combined_image, combined_peakpos: `(numpy.array, numpy.array)`
combined_image.shape = image.shape[1]
"""
assert image.shape[0] == 2
self.gain_selector.thresholds[cam_id] = threshold
waveform, gain_mask = self.gain_selector.select_gains(cam_id, waveform)
signal_mask = gain_mask.max(axis=1)
combined_image = image[0].copy()
combined_image[signal_mask] = image[1][signal_mask]
combined_peakpos = peakpos[0].copy()
combined_peakpos[signal_mask] = peakpos[1][signal_mask]
return combined_image, combined_peakpos
def combine_channels(self, event):
"""
Combine the channels for the image and peakpos arrays in the event.dl1 containers
The `event.dl1.tel[tel_id].image` and `event.dl1.tel[tel_id].peakpos` are replaced by their combined versions
Parameters
----------
event: `ctapipe.io.containers.DataContainer`
"""
for tel_id in event.r0.tels_with_data:
cam_id = event.inst.subarray.tel[tel_id].camera.cam_id
if cam_id in self.gain_thresholds:
waveform = event.r0.tel[tel_id].waveform
signals = event.dl1.tel[tel_id].image
peakpos = event.dl1.tel[tel_id].peakpos
combined_image, combined_peakpos = self.gain_selection(
waveform, signals, peakpos, cam_id,
self.gain_thresholds[cam_id])
event.dl1.tel[tel_id].image = combined_image
event.dl1.tel[tel_id].peakpos = combined_peakpos
else:
event.dl1.tel[tel_id].image = event.dl1.tel[tel_id].image[0]
event.dl1.tel[tel_id].peakpos = event.dl1.tel[tel_id].peakpos[
0]