Source code for low_comm_tools.ms_fixes.fix_scanid

from __future__ import annotations

import argparse
from pathlib import Path

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 fix_scan_numbers(ms_path: Path) -> Path:
    logger.info(f"Fixing scan numbers in {ms_path}")
    current_scan_numbers = ms_utils.get_scan_numbers_from_ms(ms_path)
    logger.info(f"There are {len(current_scan_numbers)} scan numbers")

    with table(ms_path.as_posix(), ack=False, readonly=False) as tab:
        for new_scan_number, subtab in enumerate(tab.iter("SCAN_NUMBER")):
            scan_col = subtab.getcol("SCAN_NUMBER")
            new_scan_col = new_scan_number * np.ones_like(scan_col)
            subtab.putcol("SCAN_NUMBER", value=new_scan_col)

    logger.info(f"Updated SCAN_NUMBER column in {ms_path}")
    return ms_path





def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Fix SCAN_NUMBER column 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_scan_numbers(
        ms_path=args.ms_path,
    )



if __name__ == "__main__":
    main()