Merge pull request #23 from tnakazato/refactor-filler-feed

refactor: filler for FEED table

src/nro45data/psw/ms2/filler/__init__.py

@@ -7,7 +7,7 @@
 import nro45data.psw.ms2._casa as _casa
 from .antenna import fill_antenna
 from .data_description import fill_data_description
-from .feed import _fill_feed_columns, _get_feed_columns
+from .feed import fill_feed
 from .field import _fill_field_columns, _get_field_columns
 from .observation import _fill_observation_columns, _get_observation_columns
 from .polarization import _fill_polarization_columns, _get_polarization_columns
@@ -26,11 +26,6 @@
 LOG = logging.getLogger(__name__)
 
 
-def fill_feed(msfile: str, hdu: BinTableHDU):
-    columns = _get_feed_columns(hdu)
-    _fill_feed_columns(msfile, columns)
-
-
 def fill_field(msfile: str, hdu: BinTableHDU):
     columns = _get_field_columns(hdu)
     _fill_field_columns(msfile, columns)

src/nro45data/psw/ms2/filler/feed.py

@@ -1,91 +1,119 @@
+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
+from .._casa import datestr2mjd
+from .utils import get_array_configuration, fill_ms_table
 
 if TYPE_CHECKING:
     from astropy.io.fits.hdu.BinTableHDU import BinTableHDU
 
 LOG = logging.getLogger(__name__)
 
 
-def _get_feed_columns(hdu: "BinTableHDU") -> dict:
-    array_conf = get_array_configuration(hdu)
-    num_beam = len(np.unique([v[1] for v in array_conf.values()]))
-
-    antenna_id = np.arange(num_beam, dtype=int)
-
-    feed_id = np.zeros(num_beam, dtype=int)
-
-    spw_id = np.ones(num_beam, dtype=int) * -1
+def _get_feed_row(hdu: BinTableHDU) -> Generator[dict, None, None]:
+    """Provide feed row information.
 
-    num_receptors = np.ones(num_beam, dtype=int) * 2
+    Args:
+        hdu: NRO45m psw data in the form of BinTableHDU object.
 
-    beam_id = np.zeros(num_beam, dtype=int)
-
-    history_cards = hdu.header["HISTORY"]
-    start_time_card = [x for x in history_cards if x.startswith("NEWSTAR START-TIME")]
-    start_time = float(start_time_card[0].split("=")[-1].strip(" '"))
-    end_time_card = [x for x in history_cards if x.startswith("NEWSTAR END-TIME")]
-    end_time = float(end_time_card[0].split("=")[-1].strip(" '"))
-    feed_time = np.ones(num_beam, dtype=int) * (end_time + start_time) / 2
-    feed_interval = np.ones(num_beam, dtype=int) * (end_time - start_time)
-
-    beam_offset = [np.zeros((2, 2), dtype=float) for _ in range(num_beam)]
+    Yields:
+        Dictionary containing feed row information.
+    """
+    array_conf = get_array_configuration(hdu)
+    num_beam = len(np.unique([v[1] for v in array_conf.values()]))
 
     rx_str = np.unique(hdu.data["RX"])
-    if rx_str[0][-1] in ("H", "V"):
+    poltp_str = np.unique(hdu.data["POLTP"])
+    if np.all(poltp_str == ""):
+        pol_spec = rx_str[0].strip()[-1]
+    else:
+        pol_spec = poltp_str[0].strip()[-1]
+    if pol_spec in ("H", "V"):
         _pol_type = ["X", "Y"]
-    elif rx_str[0][-1] in ("R", "L"):
+    elif pol_spec in ("R", "L"):
         _pol_type = ["R", "L"]
     else:
         # fall back to ["X", "Y"]
         _pol_type = ["X", "Y"]
