calibrate_lightning.zsh

export obsid=1333747784
export srclist=/opt/cal/srclist_pumav3_EoR0LoBES_EoR1pietro_CenA-GP_2023-11-07.yaml
export full_srclist=$srclist
# export cargo="cargo +nightly -Z unstable-options -C ~/src/hyperdrive-peel"
# export test="$cargo test --release --features=hdf5-static,cuda,gpu-single"
export hyperdrive="hyperdrive" # optionally use a different hyperdrive binary
# export hyperdrive="$cargo run --release --features=hdf5-static,cuda,gpu-single"

## Preprocessing settings
export timeres_s=2        # time resolution to average to in seconds
export freqres_khz=40     # frequency resolution to average to in kHz
export edgewidth_khz=80   # edge width to flag in kHz
export birli_args=""      # extra birli args if any
export prep_name="_${timeres_s}s_${freqres_khz}kHz"

## Srclist settings
export num_sources=4000   # number of sources to model in calibration
# export num_sources=500
export dist_cutoff=90     # distance cutoff for srclist in degrees
# export dist_cutoff=50
export veto_thresh=0.005  # minimum source brightness in Jy
# export veto_thresh=0.01
export srclist_args="-n ${num_sources} --source-dist-cutoff=${dist_cutoff} --veto-threshold ${veto_thresh}"
export srclist_name="_${num_sources}src_${dist_cutoff}deg_${veto_thresh}Jy"

## Calibration settings
export uvw_min="50l"       # minimum uvw cutoff (lamda)
# export uvw_min="30l"
export max_iters=300       # number of di-cal iterations
# export max_iters=50
export stop_thresh=1e-20   # stop threshold for di-cal
# export stop_thresh=1e-8
export dical_freq="80kHz"  # frequency resolution for di-cal
export dical_args="--uvw-min ${uvw_min} --max-iterations ${max_iters} --stop-thresh ${stop_thresh} --freq-average ${dical_freq}"
export dical_name="_${uvw_min}_${num_sources}src_${max_iters}it_${dical_freq}"

## Apply settings
export apply_time=4        # apply time averaging (seconds)
export apply_freq=80       # apply frequency averaging (kHz)
export apply_name="_${apply_time}s_${apply_freq}kHz"
export apply_args="--time-average ${apply_time}s --freq-average ${apply_freq}kHz"

export outdir=/data/dev/${obsid}
mkdir -p ${outdir}/log
mkdir -p ${outdir}/raw
export metafits=${outdir}/raw/${obsid}.metafits
export metafits_fixed=${outdir}/raw/${obsid}_fixed.metafits

mkdir -p ${outdir}/prep
export prep_ms="${outdir}/prep/birli_${obsid}${prep_name}.ms"
export prep_fixed_ms="${outdir}/prep/birli_${obsid}${prep_name}_fixed.ms"

mkdir -p ${outdir}/cal
export hyp_srclist="${outdir}/cal/srclist_${obsid}${srclist_name}.yaml"
export hyp_soln="${outdir}/cal/hyp_soln_${obsid}${dical_name}.fits"
export hyp_soln_fixed="${outdir}/cal/hyp_soln_${obsid}${dical_name}_fixed.fits"
export cal_ms="${outdir}/cal/hyp_${obsid}${dical_name}.ms"
export model_ms="${outdir}/cal/hyp_${obsid}${dical_name}.ms"

mkdir -p "${outdir}/ionoqa"
export ionosub_qa="${outdir}/ionoqa/ionoqa.json"

# if you uncomment this, then sourcing this script won't run the rest of it
return 0

# download raw files
# giant-squid submit-vis 1333747784
# cd $outdir/raw
# giant-squid download 1333747784

# ensure you are using the latest metafits, as new flags can be added.

[ -f "$metafits" ] || wget -O "$metafits" $'http://ws.mwatelescope.org/metadata/fits?obs_id='${obsid}

# preprocess raw files
[ -d ${prep_ms} ] || birli ${birli_args} \
    --avg-freq-res ${freqres_khz} --avg-time-res ${timeres_s} \
    --flag-edge-width ${edgewidth_khz} \
    -M "${prep_ms}" \
    -m "${metafits}" \
    ${outdir}/raw/${obsid}_2*.fits \
    2>&1 | tee "${outdir}/log/birli${prep_name}.log"

# debug hyperdrive srclist
[ -f "$hyp_srclist" ] || eval $hyperdrive srclist-by-beam -vv ${=srclist_args} \
    --metafits "$metafits" \
    "$full_srclist" \
    "$hyp_srclist" \
    | tee "${outdir}/log/hyp_srclist${srclist_name}.log"

