plot_illustrations.py

#!/usr/bin/env python

from glob import glob
import os
import numpy
import pyfits
import matplotlib
from matplotlib import pyplot
from matplotlib import cm
from numpy import log, log10, exp, power, pi
from numpy.polynomial.chebyshev import Chebyshev
from scipy.stats import scoreatpercentile
from scipy.special import gamma, gammaincinv
from scipy.optimize import fmin
from scipy.integrate import quad
from scipy.interpolate import interp1d
from RGBImage import *

matplotlib.rcParams.update({'font.size': 16,
                            'font.family': 'serif',
                            'font.serif': 'times',
                            'text.usetex': True})

labelsize = 16

rescolmap = matplotlib.colors.LinearSegmentedColormap.from_list('rescol', ('blue', 'black', 'white', 'red'), N=256, gamma=1.0)

bands = ['u', 'g', 'r', 'i', 'z', 'Y', 'J', 'H', 'K']

w = numpy.array([3543,4770,6231,7625,9134,10305,12483,16313,22010], numpy.float)

#zp = numpy.array([16.75,15.957,15.0,14.563,14.259,14.162,13.955,13.636,13.525])
#zpscale = 10**(-0.4*(zp-15.0))

zp = numpy.array([29.0]*9)

xlim = (2000, 23000)

sim_std = {'MAG': 1.0 + numpy.array([16.935,15.964,15.0,14.562,14.267,14.183,13.992,13.672,13.547])}
sim_A_disk = {'MAG': 1.0 + numpy.array([17.687,16.717,15.753,15.315,15.019,14.936,14.745,14.425,14.299]),
              'Re': numpy.array([18.0]*9), 'n': numpy.array([1.0]*9),
              'AR': numpy.array([0.4]*9), 'PA': numpy.array([45.0]*9)}
sim_A_bulge = {'MAG': 1.0 + numpy.array([17.687,16.717,15.753,15.315,15.019,14.936,14.745,14.425,14.299]),
              'Re': numpy.array([6.0]*9), 'n': numpy.array([4.0]*9),
              'AR': numpy.array([0.8]*9), 'PA': numpy.array([45.0]*9)}
sim_D_disk = {'MAG': 1.0 + numpy.array([17.328,16.509,15.753,15.374,15.112,15.056,14.882,14.597,14.5]),
              'Re': numpy.array([18.0]*9), 'n': numpy.array([1.0]*9),
              'AR': numpy.array([0.4]*9), 'PA': numpy.array([45.0]*9)}
sim_D_bulge = {'MAG': 1.0 + numpy.array([18.229,16.974,15.753,15.258,14.934,14.827,14.623,14.276,14.13]),
               'Re': numpy.array([6.0]*9), 'n': numpy.array([4.0]*9),
               'AR': numpy.array([0.8]*9), 'PA': numpy.array([45.0]*9)}
sim_E_disk = {'MAG': 1.0 + numpy.array([16.999,16.127,15.753,15.529,15.389,15.41,15.384,15.192,15.281]),
              'Re': numpy.array([18.0]*9), 'n': numpy.array([1.0]*9),
              'AR': numpy.array([0.8]*9), 'PA': numpy.array([45.0]*9)}
sim_E_bulge = {'MAG': 1.0 + numpy.array([18.546,17.519,15.753,15.084,14.764,14.623,14.433,14.02,13.78]),
               'Re': numpy.array([6.0]*9), 'n': numpy.array([4.0]*9),
               'AR': numpy.array([0.9]*9), 'PA': numpy.array([45.0]*9)}

marker = ['o', '^', 's', 'D', 'x', '+', '*']
linestyle = [':', '-', '-.', (0, '.-.')]

ylim_std = {'MAG': (19.05, 13.45), 'Re': (6.05, 28.95), 'n': (2.05, 5.95),
            'AR': (0.41, 0.79), 'PA': (35.05, 64.95)}

ylim_disk = {'MAG': (20.05, 14.45), 'Re': (13.05, 22.95), 'n': (0.05, 2.95),
             'AR': (0.21, 0.89), 'PA': (35.05, 64.95)}

ylim_bulge = {'MAG': (20.05, 14.45), 'Re': (1.05, 12.95), 'n': (2.05, 6.95),
              'AR': (0.61, 1.09), 'PA': (0.05, 89.95)}

#varlist_std = ('MAG', 'Re', 'n', 'AR', 'PA')
varlist_std = ('MAG', 'Re', 'n')