-    pol_type = [np.array(_pol_type) for _ in range(num_beam)]
-
-    pol_response = [np.zeros((2, 2), dtype=complex) for _ in range(num_beam)]
-
-    feed_position = [np.zeros(3, dtype=float) for _ in range(num_beam)]
-
-    receptor_angle = [np.zeros(2, dtype=float) for _ in range(num_beam)]
-
-    columns = {
-        "ANTENNA_ID": antenna_id,
-        "FEED_ID": feed_id,
-        "SPECTRAL_WINDOW_ID": spw_id,
-        "TIME": feed_time,
-        "INTERVAL": feed_interval,
-        "NUM_RECEPTORS": num_receptors,
-        "BEAM_ID": beam_id,
-        "BEAM_OFFSET": beam_offset,
-        "POLARIZATION_TYPE": pol_type,
-        "POL_RESPONSE": pol_response,
-        "POSITION": feed_position,
-        "RECEPTOR_ANGLE": receptor_angle,
-    }
-
-    return columns
-
-
-def _fill_feed_columns(msfile: str, columns: dict):
-    with open_table(msfile + "/FEED", read_only=False) as tb:
-        num_feed = len(columns["ANTENNA_ID"])
-        fix_nrow_to(num_feed, tb)
-
-        tb.putcol("ANTENNA_ID", columns["ANTENNA_ID"])
-        tb.putcol("FEED_ID", columns["FEED_ID"])
-        tb.putcol("SPECTRAL_WINDOW_ID", columns["SPECTRAL_WINDOW_ID"])
-        tb.putcol("TIME", columns["TIME"])
-        tb.putcol("INTERVAL", columns["INTERVAL"])
-        tb.putcol("NUM_RECEPTORS", columns["NUM_RECEPTORS"])
-        tb.putcol("BEAM_ID", columns["BEAM_ID"])
-
-        for i in range(num_feed):
-            tb.putcell("BEAM_OFFSET", i, columns["BEAM_OFFSET"][i])
-            tb.putcell("POLARIZATION_TYPE", i, columns["POLARIZATION_TYPE"][i])
-            tb.putcell("POL_RESPONSE", i, columns["POL_RESPONSE"][i])
-            tb.putcell("POSITION", i, columns["POSITION"][i])
-            tb.putcell("RECEPTOR_ANGLE", i, columns["RECEPTOR_ANGLE"][i])
+
+    for i in range(num_beam):
+        # ANTENNA_ID
+        antenna_id = i
+
+        # FEED_ID
+        feed_id = 0
+
+        # SPECTRAL_WINDOW_ID
+        spw_id = -1
+
+        # NUM_RECEPTORS
+        num_receptors = 2
+
+        # BEAM_ID
+        beam_id = 0
+
+        # TIME and INTERVAL
+        history_cards = hdu.header["HISTORY"]
+        start_time_card = [x for x in history_cards if x.startswith("NEWSTAR START-TIME")]
+        sstr = start_time_card[0].split("=")[-1].strip(" '")
+        LOG.debug("sstr: %s", sstr)
+        datestr = sstr[0:4] + "/" + sstr[4:6] + "/" + sstr[6:8] + " " + \
+            sstr[8:10] + ":" + sstr[10:12] + ":" + sstr[12:14]
+        LOG.debug("formatted sstr: %s", datestr)
+        start_time = datestr2mjd(datestr) - 9 * 3600
+        end_time_card = [x for x in history_cards if x.startswith("NEWSTAR END-TIME")]
+        estr = end_time_card[0].split("=")[-1].strip(" '")
+        LOG.debug("estr: %s", estr)
+        datestr = estr[0:4] + "/" + estr[4:6] + "/" + estr[6:8] + " " + \
+            estr[8:10] + ":" + estr[10:12] + ":" + estr[12:14]
+        LOG.debug("formatted estr: %s", datestr)
+        end_time = datestr2mjd(datestr) - 9 * 3600
+
+        feed_time = (end_time + start_time) / 2
+        feed_interval = end_time - start_time
+
+        # BEAM_OFFSET
+        beam_offset = np.zeros((2, 2), dtype=float)
+
+        # POLARIZATION_TYPE
+        pol_type = np.array(_pol_type)
+
+        # POL_RESPONSE
+        pol_response = np.zeros((2, 2), dtype=complex)
+
+        # POSITION
+        feed_position = np.zeros(3, dtype=float)
+
+        # RECEPTOR_ANGLE
+        receptor_angle = np.zeros(2, dtype=float)
+
+        row = {
+            "ANTENNA_ID": antenna_id,
+            "FEED_ID": feed_id,
+            "SPECTRAL_WINDOW_ID": spw_id,
+            "TIME": feed_time,
+            "INTERVAL": feed_interval,
+            "NUM_RECEPTORS": num_receptors,
+            "BEAM_ID": beam_id,
+            "BEAM_OFFSET": beam_offset,
+            "POLARIZATION_TYPE": pol_type,
+            "POL_RESPONSE": pol_response,
+            "POSITION": feed_position,
+            "RECEPTOR_ANGLE": receptor_angle,
+        }
+
+        yield row
+
+
+def fill_feed(msfile: str, hdu: BinTableHDU):
+    """Fill MS FEED table.
+
+    Args:
+        msfile: Name of MS file.
+        hdu: NRO45m psw data in the form of BinTableHDU object.
+    """
+    fill_ms_table(msfile, hdu, "FEED", _get_feed_row)

