def test_append_column_chunked_cta(cta_file, cta_config):
from aict_tools.io import read_telescope_data_chunked, read_data
from aict_tools.io import append_column_to_hdf5
new_column_name = 'foobar'
chunk_size = 125
table_name = 'telescope_events'
df = read_data(cta_file, table_name)
assert new_column_name not in df.columns
columns = cta_config.energy.columns_to_read_train
generator = read_telescope_data_chunked(cta_file, cta_config, chunk_size, columns=columns)
for df, start, stop in generator:
assert not df.empty
new_data = np.arange(start, stop, step=1)
append_column_to_hdf5(cta_file, new_data, table_name, new_column_name)
df = read_data(cta_file, table_name)
assert new_column_name in df.columns
assert np.array_equal(df.foobar, np.arange(0, len(df)))
df.set_index(
['run_id', 'array_event_id', 'telescope_id'],
drop=True,
verify_integrity=True,
inplace=True,
)
def test_append_column_chunked(hdf5_file):
from aict_tools.io import read_telescope_data_chunked, read_data
from aict_tools.io import append_column_to_hdf5
path, table_name, config = hdf5_file
new_column_name = 'foobar'
chunk_size = 125
df = read_data(path, table_name)
assert new_column_name not in df.columns
columns = config.energy.columns_to_read_train
generator = read_telescope_data_chunked(path, config, chunk_size, columns=columns)
for df, start, stop in generator:
assert not df.empty
new_data = np.arange(start, stop, step=1)
append_column_to_hdf5(path, new_data, table_name, new_column_name)
df = read_data(path, table_name)
assert new_column_name in df.columns
assert np.array_equal(df.foobar, np.arange(0, len(df)))
def test_append_column(hdf5_file):
from aict_tools.io import read_data
from aict_tools.io import append_column_to_hdf5
path, table_name, _ = hdf5_file
new_column_name = 'foobar'
df = read_data(path, table_name)
assert new_column_name not in df.columns
random_data = np.random.normal(size=len(df))
append_column_to_hdf5(path, random_data, table_name, new_column_name)
df = read_data(path, table_name)
assert new_column_name in df.columns
from aict_tools.io import (read_data, append_column_to_hdf5)
import argparse
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='PATH AND STUFF')
parser.add_argument('data_path', type=str)
parser.add_argument('column', type=str)
args = parser.parse_args()
a = read_data(args.data_path, 'array_events')
t = read_data(args.data_path,
'telescope_events')[['run_id', 'array_event_id', 'width']]
new_t = t.merge(a[['run_id', 'array_event_id', args.column]],
on=['run_id', 'array_event_id'],
how='left')
append_column_to_hdf5(args.data_path, new_t[args.column],
'telescope_events', args.column)
def main(infile):
with h5py.File(infile, mode='r') as f:
is_simulation = 'corsika_runs' in f
if is_simulation:
df = read_h5py(infile, key='events', columns=columns_sim)
obstime = None
else:
df = read_h5py(infile, key='events', columns=columns_obs)
obstime = Time(df.dragon_time, format='unix')
altaz = AltAz(obstime=obstime, location=location)
pointing = SkyCoord(
alt=u.Quantity(df.alt_tel.values, u.rad, copy=False),
az=u.Quantity(df.az_tel.values, u.rad, copy=False),
frame=altaz,
)
camera_frame = CameraFrame(telescope_pointing=pointing,
location=location,
obstime=obstime,
focal_length=28 * u.m)
prediction_cam = SkyCoord(
x=u.Quantity(df.source_x_prediction.values, u.m, copy=False),
y=u.Quantity(df.source_y_prediction.values, u.m, copy=False),
frame=camera_frame,
)
prediction_altaz = prediction_cam.transform_to(altaz)
append_column_to_hdf5(infile, prediction_altaz.alt.rad, 'events',
'source_alt_prediction')
append_column_to_hdf5(infile, prediction_altaz.az.rad, 'events',
'source_az_prediction')
if not is_simulation:
prediction_icrs = prediction_altaz.transform_to('icrs')
pointing_icrs = pointing.transform_to('icrs')
append_column_to_hdf5(infile, prediction_icrs.ra.rad, 'events',
'source_ra_prediction')
append_column_to_hdf5(infile, prediction_icrs.dec.rad, 'events',
'source_dec_prediction')
append_column_to_hdf5(infile, pointing_icrs.ra.rad, 'events',
'pointing_ra')
append_column_to_hdf5(infile, pointing_icrs.dec.rad, 'events',
'pointing_dec')