labels = {'MAG': '$m$', 'Re': '$R_e$', 'n': '$n$', 'AR': '$b/a$', 'PA': '$\\theta$'}

def ugrizYJHK_cheb(wl):
    y = numpy.array([100.000,100.058,100.217,100.541,100.924,101.307,101.631,101.790,101.848])
    fn = interp1d(w, y, 'linear', bounds_error=False)
    return fn(wl)

wlfuncs = {'A1c': numpy.log10, 'Ah1c': numpy.log10, 'A1e': ugrizYJHK_cheb}

def poster_plots():
    plot(('D1', 'D4'), 1, 'D1D4-1', 'True', ylim=ylim_bulge, sim=sim_D_bulge, varlist=('MAG', 'Re', 'n')) # and D6


def plot_all():
    plot_standard()
    plot_nonparam()

def plot_standard():
    plot(('A2', 'A1', 'A3'), 1, '01', 'True', varlist=('MAGNOSUB', 'MAG', 'Re', 'n'))
    plot(('Ah2', 'Ah1', 'Ah3'), 1, '02', 'True')
    plot(('Bh2', 'Bh1', 'Bh3'), 1, '03', 'True')
    plot(('A1e', 'A1c', 'A1'), 1, '04', 'True')  # add additional wavelength scale
    plot(('Ah1c', 'Ah1'), 1, '04h', 'True')  # add additional wavelength scale
    plot(('A1a', 'A1', 'A1b'), 1, '05', 'True', varlist=('MAG',))
    plot(('Ah1a', 'Ah1', 'Ah1b'), 1, '05h', 'True', varlist=('MAG',))
    # plot(('A1', 'A1c', 'A1d'), 1, '06', 'True')
    # plot(('Ah1', 'Ah1c', 'Ah1d'), 1, '06h', 'True')
    # illustration 7 requires a different kind of plot
    plot(('D2', 'D1', 'D3'), 1, '08', 'True', varlist=('MAG', 'Re', 'n', 'AR', 'PA'))
    plot(('D1a', 'D1'), 1, '08a', 'True', varlist=('MAG', 'Re', 'n', 'AR', 'PA'))
    plot(('A5', 'A4', 'A6'), 1, '09-1', 'True', ylim=ylim_bulge, sim=sim_A_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('A5', 'A4', 'A6'), 2, '09-2', 'True', ylim=ylim_disk, sim=sim_A_disk, varlist=('MAG', 'Re'))
    plot(('D5', 'D4', 'D6'), 1, '10-1', 'True', ylim=ylim_bulge, sim=sim_D_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('D5', 'D4', 'D6'), 2, '10-2', 'True', ylim=ylim_disk, sim=sim_D_disk, varlist=('MAG', 'Re'))
    plot(('Dh5', 'Dh4', 'Dh6'), 1, '10h-1', 'True', ylim=ylim_bulge, sim=sim_D_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('Dh5', 'Dh4', 'Dh6'), 2, '10h-2', 'True', ylim=ylim_disk, sim=sim_D_disk, varlist=('MAG', 'Re'))
    plot(('E5', 'E4', 'E6'), 1, '11-1', 'True', ylim=ylim_bulge, sim=sim_E_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('E5', 'E4', 'E6'), 2, '11-2', 'True', ylim=ylim_disk, sim=sim_E_disk, varlist=('MAG', 'Re'))
    plot(('Eh5', 'Eh4', 'Eh6'), 1, '11h-1', 'True', ylim=ylim_bulge, sim=sim_E_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('Eh5', 'Eh4', 'Eh6'), 2, '11h-2', 'True', ylim=ylim_disk, sim=sim_E_disk, varlist=('MAG', 'Re'))
    plot(('Dc5', 'Dc4', 'Dc6'), 1, '12-1', 'True', ylim=ylim_bulge, sim=sim_D_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('Dc5', 'Dc4', 'Dc6'), 2, '12-2', 'True', ylim=ylim_disk, sim=sim_D_disk, varlist=('MAG', 'Re'))

