diff --git a/src/nro45data/psw/ms2/filler/__init__.py b/src/nro45data/psw/ms2/filler/__init__.py index 5f76596..277cc86 100644 --- a/src/nro45data/psw/ms2/filler/__init__.py +++ b/src/nro45data/psw/ms2/filler/__init__.py @@ -5,8 +5,8 @@ from typing import TYPE_CHECKING import nro45data.psw.ms2._casa as _casa -from .antenna import fill_antenna # _fill_antenna_columns, _get_antenna_columns -from .data_description import _fill_data_description_columns, _get_data_description_columns +from .antenna import fill_antenna +from .data_description import fill_data_description from .feed import _fill_feed_columns, _get_feed_columns from .field import _fill_field_columns, _get_field_columns from .observation import _fill_observation_columns, _get_observation_columns @@ -26,16 +26,6 @@ LOG = logging.getLogger(__name__) -# def fill_antenna(msfile: str, hdu: BinTableHDU): -# columns = _get_antenna_columns(hdu) -# _fill_antenna_columns(msfile, columns) - - -def fill_data_description(msfile: str, hdu: BinTableHDU): - columns = _get_data_description_columns(hdu) - _fill_data_description_columns(msfile, columns) - - def fill_feed(msfile: str, hdu: BinTableHDU): columns = _get_feed_columns(hdu) _fill_feed_columns(msfile, columns) diff --git a/src/nro45data/psw/ms2/filler/antenna.py b/src/nro45data/psw/ms2/filler/antenna.py index ac37301..8e7068b 100644 --- a/src/nro45data/psw/ms2/filler/antenna.py +++ b/src/nro45data/psw/ms2/filler/antenna.py @@ -5,7 +5,6 @@ import numpy as np -from .._casa import open_table from .utils import get_array_configuration, fill_ms_table if TYPE_CHECKING: diff --git a/src/nro45data/psw/ms2/filler/data_description.py b/src/nro45data/psw/ms2/filler/data_description.py index a2f806e..4fafd75 100644 --- a/src/nro45data/psw/ms2/filler/data_description.py +++ b/src/nro45data/psw/ms2/filler/data_description.py @@ -1,10 +1,11 @@ +from __future__ import annotations + import logging -from typing import TYPE_CHECKING +from typing import Generator, TYPE_CHECKING import numpy as np -from .._casa import open_table -from .utils import fix_nrow_to, get_array_configuration, get_data_description_map +from .utils import get_array_configuration, get_data_description_map, fill_ms_table if TYPE_CHECKING: from astropy.io.fits.hdu.BinTableHDU import BinTableHDU @@ -12,28 +13,39 @@ LOG = logging.getLogger(__name__) -def _get_data_description_columns(hdu: "BinTableHDU") -> dict: +def _get_data_description_row(hdu: BinTableHDU) -> Generator[dict, None, None]: + """Provide data description row information. + + Args: + hdu: NRO45m psw data in the form of BinTableHDU object. + + Yields: + Dictionary containing data description row information. + """ array_conf = get_array_configuration(hdu) ddd, _, _, _ = get_data_description_map(array_conf) - num_dd = len(ddd) - spw_id = np.zeros(num_dd, dtype=int) - pol_id = np.zeros(num_dd, dtype=int) - for dd_id, (spw, pol) in ddd.items(): - spw_id[dd_id] = spw - pol_id[dd_id] = pol - flag_row = np.zeros(num_dd, dtype=bool) + for _, (spw, pol) in ddd.items(): + spw_id = spw + pol_id = pol + + flag_row = False - columns = {"SPECTRAL_WINDOW_ID": spw_id, "POLARIZATION_ID": pol_id, "FLAG_ROW": flag_row} - LOG.debug("columns: %s", columns) + row = { + "SPECTRAL_WINDOW_ID": spw_id, + "POLARIZATION_ID": pol_id, + "FLAG_ROW": flag_row + } + LOG.debug("row: %s", row) - return columns + yield row -def _fill_data_description_columns(msfile: str, columns: dict): - with open_table(msfile + "/DATA_DESCRIPTION", read_only=False) as tb: - num_dd = len(columns["SPECTRAL_WINDOW_ID"]) - fix_nrow_to(num_dd, tb) +def fill_data_description(msfile: str, hdu: BinTableHDU): + """Fill MS DATA_DESCRIPTION table. - tb.putcol("SPECTRAL_WINDOW_ID", columns["SPECTRAL_WINDOW_ID"]) - tb.putcol("POLARIZATION_ID", columns["POLARIZATION_ID"]) + Args: + msfile: Name of MS file. + hdu: NRO45m psw data in the form of BinTableHDU object. + """ + fill_ms_table(msfile, hdu, "DATA_DESCRIPTION", _get_data_description_row) diff --git a/tests/ms2_filler/test_forest.py b/tests/ms2_filler/test_forest.py index 65bedb3..6218266 100644 --- a/tests/ms2_filler/test_forest.py +++ b/tests/ms2_filler/test_forest.py @@ -61,6 +61,14 @@ def test_forest_ms2_structure(msfile): position = tb.getcell("POSITION", i) assert position.shape == (3,) assert np.allclose(position, np.array([-3871023.46, 3428106.87, 3724039.47])) + assert tb.getcell("FLAG_ROW", i) is False + + with open_table(os.path.join(msfile, "DATA_DESCRIPTION")) as tb: + assert tb.nrows() == num_spws + for i in range(num_spws): + assert tb.getcell("SPECTRAL_WINDOW_ID", i) == i + assert tb.getcell("POLARIZATION_ID", i) == 0 + assert tb.getcell("FLAG_ROW", i) is False with open_table(os.path.join(msfile, "STATE")) as tb: intents_map = dict((i, v) for i, v in enumerate(tb.getcol("OBS_MODE"))) diff --git a/tests/ms2_filler/test_h40.py b/tests/ms2_filler/test_h40.py index 80b8a0c..e4ada3c 100644 --- a/tests/ms2_filler/test_h40.py +++ b/tests/ms2_filler/test_h40.py @@ -63,6 +63,14 @@ def test_h40_ms2_structure(msfile): position = tb.getcell("POSITION", i) assert position.shape == (3,) assert np.allclose(position, np.array([-3871023.46, 3428106.87, 3724039.47])) + assert tb.getcell("FLAG_ROW", i) is False + + with open_table(os.path.join(msfile, "DATA_DESCRIPTION")) as tb: + assert tb.nrows() == num_spws + for i in range(num_spws): + assert tb.getcell("SPECTRAL_WINDOW_ID", i) == i + assert tb.getcell("POLARIZATION_ID", i) == 0 + assert tb.getcell("FLAG_ROW", i) is False with open_table(os.path.join(msfile, "STATE")) as tb: intents_map = dict((i, v) for i, v in enumerate(tb.getcol("OBS_MODE"))) diff --git a/tests/ms2_filler/test_z45.py b/tests/ms2_filler/test_z45.py index 8c9c939..6f15e8d 100644 --- a/tests/ms2_filler/test_z45.py +++ b/tests/ms2_filler/test_z45.py @@ -62,6 +62,14 @@ def test_z45_ms2_structure(msfile): position = tb.getcell("POSITION", i) assert position.shape == (3,) assert np.allclose(position, np.array([-3871023.46, 3428106.87, 3724039.47])) + assert tb.getcell("FLAG_ROW", i) is False + + with open_table(os.path.join(msfile, "DATA_DESCRIPTION")) as tb: + assert tb.nrows() == num_spws + for i in range(num_spws): + assert tb.getcell("SPECTRAL_WINDOW_ID", i) == i + assert tb.getcell("POLARIZATION_ID", i) == 0 + assert tb.getcell("FLAG_ROW", i) is False with open_table(os.path.join(msfile, "STATE")) as tb: intents_map = dict((i, v) for i, v in enumerate(tb.getcol("OBS_MODE")))