df_file = os.path.join(comp.paths.comp_data_dir, args.config + '_data',
'anisotropy_dataframe.hdf')
nrows = get_nrows(df_file)
data_df = get_dataframe_batch(df_file,
args.n_splits,
args.split_idx,
nrows=nrows)
times = get_random_times(df_file,
data_df.shape[0],
n_resamples=20,
nrows=nrows)
mask = np.ones(data_df.shape[0], dtype=bool)
if args.composition in ['light', 'heavy']:
mask[data_df['pred_comp'] != args.composition] = False
# # Ensure that effective area has plateaued
# mask[data_df['lap_log_energy'] < 6.4] = False
mask[data_df['lap_log_energy'] < 6.0] = False
# If specified, remove high-energy events
if args.low_energy:
mask[data_df['lap_log_energy'] > 7.0] = False
# mask[data_df['lap_log_energy'] > 6.75] = False
data_df = data_df.loc[mask, :].reset_index(drop=True)
times = times[mask]
maps = anisotropy.make_skymaps(data_df, times, n_side=args.n_side)
hp.write_map(args.outfile, maps, coord='C')
# Remove rows not in split_indices
split_mask = df.index.isin(split_indices)
df = df.loc[split_mask, :].reset_index(drop=True)
times = get_random_times(store,
split_indices,
n_events=df.shape[0],
n_resamples=20)
# If specified, remove high-energy events
if args.low_energy:
low_energy_mask = df['lap_log_energy'] <= 6.75
df = df.loc[low_energy_mask, :].reset_index(drop=True)
times = times[low_energy_mask]
data, ref, local = anisotropy.make_skymaps(df,
times,
n_side=args.n_side)
data_map += data
ref_map += ref
local_map += local
i += 1
if i > 2:
break
maps = (data_map, ref_map, local_map)
if split_idx == 0:
outfile = args.outfile_sample_0
else:
outfile = args.outfile_sample_1
hp.write_map(outfile, maps, coord='C')
n_rows = store.get_storer('dataframe').nrows
for df in store.select('dataframe', chunksize=args.chunksize):
# for df in store.select('dataframe', chunksize=args.chunksize, stop=10000):
df.reset_index(drop=True, inplace=True)
times = get_random_times(store, n_rows, args.chunksize, n_resamples=20)
# If specified, remove high-energy events
if args.low_energy:
low_energy_mask = df['lap_log_energy'] <= 6.75
df = df.loc[low_energy_mask, :].reset_index(drop=True)
times = times[low_energy_mask]
train_df, test_df = train_test_split(df, test_size=0.4)
for split_idx, split_df in enumerate([train_df, test_df]):
split_times = times[split_df.index]
split_df.reset_index(drop=True, inplace=True)
data, ref, local = anisotropy.make_skymaps(split_df,
split_times, n_side=args.n_side)
data_maps[split_idx] += data
ref_maps[split_idx] += ref
local_maps[split_idx] += local
maps_sample_0 = (data_maps[0], ref_maps[0], local_maps[0])
hp.write_map(args.outfile_sample_0, maps_sample_0, coord='C')
maps_sample_1 = (data_maps[1], ref_maps[1], local_maps[1])
hp.write_map(args.outfile_sample_1, maps_sample_1, coord='C')