def setup(self):
self.mean_pe_photon = self.config['mean_pe_per_photon']
# Getting S1 AFT maps
self.s1_aft_map = straxen.InterpolatingMap(
straxen.get_resource(
self.config['s1_aft_map'],
fmt=self._infer_map_format(self.config['s1_aft_map'])))
# Getting optical maps
self.s1_pattern_map = straxen.InterpolatingMap(
straxen.get_resource(
self.config['s1_optical_map'],
fmt=self._infer_map_format(self.config['s1_optical_map'])))
self.s2_pattern_map = straxen.InterpolatingMap(
straxen.get_resource(
self.config['s2_optical_map'],
fmt=self._infer_map_format(self.config['s2_optical_map'])))
# Getting gain model to get dead PMTs
self.to_pe = straxen.get_correction_from_cmt(self.run_id, self.config['gain_model'])
self.dead_PMTs = np.where(self.to_pe == 0)[0]
self.pmtbool = ~np.in1d(np.arange(0, self.config['n_tpc_pmts']), self.dead_PMTs)
self.pmtbool_top = self.pmtbool[:self.config['n_top_pmts']]
self.pmtbool_bottom = self.pmtbool[self.config['n_top_pmts']:self.config['n_tpc_pmts']]
def setup(self):
self.electron_drift_velocity = get_correction_from_cmt(
self.run_id, self.config['electron_drift_velocity'])
self.electron_drift_time_gate = get_correction_from_cmt(
self.run_id, self.config['electron_drift_time_gate'])
self.mean_pe_photon = self.config['mean_pe_per_photon']
# Getting S1 AFT maps
self.s1_aft_map = straxen.InterpolatingMap(
straxen.get_resource(self.config['s1_aft_map'],
fmt=self._infer_map_format(
self.config['s1_aft_map'])))
# Getting optical maps
self.s1_pattern_map = straxen.InterpolatingMap(
straxen.get_resource(self.config['s1_optical_map'],
fmt=self._infer_map_format(
self.config['s1_optical_map'])))
self.s2_pattern_map = straxen.InterpolatingMap(
straxen.get_resource(self.config['s2_optical_map'],
fmt=self._infer_map_format(
self.config['s2_optical_map'])))
# Getting S2 data-driven tensorflow models
downloader = straxen.MongoDownloader()
self.model_file = downloader.download_single(
self.config['s2_tf_model'])
with tempfile.TemporaryDirectory() as tmpdirname:
tar = tarfile.open(self.model_file, mode="r:gz")
tar.extractall(path=tmpdirname)
import tensorflow as tf
def _logl_loss(patterns_true, likelihood):
return likelihood / 10.
self.model = tf.keras.models.load_model(
tmpdirname, custom_objects={"_logl_loss": _logl_loss})
self.model_chi2 = tf.keras.Model(
self.model.inputs,
self.model.get_layer('Likelihood').output)
# Getting gain model to get dead PMTs
self.to_pe = straxen.get_correction_from_cmt(self.run_id,
self.config['gain_model'])
self.dead_PMTs = np.where(self.to_pe == 0)[0]
self.pmtbool = ~np.in1d(np.arange(0, self.config['n_tpc_pmts']),
self.dead_PMTs)
self.pmtbool_top = self.pmtbool[:self.config['n_top_pmts']]
self.pmtbool_bottom = self.pmtbool[self.config['n_top_pmts']:self.
config['n_tpc_pmts']]
def make_patternmap(map_file,
fmt=None,
method='WeightedNearestNeighbors',
pmt_mask=None):
""" This is special interpretation of the of previous make_map(), but designed
for pattern map loading with provided PMT mask. This way simplifies both S1 and S2
cases
"""
# making tests not failing, we can probably overwrite it completel
if isinstance(map_file, list):
log.warning(
f'Using dummy map with pattern mask! This has no effect here!')
assert map_file[0] == 'constant dummy', (
'Alternative file input can only be '
'("constant dummy", constant: int, shape: list')
return DummyMap(map_file[1], map_file[2])
elif isinstance(map_file, str):
if fmt is None:
fmt = parse_extension(map_file)
map_data = deepcopy(straxen.get_resource(map_file, fmt=fmt))
# XXX: straxed deals with pointers and caches resources, it means that resources are global
# what is bad, so we make own copy here and modify it locally
if 'compressed' in map_data:
compressor, dtype, shape = map_data['compressed']
map_data['map'] = np.frombuffer(
strax.io.COMPRESSORS[compressor]['decompress'](
map_data['map']),
dtype=dtype).reshape(*shape)
del map_data['compressed']
if 'quantized' in map_data:
map_data['map'] = map_data['quantized'] * map_data['map'].astype(
np.float32)
del map_data['quantized']
if not (pmt_mask is None):
assert (
map_data['map'].shape[-1] == pmt_mask.shape[0]
), "Error! Pattern map and PMT gains must have same dimensions!"
map_data['map'][..., ~pmt_mask] = 0.0
return straxen.InterpolatingMap(map_data, method=method)
else:
raise TypeError("Can't handle map_file except a string or a list")
def make_map(map_file, fmt=None, method='WeightedNearestNeighbors'):
"""Fetch and make an instance of InterpolatingMap based on map_file
Alternatively map_file can be a list of ["constant dummy", constant: int, shape: list]
return an instance of DummyMap"""
if isinstance(map_file, list):
assert map_file[0] == 'constant dummy', (
'Alternative file input can only be '
'("constant dummy", constant: int, shape: list')
return DummyMap(map_file[1], map_file[2])
elif isinstance(map_file, str):
if fmt is None:
fmt = parse_extension(map_file)
log.debug(f'Initialize map interpolator for file {map_file}')
map_data = straxen.get_resource(map_file, fmt=fmt)
return straxen.InterpolatingMap(map_data, method=method)
else:
raise TypeError("Can't handle map_file except a string or a list")
def make_map(map_file: str, fmt='text'):
map_data = straxen.get_resource(map_file, fmt)
return straxen.InterpolatingMap(map_data)
def load_corrected_positions(context,
run_id,
events,
cmt_version=None,
posrec_algos=('mlp', 'gcn', 'cnn')):
"""
Returns the corrected position for each position algorithm available,
without the need to reprocess event_basics, as the needed
information is already stored in event_basics.
:param cmt_version: CMT version to use (it can be a list of same
length as posrec_algos, if different versions are required for
different posrec algorithms, default 'ONLINE')
:param posrec_algos: list of position reconstruction algorithms to
use (default ['mlp', 'gcn', 'cnn'])
"""
posrec_algos = strax.to_str_tuple(posrec_algos)
if cmt_version is None:
fdc_config = None
try:
fdc_config = context.get_single_plugin(
run_id, 'event_positions').config['fdc_map']
cmt_version = fdc_config[1][1]
except IndexError as e:
raise ValueError(
f'CMT is not set? Your fdc config is {fdc_config}') from e
if hasattr(cmt_version, '__len__') and not isinstance(
cmt_version, str) and len(cmt_version) != len(posrec_algos):
raise TypeError(
f"cmt_version is a list but does not match the posrec_algos ({posrec_algos}) length."
)
cmt_version = (cmt_version, ) * len(posrec_algos) if isinstance(
cmt_version, str) else cmt_version
dtype = []
for algo in posrec_algos:
for xyzr in 'x y z r'.split():
dtype += [
((f'Interaction {xyzr}-position, field-distortion corrected (cm) - '
f'{algo.upper()} posrec algorithm', f'{xyzr}_{algo}'),
np.float32),
]
dtype += [((
f'Interaction r-position using observed S2 positions directly (cm) -'
f' {algo.upper()} posrec algorithm', f'r_naive_{algo}'),
np.float32),
((f'Correction added to r_naive for field distortion (cm) - '
f'{algo.upper()} posrec algorithm',
f'r_field_distortion_correction_{algo}'), np.float32),
((f'Interaction angular position (radians) - {algo.upper()} '
f'posrec algorithm', f'theta_{algo}'), np.float32)]
dtype += [(('Interaction z-position using mean drift velocity only (cm)',
'z_naive'), np.float32)]
result = np.zeros(len(events), dtype=dtype)
drift_speed = context.get_single_plugin(
run_id, 'event_positions').config['electron_drift_velocity']
z_obs = -drift_speed * events['drift_time']
for algo, v_cmt in zip(posrec_algos, cmt_version):
fdc_tmp = ('CMT_model', (f'fdc_map_{algo}', v_cmt), True)
map_tmp = straxen.get_config_from_cmt(run_id, fdc_tmp)
itp_tmp = straxen.InterpolatingMap(
straxen.common.get_resource(map_tmp, fmt='binary'))
itp_tmp.scale_coordinates([1., 1., -drift_speed])
orig_pos = np.vstack(
[events[f's2_x_{algo}'], events[f's2_y_{algo}'], z_obs]).T
r_obs = np.linalg.norm(orig_pos[:, :2], axis=1)
delta_r = itp_tmp(orig_pos)
# apply radial correction
with np.errstate(invalid='ignore', divide='ignore'):
r_cor = r_obs + delta_r
scale = r_cor / r_obs
with np.errstate(invalid='ignore'):
z_cor = -(z_obs**2 - delta_r**2)**0.5
invalid = np.abs(z_obs) < np.abs(delta_r)
z_cor[invalid] = z_obs[invalid]
result[f'x_{algo}'] = orig_pos[:, 0] * scale
result[f'y_{algo}'] = orig_pos[:, 1] * scale
result[f'r_{algo}'] = r_cor
result[f'r_naive_{algo}'] = r_obs
result[f'r_field_distortion_correction_{algo}'] = delta_r
result[f'theta_{algo}'] = np.arctan2(orig_pos[:, 1], orig_pos[:, 0])
result[f'z_{algo}'] = z_cor
result['z_naive'] = z_obs
return result
def load_corrected_positions(context,
run_id,
events,
alt_s1=False,
alt_s2=False,
cmt_version=None,
posrec_algos=('mlp', 'gcn', 'cnn')):
"""
Returns the corrected position for each position algorithm available,
without the need to reprocess event_basics, as the needed
information is already stored in event_basics.
:param alt_s1: False by default, if True it uses alternative S1 as main one
:param alt_s2: False by default, if True it uses alternative S2 as main one
:param cmt_version: CMT version to use (it can be a list of same
length as posrec_algos, if different versions are required for
different posrec algorithms, default 'local_ONLINE')
:param posrec_algos: list of position reconstruction algorithms to
use (default ['mlp', 'gcn', 'cnn'])
"""
posrec_algos = strax.to_str_tuple(posrec_algos)
if cmt_version is None:
fdc_config = context.get_single_plugin(
run_id, 'event_positions').config['fdc_map']
if isinstance(fdc_config, str) and 'cmt://' in fdc_config:
cmt_version = straxen.URLConfig.split_url_kwargs(fdc_config)
elif straxen.is_cmt_option(fdc_config):
cmt_version = fdc_config[1]
else:
raise ValueError(
'FDC map is not a CMT option, cannot infer cmt version.')
if (isinstance(cmt_version, (tuple, list))
and len(cmt_version) != len(posrec_algos)):
raise TypeError(f"cmt_version is a list but does not match the "
f"posrec_algos ({posrec_algos}) length.")
cmt_version = ((cmt_version, ) * len(posrec_algos) if isinstance(
cmt_version, str) else cmt_version)
# Get drift from CMT
ep = context.get_single_plugin(run_id, 'event_positions')
drift_speed = ep.electron_drift_velocity
drift_time_gate = ep.electron_drift_time_gate
dtype = load_dtypes(posrec_algos)
result = np.zeros(len(events), dtype=dtype)
s1_pre = 'alt_' if alt_s1 else ''
s2_pre = 'alt_' if alt_s2 else ''
drift_time = events['drift_time'] if not (alt_s1 or alt_s2) else events[
s2_pre + 's2_center_time'] - events[s1_pre + 's1_center_time']
z_obs = -drift_speed * drift_time
for algo, v_cmt in zip(posrec_algos, cmt_version):
fdc_tmp = (f'fdc_map_{algo}', v_cmt, True)
map_tmp = straxen.get_correction_from_cmt(run_id, fdc_tmp)
itp_tmp = straxen.InterpolatingMap(
straxen.common.get_resource(map_tmp, fmt='binary'))
itp_tmp.scale_coordinates([1., 1., -drift_speed])
orig_pos = np.vstack([
events[f'{s2_pre}s2_x_{algo}'], events[f'{s2_pre}s2_y_{algo}'],
z_obs
]).T
r_obs = np.linalg.norm(orig_pos[:, :2], axis=1)
delta_r = itp_tmp(orig_pos)
z_obs = z_obs + drift_speed * drift_time_gate
# apply radial correction
with np.errstate(invalid='ignore', divide='ignore'):
r_cor = r_obs + delta_r
scale = r_cor / r_obs
with np.errstate(invalid='ignore'):
z_cor = -(z_obs**2 - delta_r**2)**0.5
invalid = np.abs(z_obs) < np.abs(delta_r)
z_cor[invalid] = z_obs[invalid]
result[f'x_{algo}'] = orig_pos[:, 0] * scale
result[f'y_{algo}'] = orig_pos[:, 1] * scale
result[f'r_{algo}'] = r_cor
result[f'r_naive_{algo}'] = r_obs
result[f'r_field_distortion_correction_{algo}'] = delta_r
result[f'theta_{algo}'] = np.arctan2(orig_pos[:, 1], orig_pos[:, 0])
result[f'z_{algo}'] = z_cor
result['z_naive'] = z_obs
return result
def load_map(some_map, method='WeightedNearestNeighbors', **kwargs):
"""Make an InterpolatingMap"""
return straxen.InterpolatingMap(some_map, method=method, **kwargs)
