montage-workflow.py

#!/usr/bin/env python3

'''
Pegasus wf generator for the Montage toolkit. The generated
workflow will support multiple bands and colors to produce
a color image.

#  Copyright 2020 University Of Southern California
#
#  Licensed under the Apache License, Version 2.0 (the 'License');
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing,
#  software distributed under the License is distributed on an 'AS IS' BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

'''

import os
import argparse
import re
import subprocess
import sys

#Insert this directory in our search path
os.sys.path.insert(0, os.getcwd())

from astropy.io import ascii
from Pegasus.api import *

common_files = {}


def which(file):
    for path in os.environ['PATH'].split(os.pathsep):
        if os.path.exists(os.path.join(path, file)):
            return os.path.join(path, file)
    return None


def repl_base_dir():
    '''
    In general, this is os.getcwd(), but for containers we might have to
    override the replica base dir 
    '''
    pass


def build_transformation_catalog(tc_target, wf):
    '''
    Some transformations in Montage uses multiple executables
    '''

    tc = TransformationCatalog()
    trans = {}

    exes = {}
    full_path = which('mProject')
    if full_path is None:
        raise RuntimeError('mProject is not in the $PATH')
    base_dir = os.path.dirname(full_path)
   
    container = None
    if tc_target == 'container':
        container = Container('montage',
            Container.SINGULARITY,
            'https://data.isi.edu/montage/images/montage-workflow-v3.sif'
            ).add_env(MONTAGE_HOME='/opt/Montage')
        tc.add_containers(container)

    for fname in os.listdir(base_dir):
        transformation = None
        if fname[0] == '.':
            continue

        if tc_target == 'regular':
            transformation = Transformation(fname, 
                                            site='local',
                                            pfn=os.path.join(base_dir, fname), 
                                            is_stageable=True)
            transformation.add_env(PATH='/usr/bin:/bin:.')
        else:
            # container
            transformation = Transformation(fname,
                                            site='insidecontainer',
                                            pfn=os.path.join(base_dir, fname),
                                            container=container,
                                            is_stageable=False)

        # memory requirement
        transformation.add_profiles(Namespace.CONDOR, 'request_memory', '1 GB')

        # some transformations can be clustered for effiency
        if fname in ['gmProject', 'mDiff', 'mDiffFit', 'mBackground']:
            transformation.add_profiles(Namespace.PEGASUS, 'clusters.size', '3')

        # keep a handle to added ones, for use later
        trans[fname] = transformation
            
        tc.add_transformations(transformation)

    # some Montage tools depend on other tools
    trans['mDiffFit'].add_requirement(trans['mDiff'])
    trans['mDiffFit'].add_requirement(trans['mFitplane'])

    wf.add_transformation_catalog(tc)


def generate_region_hdr(wf, rc, center, degrees):

    global common_files

    (crval1, crval2) = center.split()
    crval1 = float(crval1)
    crval2 = float(crval2)

    cdelt = 0.000277778
    naxis = int((float(degrees) / cdelt) + 0.5)
    crpix = (naxis + 1) / 2.0

    f = open('data/region.hdr', 'w')
    f.write('SIMPLE  = T\n')
    f.write('BITPIX  = -64\n')
    f.write('NAXIS   = 2\n')
    f.write('NAXIS1  = %d\n' %(naxis))
    f.write('NAXIS2  = %d\n' %(naxis))
    f.write('CTYPE1  = \'RA---TAN\'\n')
    f.write('CTYPE2  = \'DEC--TAN\'\n')
    f.write('CRVAL1  = %.6f\n' %(crval1))
    f.write('CRVAL2  = %.6f\n' %(crval2))
    f.write('CRPIX1  = %.6f\n' %(crpix))
    f.write('CRPIX2  = %.6f\n' %(crpix))
    f.write('CDELT1  = %.9f\n' %(-cdelt))
    f.write('CDELT2  = %.9f\n' %(cdelt))
    f.write('CROTA2  = %.6f\n' %(0.0))
    f.write('EQUINOX = %d\n' %(2000))
    f.write('END\n')
    f.close()

    common_files['region.hdr'] = File('region.hdr')
    rc.add_replica('local', common_files['region.hdr'], 'file://' + os.getcwd() + '/data/region.hdr')

    # we also need an oversized region which will be used in the first part of the 
    # workflow to get the background correction correct
    f = open('data/region-oversized.hdr', 'w')
    f.write('SIMPLE  = T\n')
    f.write('BITPIX  = -64\n')
    f.write('NAXIS   = 2\n')
    f.write('NAXIS1  = %d\n' %(naxis + 3000))
    f.write('NAXIS2  = %d\n' %(naxis + 3000))
    f.write('CTYPE1  = \'RA---TAN\'\n')
    f.write('CTYPE2  = \'DEC--TAN\'\n')
    f.write('CRVAL1  = %.6f\n' %(crval1))
    f.write('CRVAL2  = %.6f\n' %(crval2))
    f.write('CRPIX1  = %.6f\n' %(crpix + 1500))
    f.write('CRPIX2  = %.6f\n' %(crpix + 1500))
    f.write('CDELT1  = %.9f\n' %(-cdelt))
    f.write('CDELT2  = %.9f\n' %(cdelt))
    f.write('CROTA2  = %.6f\n' %(0.0))
    f.write('EQUINOX = %d\n' %(2000))
    f.write('END\n')
    f.close()

    common_files['region-oversized.hdr'] = File('region-oversized.hdr')
    rc.add_replica('local', common_files['region-oversized.hdr'], 'file://' + os.getcwd() + '/data/region-oversized.hdr')


