From aa8d6cf4b94eb64d6269f8ef98db304516098628 Mon Sep 17 00:00:00 2001 From: Jonathan Kenyon Date: Fri, 8 Nov 2024 13:57:56 +0200 Subject: [PATCH] Improve removal of coords/attts added by quartical during writes. --- quartical/data_handling/ms_handler.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/quartical/data_handling/ms_handler.py b/quartical/data_handling/ms_handler.py index 3566deff..adc77024 100644 --- a/quartical/data_handling/ms_handler.py +++ b/quartical/data_handling/ms_handler.py @@ -15,6 +15,13 @@ from quartical.data_handling.selection import filter_xds_list from quartical.data_handling.angles import apply_parangles +DASKMS_ATTRS = { + "__daskms_partition_schema__", + "SCAN_NUMBER", + "FIELD_ID", + "DATA_DESC_ID" +} + def read_xds_list(model_columns, ms_opts): """Reads a measurement set and generates a list of xarray data sets. @@ -292,14 +299,18 @@ def write_xds_list(xds_list, ref_xds_list, ms_path, output_opts): logger.info("Outputs will be written to {}.", ", ".join(output_cols)) + # Select only the output columns to simplify datasets. + xds_list = [xds[list(output_cols)] for xds in xds_list] + + # Remove all coords bar ROWID so that they do not get written. + xds_list = [ + xds.drop_vars(set(xds.coords.keys()) - {"ROWID"}, errors='ignore') + for xds in xds_list + ] + # Remove attrs added by QuartiCal so that they do not get written. for xds in xds_list: - xds.attrs.pop("UTIME_CHUNKS", None) - xds.attrs.pop("FIELD_NAME", None) - - # Remove coords added by QuartiCal so that they do not get written. - xds_list = [xds.drop_vars(["chan", "corr"], errors='ignore') - for xds in xds_list] + xds.attrs = {k: v for k, v in xds.attrs.items() if k in DASKMS_ATTRS} with warnings.catch_warnings(): # We anticipate spurious warnings. warnings.simplefilter("ignore")