def plot_nonparam():
    plot(('NA1n', 'NA1'), 1, 'N01', 'True')
    plot(('NA2n', 'NA2'), 1, 'N02', 'True')
    plot(('NA4n', 'NA4'), 1, 'N03-1', 'True', ylim=ylim_bulge, sim=sim_A_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('NA4n', 'NA4'), 2, 'N03-2', 'True', ylim=ylim_disk, sim=sim_A_disk, varlist=('MAG', 'Re'))
    plot(('NB4n', 'NB4'), 1, 'N04-1', 'True', ylim=ylim_bulge, sim=sim_A_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('NB4n', 'NB4'), 2, 'N04-2', 'True', ylim=ylim_disk, sim=sim_A_disk, varlist=('MAG', 'Re'))
    plot(('NC4n', 'NC4', 'NC4m'), 1, 'N05-1', 'True', ylim=ylim_bulge, sim=sim_A_bulge, varlist=('MAG', 'Re', 'n'))
    plot(('NC4n', 'NC4', 'NC4m'), 2, 'N05-2', 'True', ylim=ylim_disk, sim=sim_A_disk, varlist=('MAG', 'Re'))


def plot(id=('A2', 'A1'), compno=1, name='0', show_func=False,
         varlist=varlist_std, ylim=ylim_std, sim=sim_std, submag=True,
         legends=None):
    print name, ':', id
    res = [fit_results(i) for i in id]
    if show_func:
        func = [fit_func(i) for i in id]
    else:
        func = None
    nvar = len(varlist)
    fig = pyplot.figure(figsize=(5, 15))
    fig.subplots_adjust(bottom=0.1, top=0.9, left=0.2, right=0.95, hspace=0.075)
    for i, v in enumerate(varlist):
        if v == 'MAGNOSUB':
            v = 'MAG'
            vsubmag = False
        else:
            vsubmag = submag
        ax = make_bands_plot(fig, (5, 1, i+1), labels[v], i==0, i==nvar-1)
        sub = norm = None
        if v in sim.keys():
            sv = sim[v]
            if vsubmag and v == 'MAG':
                sub = interp1d(w, sim[v], 'cubic', bounds_error=False, fill_value='extrapolate')
            if sub is not None:
                sv = sim[v] - sub(w)
                ax.set_ylabel('$\Delta ' + labels[v][1:])
            pyplot.plot(w, sv, '-k', alpha=0.75)
        plotres(res, id, 'COMP%i_%s'%(compno, v), func, sub=sub, norm=norm, legends=None)
        if v in sim.keys():
            pyplot.plot(w, sv, 'xk', markersize=10.0, alpha=0.75)
        if sub is None:
            pyplot.ylim(ylim[v])
        else:
            pyplot.ylim(numpy.subtract(ylim[v], numpy.mean(ylim[v], 0))/3.0)
            #pyplot.legend(loc='lower right', numpoints=1, prop={'size': 16})
        if compno == 1 and ('-' not in name):
            if i == nvar-1:
                pyplot.legend(loc='upper left', numpoints=1, prop={'size': 16},
                              bbox_to_anchor=(0., -.35, 1., .1),
                              ncol=4, mode="expand", borderaxespad=0.)
        elif compno == 1:
            if i == nvar-1:
                pyplot.text(0.5, -0.25, 'bulge',
                            horizontalalignment='center',
                            verticalalignment='top',
                            transform = ax.transAxes)
        elif compno == 2:
            if i == nvar-1:
                pyplot.text(0.5, -0.25, 'disc',
                            horizontalalignment='center',
                            verticalalignment='top',
                            transform = ax.transAxes)
            if i == 0:
                pyplot.legend(loc='lower left', numpoints=1, prop={'size': 16},
                              bbox_to_anchor=(0., 1.45, 1., .1),
                              ncol=4, mode="expand", borderaxespad=0.)
        if i == nvar-1:
            # invisible text to ensure plots line up after cropping
            pyplot.text(0.5, -0.30, '\_', color='white',
                        horizontalalignment='center',
                        verticalalignment='top',
                        transform = ax.transAxes)
    fig.savefig('plots/illustration_%s.pdf'%name)
    pyplot.close('all')
    if compno==1:
        plotimg(id, name)
        plotcolimg(id, name)
        plotprof(id, name)
        plotcolprof(id, name)
        npid = [j for j in id if j[0] == 'N' and j[-1] in 'nm']
        if len(npid) > 0:
            plotnonparamcolimg(npid, name)