tests/ms2_filler/test_forest.py

@@ -70,6 +70,40 @@ def test_forest_ms2_structure(msfile):
             assert tb.getcell("POLARIZATION_ID", i) == 0
             assert tb.getcell("FLAG_ROW", i) is False
 
+    with open_table(os.path.join(msfile, "FEED")) as tb:
+        assert tb.nrows() == num_beams
+        # start time: 2024/09/30 17:49:19
+        # end time: 2024/09/30 17:57:32
+        # ---> mid-time 2024/09/30 17:53:25.5
+        #      interval 8min 13sec = 493sec
+        time_expected = datetime.datetime(2024, 9, 30, 17, 53, 25, 500000, tzinfo=datetime.timezone.utc)
+        interval_expected = 493.0  # sec
+        for i in range(num_beams):
+            assert tb.getcell("ANTENNA_ID", i) == i
+            assert tb.getcell("FEED_ID", i) == 0
+            assert tb.getcell("SPECTRAL_WINDOW_ID", i) == -1
+            assert tb.getcell("NUM_RECEPTORS", i) == 2
+            assert tb.getcell("BEAM_ID", i) == 0
+            feed_time = mjd2datetime(tb.getcell("TIME", i))
+            assert abs((feed_time - time_expected).total_seconds()) < 1e-3
+            assert tb.getcell("INTERVAL", i) == interval_expected
+            beam_offset = np.array(tb.getcell("BEAM_OFFSET", i))
+            assert beam_offset.shape == (2, 2)
+            assert np.all(beam_offset == 0.0)
+            pol_type = np.array(tb.getcell("POLARIZATION_TYPE", i))
+            assert pol_type.shape == (2,)
+            assert pol_type[0] == "X"
+            assert pol_type[1] == "Y"
+            pol_response = np.array(tb.getcell("POL_RESPONSE", i))
+            assert pol_response.shape == (2, 2)
+            assert np.all(pol_response == 0.0)
+            feed_position = np.array(tb.getcell("POSITION", i))
+            assert feed_position.shape == (3,)
+            assert np.all(feed_position == 0.0)
+            receptor_angle = np.array(tb.getcell("RECEPTOR_ANGLE", i))
+            assert receptor_angle.shape == (2,)
+            assert np.all(receptor_angle == 0.0)
+
     with open_table(os.path.join(msfile, "STATE")) as tb:
         intents_map = dict((i, v) for i, v in enumerate(tb.getcol("OBS_MODE")))
 

tests/ms2_filler/test_h40.py

@@ -72,6 +72,40 @@ def test_h40_ms2_structure(msfile):
             assert tb.getcell("POLARIZATION_ID", i) == 0
             assert tb.getcell("FLAG_ROW", i) is False
 
