given a set of MWA raw files on ch137, and their metafits, crop 40kHz wide around the signal and write to disk

ssins_starlink_175.py

#!/usr/bin/env python

# TODO: FREQUENCY BROADCAST

import os
import json

from SSINS import Catalog_Plot as cp, util as ssins_util
from SSINS import SS, INS, MF, plot_lib
from matplotlib import pyplot as plt, cm
import numpy as np
import pylab
from pyuvdata import UVData, UVFlag, utils as uvutils
import sys
import shlex
from itertools import groupby
import traceback

def get_parser():
    import argparse

    parser = argparse.ArgumentParser(
        description="Run SSINS on data visibilities.")

    parser.add_argument('--data', nargs='+',
                        help='files to flag')

    plot_group = parser.add_argument_group('PLOTTING OPTIONS')
    plot_group.add_argument('--output_prefix',
                            help='Prefix of output plot files', default=''
                            )
    plot_group.add_argument('--plot_title', default=None,
                            help="Optional title for the plot")
    plot_group.add_argument('--sel_ants', default=[], nargs='*', type=int,
                            help="antenna indices to select")

    return parser


"""
example:

```bash
export script=ssins_starlink_175.py
[ -f $script ] || wget -O $script 'https://gist.githubusercontent.com/d3v-null/a1e5c8dc5fa74f2ab07f752d5c719d0f/raw/46e302748566dd4b5aa11389411ea6ce6576ad83/ssins_starlink_175.py'
singularity exec \
    --bind ${PWD} --cleanenv --home ${HOME}/mplhome \
    docker://d3vnull0/ssins:latest python \
    $script \
    --plot_title="test" \
    --output_prefix="test" \
    --data \
      1390650472_20240130114734_ch137_000.fits 1390650472.metafits \
      1390650592_20240130114934_ch137_000.fits 1390650592.metafits # ...
"""

def ss_combine(ss1, ss2):
    """
    try and combine two SS objects, no biggie if it fails
    """
    try:
        # ss1.fast_concat(
        #     ss2,
        #     inplace=True,
        #     axis='blt',
        #     # verbose_history=True,
        #     run_check=False,
        #     check_extra=False,
        #     run_check_acceptability=False,
        #     strict_uvw_antpos_check=False,
        #     ignore_name=True,
        # )
        ss1.__add__(
            ss2,
            inplace=True,
            # verbose_history=True,
            run_check=False,
            check_extra=False,
            run_check_acceptability=False,
            strict_uvw_antpos_check=False,
            ignore_name=True,
        )
    except ValueError as exc:
        traceback.print_exc()
    return ss1