def plotimg(id, name='0', asinh=True):
    cmap_img = pyplot.cm.gray
    cmap_res = rescolmap
    norm_res = None
    nbands = len(bands)
    nid = len(id)
    fig = pyplot.figure(figsize=(15.0/nbands * (1+nid*2), 15))
    fig.subplots_adjust(bottom=0.05, top=0.95, left=0.05, right=0.95, hspace=0.0, wspace=0.0)
    for i, iid in enumerate(id):
        img = fit_images(iid)
        if asinh:
            for j, jimg in enumerate(img):
                for iimg in jimg:
                    if j != 2:
                        iimg[iimg <= 0] = 0.0
                    else:
                        iimg /= 10.0
                    iimg *= 0.1
                    iimg[:] = numpy.arcsinh(iimg)
        if i == 0:
            vmin = []
            vmax = []
            vrange = []
            for ib, b in enumerate(bands):
                ax = fig.add_subplot(nbands, 1+2*len(id), 1+ib*(1+nid*2)+i*2)
                if ib==nbands-1:
                    ax.set_xlabel('image', fontsize=labelsize)
                ticksoff(ax)
                vmin.append(scoreatpercentile(img[0][ib].ravel(), 50))
                vmax.append(scoreatpercentile(img[0][ib].ravel(), 99.9) * 1.1)
                vrange.append((scoreatpercentile(img[2][ib].ravel(), 99) - scoreatpercentile(img[2][ib].ravel(), 1)) * 1.5)
                pyplot.imshow(img[0][ib][::-1], cmap=cmap_img, vmin=vmin[ib], vmax=vmax[ib], interpolation='nearest')
                ax.set_ylabel('$%s$'%b, fontsize=labelsize)
        for ib, b in enumerate(bands):
            ax = fig.add_subplot(nbands, 1+2*len(id), 2+ib*(1+nid*2)+i*2)
            if ib==nbands-1:
                ax.set_xlabel('model %s'%iid, fontsize=labelsize)
            ticksoff(ax)
            pyplot.imshow(img[1][ib][::-1], cmap=cmap_img, vmin=vmin[ib], vmax=vmax[ib], interpolation='nearest')
        for ib, b in enumerate(bands):
            ax = fig.add_subplot(nbands, 1+2*len(id), 3+ib*(1+nid*2)+i*2)
            if ib==nbands-1:
                ax.set_xlabel('residual %s'%iid, fontsize=labelsize)
            ticksoff(ax)
            pyplot.imshow(img[2][ib][::-1], cmap=cmap_res, norm=norm_res, vmin=-vrange[ib], vmax=vrange[ib], interpolation='nearest')
    fig.savefig('plots/images_%s.pdf'%name, dpi=300)
    pyplot.close('all')


def plotcolimg(id, name='0', rgb='Hzg', desaturate=True, pedestal=0):
    nbands = len(bands)
    nid = len(id)
    beta = 2.5
    scales = numpy.array((0.04, 0.055, 0.2))
    # offsets not so necessary now have nice desaturation feature working
    offsets = numpy.array([75.0, 40.0, 8.0]) * 0.5
    fig = pyplot.figure(figsize=(15.0/nbands * (1+nid*2), 15))
    fig.subplots_adjust(bottom=0.05, top=0.95, left=0.05, right=0.95, hspace=0.0, wspace=0.0)
    for i, iid in enumerate(id):
        img = fit_images(iid, rgb)
        img[0] = [img[0][j] - offsets[j] for j in range(3)]
        img[1] = [img[1][j] - offsets[j] for j in range(3)]
        img[2] = [img[2][j] + scales[j]*2*offsets.mean() for j in range(3)]
        if i == 0:
            ax = fig.add_subplot(nbands, 1+2*nid, 1+i*2)
            ticksoff(ax)
            ax.set_xlabel('image', fontsize=labelsize)
            colimg = RGBImage(*img[0], scales=scales, beta=beta,
                              desaturate=desaturate, pedestal=pedestal).img
            pyplot.imshow(colimg, interpolation='nearest', origin='lower')
        ax = fig.add_subplot(nbands, 1+2*nid, 2+i*2)
        ticksoff(ax)
        ax.set_xlabel('model %s'%iid, fontsize=labelsize)
        colimg = RGBImage(*img[1], scales=scales, beta=beta,
                          desaturate=False, pedestal=pedestal).img
        pyplot.imshow(colimg, interpolation='nearest', origin='lower')
        ax = fig.add_subplot(nbands, 1+2*nid, 3+i*2)
        ticksoff(ax)
        ax.set_xlabel('residual %s'%iid, fontsize=labelsize)
        colimg = RGBImage(*img[2], scales=scales, beta=beta,
                          desaturate=desaturate).img
        pyplot.imshow(colimg, interpolation='nearest', origin='lower')
    fig.savefig('plots/colimages_%s.pdf'%name, dpi=300)
    pyplot.close('all')


