Source code for low_comm_tools.flag_rfi

from __future__ import annotations

import argparse
import subprocess as sp
from importlib import resources
from pathlib import Path

from low_comm_tools import data



STRATEGY_FILES = {
    "comprehensive": "SKA-Low_sharpRFI_comprehensive.lua",
    "fast": "SKA-Low_sharpRFI_fast.lua",
    "extended": "SKA-Low_extendedRFI.lua",
}





def run_aoflagger(
    vis_path: Path,
    strategy: str | None = None,
    chunk_size: int = 1800,
) -> Path:
    if strategy is None:
        sp.run(
            [
                "aoflagger",
                "-chunk-size",
                str(chunk_size),
                vis_path.as_posix(),
            ],
            check=True,
        )
        return vis_path

    strategy_file = STRATEGY_FILES[strategy]
    strategy_resource = resources.files(data) / "rfi" / strategy_file

    with resources.as_file(strategy_resource) as strategy_path:
        sp.run(
            [
                "aoflagger",
                "-chunk-size",
                str(chunk_size),
                "--strategy",
                str(strategy_path),
                vis_path.as_posix(),
            ],
            check=True,
        )

    return vis_path





def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Run AOFlagger on a Measurement Set (MS).",
    )
    parser.add_argument("ms_path", type=Path, help="Path to the Measurement Set")
    parser.add_argument(
        "--rfi-strategy",
        choices=["comprehensive", "fast", "extended"],
        default=None,
        help="Optional packaged AOFlagger strategy to use.",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=1800,
        help="AOFlagger chunk size in timesteps. Default: 1800.",
    )
    return parser





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

    run_aoflagger(
        vis_path=args.ms_path,
        strategy=args.rfi_strategy,
        chunk_size=args.chunk_size,
    )
    return 0



if __name__ == "__main__":
    raise SystemExit(main())