Source code for low_comm_tools.ms_fixes.fix_interval

from __future__ import annotations

import argparse
from pathlib import Path

import astropy.units as u
import numpy as np
from casacore.tables import table

from low_comm_tools import ms_utils
from low_comm_tools.log_config import logger




def get_interval_from_time(ms_path: Path) -> u.Quantity[u.s]:
    current_interval = ms_utils.get_interval_from_ms(
        ms_path,
    )

    logger.info(f"Current INTERVAL column has {current_interval=}")

    with table(ms_path.as_posix(), ack=False, readonly=True) as tab:
        time_mjds = tab.getcol("TIME") * u.s

    unique_time_mjds, time_index = np.unique(time_mjds, return_inverse=True)
    diff_s = np.diff(unique_time_mjds)

    diff_s_padded = np.append(diff_s, diff_s[-1])

    new_interval = current_interval.copy()
    new_interval[:] = diff_s_padded[time_index]

    return new_interval





def fix_interval_from_time(ms_path: Path) -> Path:
    logger.info(f"Fixing INTERVAL column in {ms_path}")
    new_interval = get_interval_from_time(ms_path)
    logger.info(f"New INTERVAL column will be {new_interval=}")

    with table(ms_path.as_posix(), ack=False, readonly=False) as tab:
        tab.putcol(
            "INTERVAL",
            value=new_interval.to(u.s).value,
        )
    logger.info(f"Fixed INTERVAL column in {ms_path}")
    return ms_path





def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Fix INTERVAL column using TIME in MS",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("ms_path", help="Path to target MeasurementSet", type=Path)
    return parser





def main() -> None:
    parser = get_parser()
    args = parser.parse_args()

    _ = fix_interval_from_time(
        ms_path=args.ms_path,
    )



if __name__ == "__main__":
    main()