# first round of calibration
[ -f "$hyp_soln" ] || eval $hyperdrive di-cal -vv ${=srclist_args} ${=dical_args} \
    --data "$metafits" "$prep_ms" \
    --source-list "$hyp_srclist" \
    --outputs "$hyp_soln" \
    --model-filenames "$cal_model" \
    --tile-flags LBC2 \
    | tee "${outdir}/log/hyp_dical${dical_name}.log"

# plot solutions
[ -f "${hyp_soln%%.fits}_phases.png" ] || eval $hyperdrive solutions-plot \
    -m "$metafits" "${hyp_soln}" --ref-tile 7 \
    --output-directory "${outdir}/cal" \
    "$hyp_soln"

# fit cable lengths from solution
[ -f "$metafits_fixed" ] || docker run -it --rm -v ${outdir}:${outdir} -w ${outdir}/cal --entrypoint python \
    d3vnull0/mwax_mover:latest \
    /app/scripts/cal_analysis.py \
    --name "hyp${dical_name}" \
    --metafits "$metafits" --solns "$hyp_soln" \
    --phase-diff-path=/app/phase_diff.txt

# fix metafits with fitted cable lengths
[ -f "$metafits_fixed" ] || docker run -it --rm -v ${outdir}:${outdir} -w ${outdir}/cal --entrypoint python \
    d3vnull0/mwax_mover:latest \
    /app/scripts/update_cable_lengths.py \
    --name "hyp${dical_name}" \
    --metafits "$metafits"

# flag LBC2 in fixed metafits
python -c "import sys; from astropy.io import fits; p=sys.argv[-1]; h=fits.open(p); t=h[1].data; t['Flag'][t['TileName']=='LBC2']=1; h.writeto(p, overwrite=True)" $metafits_fixed

# repreprocess with fixed metafits
[ -d "$prep_fixed_ms" ] || birli ${birli_args} \
    --avg-freq-res ${freqres_khz} --avg-time-res ${timeres_s} \
    --flag-edge-width ${edgewidth_khz} \
    -M "${prep_fixed_ms}" \
    -m "${metafits_fixed}" \
    ${outdir}/raw/${obsid}_2*.fits \
    | tee "${outdir}/log/birli${prep_name}_fixed.log"

# dical with fixed metafits
[ -f "$hyp_soln_fixed" ] || eval $hyperdrive di-cal -vv ${=srclist_args} ${=dical_args} \
    --data "$metafits_fixed" "$prep_fixed_ms" \
    --source-list "$hyp_srclist" \
    --outputs "$hyp_soln_fixed" \
    | tee "${outdir}/log/hyp_dical${dical_name}_fixed.log"

# plot fixed solutions
[ -f "${hyp_soln_fixed%%.fits}_phases.png" ] || eval $hyperdrive solutions-plot \
    -m "$metafits" "${hyp_soln_fixed}" --no-ref-tile \
    --output-directory "${outdir}/cal" \
    "$hyp_soln"

# apply
[ -d "$cal_ms" ] || eval $hyperdrive apply -vv ${=apply_args} \
    --data "$metafits_fixed" "$prep_fixed_ms" \
    --solutions "$hyp_soln_fixed" \
    --outputs "$cal_ms" \
    | tee "${outdir}/log/hyp_apply${apply_name}_fixed.log"

# image
export img_base=""
unset imgs
if [[ -n "$ZSH_VERSION" ]]; then
    typeset -A imgs
else
    declare -A imgs
fi

# imgs[uni]="$img_base -weight uniform -size 4096 4096 -scale 40asec -niter 10000"
imgs[uni_24ch]="$img_base -weight uniform -size 4096 4096 -scale 40asec -niter 1000 -channels-out 24 -join-channels"

for img_name img_args in ${(kv)imgs}; do
    mkdir -p "${outdir}/img-${img_name}"
    pushd "${outdir}/img-${img_name}"
        wsclean ${=img_args} \
            -name wsclean_prep_${obsid}_${dical_name}_${apply_name}_${img_name} \
            "${prep_ms}"
        wsclean ${=img_args} \
            -name wsclean_hyp_${obsid}_${dical_name}_${apply_name}_${img_name} \
            "${cal_ms}"
        wsclean ${=img_args} \
            -name wsclean_hyp_model_${obsid}_${dical_name}_${apply_name}_${img_name} \
            "${cal_model}"
    popd
done