You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
antab/antab_editor.py

183 lines
8.2 KiB

from antab_editor_lib.meta_data import (
MetaData, get_setup_stations, get_station_scans, get_channel_map,
header2channels)
from antab_editor_lib.station_widget import create_station_widget
from antab_editor_lib.poly_plot import PolyPlotWidget
from antab_editor_lib.general import GeneralWidget
from antab_editor_lib import antab_file
from PyQt5.QtWidgets import QApplication, QTabWidget
import sys
import argparse
import os
import copy
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GUI to edit ANTAB files. "
"Documentation available at "
"https://code.jive.eu/eldering/antab")
parser.add_argument("-e", "--experiment", default=None,
help="Experiment to handle, defaults to the current "
"directory name.")
parser.add_argument("-s", "--setupstation", default=None,
help="Force setup station, required if setup station "
"cannot be determined from correlator database.")
parser.add_argument("-a", "--addfits", default=[], nargs="+",
metavar="EXPERIMENT",
help="Additionally read FITS files for given "
"experiment(s).")
group = parser.add_mutually_exclusive_group()
group.add_argument("-l", "--spectralline",
action="store_true", default=False,
help="Generate two ANTAB files per station, one for "
"the continuum pass (from *_1_1.IDI* files) and one "
"for the spectral line pass (from *_2_1.IDI* files).")
group.add_argument("-p", "--passes",
default=None, nargs="+", type=int,
help="Only read FITS files for given passes "
"(files named *_<n>_1.IDI*).")
args = parser.parse_args()
experiment = args.experiment
if experiment is None:
experiment = os.path.split(os.getcwd())[1]
setup_station = args.setupstation
if setup_station is None:
setup_stations = get_setup_stations(experiment)
if len(setup_stations) != 1:
print("Cannot determine setup station automatically, "
"specify it using the -s/--setupstation option.")
sys.exit(1)
setup_station = setup_stations[0][0]
print(f"Using {setup_station} as setup station.")
app = QApplication(sys.argv)
if args.passes is not None:
pass_filter = set(int(pass_index) for pass_index in args.passes)
else:
pass_filter = None
meta_data = MetaData(experiment, setup_station,
args.spectralline, args.addfits, pass_filter)
vex = meta_data.get_vex()
main_widget = QTabWidget()
station_widgets = {}
scans = vex["SCHED"].values()
for station in sorted(vex["STATION"].keys()):
if len(get_station_scans(scans, station)) == 0:
print(f"Warning, ignoring {station} because it is not in any VEX "
"scan.")
continue
if len(meta_data.fits_cache.station_time_ranges[station]) == 0:
print(f"Warning, ignoring {station} because it has no data in any "
"FITS file.")
continue
station_widget = create_station_widget(
station, experiment, meta_data)
main_widget.addTab(station_widget, station)
station_widgets[station] = station_widget
poly_widget = PolyPlotWidget(station_widgets, meta_data)
main_widget.insertTab(0, poly_widget, "Polys")
general_widget = GeneralWidget(station_widgets, meta_data)
main_widget.insertTab(0, general_widget, "Overview")
def activate_station_tab(station):
for tab_index in range(main_widget.count()):
if main_widget.tabText(tab_index) == station:
main_widget.setCurrentIndex(tab_index)
break
general_widget.station_clicked.connect(activate_station_tab)
def write_antab():
assert((not args.spectralline) or
(len(meta_data.fits_cache.pass_time_ranges) == 1))
no_pass = ((not args.spectralline) and
(len(meta_data.fits_cache.pass_time_ranges) == 1))
for pass_, time_ranges in meta_data.fits_cache.pass_time_ranges.items():
pass_ext = "" if no_pass else f"_{pass_}"
with open(f"{experiment.lower()}{pass_ext}.antab", "w") as fp:
for w in station_widgets.values():
if no_pass:
w.save_to(fp)
else:
# FIX too strict filter, add times in gaps to time_ranges
w.save_to_filtered(fp, time_ranges)
if args.spectralline:
# map the spectral line correlator output channels to
# the continuum correlator output channels and
# change the INDEX field accordingly
# assertion guarantees that there is only one pass
time_ranges = next(iter(
meta_data.fits_cache.pass_time_ranges.values()))
def get_channels(timed_channels):
ret = None
for channel_time_range, channels in timed_channels.items():
overlap = any(((time_range[1] > channel_time_range[0]) and
(time_range[0] < channel_time_range[1]))
for time_range in time_ranges)
if overlap:
if ret is None:
ret = channels
elif channels != ret:
raise RuntimeError(
"Multiple frequency setups found, "
"this is not supported in spectral line mode.")
if ret is None:
raise RuntimeError(
"Failed to find any frequency setup in time range "
f"{time_range[0]} - {time_range[1]}.")
return ret
channel_map = get_channel_map(
get_channels(meta_data.fits_cache.spectral_frequency),
get_channels(meta_data.fits_cache.time_frequency))
with open(f"{experiment.lower()}_2.antab", "w") as fp:
for station in sorted(station_widgets.keys()):
widget = station_widgets[station]
for station_data in widget.get_station_data_sections():
rename_map = {}
for continuum_column_header in \
station_data.data.columns:
old_channels = set(
header2channels(continuum_column_header))
# the new INDEX column header is composed of all
# spectral channels which map to a continuum channel
# in the old INDEX column header
new_channels = [k for k, v in channel_map.items()
if v in old_channels]
if len(new_channels) > 0:
new_value = "|".join(new_channels)
else:
new_value = "X"
rename_map[continuum_column_header] = new_value
data = station_data.data.rename(columns=rename_map,
copy=False,
inplace=False,
errors="raise")
antab_file.write(fp, station_data.header, data)
general_widget.generate_clicked.connect(write_antab)
def save_all():
for w in station_widgets.values():
w.save()
general_widget.save_all_clicked.connect(save_all)
main_widget.setCurrentIndex(0)
main_widget.show()
desktop = QApplication.desktop()
main_widget.resize(main_widget.minimumSizeHint().expandedTo(
desktop.screenGeometry(desktop.screenNumber(main_widget)).size()/2))
app.exec_()