def plotnonparamcolimg(id, name='0', rgb='Hzg', desaturate=True, pedestal=0):
    nbands = len(bands)
    nid = len(id)
    beta = 2.5
    scales = numpy.array((0.04, 0.055, 0.2))
    # offsets not so necessary now have nice desaturation feature working
    offsets = numpy.array([75.0, 40.0, 8.0]) * 0.5
    fig = pyplot.figure(figsize=(15.0/nbands * (1+nid*2), 15))
    fig.subplots_adjust(bottom=0.05, top=0.95, left=0.05, right=0.95, hspace=0.0, wspace=0.0)
    original_iid = None
    for i, iid in enumerate(id):
        # First row, results without nonparam
        if original_iid != iid[:-1]:
            original_iid = iid[:-1]
            img = fit_images(original_iid, rgb)
            img[0] = [img[0][j] - offsets[j] for j in range(3)]
            img[1] = [img[1][j] - offsets[j] for j in range(3)]
            img[2] = [img[2][j] + scales[j]*2*offsets.mean() for j in range(3)]
            if i == 0:
                ax = fig.add_subplot(nbands, 1+2*nid, 1+i*2)
                ticksoff(ax)
                ax.set_title('image', fontsize=labelsize)
                colimg = RGBImage(*img[0], scales=scales, beta=beta,
                                  desaturate=desaturate, pedestal=pedestal).img
                pyplot.imshow(colimg, interpolation='nearest', origin='lower')
            ax = fig.add_subplot(nbands, 1+2*nid, 2+i*2)
            ticksoff(ax)
            ax.set_title('model %s'%original_iid, fontsize=labelsize)
            colimg = RGBImage(*img[1], scales=scales, beta=beta,
                              desaturate=desaturate, pedestal=pedestal).img
            pyplot.imshow(colimg, interpolation='nearest', origin='lower')
            ax = fig.add_subplot(nbands, 1+2*nid, 3+i*2)
            ticksoff(ax)
            ax.set_title('residual %s'%original_iid, fontsize=labelsize)
            colimg = RGBImage(*img[2], scales=scales, beta=beta,
                              desaturate=desaturate).img
            pyplot.imshow(colimg, interpolation='nearest', origin='lower')
        # Second row, results with nonparam
        img = fit_images(iid, rgb)
        img[0] = [img[0][j] - offsets[j] for j in range(3)]
        img[1] = [img[1][j] - offsets[j] for j in range(3)]
        img[2] = [img[2][j] + scales[j]*2*offsets.mean() for j in range(3)]
        ax = fig.add_subplot(nbands, 1+2*nid, 1+2*nid+2+i*2)
        ticksoff(ax)
        colimg = RGBImage(*img[1], scales=scales, beta=beta,
                          desaturate=False, pedestal=pedestal).img
        pyplot.imshow(colimg, interpolation='nearest', origin='lower')
        ax = fig.add_subplot(nbands, 1+2*nid, 1+2*nid+3+i*2)
        ticksoff(ax)
        colimg = RGBImage(*img[2], scales=scales, beta=beta,
                          desaturate=desaturate).img
        pyplot.imshow(colimg, interpolation='nearest', origin='lower')
        # Third row, nonparam diagnostics
        nonparam = nonparam_images(iid, rgb)
        datasub = [img[0][j] - nonparam[j] for j in range(3)]
        nonparam = [nonparam[j] - offsets[j] for j in range(3)]
        datasub = [datasub[j] - offsets[j] for j in range(3)]
        ax = fig.add_subplot(nbands, 1+2*nid, 2+4*nid+2+i*2)
        ticksoff(ax)
        ax.set_xlabel('nonparam %s'%iid, fontsize=labelsize)
        colimg = RGBImage(*nonparam, scales=scales, beta=beta,
                          desaturate=desaturate).img
        pyplot.imshow(colimg, interpolation='nearest', origin='lower')
        ax = fig.add_subplot(nbands, 1+2*nid, 2+4*nid+3+i*2)
        ticksoff(ax)
        ax.set_xlabel('datasub %s'%iid, fontsize=labelsize)
        colimg = RGBImage(*datasub, scales=scales, beta=beta,
                          desaturate=desaturate).img
        pyplot.imshow(colimg, interpolation='nearest', origin='lower')
    fig.savefig('plots/nonparamcolimages_%s.pdf'%name, dpi=300)
    pyplot.close('all')


def ticksoff(ax):
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.set_xticks([])
    ax.set_yticks([])