def main():

    parser = get_parser()

    if len(sys.argv) > 1:
        args = parser.parse_args()
    else:
        # is being called directly from nextflow
        args = parser.parse_args(shlex.split("""${args}"""))

    print(vars(args))
    ins_plot_args = {
        "file_ext": "png",
        "title": args.plot_title,
        # "extent_time_format": "lst"
    }
    shape_dict = {}
    sig_thresh = {
        'narrow': 5,
        'streak': 8,
    }

    # ss = None
    ins_cross = None
    read_kwargs = {
        # "times": times,
        # "correct_coarse_band": False
        "remove_coarse_band": False,
    }
    if args.sel_ants:
        read_kwargs["antenna_nums"] = args.sel_ants
    groups = groupby(args.data, lambda x: x.split('/')[-1][:10])
    groups = sorted([(k, [*v]) for k, v in groups], key=lambda x: x[0])
    # groups = groups[2:4] # deleteme
    sss = {}
    for k, v in groups:
        v = [*v]
        if len(v) < 2:
            continue
        print(f'processing {k}')
        ss = SS()
        ss.read(v, read_data=False, **read_kwargs)
        # only look at flags 174.8-175.2 MHz
        # print(ss.freq_array)
        freq_idxs = np.where(np.logical_and(ss.freq_array > 174.95e6, ss.freq_array < 175.05e6))[1]
        print(f"{freq_idxs=}")
        ss.read(v, read_data=True, diff=True, freq_chans=freq_idxs, **read_kwargs)
        ss.apply_flags(flag_choice='original')


        ins_cross_ = INS(ss, spectrum_type='cross')
        sss[k] = {
            'ss': ss,
            'files': v,
            'ins_cross': ins_cross_,
        }

        if ins_cross is None:
            ins_cross = ins_cross_
        else:
            ins_cross += ins_cross_


    cp.INS_plot(ins_cross, f'{args.output_prefix}cross', **ins_plot_args)

    mf = MF(ins_cross.freq_array, sig_thresh,
            shape_dict=shape_dict, streak=True, narrow=True)
    ins_cross.metric_array[ins_cross.metric_array == 0] = np.ma.masked
    ins_cross.metric_ms = ins_cross.mean_subtract()
    ins_cross.sig_array = np.ma.copy(ins_cross.metric_ms)
    mf.apply_match_test(ins_cross)

    cp.INS_plot(ins_cross, f'{args.output_prefix}flagged', **ins_plot_args)
    occ_dict = ssins_util.calc_occ(ins_cross, mf, 0)
    with open(f'{args.output_prefix}ssins_occ.json', "w") as json_file:
        json.dump(occ_dict, json_file, indent=4)

    # with open(f'{args.output_prefix}ssins_events.json', "w") as events_file:
    #     events_obj = [
    #         {
    #             'time_bounds': [int(event[0].start), int(event[0].stop)],
    #             'freq_bounds': [int(event[1].start), int(event[1].stop)],
    #             'shape': event[2],
    #             'sig': event[3] if event[3] is None else float(event[3])
    #         } for event in ins_cross.match_events
    #     ]
    #     json.dump(events_obj, events_file, indent=4)

    ins_cross.write(f'{args.output_prefix}', output_type='mask', clobber=True)

    flags = ins_cross.mask_to_flags()[:, :, 0]
    freq_array = ins_cross.freq_array
    time_array = ins_cross.time_array

    # only look at flags 174.8-175.2 MHz
    freq_idxs = np.where(np.logical_and(freq_array > 174.8e6, freq_array < 175.2e6))[0]
    freq_array = freq_array[freq_idxs]
    flags = flags[:, freq_idxs]

    # time idxs that are not completely flagged
    partial_time_idxs = np.where(np.any(np.logical_not(flags), axis=1))[0]
    partial_time_idxs = partial_time_idxs[partial_time_idxs < len(time_array)]
    print(f"{partial_time_idxs=}")
    time_array = time_array[partial_time_idxs]
    flags = flags[partial_time_idxs, :]

    # time idxs that are completely unflagged
    unflagged_time_idxs = np.where(np.all(np.logical_not(flags), axis=1))[0]
    print(f"{unflagged_time_idxs=}")
    unflagged_time_array = time_array[unflagged_time_idxs]
    print(f"{unflagged_time_array=}")

    # get closest freq to 175MHz
    freq_idx175 = np.argmin(np.abs(freq_array - 175e6))
    print(f"{freq_idx175=}")
    time_idxs_175 = np.where(flags[:, freq_idx175])[0]
    print(f"{time_idxs_175=}")
    time_array_175 = time_array[time_idxs_175]

    for k, meta in sss.items():
        uvd = UVData()
        uvd.read(meta['files'], read_data=False)
        time_min = np.min(uvd.time_array)
        time_max = np.max(uvd.time_array)
        n_unflagged = np.sum(np.logical_and(unflagged_time_array >= time_min, unflagged_time_array <= time_max))
        n_175 = np.sum(np.logical_and(time_array_175 >= time_min, time_array_175 <= time_max))
        print(f"{k=}: {time_min=}, {time_max=}, {n_unflagged=}, {n_175=}")
        if n_unflagged < 2:
            print(f"{k=}: not enough intersection with unflagged time array: {n_unflagged}")
            continue
        if n_175 < 2:
            print(f"{k=}: not enough intersection with 175MHz signal: {n_175}")
            continue
        uvd.read(meta['files'], frequencies=freq_array, **read_kwargs)
        path = f'prep/{k}.uvfits'
        print(f"writing to {path=}")
        uvd.write_uvfits(path, force_phase=True)


if __name__ == '__main__':
    main()