def add_band(wf, rc, band_id, center, degrees, survey, band, color):

    band_id = str(band_id)

    print('\nAdding band %s (%s %s -> %s)' %(band_id, survey, band, color))

    # data find - go a little bit outside the box - see mExec implentation
    degrees_datafind = str(float(degrees) * 1.42)
    cmd = 'mArchiveList %s %s \'%s\' %s %s data/%s-images.tbl' \
          %(survey, band, center, degrees_datafind, degrees_datafind, band_id)
    print('Running sub command: ' + cmd)
    if subprocess.call(cmd, shell=True) != 0:
        print('Command failed!')
        sys.exit(1)
    rc.add_replica('local', '%s-images.tbl' %(band_id), 'file://' + os.getcwd() + '/data/%s-images.tbl' %(band_id))

    # image tables
    raw_tbl = File('%s-raw.tbl' %(band_id))
    rc.add_replica('local', raw_tbl, 'file://' + os.getcwd() + '/data/' + raw_tbl.lfn)
    projected_tbl = File('%s-projected.tbl' %(band_id))
    rc.add_replica('local', projected_tbl, 'file://' + os.getcwd() + '/data/' + projected_tbl.lfn)
    corrected_tbl = File('%s-corrected.tbl' %(band_id))
    rc.add_replica('local', corrected_tbl, 'file://' + os.getcwd() + '/data/' + corrected_tbl.lfn)
    cmd = 'cd data && mDAGTbls %s-images.tbl region-oversized.hdr %s %s %s' \
          %(band_id, raw_tbl.lfn, projected_tbl.lfn, corrected_tbl.lfn)
    print('Running sub command: ' + cmd)
    if subprocess.call(cmd, shell=True) != 0:
        print('Command failed!')
        sys.exit(1)
    
    # diff table
    cmd = 'cd data && mOverlaps %s-raw.tbl %s-diffs.tbl' \
          %(band_id, band_id)
    print('Running sub command: ' + cmd)
    if subprocess.call(cmd, shell=True) != 0:
        print('Command failed!')
        sys.exit(1)

    # statfile table
    t = ascii.read('data/%s-diffs.tbl' %(band_id))
    # make sure we have a wide enough column
    t['stat'] = '                                                                  '
    for row in t:
        base_name = re.sub('(diff\.|\.fits.*)', '', row['diff'])
        row['stat'] = '%s-fit.%s.txt' %(band_id, base_name)
    ascii.write(t, 'data/%s-stat.tbl' %(band_id), format='ipac')
    rc.add_replica('local', '%s-stat.tbl' %(band_id), 'file://' + os.getcwd() + '/data/%s-stat.tbl' %(band_id))

    # for all the input images in this band, and them to the rc, and
    # add reproject tasks
    data = ascii.read('data/%s-images.tbl' %(band_id))  
    for row in data:
        
        base_name = re.sub('\.fits.*', '', row['file'])

        # add an entry to the replica catalog
        rc.add_replica('ipac', base_name + '.fits', row['URL'])

        # projection job
        j = Job('mProject')
        in_fits = File(base_name + '.fits')
        projected_fits = File('p' + base_name + '.fits')
        area_fits = File('p' + base_name + '_area.fits')
        j.add_inputs(common_files['region-oversized.hdr'], in_fits)
        j.add_outputs(projected_fits, area_fits, stage_out=False)
        j.add_args('-X', in_fits, projected_fits, common_files['region-oversized.hdr'])
        wf.add_jobs(j)

    fit_txts = []
    data = ascii.read('data/%s-diffs.tbl' %(band_id))
    for row in data:
        
        base_name = re.sub('(diff\.|\.fits.*)', '', row['diff'])

        # mDiffFit job
        j = Job('mDiffFit')
        plus = File('p' + row['plus'])
        plus_area = File(re.sub('\.fits', '_area.fits', plus.lfn))
        minus = File('p' + row['minus'])
        minus_area = File(re.sub('\.fits', '_area.fits', minus.lfn))
        fit_txt = File('%s-fit.%s.txt' %(band_id, base_name))
        diff_fits = File('%s-diff.%s.fits' %(band_id, base_name))
        j.add_inputs(plus, plus_area, minus, minus_area, common_files['region-oversized.hdr'])
        j.add_outputs(fit_txt, stage_out=False)
        j.add_args('-d', '-s', fit_txt, plus, minus, diff_fits, common_files['region-oversized.hdr'])
        wf.add_jobs(j)
        fit_txts.append(fit_txt)

    # mConcatFit
    j = Job('mConcatFit')
    stat_tbl = File('%s-stat.tbl' %(band_id))
    j.add_inputs(stat_tbl)
    for fit_txt in fit_txts:
        j.add_inputs(fit_txt)
    fits_tbl = File('%s-fits.tbl' %(band_id))
    j.add_outputs(fits_tbl, stage_out=False)
    j.add_args(stat_tbl, fits_tbl, '.')
    wf.add_jobs(j)

    # mBgModel
    j = Job('mBgModel')
    images_tbl = File('%s-images.tbl' %(band_id))
    corrections_tbl = File('%s-corrections.tbl' %(band_id))
    j.add_inputs(images_tbl, fits_tbl)
    j.add_outputs(corrections_tbl, stage_out=False)
    j.add_args('-i', '100000', images_tbl, fits_tbl, corrections_tbl)
    wf.add_jobs(j)

    # mBackground
    data = ascii.read('data/%s-raw.tbl' %(band_id))  
    for row in data:
        base_name = re.sub('(diff\.|\.fits.*)', '', row['file'])

        # mBackground job
        j = Job('mBackground')
        projected_fits = File('p' + base_name + '.fits')
        projected_area = File('p' + base_name + '_area.fits')
        corrected_fits = File('c' + base_name + '.fits')
        corrected_area = File('c' + base_name + '_area.fits')
        j.add_inputs(projected_fits, projected_area, projected_tbl, corrections_tbl)
        j.add_outputs(corrected_fits, corrected_area, stage_out=False)
        j.add_args('-t', projected_fits, corrected_fits, projected_tbl, corrections_tbl)
        wf.add_jobs(j)

    # mImgtbl - we need an updated corrected images table because the pixel offsets and sizes need
    # to be exactly right and the original is only an approximation
    j = Job('mImgtbl')
    updated_corrected_tbl = File('%s-updated-corrected.tbl' %(band_id))
    j.add_inputs(corrected_tbl)
    j.add_outputs(updated_corrected_tbl, stage_out=False)
    j.add_args('.', '-t', corrected_tbl, updated_corrected_tbl)
    data = ascii.read('data/%s-corrected.tbl' %(band_id))  
    for row in data:
        base_name = re.sub('(diff\.|\.fits.*)', '', row['file'])
        projected_fits = File(base_name + '.fits')
        j.add_inputs(projected_fits)
    wf.add_jobs(j)

    # mAdd
    j = Job('mAdd')
    mosaic_fits = File('%s-mosaic.fits' %(band_id))
    mosaic_area = File('%s-mosaic_area.fits' %(band_id))
    j.add_inputs(updated_corrected_tbl, common_files['region.hdr'])
    j.add_outputs(mosaic_fits, mosaic_area, stage_out=True)
    j.add_args('-e', updated_corrected_tbl, common_files['region.hdr'], mosaic_fits)
    data = ascii.read('data/%s-corrected.tbl' %(band_id))  
    for row in data:
        base_name = re.sub('(diff\.|\.fits.*)', '', row['file'])
        corrected_fits = File(base_name + '.fits')
        corrected_area = File(base_name + '_area.fits')
        j.add_inputs(corrected_fits, corrected_area)
    wf.add_jobs(j)

    # mViewer - Make the JPEG for this channel
    j = Job('mViewer')
    mosaic_png = File('%s-mosaic.png' %(band_id))
    j.add_inputs(mosaic_fits)
    j.add_outputs(mosaic_png, stage_out=True)
    j.add_args('-ct', '1', '-gray', mosaic_fits, '-1s', 'max', 'gaussian', \
               '-png', mosaic_png)
    wf.add_jobs(j)


