def load_clsim_table_minimal(fpath, mmap=False, include_overflow=False):
"""Load a CLSim table from disk (optionally compressed with zstd).
Similar to the `load_clsim_table` function but the full table, including
under/overflow bins, is kept and no normalization or further processing is
performed on the table data besides populating the ouptput OrderedDict.
Parameters
----------
fpath : string
Path to file to be loaded. If the file has extension 'zst', 'zstd', or
'zstandard', the file will be decompressed using the `python-zstandard`
Python library before passing to `fits` for interpreting.
mmap : bool, optional
Whether to memory map the table
include_overflow : bool, optional
By default, overflow bins (if present) are removed
Returns
-------
table : OrderedDict
"""
t0 = time()
table = OrderedDict()
fpath = expand(fpath)
if DEBUG:
wstderr('Loading table from {} ...\n'.format(fpath))
if isdir(fpath):
indir = fpath
if mmap:
mmap_mode = 'r'
else:
mmap_mode = None
for rel_fpath in listdir(indir):
key, ext = splitext(rel_fpath)
abs_fpath = join(indir, rel_fpath)
if not (isfile(abs_fpath) and ext == '.npy'):
continue
if DEBUG:
wstderr(' loading {} from "{}" ...'.format(key, abs_fpath))
t1 = time()
val = np.load(abs_fpath, mmap_mode=mmap_mode)
# Pull "small" things (less than 10 MiB) into memory so we don't
# have too many file handles open due to memory mapping
if mmap and val.nbytes < 10 * 1024**2:
val = np.copy(val)
table[key] = val
if DEBUG:
wstderr(' ({} ms)\n'.format(np.round((time() - t1) * 1e3, 3)))
elif isfile(fpath):
from astropy.io import fits
fobj = get_decompressd_fobj(fpath)
pf_table = None
try:
pf_table = fits.open(fobj, mode='readonly', memmap=mmap)
header = pf_table[0].header # pylint: disable=no-member
table['table_shape'] = np.array(pf_table[0].data.shape, dtype=int) # pylint: disable=no-member
table['group_refractive_index'] = set_explicit_dtype(
force_little_endian(header['_i3_n_group']))
table['phase_refractive_index'] = set_explicit_dtype(
force_little_endian(header['_i3_n_phase']))
n_dims = len(table['table_shape'])
new_style = False
axnames = [None] * n_dims
binning = [None] * n_dims
for key in header.keys():
if not key.startswith('_i3_ax_'):
continue
new_style = True
axnum = header[key]
axname = key[len('_i3_ax_'):]
be0 = header['_i3_{}_min'.format(axname)]
be1 = header['_i3_{}_max'.format(axname)]
n_bins = header['_i3_{}_n_bins'.format(axname)]
power = header.get('_i3_{}_power'.format(axname), 1)
bin_edges = force_little_endian(pf_table[axnum + 1].data) # pylint: disable=no-member
assert np.isclose(bin_edges[0],
be0), '%f .. %f' % (be0, bin_edges[0])
assert np.isclose(bin_edges[-1],
be1), '%f .. %f' % (be1, bin_edges[-1])
assert len(bin_edges) == n_bins + 1, '%d vs. %d' % (
len(bin_edges), n_bins + 1)
assert np.allclose(
bin_edges,
powerspace(start=be0,
stop=be1,
num=n_bins + 1,
power=power),
)
axnames[axnum] = axname
binning[axnum] = bin_edges
if not new_style:
if n_dims == 5:
axnames = [
'r', 'costheta', 't', 'costhetadir', 'deltaphidir'
]
elif n_dims == 6:
axnames = [
'r', 'costheta', 'phi', 't', 'costhetadir',
'deltaphidir'
]
else:
raise NotImplementedError(
'{}-dimensional table not handled for old-style CLSim'
' tables'.format(n_dims))
binning = [
force_little_endian(pf_table[i + 1].data).flat
for i in range(len(axnames))
] # pylint: disable=no-member
for axnum, (axname, bin_edges) in enumerate(zip(axnames, binning)):
assert axname is not None, 'missing axis %d name' % axnum
assert bin_edges is not None, 'missing axis %d binning' % axnum
dtype = np.dtype([(axname, np.float64, dim.size)
for axname, dim in zip(axnames, binning)])
table['binning'] = np.array(tuple(binning), dtype=dtype)
for keyroot in GENERIC_KEYS:
keyname = '_i3_' + keyroot
if keyname in header:
val = force_little_endian(header[keyname])
if keyroot in (
't_is_residual_time',
'disable_tilt',
'disable_anisotropy',
):
val = np.bool8(val)
else:
val = set_explicit_dtype(val)
table[keyroot] = val
# Get string values from keys that have a prefix preceded by the
# value all in the key (I3 software had issues saving strings as
# values in the header "dict" so the workaround was to store the
# string value in this way)
for infix in INFIX_KEYS:
keyroot = '_i3_' + infix + '_'
for keyname in header.keys():
if not keyname.startswith(keyroot):
continue
val = keyname[len(keyroot):]
table[infix] = np.string0(val)
if include_overflow:
slicer = (slice(None), ) * n_dims
else:
slicer = (slice(1, -1), ) * n_dims
table['table'] = force_little_endian(pf_table[0].data[slicer]) # pylint: disable=no-member
wstderr(' (load took {} s)\n'.format(np.round(time() - t0, 3)))
except:
wstderr('ERROR: Failed to load "{}"\n'.format(fpath))
raise
finally:
del pf_table
if hasattr(fobj, 'close'):
fobj.close()
del fobj
else: # fpath is neither dir nor file
raise ValueError('Table does not exist at path "{}"'.format(fpath))
if 'step_length' not in table:
table['step_length'] = 1
if 't_is_residual_time' not in table:
table['t_is_residual_time'] = True
if DEBUG:
wstderr(' Total time to load: {} s\n'.format(np.round(time() - t0,
3)))
return table
def load_t_r_theta_table(fpath,
depth_idx,
scale=1,
exponent=1,
photon_info=None):
"""Extract info from a file containing a (t, r, theta)-binned Retro table.
Parameters
----------
fpath : string
Path to FITS file corresponding to the passed ``depth_idx``.
depth_idx : int
Depth index (e.g. from 0 to 59)
scale : float
Scaling factor to apply to the photon survival probability from the
table, e.g. for quantum efficiency. This is applied _before_
`exponent`. See `Notes` for more info.
exponent : float >= 0, optional
Modify probabilties in the table by ``prob = 1 - (1 - prob)**exponent``
to allow for up- and down-scaling the efficiency of the DOMs. This is
applied to each DOM's table _after_ `scale`. See `Notes` for more
info.
photon_info : None or RetroPhotonInfo namedtuple of dicts
If None, creates a new RetroPhotonInfo namedtuple with empty dicts to
fill. If one is provided, the existing component dictionaries are
updated.
Returns
-------
photon_info : RetroPhotonInfo namedtuple of dicts
Tuple fields are 'survival_prob', 'theta', 'phi', and 'length'. Each
dict is keyed by `depth_idx` and values are the arrays loaded
from the FITS file.
bin_edges : TimeSphCoord namedtuple
Each element of the tuple is an array of bin edges.
Notes
-----
The parameters `scale` and `exponent` modify a table's probability `P` by::
P = 1 - (1 - P*scale)**exponent
This allows for `scale` (which must be from 0 to 1) to be used for e.g.
quantum efficiency--which always reduces the detection probability--and
`exponent` (which must be 0 or greater) to be used as a systematic that
modifies the post-`scale` probabilities up and down while keeping them
valid (i.e., between 0 and 1). Larger values of `scale` (i.e., closer to 1)
indicate a more efficient DOM. Likewise, values of `exponent` greater than
one scale up the DOM efficiency, while values of `exponent` between 0 and 1
scale the efficiency down.
"""
# pylint: disable=no-member
from astropy.io import fits
assert 0 <= scale <= 1
assert exponent >= 0
if photon_info is None:
empty_dicts = []
for _ in RetroPhotonInfo._fields:
empty_dicts.append({})
photon_info = RetroPhotonInfo(*empty_dicts)
with fits.open(expand(fpath)) as table:
data = force_little_endian(table[0].data)
if scale == exponent == 1:
photon_info.survival_prob[depth_idx] = data
else:
photon_info.survival_prob[depth_idx] = (
1 - (1 - data * scale)**exponent)
photon_info.theta[depth_idx] = force_little_endian(table[1].data)
photon_info.deltaphi[depth_idx] = force_little_endian(table[2].data)
photon_info.length[depth_idx] = force_little_endian(table[3].data)
# Note that we invert (reverse and multiply by -1) time edges; also,
# no phi edges are defined in these tables.
data = force_little_endian(table[4].data)
t = -data[::-1]
r = force_little_endian(table[5].data)
# Previously used the following to get "agreement" w/ raw photon sim
#r_volumes = np.square(0.5 * (r[1:] + r[:-1]))
#r_volumes = (0.5 * (r[1:] + r[:-1]))**2 * (r[1:] - r[:-1])
r_volumes = 0.25 * (r[1:]**3 - r[:-1]**3)
photon_info.survival_prob[depth_idx] /= r_volumes[np.newaxis, :,
np.newaxis]
photon_info.time_indep_survival_prob[depth_idx] = np.sum(
photon_info.survival_prob[depth_idx], axis=0)
theta = force_little_endian(table[6].data)
bin_edges = TimeSphCoord(t=t,
r=r,
theta=theta,
phi=np.array([], dtype=t.dtype))
return photon_info, bin_edges
def load_tables(self, force_reload=False):
"""Load all tables that match the `proto_tile_hash`; if multiple tables
match, then stitch these together into one large TDI table."""
if self.tables_loaded and not force_reload:
return
from astropy.io import fits
t0 = time()
x_ref = self.proto_meta['x_min']
y_ref = self.proto_meta['y_min']
z_ref = self.proto_meta['z_min']
must_match = [
'binmap_hash',
'geom_hash',
'dom_tables_hash',
'time_indices',
'binwidth',
'anisotropy',
# For simplicity, assume all tiles have equal widths. If there's a
# compelling reason to use something more complicated, we could
# implement it... but I see no reason to do so now
'x_width',
'y_width',
'z_width'
]
# Work with "survival_prob" table filepaths, which generalizes to all
# table filepaths (so long as they exist)
fpaths = glob(
join(expand(self.tables_dir),
'retro_tdi_table_*survival_prob.fits'))
lowermost_corner = np.array([np.inf] * 3)
uppermost_corner = np.array([-np.inf] * 3)
to_load_meta = {}
for fpath in fpaths:
meta = self.get_table_metadata(fpath)
if meta is None:
continue
is_match = True
for key in must_match:
if meta[key] != self.proto_meta[key]:
is_match = False
if not is_match:
continue
# Make sure that the corner falls on the reference grid (within
# micrometer precision)
x_float_idx = (meta['x_min'] - x_ref) / self.x_tile_width
y_float_idx = (meta['y_min'] - y_ref) / self.y_tile_width
z_float_idx = (meta['z_min'] - z_ref) / self.x_tile_width
indices_widths = ([x_float_idx, y_float_idx, z_float_idx], [
self.x_tile_width, self.y_tile_width, self.z_tile_width
])
for float_idx, tile_width in zip(indices_widths):
if abs(np.round(float_idx) - float_idx) * tile_width >= 1e-6:
continue
# Extend the limits of the tiled volume to include this tile
lower_corner = [meta['x_min'], meta['y_min'], meta['z_min']]
upper_corner = [meta['x_max'], meta['y_max'], meta['z_max']]
lowermost_corner = np.min([lowermost_corner, lower_corner], axis=0)
uppermost_corner = np.max([uppermost_corner, upper_corner], axis=0)
# Store the metadata by relative tile index
rel_idx = tuple(
int(np.round(i))
for i in (x_float_idx, y_float_idx, z_float_idx))
to_load_meta[rel_idx] = meta
x_min, y_min, z_min = lowermost_corner
x_max, y_max, z_max = uppermost_corner
# Figure out how many tiles we _should_ have
nx_tiles = int(np.round((x_max - x_min) / self.x_tile_width))
ny_tiles = int(np.round((y_max - y_min) / self.y_tile_width))
nz_tiles = int(np.round((z_max - z_min) / self.z_tile_width))
n_tiles = nx_tiles * ny_tiles * nz_tiles
if len(to_load_meta) < n_tiles:
raise ValueError('Not enough tiles found! Cannot fill the extents'
' of the outermost extents of the volume defined'
' by the tiles found.')
elif len(to_load_meta) > n_tiles:
print(self.proto_meta['tdi_hash'])
print('x:', self.proto_meta['x_min'], self.proto_meta['x_max'],
self.proto_meta['x_width'])
print('y:', self.proto_meta['y_min'], self.proto_meta['y_max'],
self.proto_meta['y_width'])
print('z:', self.proto_meta['z_min'], self.proto_meta['z_max'],
self.proto_meta['z_width'])
print('')
for v in to_load_meta.values():
print(v['tdi_hash'])
print('x:', v['x_min'], v['x_max'], v['x_width'])
print('y:', v['y_min'], v['y_max'], v['y_width'])
print('z:', v['z_min'], v['z_max'], v['z_width'])
print('')
raise ValueError(
'WTF? How did we get here? to_load_meta = %d, n_tiles = %d' %
(len(to_load_meta), n_tiles))
# Figure out how many bins in each dimension fill the volume
nx = int(np.round(nx_tiles * self.x_tile_width / self.binwidth))
ny = int(np.round(ny_tiles * self.y_tile_width / self.binwidth))
nz = int(np.round(nz_tiles * self.z_tile_width / self.binwidth))
# Number of bins per dimension in the tile
nx_per_tile = int(np.round(self.x_tile_width / self.binwidth))
ny_per_tile = int(np.round(self.y_tile_width / self.binwidth))
nz_per_tile = int(np.round(self.z_tile_width / self.binwidth))
# Create empty arrays to fill
survival_prob = np.empty((nx, ny, nz), dtype=np.float32)
if self.use_directionality:
avg_photon_x = np.empty((nx, ny, nz), dtype=np.float32)
avg_photon_y = np.empty((nx, ny, nz), dtype=np.float32)
avg_photon_z = np.empty((nx, ny, nz), dtype=np.float32)
else:
avg_photon_x, avg_photon_y, avg_photon_z = None, None, None
anisotropy_str = generate_anisotropy_str(self.anisotropy)
tables_meta = {} #[[[None]*nz_tiles]*ny_tiles]*nx_tiles
for meta in to_load_meta.values():
tile_x_idx = int(
np.round((meta['x_min'] - x_min) / self.x_tile_width))
tile_y_idx = int(
np.round((meta['y_min'] - y_min) / self.y_tile_width))
tile_z_idx = int(
np.round((meta['z_min'] - z_min) / self.z_tile_width))
x0_idx = int(np.round((meta['x_min'] - x_min) / self.binwidth))
y0_idx = int(np.round((meta['y_min'] - y_min) / self.binwidth))
z0_idx = int(np.round((meta['z_min'] - z_min) / self.binwidth))
bin_idx_range = (slice(x0_idx, x0_idx + nx_per_tile),
slice(y0_idx, y0_idx + ny_per_tile),
slice(z0_idx, z0_idx + nz_per_tile))
kwargs = deepcopy(meta)
kwargs.pop('table_name')
to_fill = [('survival_prob', survival_prob)]
if self.use_directionality:
to_fill.extend([('avg_photon_x', avg_photon_x),
('avg_photon_y', avg_photon_y),
('avg_photon_z', avg_photon_z)])
for table_name, table in to_fill:
fpath = join(
self.tables_dir, TDI_TABLE_FNAME_PROTO[-1].format(
table_name=table_name,
anisotropy_str=anisotropy_str,
**kwargs).lower())
with fits.open(fpath) as fits_table:
data = force_little_endian(fits_table[0].data) # pylint: disable=no-member
if self.scale != 1 and table_name == 'survival_prob':
data = 1 - (1 - data)**self.scale
table[bin_idx_range] = data
tables_meta[(tile_x_idx, tile_y_idx, tile_z_idx)] = meta
# Since we have made it to the end successfully, it is now safe to
# store the above-computed info to the object for later use
self.nx, self.ny, self.nz = nx, ny, nz
self.nx_tiles = nx_tiles
self.ny_tiles = ny_tiles
self.nz_tiles = nz_tiles
self.n_bins = self.nx * self.ny * self.nz
self.n_tiles = self.nx_tiles * self.ny_tiles * self.nz_tiles
self.x_min, self.y_min, self.z_min = x_min, y_min, z_min
self.x_max, self.y_max, self.z_max = x_max, y_max, z_max
self.survival_prob = survival_prob
self.avg_photon_x = avg_photon_x
self.avg_photon_y = avg_photon_y
self.avg_photon_z = avg_photon_z
self.tables_meta = tables_meta
self.tables_loaded = True
if self.n_tiles == 1:
tstr = 'tile'
else:
tstr = 'tiles'
print('Loaded %d %s spanning'
' x ∈ [%.2f, %.2f) m,'
' y ∈ [%.2f, %.2f) m,'
' z ∈ [%.2f, %.2f) m;'
' bins are (%.3f m)³' %
(self.n_tiles, tstr, self.x_min, self.x_max, self.y_min,
self.y_max, self.z_min, self.z_max, self.binwidth))
print('Time to load: {} s'.format(np.round(time() - t0, 3)))
'_{hash}_tile_{tile}_string_{string}_dom_{dom}'
'_seed_{seed}_n_{n_events}.fits'
.format(hash=TILESET_HASH, **spec)
)
tile_fits = fits.open(
join(tdi_tile_dir, tile_fname),
memmap=True,
)
header = tile_fits[0].header
tile = tile_fits[0].data
# Get rid of underflow and overflow bins
tile = tile[(slice(1, -1),) * tile.ndim]
# Convert to little-endian for numba codes (plus native format on Intel)
tile = force_little_endian(tile)
if np.product(tile.shape) == 0:
wstdout(' tile failed!.\n')
continue
n_photons = header['_i3_n_photons']
n_phase = header['_i3_n_phase']
cos_ckv = 1 / (n_phase * BETA)
sin_ckv = np.sin(np.arccos(cos_ckv))
if avg_angsens is None:
angsens_model = [
k.replace('_i3_angsens_', '')
for k in header.keys() if k.startswith('_i3_angsens_')
][0]
def load_clsim_table_minimal(fpath, step_length=None, mmap=False):
"""Load a CLSim table from disk (optionally compressed with zstd).
Similar to the `load_clsim_table` function but the full table, including
under/overflow bins, is kept and no normalization or further processing is
performed on the table data besides populating the ouptput OrderedDict.
Parameters
----------
fpath : string
Path to file to be loaded. If the file has extension 'zst', 'zstd', or
'zstandard', the file will be decompressed using the `python-zstandard`
Python library before passing to `pyfits` for interpreting.
mmap : bool, optional
Whether to memory map the table (if it's stored in a directory
containing .npy files).
Returns
-------
table : OrderedDict
Items include
- 'table_shape' : tuple of int
- 'table' : np.ndarray
- 't_indep_table' : np.ndarray (if available)
- 'n_photons' :
- 'phase_refractive_index' :
- 'r_bin_edges' :
- 'costheta_bin_edges' :
- 't_bin_edges' :
- 'costhetadir_bin_edges' :
- 'deltaphidir_bin_edges' :
"""
table = OrderedDict()
fpath = expand(fpath)
if DEBUG:
wstderr('Loading table from {} ...\n'.format(fpath))
if isdir(fpath):
t0 = time()
indir = fpath
if mmap:
mmap_mode = 'r'
else:
mmap_mode = None
for key in MY_CLSIM_TABLE_KEYS + ['t_indep_table']:
fpath = join(indir, key + '.npy')
if DEBUG:
wstderr(' loading {} from "{}" ...'.format(key, fpath))
t1 = time()
if isfile(fpath):
table[key] = np.load(fpath, mmap_mode=mmap_mode)
elif key != 't_indep_table':
raise ValueError(
'Could not find file "{}" for loading table key "{}"'
.format(fpath, key)
)
if DEBUG:
wstderr(' ({} ms)\n'.format(np.round((time() - t1)*1e3, 3)))
if step_length is not None and 'step_length' in table:
assert step_length == table['step_length']
if DEBUG:
wstderr(' Total time to load: {} s\n'.format(np.round(time() - t0, 3)))
return table
if not isfile(fpath):
raise ValueError('Table does not exist at path "{}"'.format(fpath))
if mmap:
print('WARNING: Cannot memory map a fits or compressed fits file;'
' ignoring `mmap=True`.')
import pyfits
t0 = time()
fobj = get_decompressd_fobj(fpath)
try:
pf_table = pyfits.open(fobj)
table['table_shape'] = pf_table[0].data.shape # pylint: disable=no-member
table['n_photons'] = force_little_endian(
pf_table[0].header['_i3_n_photons'] # pylint: disable=no-member
)
table['group_refractive_index'] = force_little_endian(
pf_table[0].header['_i3_n_group'] # pylint: disable=no-member
)
table['phase_refractive_index'] = force_little_endian(
pf_table[0].header['_i3_n_phase'] # pylint: disable=no-member
)
if step_length is not None:
table['step_length'] = step_length
n_dims = len(table['table_shape'])
if n_dims == 5:
# Space-time dimensions
table['r_bin_edges'] = force_little_endian(
pf_table[1].data # meters # pylint: disable=no-member
)
table['costheta_bin_edges'] = force_little_endian(
pf_table[2].data # pylint: disable=no-member
)
table['t_bin_edges'] = force_little_endian(
pf_table[3].data # nanoseconds # pylint: disable=no-member
)
# Photon directionality
table['costhetadir_bin_edges'] = force_little_endian(
pf_table[4].data # pylint: disable=no-member
)
table['deltaphidir_bin_edges'] = force_little_endian(
pf_table[5].data # pylint: disable=no-member
)
else:
raise NotImplementedError(
'{}-dimensional table not handled'.format(n_dims)
)
table['table'] = force_little_endian(pf_table[0].data) # pylint: disable=no-member
wstderr(' (load took {} s)\n'.format(np.round(time() - t0, 3)))
finally:
del pf_table
if hasattr(fobj, 'close'):
fobj.close()
del fobj
return table