+    with open_table(os.path.join(msfile, "FEED")) as tb:
+        assert tb.nrows() == num_beams
+        # start time: 2024/09/25 15:58:54
+        # end time: 2024/09/25 16:03:50
+        # ---> mid-time 2024/09/25 16:01:22
+        #      interval 4min 56sec = 296sec
+        time_expected = datetime.datetime(2024, 9, 25, 16, 1, 22, tzinfo=datetime.timezone.utc)
+        interval_expected = 296.0  # sec
+        for i in range(num_beams):
+            assert tb.getcell("ANTENNA_ID", i) == i
+            assert tb.getcell("FEED_ID", i) == 0
+            assert tb.getcell("SPECTRAL_WINDOW_ID", i) == -1
+            assert tb.getcell("NUM_RECEPTORS", i) == 2
+            assert tb.getcell("BEAM_ID", i) == 0
+            feed_time = mjd2datetime(tb.getcell("TIME", i))
+            assert abs((feed_time - time_expected).total_seconds()) < 1e-3
+            assert tb.getcell("INTERVAL", i) == interval_expected
+            beam_offset = np.array(tb.getcell("BEAM_OFFSET", i))
+            assert beam_offset.shape == (2, 2)
+            assert np.all(beam_offset == 0.0)
+            pol_type = np.array(tb.getcell("POLARIZATION_TYPE", i))
+            assert pol_type.shape == (2,)
+            assert pol_type[0] == "R"
+            assert pol_type[1] == "L"
+            pol_response = np.array(tb.getcell("POL_RESPONSE", i))
+            assert pol_response.shape == (2, 2)
+            assert np.all(pol_response == 0.0)
+            feed_position = np.array(tb.getcell("POSITION", i))
+            assert feed_position.shape == (3,)
+            assert np.all(feed_position == 0.0)
+            receptor_angle = np.array(tb.getcell("RECEPTOR_ANGLE", i))
+            assert receptor_angle.shape == (2,)
+            assert np.all(receptor_angle == 0.0)
+
     with open_table(os.path.join(msfile, "STATE")) as tb:
         intents_map = dict((i, v) for i, v in enumerate(tb.getcol("OBS_MODE")))
 

tests/ms2_filler/test_z45.py

@@ -71,14 +71,48 @@ def test_z45_ms2_structure(msfile):
             assert tb.getcell("POLARIZATION_ID", i) == 0
             assert tb.getcell("FLAG_ROW", i) is False
 
+    with open_table(os.path.join(msfile, "FEED")) as tb:
+        assert tb.nrows() == num_beams
+        # start time: 2024/09/30 18:00:32
+        # end time: 2024/09/30 18:06:33
+        # ---> mid-time 2024/09/30 18:03:32.5
+        #      interval 6min 1sec = 361sec
+        time_expected = datetime.datetime(2024, 9, 30, 18, 3, 32, 500000, tzinfo=datetime.timezone.utc)
+        interval_expected = 361.0  # sec
+        for i in range(num_beams):
+            assert tb.getcell("ANTENNA_ID", i) == i
+            assert tb.getcell("FEED_ID", i) == 0
+            assert tb.getcell("SPECTRAL_WINDOW_ID", i) == -1
+            assert tb.getcell("NUM_RECEPTORS", i) == 2
+            assert tb.getcell("BEAM_ID", i) == 0
+            feed_time = mjd2datetime(tb.getcell("TIME", i))
+            assert abs((feed_time - time_expected).total_seconds()) < 1e-3
+            assert tb.getcell("INTERVAL", i) == interval_expected
+            beam_offset = np.array(tb.getcell("BEAM_OFFSET", i))
+            assert beam_offset.shape == (2, 2)
+            assert np.all(beam_offset == 0.0)
+            pol_type = np.array(tb.getcell("POLARIZATION_TYPE", i))
+            assert pol_type.shape == (2,)
+            assert pol_type[0] == "X"
+            assert pol_type[1] == "Y"
+            pol_response = np.array(tb.getcell("POL_RESPONSE", i))
+            assert pol_response.shape == (2, 2)
+            assert np.all(pol_response == 0.0)
+            feed_position = np.array(tb.getcell("POSITION", i))
+            assert feed_position.shape == (3,)
+            assert np.all(feed_position == 0.0)
+            receptor_angle = np.array(tb.getcell("RECEPTOR_ANGLE", i))
+            assert receptor_angle.shape == (2,)
+            assert np.all(receptor_angle == 0.0)
+
     with open_table(os.path.join(msfile, "STATE")) as tb:
         intents_map = dict((i, v) for i, v in enumerate(tb.getcol("OBS_MODE")))
 
     with open_table(os.path.join(msfile, "OBSERVATION")) as tb:
         assert tb.nrows() == 1
         time_range = tb.getcell("TIME_RANGE", 0)
-        # start time: 2024/09/25 15:58:54
-        # end time: 2024/09/25 16:03:50
+        # start time: 2024/09/30 18:00:32
+        # end time: 2024/09/30 18:06:33
         start_expected = datetime.datetime(2024, 9, 30, 18, 00, 32, tzinfo=datetime.timezone.utc)
         start_time = mjd2datetime(time_range[0])
         # impose msec accuracy