def color_png(wf, rc, red_id, green_id, blue_id):

    red_id = str(red_id)
    green_id = str(green_id)
    blue_id = str(blue_id)

    # mJPEG - Make the JPEG for this channel
    j = Job('mViewer')
    mosaic_png = File('mosaic-color.png')
    red_fits = File('%s-mosaic.fits' %(red_id))
    green_fits = File('%s-mosaic.fits' %(green_id))
    blue_fits = File('%s-mosaic.fits' %(blue_id))
    j.add_inputs(red_fits, green_fits, blue_fits)
    j.add_outputs(mosaic_png, stage_out=True)
    j.add_args( \
            '-red', red_fits, '-0.5s', 'max', 'gaussian-log', \
            '-green', green_fits, '-0.5s', 'max', 'gaussian-log', \
            '-blue', blue_fits, '-0.5s', 'max', 'gaussian-log', \
            '-png', mosaic_png)
    wf.add_jobs(j)


def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('--work-dir', action = 'store', dest = 'work_dir',
                        help = 'Work directory to chdir to')
    parser.add_argument('--center', action = 'store', dest = 'center',
                        help = 'Center of the output, for example M17 or 56.5 23.75')
    parser.add_argument('--degrees', action = 'store', dest = 'degrees',
                        help = 'Number of degrees of side of the output')
    parser.add_argument('--band', action = 'append', dest = 'bands',
                        help = 'Band definition. Example: dss:DSS2B:red')
    parser.add_argument('--tc-target', action = 'store', dest = 'tc_target',
                        help = 'Transformation catalog: regular or container')
    args = parser.parse_args()

    if args.work_dir:
        os.chdir(args.work_dir)

    if args.tc_target is None:
        args.tc_target = 'regular'

    if os.path.exists('data'):
        print('data/ directory already exists')
        sys.exit(1)
    os.mkdir('data')

    wf = Workflow('montage', infer_dependencies=True)
    rc = ReplicaCatalog()

    # FIXME
    # email notificiations for when the state of the workflow changes
    #share_dir = subprocess.Popen('pegasus-config --sh-dump | grep ^PEGASUS_SHARE_DIR= | sed -e 's/.*=//' -e 's/\'//g'',
    #                             shell=True,
    #                             stdout=subprocess.PIPE).communicate()[0]
    #share_dir = share_dir.strip()
    #wf.invoke('start', share_dir + '/notification/email')
    #wf.invoke('on_error', share_dir + '/notification/email')
    #wf.invoke('on_success', share_dir + '/notification/email --report=pegasus-statistics')

    build_transformation_catalog(args.tc_target, wf)

    # region.hdr is the template for the ouput area
    generate_region_hdr(wf, rc, args.center, args.degrees)

    band_id = 0
    color_band = {}
    for band_def in args.bands:
        band_id += 1
        (survey, band, color) = band_def.split(':')
        add_band(wf, rc, band_id, args.center, args.degrees, survey, band, color)
        color_band[color] = band_id

    # if we have 3 bands in red, blue, green, try to create a color jpeg
    if 'red' in color_band and 'green' in color_band and 'blue' in color_band:
        color_png(wf, rc, color_band['red'], color_band['green'], color_band['blue'])

    # write out the workflow and catalogs
    wf.add_replica_catalog(rc)
    wf.write('data/montage-workflow.yml')


if __name__ == '__main__':
    main()