Tnakefile

# https://stackoverflow.com/questions/66466082/how-do-i-write-a-snakemake-input-when-not-all-jobs-successfully-output-files-fro

configfile: "config.yaml"
FIELD = config['FIELD']
CUBE = str(config['CUBE'])

beams, = glob_wildcards(FIELD+"/HI_B0{xx}_cube"+CUBE+"_image.fits")

wildcard_constraints:
    beam="\d+"

rule all:
    input:
        expand(FIELD + "/HI_B0{beam}_cube"+CUBE+".txt", beam=beams)

rule make_bin_mask:
    input:
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask.fits"
    output:
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask_bin.fits"
    threads:
        18
    resources:
        mem_mb = 5000
    shell:
        "python3 /mnt/scratch/stuff/aper_sf2/src/binary_mask.py -t "+FIELD+" -c "+CUBE

rule dilate_mask:
    input:
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask_bin.fits"
    output:
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask_bin_dil.fits"
    threads:
        18
    resources:
        mem_mb = 5000
    shell:
        "python3 /mnt/scratch/stuff/aper_sf2/src/dilate_binmask.py -t "+FIELD+" -c "+CUBE

rule regrid_masks:
    input:
        FIELD+"/HI_B0{bm}_cube"+CUBE+"_image.fits",
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask_bin_dil.fits"
    output:
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask_bin_dil{bm}_regrid.fits"
#    threads:
#        #5  # Normal
#        3  # MDS fields (should really be doing this with memory allocation not threads
    resources:
        mem_mb = 55000
    shell:
        "python3 /mnt/scratch/stuff/aper_sf2/src/regridmask.py {input[0]} {input[1]}"

rule generate_new_pb:
    input:
        FIELD+"/HI_B0{bm}_cube"+CUBE+"_image.fits",
        "/mnt/scratch/apertif/cbeams/{bm}_gp_avg_orig.fits"
    output:
        FIELD+"/HI_B0{bm}_cube"+CUBE+"_spline_clean_pb.fits"
    threads:
        1
    run:
        shell("python3 /mnt/scratch/stuff/aper_cube_stack/modules/regrid_aperpb.py -t {input[0]} -b {wildcards.bm} -p /mnt/scratch/apertif/cbeams/")
        shell("mv "+FIELD+"/HI_B0{wildcards.bm}_cube"+CUBE+"_pb.fits {output}")

rule stack_psf:
    input:
        FIELD+"/HI_B0{bm}_cube"+CUBE+"_image.fits"
    output:
        FIELD+"/HI_B0{bm}_cube"+CUBE+"_psf.fits"
    threads:
        4   # This is a trick to not exceed the total memory (see below; sad face)
    resources:
        mem_mb = 30000 # For AWES;  MEMORY PER CORE NOT TOTAL!!!
      #  mem_mb = 185000  # For MDS 
    shell:
        "python3 /mnt/scratch/stuff/aper_cube_stack/psf_stack.py -f "+FIELD+" -b {wildcards.bm} -c "+CUBE

#  NEEDS TO BE IN A CONTAINER!!
checkpoint f:
# rule f:
    input:
        FIELD+"/HI_B0{beam}_cube"+CUBE+"_image.fits",
        FIELD+"/HI_B0{beam}_cube"+CUBE+"_psf.fits",
        FIELD + "/HI_B0{beam}_cube"+CUBE+"_spline_clean_pb.fits",  # not essential to this step, but prepares rest by-hand
        "mos_"+FIELD+"/"+FIELD+"_HIcube"+CUBE+"_image_sofiaFS_mask_bin_dil{beam}_regrid.fits"
    output:
        FIELD + "/HI_B0{beam}_cube"+CUBE+".txt"
        # protected(FIELD+"/HI_B0{beam}_cube"+CUBE+"_spline.fits"),
        # protected(FIELD+"/HI_B0{beam}_cube"+CUBE+"_spline_clean_image.fits")
    singularity:
        "/mnt/software/containers/apercal.sif"
    threads:
        11
    shell:
        "touch {output}"