def plotres(res, id, field, func=None, sub=None, norm=None, legends=None):
    nid = len(id)
    #mec = ['black', None] * (1+nid//2)
    #mfc = ['white', 'black'] * (1+nid//2)
    #color = ['grey', 'grey'] * (1+nid//2)
    mec_func = ['DeepSkyBlue' ,'DarkGreen', 'Orange']
    mfc_func = ['DeepSkyBlue', 'white', 'white']
    color = ['DeepSkyBlue', 'DarkGreen', 'Orange']
    mec_nofunc = ['MediumPurple', 'MediumSeaGreen', 'Pink']
    mfc_nofunc = ['MediumPurple', 'MediumSeaGreen', 'Pink']
    ymin, ymax = (1e99, -1e99)
    for i, iid in enumerate(id):
        if nid%2 == 0:
            x = w + 150 * [-1, 1, -2, 2][i]
        else:
            x = w + 300 * [-1, 0, 1, -2, 2][i]
        if func is not None and func[i] is not None:
            mec = mec_func
            mfc = mfc_func
            f = func[i][field]
            plotfunc(f, wlfunc=wlfuncs.get(iid), color=color[i],
                     sub=sub, norm=norm)
        else:
            mec = mec_nofunc
            mfc = mfc_nofunc
        r = res[i][field]
        rerr = res[i][field+'_ERR']
        if sub is not None:
            r = r - sub(w)
        if norm is not None:
            r = r / norm(w)
            rerr = rerr / norm(w)
        #alpha = 1 - i * 0.2
        alpha = 1
        if legends is not None:
            label = legends[i]
        else:
            label = iid
        pyplot.errorbar(x, r, rerr, color=mec[i],
                        marker=marker[i//2], mec=mec[i],
                        markerfacecolor=mfc[i], linestyle='',
                        label=label, alpha=alpha)
        ymin = min(ymin, (r-rerr).min())
        ymax = max(ymax, (r+rerr).max())
    yrange = ymax - ymin
    ymin -= 0.05 * yrange
    ymax += 0.05 * yrange
    pyplot.ylim(ymin, ymax)

def plotfunc(func, wlfunc=None, color='red', label='', sub=None, norm=None):
    #dx = (xlim[1] - xlim[0]) / 1000.0
    #x = numpy.arange(xlim[0], xlim[1]+dx/2.0, dx)
    dx = (w[-1] - w[0]) / 1000.0
    x = numpy.arange(w[0], w[-1]+dx/2.0, dx)
    if wlfunc is None:
        xfunc = x
    else:
        xfunc = wlfunc(x)
    y = func(xfunc)
    if sub is not None:
        y = y - sub(x)
    if norm is not None:
        y = y / norm(x)
    return pyplot.plot(x, y, ':', color=color, label=label, alpha=0.5)

def fit_results(f):
    fn = 'fits/%s/fit%s.fits'%(f,f)
    r = None
    if os.path.exists(fn):
        r = pyfits.getdata(fn, 'final_band')
    else:
        r = numpy.concatenate([pyfits.getdata('fits/%s/fit%s%s.fits'%(f,f, b), 'final_band')
                              for b in bands])
    return r


def fit_images(f, bands=bands, zoom=2.0,
               extensions=['input', 'model', 'residual']):
    """Get the images from the specified galfit(m) filename(s)

    `bands` should be a list of band id strings.
    If fn includes a format placeholder `{}`, this will be replaced by each
    band id in turn. Otherwise all bands are assumed to be in a single file.
    """
    fn = 'fits/%s/fit%s.fits'%(f,f)
    if not os.path.exists(fn):
        fn = fn.replace('.fits', '{}.fits')
    out = []
    # auto-discovery of bands for galfitm files
    multiband = True
    if '{}' not in fn:
        with pyfits.open(fn) as p:
            names = [x.name for x in p]
        ext = extensions[0].upper()
        if ext in names:
            multiband = False
        else:
            ext = ext + '_'
            found_bands = [x.name.replace(ext, '') for x in p
                           if x.name.startswith(ext)]
            if bands is None:
                bands = found_bands
    if not multiband:
        for ext in extensions:
            try:
                hdu = [pyfits.getdata(fn, ext)]
            except KeyError:
                hdu = None
            out.append(hdu)
    else:
        for ext in extensions:
            try:
                if '{}' in fn:
                    hdu = [pyfits.getdata(fn.format(b), ext)
                           for b in bands]
                else:
                    hdu = [pyfits.getdata(fn, '{}_{}'.format(ext, b))
                           for b in bands]
            except KeyError:
                hdu = None
            out.append(hdu)
    if zoom is not None:
        for ib, b in enumerate(bands):
            for i, xx in enumerate(out):
                if xx is not None:
                    shape = numpy.array(out[i][ib].shape)
                    crop = shape * (1 - 1 / zoom) / 2
                    crop = crop.round().astype(numpy.int)
                    crop = crop.clip(0, shape // 2 - 1)
                    icrop = [crop[0]] * 2
                    jcrop = [crop[1]] * 2
                    if xx is not None:
                        out[i][ib] = out[i][ib][icrop[0]:-icrop[1],
                                                jcrop[0]:-jcrop[1]]
    return out


def nonparam_images(f, bands=bands):
    return fit_images(f, bands, extensions=['nonparam'])[0]


def fit_func(f):
    fn = 'fits/%s/fit%s.fits'%(f,f)
    if os.path.exists(fn):
        r = {}
        d = pyfits.getdata('fits/%s/fit%s.fits'%(f,f), 'fit_info')[0]
        ref = d.field('refwlband')
        low = d.field('lowdwlband') + ref
        high = d.field('highdwlband') + ref
        d = pyfits.getdata('fits/%s/fit%s.fits'%(f,f), 'final_cheb')
        for n in d.names:
            r[n] = Chebyshev(d.field(n), (low, high))
    else:
        r = None
    return r


def make_bands_plot(fig, subplot=111, ylabel='', top=True, bottom=True):
    ax1 = fig.add_subplot(*subplot)
    ax2 = ax1.twiny()
    ax1.set_ylabel(ylabel)
    ax1.set_xlim(xlim)
    if top:
        ax2.set_xlabel('wavelength, \AA')
    else:
        ax2.set_xticklabels([])
    ax2.set_xlim(xlim)
    ax1.set_xticks(w)
    if bottom:
        ax1.set_xticklabels(['$'+i+'$' for i in bands])
    else:
        ax1.set_xticklabels([])
    pyplot.setp(ax1.get_xticklabels(), va='baseline')
    pyplot.setp(ax1.get_xaxis().get_major_ticks(), pad=20.)
    pyplot.setp(ax1.get_yaxis().get_major_ticks(), pad=8.)
    ax2.xaxis.labelpad = 12
    return ax1


class Sersic:
    # currently doesn't handle uncertainties
    def __init__(self, mag, re, n, ar=1.0, pa=0.0,
                 mag_err=None, re_err=None, n_err=None, ar_err=None, pa_err=None, xc_err=None, yc_err=None):
        self.mag = mag
        self.re = re
        self.n = n
        self.ar = ar
        self.pa = pa
        self.mag_err = mag_err
        self.re_err = re_err
        self.n_err = n_err
        self.ar_err = ar_err
        self.pa_err = pa_err
    def __call__(self, r):
        return self.mu_r(r)
    def mu_r(self, r):
        # Returns the surface brightess at specified major axis radius,
        # within annular ellipses corresponding to the shape of each component individualy
        # Taking, e.g. colours, this currently assumes major axes of components align
        # to be more generally correct need to account for AR, PA, XC, YC,
        # and either select specific vector, or properly compute azimuthal average
        mag = self.mag
        re = self.re
        n = self.n
        bn = self.bn()
        mue = mag + 5.0*log10(re) + 2.5*log10(2.0*pi*n*gamma(2.0*n)*exp(bn)/power(bn, 2.0*n))
        mu = mue + 2.5 * bn / log(10) * (power(r/re, 1.0/n) - 1.0)
        return mu
    def bn(self):
        return gammaincinv(2.0*self.n, 0.5)
    # These need testing
    def I_el(self, r_m, ar_m, pa_m=0):
        return quad(self.I_el_theta, 0, 2*pi, args=(r_m, ar_m, pa_m))[0] / (2*pi)
    def mu_el(self, r_m, ar_m, pa_m=0):
        return -2.5*numpy.log10(self.I_el(r_m, ar_m, pa_m))
    def mu_el_theta(self, theta, r_m, ar_m, pa_m=0):
        x = r_m * numpy.cos(theta - pa_m)
        y = ar_m * r_m * numpy.sin(theta - pa_m)
        r_c = numpy.sqrt(x**2 + self.ar**2 * y**2)
        return self.mu_r(r_c)
    def I_el_theta(self, theta, r_m, ar_m, pa_m=0):
        return 10**(-0.4*self.mu_el_theta(theta, r_m, ar_m, pa_m))


def plotprof(id=('A1', 'A2'), name='0'):
    print name, ':', id
    color = [cm.gist_rainbow(i) for i in numpy.linspace(1.0, 0.0, 9)]
    func, remax = make_funcs(id)
    fig = pyplot.figure(figsize=(5, 5))
    fig.subplots_adjust(bottom=0.15, top=0.95, left=0.15, right=0.95, hspace=0.0, wspace=0.0)
    rmax = remax*3.0001
    r = numpy.arange(rmax/10000.0, rmax, rmax/100.0)
    for i, iid in enumerate(id):
        for j in range(len(func[i])):
            for k, band in enumerate(bands):
                if k == 0:
                    label = "%s\_%i"%(iid, j)
                else:
                    label = ""
                pyplot.plot(r, func[i][j][k](r), linestyle=linestyle[i],
                            marker=None, color=color[k], label=label)
    pyplot.legend(loc='upper right', numpoints=1, prop={'size': 16})
    pyplot.xlabel('$r_{\mathrm{e}}$')
    pyplot.ylabel('$\mu$')
    #fig.gca().invert_yaxis()
    pyplot.xlim(0.0, rmax)
    pyplot.ylim(26, 16)
    fig.savefig('plots/profiles_%s.pdf'%name)


def plotcolprof(id=('A1', 'A2'), name='0'):
    # need to decide and implement consistent annuli in which to determine colour
    # would be nice to plot lines for input model too
    # normalised at remax and offset for display purposes
    print name, ':', id
    offset = 0.5
    color = [cm.gist_rainbow(i) for i in numpy.linspace(1.0, 0.0, 9)]
    func, remax = make_funcs(id)
    fig = pyplot.figure(figsize=(5, 5))
    fig.subplots_adjust(bottom=0.15, top=0.95, left=0.15, right=0.95, hspace=0.0, wspace=0.0)
    rmax = remax*3.0001
    r = numpy.arange(rmax/10000.0, rmax, rmax/100.0)
    for i, iid in enumerate(id):
        for k in range(len(bands)-1):
            f1 = f2 = f1max = f2max = 0.0
            for j in range(len(func[i])):
                if k == 0:
                    #label = "%s_%i_%s-%s"%(iid, j, bands[k], bands[k+1])
                    #label = "%s_%i"%(iid, j)
                    label = "%s"%iid
                else:
                    label = ""
                # to use elliptically averaged surface brightnesses will need
                # to supply multi-component fits with single-Sersic info
                f1 += 10**(-0.4*func[i][j][k](r))
                f2 += 10**(-0.4*func[i][j][k+1](r))
                f1max += 10**(-0.4*func[i][j][k](remax))
                f2max += 10**(-0.4*func[i][j][k+1](remax))
            colour = -2.5*numpy.log10(f1/f2)
            colour_remax = -2.5*numpy.log10(f1max/f2max)
            colour -= colour_remax
            colour += offset*k
            pyplot.hlines([offset*k], 0.0, rmax, colors='grey')
            pyplot.plot(r, colour, linestyle=linestyle[i],
                        marker=None, color=color[k], label=label)
    pyplot.legend(loc='upper right', numpoints=1, prop={'size': 16})
    pyplot.xlabel('$r_{\mathrm{e}}$')
    pyplot.ylabel('Colour')
    pyplot.ylim(-1*offset, (len(bands)+1) * offset)
    pyplot.xlim(0.0, rmax)
    fig.savefig('plots/colprofiles_%s.pdf'%name)


def make_funcs(id):
    res = [fit_results(i) for i in id]
    func = []
    for i, iid in enumerate(id):
        func.append([])
        remax = 0
        compno = 0
        while True:
            compno += 1
            field = 'COMP%i_MAG'%compno
            if field not in res[i].dtype.names:
                break
            mag = res[i][field]
            re, n, ar, pa, xc, yc =  [res[i]['COMP%i_%s'%(compno, par)] for par in
                                      ('Re', 'n', 'AR', 'PA', 'XC', 'YC')]
            func[i].append([])
            for k, band in enumerate(bands):
                func[i][compno-1].append(Sersic(mag[k], re[k], n[k], ar[k], pa[k], xc[k], yc[k]))
            remax = max(remax, re.max())
    return func, remax

if __name__ =='__main__':
    plot_all()