make_uwish_mosaic.py

import bs4
import requests
import os
import astropy.utils.data as aud
import glob
import montage
from astropy.io import fits
import shutil

# These are URLs corresponding to the region that overlaps l=49 to l=50
urls = """
http://astro.kent.ac.uk/uwish2/fields/w20110730_00166.html
http://astro.kent.ac.uk/uwish2/fields/w20100714_00710.html
http://astro.kent.ac.uk/uwish2/fields/w20100819_00361.html
http://astro.kent.ac.uk/uwish2/fields/w20110730_00225.html
http://astro.kent.ac.uk/uwish2/fields/w20110730_00274.html
http://astro.kent.ac.uk/uwish2/fields/w20100817_00443.html
http://astro.kent.ac.uk/uwish2/fields/w20100816_00405.html
http://astro.kent.ac.uk/uwish2/fields/w20110730_00323.html
http://astro.kent.ac.uk/uwish2/fields/w20110731_00116.html
"""

# H2: http://surveys.roe.ac.uk/wsa/cgi-bin/fits_download.cgi?file=/disk40/wsa/ingest/fits/20110731_v1/w20110731_00116_sf_st.fit&MFID=4690635&rID=561

def get_link(pg):
    """
    The "main" URL has frames; this digs into the frames and finds the URL with
    the actual FITS links
    """
    soup = bs4.BeautifulSoup(pg)
    frames = soup.findAll('frame')
    for f in frames:
        if f.attrs['name'] == 'links':
            return f.attrs['src']

    raise ValueError("No match.")

def get_right_url(url):
    """
    Take the base URL and append the frame's URL to it
    """
    R = requests.get(url)
    rellink = get_link(R.content)
    left = os.path.split(url)[0]
    return os.path.join(left,rellink)

def get_target_urls(url, type='H2-K'):
    """
    Acquire the URL of the FITS files
    """
    right_url = get_right_url(url)

    R = requests.get(right_url)

    soup = bs4.BeautifulSoup(R.content)
    if type == 'H2-K':
        # Get the links for the H2 - K "difference" images
        links = [x.attrs['href'] for x in soup.findAll('a')
                 if '_diff_' in x.attrs['href']]
    elif type == 'H2':
        table = soup.find('table')
        rows = table.findAll('tr')
        links = []
        for tr in rows:
            cols = tr.findAll('td')
            # Only get the first column, which is H2
            # cols[1]=J, cols[2]=H, cols[3]=K
            if cols[0].a:
                links.append(cols[0].a.attrs['href'])
    else:
        raise ValueError("Invalid type %s" % type)

    return links

def download_files(url, timeout=20, **kwargs):
    """
    Wrapper: Download all the FITS files at a given URL
    """
    links = get_target_urls(url, **kwargs)

    results = []
    for U in links:
        try:
            x = aud.download_file(U, cache=True, show_progress=True, timeout=timeout)
            filename = os.path.split(U)[1].split("&")[0]
            results.append((filename,x))
        except Exception as ex:
            print ex
            if isinstance(ex,KeyboardInterrupt):
                raise ex

    return results

def redownload(type):
    """
    Actually download all files of a given type, AND hardlink them from the
    astropy cache directory to their local filenames
    """
    for url in urls.split():
        results = download_files(url, type=type)
        for fn, savename in results:
            try:
                os.link(savename,fn)
            except Exception as ex:
                print fn,":",ex
                if isinstance(ex,KeyboardInterrupt):
                    raise ex

def extract_image_hdus(clobber=False, skip=False):
    """
    Extract image HDUs from all FIT files in pwd
    """
    files = glob.glob("w*.fit")
    for fn in files:
        f = fits.open(fn)
        print f
        for ii,hdu in enumerate(f):
            print hdu
            if hasattr(hdu.data,'shape') and hdu.data.ndim == 2:
                outfn = 'hdu%i_%s' % (ii,fn)
                if os.path.exists(outfn) and skip:
                    continue
                hdu.writeto(outfn, clobber=clobber)
        

def mosaic_h2mK(redl=False):
    """
    Wrapper: Make an H2 - K mosaic
    """
    if redl:
        redownload(type='H2-K')
    files = glob.glob("h2mk/*diff*fit*")
    montage.wrapper(files, outfile='UWISH2_minusK_W51_mosaic.fits', header='header.hdr', background_match=True)

def mosaic_h2(redl=False):
    """
    Wrapper: Make an H2 mosaic
    """
    if not os.path.exists('h2'):
        os.mkdir('h2')
    os.chdir('h2')
    if redl:
        redownload(type='H2')

    # extract the USEFUL headers...
    extract_image_hdus()
    if not os.path.exists('raw'):
        os.mkdir('raw')
    wfiles = glob.glob("w*fit")
    for fn in wfiles:
        shutil.move(fn, 'raw/')

    files = glob.glob("*.fit")
    for fn in files:
        if 'hdu' not in fn:
            raise ValueError("FIT file without hdu prefix.")

    if not os.path.exists('header.hdr'):
        os.link('../header.hdr','header.hdr')

    montage.wrapper(files, outfile='UWISH2_W51_mosaic.fits', header='header.hdr', background_match=True)
    os.chdir('../')