<reponame>rsiverd/ultracool<gh_stars>0
#!/usr/bin/env python
# vim: set fileencoding=utf-8 ts=4 sts=4 sw=4 et tw=80 :
#
# Compare <NAME>' 5th-order solution to what Spitzer currently provides.
#
# <NAME>
# Created:       2019-09-04
# Last modified: 2019-11-11
#--------------------------------------------------------------------------
#**************************************************************************
#--------------------------------------------------------------------------

## Current version:
__version__ = "0.1.5"

## Python version-agnostic module reloading:
try:
    reload                              # Python 2.7
except NameError:
    try:
        from importlib import reload    # Python 3.4+
    except ImportError:
        from imp import reload          # Python 3.0 - 3.3

## Modules:
import argparse
#import shutil
#import resource
import signal
#import glob
#import gc
import os
import sys
import time
import numpy as np
#from numpy.lib.recfunctions import append_fields
#import datetime as dt
#from dateutil import parser as dtp
#import scipy.linalg as sla
#import scipy.signal as ssig
#import scipy.ndimage as ndi
#import scipy.optimize as opti
#import scipy.interpolate as stp
#import scipy.spatial.distance as ssd
import matplotlib.pyplot as plt
#import matplotlib.cm as cm
#import matplotlib.ticker as mt
#import matplotlib._pylab_helpers as hlp
#from matplotlib.colors import LogNorm
#from matplotlib import colors
#import matplotlib.colors as mplcolors
#import matplotlib.gridspec as gridspec
#from functools import partial
#from collections import OrderedDict
#import multiprocessing as mp
#np.set_printoptions(suppress=True, linewidth=160)
#import pandas as pd
import statsmodels.api as sm
#import statsmodels.formula.api as smf
#from statsmodels.regression.quantile_regression import QuantReg
#import PIL.Image as pli
#import seaborn as sns
#import cmocean
#import theil_sen as ts
#import window_filter as wf
#import itertools as itt
_have_np_vers = float('.'.join(np.__version__.split('.')[:2]))

## Because obviously:
#import warnings
#if not sys.warnoptions:
#    warnings.simplefilter("ignore", category=DeprecationWarning)
#    warnings.simplefilter("ignore", category=UserWarning)
#    warnings.simplefilter("ignore")
#with warnings.catch_warnings():
#    some_risky_activity()
#with warnings.catch_warnings():
#    warnings.filterwarnings("ignore", category=DeprecationWarning)
#    import problem_child1, problem_child2

## Attempt Kraus solution import:
try:
    import akspoly
    reload(akspoly)
    akp = akspoly.AKSPoly()
except ImportError:
    sys.stderr.write("Failed to import akspoly, please fix!\n")
    sys.exit(1)

##--------------------------------------------------------------------------##
## Projections with cartopy:
#try:
#    import cartopy.crs as ccrs
#    from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
#    from cartopy.feature.nightshade import Nightshade
#    #from cartopy import config as cartoconfig
#except ImportError:
#    sys.stderr.write("Error: cartopy module not found!\n")
#    sys.exit(1)

##--------------------------------------------------------------------------##
## Disable buffering on stdout/stderr:
class Unbuffered(object):
   def __init__(self, stream):
       self.stream = stream
   def write(self, data):
       self.stream.write(data)
       self.stream.flush()
   def __getattr__(self, attr):
       return getattr(self.stream, attr)

sys.stdout = Unbuffered(sys.stdout)
sys.stderr = Unbuffered(sys.stderr)

##--------------------------------------------------------------------------##

## Home-brew robust statistics:
#try:
#    import robust_stats
#    reload(robust_stats)
#    rs = robust_stats
#except ImportError:
#    sys.stderr.write("\nError!  robust_stats module not found!\n"
#           "Please install and try again ...\n\n")
#    sys.exit(1)

## Fast FITS I/O:
#try:
#    import fitsio
#except ImportError:
#    sys.stderr.write("\nError: fitsio module not found!\n")
#    sys.exit(1)

## FITS I/O:
try:
    import astropy.io.fits as pf
except ImportError:
    try:
       import pyfits as pf
    except ImportError:
        sys.stderr.write("\nError!  No FITS I/O module found!\n"
               "Install either astropy.io.fits or pyfits and try again!\n\n")
        sys.exit(1)

## Various from astropy:
try:
#    import astropy.io.ascii as aia
#    import astropy.table as apt
#    import astropy.time as astt
    import astropy.wcs as awcs
#    from astropy import coordinates as coord
#    from astropy import units as uu
except ImportError:
    sys.stderr.write("\nError: astropy module not found!\n")
    sys.exit(1)

## Star extraction:
#try:
#    import easy_sep
#    reload(easy_sep)
#except ImportError:
#    sys.stderr.write("Error: easy_sep module not found!\n\n")
#    sys.exit(1)
#pse = easy_sep.EasySEP()

##--------------------------------------------------------------------------##
## Colors for fancy terminal output:
NRED    = '\033[0;31m'   ;  BRED    = '\033[1;31m'
NGREEN  = '\033[0;32m'   ;  BGREEN  = '\033[1;32m'
NYELLOW = '\033[0;33m'   ;  BYELLOW = '\033[1;33m'
NBLUE   = '\033[0;34m'   ;  BBLUE   = '\033[1;34m'
NMAG    = '\033[0;35m'   ;  BMAG    = '\033[1;35m'
NCYAN   = '\033[0;36m'   ;  BCYAN   = '\033[1;36m'
NWHITE  = '\033[0;37m'   ;  BWHITE  = '\033[1;37m'
ENDC    = '\033[0m'

## Suppress colors in cron jobs:
if (os.getenv('FUNCDEF') == '--nocolors'):
    NRED    = ''   ;  BRED    = ''
    NGREEN  = ''   ;  BGREEN  = ''
    NYELLOW = ''   ;  BYELLOW = ''
    NBLUE   = ''   ;  BBLUE   = ''
    NMAG    = ''   ;  BMAG    = ''
    NCYAN   = ''   ;  BCYAN   = ''
    NWHITE  = ''   ;  BWHITE  = ''
    ENDC    = ''

## Fancy text:
degree_sign = u'\N{DEGREE SIGN}'

## Dividers:
halfdiv = '-' * 40
fulldiv = '-' * 80

##--------------------------------------------------------------------------##
## Save FITS image with clobber (astropy / pyfits):
def qsave(iname, idata, header=None, padkeys=1000, **kwargs):
    this_func = sys._getframe().f_code.co_name
    parent_func = sys._getframe(1).f_code.co_name
    sys.stderr.write("Writing to '%s' ... " % iname)
    if header:
        while (len(header) < padkeys):
            header.append() # pad header
    if os.path.isfile(iname):
        os.remove(iname)
    pf.writeto(iname, idata, header=header, **kwargs)
    sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
## Save FITS image with clobber (fitsio):
#def qsave(iname, idata, header=None, **kwargs):
#    this_func = sys._getframe().f_code.co_name
#    parent_func = sys._getframe(1).f_code.co_name
#    sys.stderr.write("Writing to '%s' ... " % iname)
#    #if os.path.isfile(iname):
#    #    os.remove(iname)
#    fitsio.write(iname, idata, clobber=True, header=header, **kwargs)
#    sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
def ldmap(things):
    return dict(zip(things, range(len(things))))

def argnear(vec, val):
    return (np.abs(vec - val)).argmin()

##--------------------------------------------------------------------------##
## Rotation matrix (radians):
def make_rmat_rad(theta):
    return np.matrix([[np.cos(theta), -np.sin(theta)],
                        [np.sin(theta), np.cos(theta)]])

## Rotation matrix (degrees):
def make_rmat_deg(theta):
    return make_rmat_rad(np.radians(theta))


##--------------------------------------------------------------------------##
##------------------         Parse Command Line             ----------------##
##--------------------------------------------------------------------------##

## Parse arguments and run script:
class MyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)

## Enable raw text AND display of defaults:
class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                        argparse.RawDescriptionHelpFormatter):
    pass

## Parse the command line:
if __name__ == '__main__':

    # ------------------------------------------------------------------
    prog_name = os.path.basename(__file__)
    descr_txt = """
    Spitzer coordinate distortion comparison.
    
    Version: %s
    """ % __version__
    parser = MyParser(prog=prog_name, description=descr_txt,
                          formatter_class=argparse.RawTextHelpFormatter)
    # ------------------------------------------------------------------
    #parser.set_defaults(thing1='value1', thing2='value2')
    # ------------------------------------------------------------------
    #parser.add_argument('firstpos', help='first positional argument')
    #parser.add_argument('-w', '--whatever', required=False, default=5.0,
    #        help='some option with default [def: %(default)s]', type=float)
    #parser.add_argument('-s', '--site',
    #        help='Site to retrieve data for', required=True)
    #parser.add_argument('-n', '--number_of_days', default=1,
    #        help='Number of days of data to retrieve.')
    #parser.add_argument('-o', '--output_file', 
    #        default='observations.csv', help='Output filename.')
    #parser.add_argument('remainder', help='other stuff', nargs='*')
    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    iogroup = parser.add_argument_group('File I/O')
    iogroup.add_argument('-i', '--input_image', default=None, required=True,
            help='Modern Spitzer cbcd FITS image', type=str)
    #iogroup.add_argument('-o', '--output_file', default=None, required=True,
    #        help='Output filename', type=str)
    #iogroup.add_argument('-R', '--ref_image', default=None, required=True,
    #        help='KELT image with WCS')
    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    # Miscellany:
    miscgroup = parser.add_argument_group('Miscellany')
    miscgroup.add_argument('--debug', dest='debug', default=False,
            help='Enable extra debugging messages', action='store_true')
    miscgroup.add_argument('-q', '--quiet', action='count', default=0,
            help='less progress/status reporting')
    miscgroup.add_argument('-v', '--verbose', action='count', default=0,
            help='more progress/status reporting')
    # ------------------------------------------------------------------

    context = parser.parse_args()
    context.vlevel = 99 if context.debug else (context.verbose-context.quiet)
    context.prog_name = prog_name

##--------------------------------------------------------------------------##
## Load Spitzer image:
if not os.path.isfile(context.input_image):
    sys.stderr.write("Error: file not found: %s\n" % context.input_image)
    sys.exit(1)

sys.stderr.write("Loading %s ... " % context.input_image)
sdata, skeys = pf.getdata(context.input_image, header=True)
sys.stderr.write("done.\n")

## Check channel:
irac_channel = skeys['CHNLNUM']
sys.stderr.write("Detected IRAC channel %d.\n" % irac_channel)

##--------------------------------------------------------------------------##
## Load WCS from FITS header:
imwcs = awcs.WCS(skeys)
imorigin = 1.0

##--------------------------------------------------------------------------##
## Spitzer CCD coordinate grids:
ny, nx = sdata.shape
#x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered)
#y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered)
#xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered)
#xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered)
#xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute
#yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute
#yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute
#yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array)
yy, xx = np.mgrid[1:ny+1, 1:nx+1].astype('float32')  # absolute (pixel)

## Relative coordinates:
rx = xx - skeys['CRPIX1']
ry = yy - skeys['CRPIX2']

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##

## Intermediate coords with AK solution:
#ak_xx, ak_yy = akp.xform_xy(xx.flatten(), yy.flatten(), irac_channel)
#ak_xx, ak_yy = akp.xform_xy(xx, yy, irac_channel)
akfpx, akfpy = akp.xy2focal(xx, yy, irac_channel)

## Corresponding corrections relative to naive grid:
ak_dx = akfpx - rx
ak_dy = akfpy - ry

#fiverows_xx = xx[:5]
#fiverows_yy = yy[:5]
#fr_xr = fiverows_xx - 128.
#fr_yr = fiverows_yy - 128.
#akx5r, aky5r = akp.xy2focal(fiverows_xx, fiverows_yy, irac_channel)
#xpar = akp._a_pars[irac_channel]
#ypar = akp._b_pars[irac_channel]
#
#xresult = np.zeros_like(fr_xr)
#xresult += xpar[0]
#
#yresult = np.zeros_like(fr_xr)
#yresult += ypar[0]

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##

## Focal plane coordinates from WCS (de-distort using SIP terms):
imfpx, imfpy = imwcs.sip_pix2foc(xx, yy, imorigin)

## Corresponding deltas:
im_dx = imfpx - rx
im_dy = imfpy - ry

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## How to compare dx and dy vs. rotation angle:
def rotate_coords(xx, yy, angle_deg):
    prev_shape = xx.shape
    stacked_xy = np.vstack((xx.flatten(), yy.flatten()))
    rot_matrix = make_rmat_deg(angle_deg)
    rx, ry = np.dot(rot_matrix, stacked_xy)
    things = []
    for xypair in zip(xx.flatten(), yy.flatten()):
        things.append(np.dot(rot_matrix, np.array(xypair)))
    return rx.reshape(prev_shape), ry.reshape(prev_shape), things


##--------------------------------------------------------------------------##
## Quick FITS I/O:
#data_file = 'image.fits'
#img_vals = pf.getdata(data_file)
#hdr_keys = pf.getheader(data_file)
#img_vals, hdr_keys = pf.getdata(data_file, header=True)
#img_vals, hdr_keys = pf.getdata(data_file, header=True, uint=True) # USHORT
#img_vals, hdr_keys = fitsio.read(data_file, header=True)

#date_obs = hdr_keys['DATE-OBS']
#site_lat = hdr_keys['LATITUDE']
#site_lon = hdr_keys['LONGITUD']

## Initialize time:
##img_time = astt.Time(hdr_keys['DATE-OBS'], scale='utc', format='isot')
##img_time += astt.TimeDelta(0.5 * hdr_keys['EXPTIME'], format='sec')
##jd_image = img_time.jd
#
### Initialize location:
##observer = ephem.Observer()
##observer.lat = np.radians(site_lat)
##observer.lon = np.radians(site_lon)
##observer.date = img_time.datetime
#
##pf.writeto('new.fits', img_vals)
##qsave('new.fits', img_vals)
##qsave('new.fits', img_vals, header=hdr_keys)

## Star extraction:
#pse.set_image(img_vals, gain=3.6)
#objlist = pse.analyze(sigthresh=5.0)

##--------------------------------------------------------------------------##
## Misc:
#def log_10_product(x, pos):
#   """The two args are the value and tick position.
#   Label ticks with the product of the exponentiation."""
#   return '%.2f' % (x)  # floating-point
#
#formatter = plt.FuncFormatter(log_10_product) # wrap function for use

## Convenient, percentile-based plot limits:
#def nice_limits(vec, pctiles=[1,99], pad=1.2):
#    ends = np.percentile(vec[~np.isnan(vec)], pctiles)
#    middle = np.average(ends)
#    return (middle + pad * (ends - middle))

## Convenient plot limits for datetime/astropy.Time content:
#def nice_time_limits(tvec, buffer=0.05):
#    lower = tvec.min()
#    upper = tvec.max()
#    ndays = upper - lower
#    return ((lower - 0.05*ndays).datetime, (upper + 0.05*ndays).datetime)

## Convenient limits for datetime objects:
#def dt_limits(vec, pad=0.1):
#    tstart, tstop = vec.min(), vec.max()
#    trange = (tstop - tstart).total_seconds()
#    tpad = dt.timedelta(seconds=pad*trange)
#    return (tstart - tpad, tstop + tpad)

##--------------------------------------------------------------------------##
## Solve prep:
#ny, nx = img_vals.shape
#x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered)
#y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered)
#xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered)
#xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered)
#xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute
#yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute
#yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute
#yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array)
#yy, xx = np.mgrid[1:ny+1, 1:nx+1].astype('uint16')   # absolute (pixel)

## 1-D vectors:
#x_pix, y_pix, ivals = xx.flatten(), yy.flatten(), img_vals.flatten()
#w_vec = np.ones_like(ivals)            # start with uniform weights
#design_matrix = np.column_stack((np.ones(x_pix.size), x_pix, y_pix))

## Image fitting (statsmodels etc.):
#data = sm.datasets.stackloss.load()
#ols_res = sm.OLS(ivals, design_matrix).fit()
#rlm_res = sm.RLM(ivals, design_matrix).fit()
#rlm_model = sm.RLM(ivals, design_matrix, M=sm.robust.norms.HuberT())
#rlm_res = rlm_model.fit()
#data = pd.DataFrame({'xpix':x_pix, 'ypix':y_pix})
#rlm_model = sm.RLM.from_formula("ivals ~ xpix + ypix", data)

##--------------------------------------------------------------------------##
##------------------         Fit for im-ak Difference       ----------------##
##--------------------------------------------------------------------------##

### 1-D vectors:
#xx_1d = xx.flatten()
#yy_1d = yy.flatten()
#xdiff = (im_dx - ak_dx).flatten()
#ydiff = (im_dy - ak_dy).flatten()
#
#design_matrix = np.column_stack((np.ones_like(xdiff), xx_1d, yy_1d))
#dx_ols = sm.OLS(xdiff, design_matrix).fit()
#dy_ols = sm.OLS(ydiff, design_matrix).fit()


##--------------------------------------------------------------------------##
## Theil-Sen line-fitting (linear):
#model = ts.linefit(xvals, yvals)
#icept, slope = ts.linefit(xvals, yvals)

## Theil-Sen line-fitting (loglog):
#xvals, yvals = np.log10(original_xvals), np.log10(original_yvals)
#xvals, yvals = np.log10(df['x'].values), np.log10(df['y'].values)
#llmodel = ts.linefit(np.log10(xvals), np.log10(yvals))
#icept, slope = ts.linefit(xvals, yvals)
#fit_exponent = slope
#fit_multiplier = 10**icept
#bestfit_x = np.arange(5000)
#bestfit_y = fit_multiplier * bestfit_x**fit_exponent

## Log-log evaluator:
#def loglog_eval(xvals, model):
#    icept, slope = model
#    return 10**icept * xvals**slope
#def loglog_eval(xvals, icept, slope):
#    return 10**icept * xvals**slope

##--------------------------------------------------------------------------##
## KDE:
#kde_pnts, kde_vals = mk.go(data_vec)

##--------------------------------------------------------------------------##
## Vaex plotting:
#ds = vaex.open('big_file.hdf5')
#ds = vaex.from_arrays(x=x, y=y)     # load from arrays
#ds = vaex.from_csv('mydata.csv')

## Stats:
#ds.mean("x"), ds.std("x"), ds.correlation("vx**2+vy**2+vz**2", "E")
#ds.plot(....)
#http://vaex.astro.rug.nl/latest/tutorial_ipython_notebook.html

##--------------------------------------------------------------------------##
## Plot config:

# gridspec examples:
# https://matplotlib.org/users/gridspec.html

#gs1 = gridspec.GridSpec(4, 4)
#gs1.update(wspace=0.025, hspace=0.05)  # set axis spacing

#ax1 = plt.subplot2grid((3, 3), (0, 0), colspan=3) # top-left + center + right
#ax2 = plt.subplot2grid((3, 3), (1, 0), colspan=2) # mid-left + mid-center
#ax3 = plt.subplot2grid((3, 3), (1, 2), rowspan=2) # mid-right + bot-right
#ax4 = plt.subplot2grid((3, 3), (2, 0))            # bot-left
#ax5 = plt.subplot2grid((3, 3), (2, 1))            # bot-center


##--------------------------------------------------------------------------##
#plt.style.use('bmh')   # Bayesian Methods for Hackers style
fig_dims = (12, 10)
fig = plt.figure(1, figsize=fig_dims)
plt.gcf().clf()
#fig, axs = plt.subplots(2, 2, sharex=True, figsize=fig_dims, num=1)
# sharex='col' | sharex='row'
#fig.frameon = False # disable figure frame drawing
#fig.subplots_adjust(left=0.07, right=0.95)
#ax1 = plt.subplot(gs[0, 0])
ax1 = fig.add_subplot(111)
#ax1 = fig.add_axes([0, 0, 1, 1])
#ax1.patch.set_facecolor((0.8, 0.8, 0.8))
ax1.grid(True)
#ax1.axis('off')

## Disable axis offsets:
#ax1.xaxis.get_major_formatter().set_useOffset(False)
#ax1.yaxis.get_major_formatter().set_useOffset(False)
stride = slice(None, None, 10)
ss = (stride, stride)
#ax1.quiver(xx[ss], yy[ss], im_dx[ss], im_dy[ss])
ax1.quiver(xx[ss], yy[ss], (im_dx - ak_dx)[ss], (im_dy - ak_dy)[ss])

#ax1.plot(kde_pnts, kde_vals)

#blurb = "some text"
#ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes)
#ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes,
#      va='top', ha='left', bbox=dict(facecolor='white', pad=10.0))
#      fontdict={'family':'monospace'}) # fixed-width

#colors = cm.rainbow(np.linspace(0, 1, len(plot_list)))
#for camid, c in zip(plot_list, colors):
#    cam_data = subsets[camid]
#    xvalue = cam_data['CCDATEMP']
#    yvalue = cam_data['PIX_MED']
#    yvalue = cam_data['IMEAN']
#    ax1.scatter(xvalue, yvalue, color=c, lw=0, label=camid)

#mtickpos = [2,5,7]
#ndecades = 1.0   # for symlog, set width of linear portion in units of dex
#nonposx='mask' | nonposx='clip' | nonposy='mask' | nonposy='clip'
#ax1.set_xscale('log', basex=10, nonposx='mask', subsx=mtickpos)
#ax1.set_xscale('log', nonposx='clip', subsx=[3])
#ax1.set_yscale('symlog', basey=10, linthreshy=0.1, linscaley=ndecades)
#ax1.xaxis.set_major_formatter(formatter) # re-format x ticks
#ax1.set_ylim(ax1.get_ylim()[::-1])
#ax1.set_xlabel('whatever', labelpad=30)  # push X label down 

#ax1.set_xticks([1.0, 3.0, 10.0, 30.0, 100.0])
#ax1.set_xticks([1, 2, 3], ['Jan', 'Feb', 'Mar'])
#for label in ax1.get_xticklabels():
#    label.set_rotation(30)
#    label.set_fontsize(14) 

#ax1.xaxis.label.set_fontsize(18)
#ax1.yaxis.label.set_fontsize(18)

#ax1.set_xlim(nice_limits(xvec, pctiles=[1,99], pad=1.2))
#ax1.set_ylim(nice_limits(yvec, pctiles=[1,99], pad=1.2))

#spts = ax1.scatter(x, y, lw=0, s=5)
#cbar = fig.colorbar(spts, orientation='vertical')
#cbar.formatter.set_useOffset(False)
#cbar.update_ticks()

plot_name = 'im_minus_ak.png'
fig.tight_layout() # adjust boundaries sensibly, matplotlib v1.1+
plt.draw()
fig.savefig(plot_name, bbox_inches='tight')

# cyclical colormap ... cmocean.cm.phase
# cmocean: https://matplotlib.org/cmocean/




######################################################################
# CHANGELOG (compare_distortion.py):
#---------------------------------------------------------------------
#
#  2019-09-04:
#     -- Increased __version__ to 0.1.0.
#     -- First created compare_distortion.py.
#

#!/usr/bin/env python
# vim: set fileencoding=utf-8 ts=4 sts=4 sw=4 et tw=80 :
#
# Compare <NAME>' 5th-order solution to what Spitzer currently provides.
#
# <NAME>
# Created:       2019-09-04
# Last modified: 2019-11-11
#--------------------------------------------------------------------------
#**************************************************************************
#--------------------------------------------------------------------------

## Current version:
__version__ = "0.1.5"

## Python version-agnostic module reloading:
try:
    reload                              # Python 2.7
except NameError:
    try:
        from importlib import reload    # Python 3.4+
    except ImportError:
        from imp import reload          # Python 3.0 - 3.3

## Modules:
import argparse
#import shutil
#import resource
import signal
#import glob
#import gc
import os
import sys
import time
import numpy as np
#from numpy.lib.recfunctions import append_fields
#import datetime as dt
#from dateutil import parser as dtp
#import scipy.linalg as sla
#import scipy.signal as ssig
#import scipy.ndimage as ndi
#import scipy.optimize as opti
#import scipy.interpolate as stp
#import scipy.spatial.distance as ssd
import matplotlib.pyplot as plt
#import matplotlib.cm as cm
#import matplotlib.ticker as mt
#import matplotlib._pylab_helpers as hlp
#from matplotlib.colors import LogNorm
#from matplotlib import colors
#import matplotlib.colors as mplcolors
#import matplotlib.gridspec as gridspec
#from functools import partial
#from collections import OrderedDict
#import multiprocessing as mp
#np.set_printoptions(suppress=True, linewidth=160)
#import pandas as pd
import statsmodels.api as sm
#import statsmodels.formula.api as smf
#from statsmodels.regression.quantile_regression import QuantReg
#import PIL.Image as pli
#import seaborn as sns
#import cmocean
#import theil_sen as ts
#import window_filter as wf
#import itertools as itt
_have_np_vers = float('.'.join(np.__version__.split('.')[:2]))

## Because obviously:
#import warnings
#if not sys.warnoptions:
#    warnings.simplefilter("ignore", category=DeprecationWarning)
#    warnings.simplefilter("ignore", category=UserWarning)
#    warnings.simplefilter("ignore")
#with warnings.catch_warnings():
#    some_risky_activity()
#with warnings.catch_warnings():
#    warnings.filterwarnings("ignore", category=DeprecationWarning)
#    import problem_child1, problem_child2

## Attempt Kraus solution import:
try:
    import akspoly
    reload(akspoly)
    akp = akspoly.AKSPoly()
except ImportError:
    sys.stderr.write("Failed to import akspoly, please fix!\n")
    sys.exit(1)

##--------------------------------------------------------------------------##
## Projections with cartopy:
#try:
#    import cartopy.crs as ccrs
#    from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
#    from cartopy.feature.nightshade import Nightshade
#    #from cartopy import config as cartoconfig
#except ImportError:
#    sys.stderr.write("Error: cartopy module not found!\n")
#    sys.exit(1)

##--------------------------------------------------------------------------##
## Disable buffering on stdout/stderr:
class Unbuffered(object):
   def __init__(self, stream):
       self.stream = stream
   def write(self, data):
       self.stream.write(data)
       self.stream.flush()
   def __getattr__(self, attr):
       return getattr(self.stream, attr)

sys.stdout = Unbuffered(sys.stdout)
sys.stderr = Unbuffered(sys.stderr)

##--------------------------------------------------------------------------##

## Home-brew robust statistics:
#try:
#    import robust_stats
#    reload(robust_stats)
#    rs = robust_stats
#except ImportError:
#    sys.stderr.write("\nError!  robust_stats module not found!\n"
#           "Please install and try again ...\n\n")
#    sys.exit(1)

## Fast FITS I/O:
#try:
#    import fitsio
#except ImportError:
#    sys.stderr.write("\nError: fitsio module not found!\n")
#    sys.exit(1)

## FITS I/O:
try:
    import astropy.io.fits as pf
except ImportError:
    try:
       import pyfits as pf
    except ImportError:
        sys.stderr.write("\nError!  No FITS I/O module found!\n"
               "Install either astropy.io.fits or pyfits and try again!\n\n")
        sys.exit(1)

## Various from astropy:
try:
#    import astropy.io.ascii as aia
#    import astropy.table as apt
#    import astropy.time as astt
    import astropy.wcs as awcs
#    from astropy import coordinates as coord
#    from astropy import units as uu
except ImportError:
    sys.stderr.write("\nError: astropy module not found!\n")
    sys.exit(1)

## Star extraction:
#try:
#    import easy_sep
#    reload(easy_sep)
#except ImportError:
#    sys.stderr.write("Error: easy_sep module not found!\n\n")
#    sys.exit(1)
#pse = easy_sep.EasySEP()

##--------------------------------------------------------------------------##
## Colors for fancy terminal output:
NRED    = '\033[0;31m'   ;  BRED    = '\033[1;31m'
NGREEN  = '\033[0;32m'   ;  BGREEN  = '\033[1;32m'
NYELLOW = '\033[0;33m'   ;  BYELLOW = '\033[1;33m'
NBLUE   = '\033[0;34m'   ;  BBLUE   = '\033[1;34m'
NMAG    = '\033[0;35m'   ;  BMAG    = '\033[1;35m'
NCYAN   = '\033[0;36m'   ;  BCYAN   = '\033[1;36m'
NWHITE  = '\033[0;37m'   ;  BWHITE  = '\033[1;37m'
ENDC    = '\033[0m'

## Suppress colors in cron jobs:
if (os.getenv('FUNCDEF') == '--nocolors'):
    NRED    = ''   ;  BRED    = ''
    NGREEN  = ''   ;  BGREEN  = ''
    NYELLOW = ''   ;  BYELLOW = ''
    NBLUE   = ''   ;  BBLUE   = ''
    NMAG    = ''   ;  BMAG    = ''
    NCYAN   = ''   ;  BCYAN   = ''
    NWHITE  = ''   ;  BWHITE  = ''
    ENDC    = ''

## Fancy text:
degree_sign = u'\N{DEGREE SIGN}'

## Dividers:
halfdiv = '-' * 40
fulldiv = '-' * 80

##--------------------------------------------------------------------------##
## Save FITS image with clobber (astropy / pyfits):
def qsave(iname, idata, header=None, padkeys=1000, **kwargs):
    this_func = sys._getframe().f_code.co_name
    parent_func = sys._getframe(1).f_code.co_name
    sys.stderr.write("Writing to '%s' ... " % iname)
    if header:
        while (len(header) < padkeys):
            header.append() # pad header
    if os.path.isfile(iname):
        os.remove(iname)
    pf.writeto(iname, idata, header=header, **kwargs)
    sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
## Save FITS image with clobber (fitsio):
#def qsave(iname, idata, header=None, **kwargs):
#    this_func = sys._getframe().f_code.co_name
#    parent_func = sys._getframe(1).f_code.co_name
#    sys.stderr.write("Writing to '%s' ... " % iname)
#    #if os.path.isfile(iname):
#    #    os.remove(iname)
#    fitsio.write(iname, idata, clobber=True, header=header, **kwargs)
#    sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
def ldmap(things):
    return dict(zip(things, range(len(things))))

def argnear(vec, val):
    return (np.abs(vec - val)).argmin()

##--------------------------------------------------------------------------##
## Rotation matrix (radians):
def make_rmat_rad(theta):
    return np.matrix([[np.cos(theta), -np.sin(theta)],
                        [np.sin(theta), np.cos(theta)]])

## Rotation matrix (degrees):
def make_rmat_deg(theta):
    return make_rmat_rad(np.radians(theta))


##--------------------------------------------------------------------------##
##------------------         Parse Command Line             ----------------##
##--------------------------------------------------------------------------##

## Parse arguments and run script:
class MyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)

## Enable raw text AND display of defaults:
class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                        argparse.RawDescriptionHelpFormatter):
    pass

## Parse the command line:
if __name__ == '__main__':

    # ------------------------------------------------------------------
    prog_name = os.path.basename(__file__)
    descr_txt = """
    Spitzer coordinate distortion comparison.
    
    Version: %s
    """ % __version__
    parser = MyParser(prog=prog_name, description=descr_txt,
                          formatter_class=argparse.RawTextHelpFormatter)
    # ------------------------------------------------------------------
    #parser.set_defaults(thing1='value1', thing2='value2')
    # ------------------------------------------------------------------
    #parser.add_argument('firstpos', help='first positional argument')
    #parser.add_argument('-w', '--whatever', required=False, default=5.0,
    #        help='some option with default [def: %(default)s]', type=float)
    #parser.add_argument('-s', '--site',
    #        help='Site to retrieve data for', required=True)
    #parser.add_argument('-n', '--number_of_days', default=1,
    #        help='Number of days of data to retrieve.')
    #parser.add_argument('-o', '--output_file', 
    #        default='observations.csv', help='Output filename.')
    #parser.add_argument('remainder', help='other stuff', nargs='*')
    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    iogroup = parser.add_argument_group('File I/O')
    iogroup.add_argument('-i', '--input_image', default=None, required=True,
            help='Modern Spitzer cbcd FITS image', type=str)
    #iogroup.add_argument('-o', '--output_file', default=None, required=True,
    #        help='Output filename', type=str)
    #iogroup.add_argument('-R', '--ref_image', default=None, required=True,
    #        help='KELT image with WCS')
    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    # Miscellany:
    miscgroup = parser.add_argument_group('Miscellany')
    miscgroup.add_argument('--debug', dest='debug', default=False,
            help='Enable extra debugging messages', action='store_true')
    miscgroup.add_argument('-q', '--quiet', action='count', default=0,
            help='less progress/status reporting')
    miscgroup.add_argument('-v', '--verbose', action='count', default=0,
            help='more progress/status reporting')
    # ------------------------------------------------------------------

    context = parser.parse_args()
    context.vlevel = 99 if context.debug else (context.verbose-context.quiet)
    context.prog_name = prog_name

##--------------------------------------------------------------------------##
## Load Spitzer image:
if not os.path.isfile(context.input_image):
    sys.stderr.write("Error: file not found: %s\n" % context.input_image)
    sys.exit(1)

sys.stderr.write("Loading %s ... " % context.input_image)
sdata, skeys = pf.getdata(context.input_image, header=True)
sys.stderr.write("done.\n")

## Check channel:
irac_channel = skeys['CHNLNUM']
sys.stderr.write("Detected IRAC channel %d.\n" % irac_channel)

##--------------------------------------------------------------------------##
## Load WCS from FITS header:
imwcs = awcs.WCS(skeys)
imorigin = 1.0

##--------------------------------------------------------------------------##
## Spitzer CCD coordinate grids:
ny, nx = sdata.shape
#x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered)
#y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered)
#xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered)
#xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered)
#xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute
#yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute
#yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute
#yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array)
yy, xx = np.mgrid[1:ny+1, 1:nx+1].astype('float32')  # absolute (pixel)

## Relative coordinates:
rx = xx - skeys['CRPIX1']
ry = yy - skeys['CRPIX2']

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##

## Intermediate coords with AK solution:
#ak_xx, ak_yy = akp.xform_xy(xx.flatten(), yy.flatten(), irac_channel)
#ak_xx, ak_yy = akp.xform_xy(xx, yy, irac_channel)
akfpx, akfpy = akp.xy2focal(xx, yy, irac_channel)

## Corresponding corrections relative to naive grid:
ak_dx = akfpx - rx
ak_dy = akfpy - ry

#fiverows_xx = xx[:5]
#fiverows_yy = yy[:5]
#fr_xr = fiverows_xx - 128.
#fr_yr = fiverows_yy - 128.
#akx5r, aky5r = akp.xy2focal(fiverows_xx, fiverows_yy, irac_channel)
#xpar = akp._a_pars[irac_channel]
#ypar = akp._b_pars[irac_channel]
#
#xresult = np.zeros_like(fr_xr)
#xresult += xpar[0]
#
#yresult = np.zeros_like(fr_xr)
#yresult += ypar[0]

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##

## Focal plane coordinates from WCS (de-distort using SIP terms):
imfpx, imfpy = imwcs.sip_pix2foc(xx, yy, imorigin)

## Corresponding deltas:
im_dx = imfpx - rx
im_dy = imfpy - ry

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## How to compare dx and dy vs. rotation angle:
def rotate_coords(xx, yy, angle_deg):
    prev_shape = xx.shape
    stacked_xy = np.vstack((xx.flatten(), yy.flatten()))
    rot_matrix = make_rmat_deg(angle_deg)
    rx, ry = np.dot(rot_matrix, stacked_xy)
    things = []
    for xypair in zip(xx.flatten(), yy.flatten()):
        things.append(np.dot(rot_matrix, np.array(xypair)))
    return rx.reshape(prev_shape), ry.reshape(prev_shape), things


##--------------------------------------------------------------------------##
## Quick FITS I/O:
#data_file = 'image.fits'
#img_vals = pf.getdata(data_file)
#hdr_keys = pf.getheader(data_file)
#img_vals, hdr_keys = pf.getdata(data_file, header=True)
#img_vals, hdr_keys = pf.getdata(data_file, header=True, uint=True) # USHORT
#img_vals, hdr_keys = fitsio.read(data_file, header=True)

#date_obs = hdr_keys['DATE-OBS']
#site_lat = hdr_keys['LATITUDE']
#site_lon = hdr_keys['LONGITUD']

## Initialize time:
##img_time = astt.Time(hdr_keys['DATE-OBS'], scale='utc', format='isot')
##img_time += astt.TimeDelta(0.5 * hdr_keys['EXPTIME'], format='sec')
##jd_image = img_time.jd
#
### Initialize location:
##observer = ephem.Observer()
##observer.lat = np.radians(site_lat)
##observer.lon = np.radians(site_lon)
##observer.date = img_time.datetime
#
##pf.writeto('new.fits', img_vals)
##qsave('new.fits', img_vals)
##qsave('new.fits', img_vals, header=hdr_keys)

## Star extraction:
#pse.set_image(img_vals, gain=3.6)
#objlist = pse.analyze(sigthresh=5.0)

##--------------------------------------------------------------------------##
## Misc:
#def log_10_product(x, pos):
#   """The two args are the value and tick position.
#   Label ticks with the product of the exponentiation."""
#   return '%.2f' % (x)  # floating-point
#
#formatter = plt.FuncFormatter(log_10_product) # wrap function for use

## Convenient, percentile-based plot limits:
#def nice_limits(vec, pctiles=[1,99], pad=1.2):
#    ends = np.percentile(vec[~np.isnan(vec)], pctiles)
#    middle = np.average(ends)
#    return (middle + pad * (ends - middle))

## Convenient plot limits for datetime/astropy.Time content:
#def nice_time_limits(tvec, buffer=0.05):
#    lower = tvec.min()
#    upper = tvec.max()
#    ndays = upper - lower
#    return ((lower - 0.05*ndays).datetime, (upper + 0.05*ndays).datetime)

## Convenient limits for datetime objects:
#def dt_limits(vec, pad=0.1):
#    tstart, tstop = vec.min(), vec.max()
#    trange = (tstop - tstart).total_seconds()
#    tpad = dt.timedelta(seconds=pad*trange)
#    return (tstart - tpad, tstop + tpad)

##--------------------------------------------------------------------------##
## Solve prep:
#ny, nx = img_vals.shape
#x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered)
#y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered)
#xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered)
#xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered)
#xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute
#yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute
#yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute
#yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array)
#yy, xx = np.mgrid[1:ny+1, 1:nx+1].astype('uint16')   # absolute (pixel)

## 1-D vectors:
#x_pix, y_pix, ivals = xx.flatten(), yy.flatten(), img_vals.flatten()
#w_vec = np.ones_like(ivals)            # start with uniform weights
#design_matrix = np.column_stack((np.ones(x_pix.size), x_pix, y_pix))

## Image fitting (statsmodels etc.):
#data = sm.datasets.stackloss.load()
#ols_res = sm.OLS(ivals, design_matrix).fit()
#rlm_res = sm.RLM(ivals, design_matrix).fit()
#rlm_model = sm.RLM(ivals, design_matrix, M=sm.robust.norms.HuberT())
#rlm_res = rlm_model.fit()
#data = pd.DataFrame({'xpix':x_pix, 'ypix':y_pix})
#rlm_model = sm.RLM.from_formula("ivals ~ xpix + ypix", data)

##--------------------------------------------------------------------------##
##------------------         Fit for im-ak Difference       ----------------##
##--------------------------------------------------------------------------##

### 1-D vectors:
#xx_1d = xx.flatten()
#yy_1d = yy.flatten()
#xdiff = (im_dx - ak_dx).flatten()
#ydiff = (im_dy - ak_dy).flatten()
#
#design_matrix = np.column_stack((np.ones_like(xdiff), xx_1d, yy_1d))
#dx_ols = sm.OLS(xdiff, design_matrix).fit()
#dy_ols = sm.OLS(ydiff, design_matrix).fit()


##--------------------------------------------------------------------------##
## Theil-Sen line-fitting (linear):
#model = ts.linefit(xvals, yvals)
#icept, slope = ts.linefit(xvals, yvals)

## Theil-Sen line-fitting (loglog):
#xvals, yvals = np.log10(original_xvals), np.log10(original_yvals)
#xvals, yvals = np.log10(df['x'].values), np.log10(df['y'].values)
#llmodel = ts.linefit(np.log10(xvals), np.log10(yvals))
#icept, slope = ts.linefit(xvals, yvals)
#fit_exponent = slope
#fit_multiplier = 10**icept
#bestfit_x = np.arange(5000)
#bestfit_y = fit_multiplier * bestfit_x**fit_exponent

## Log-log evaluator:
#def loglog_eval(xvals, model):
#    icept, slope = model
#    return 10**icept * xvals**slope
#def loglog_eval(xvals, icept, slope):
#    return 10**icept * xvals**slope

##--------------------------------------------------------------------------##
## KDE:
#kde_pnts, kde_vals = mk.go(data_vec)

##--------------------------------------------------------------------------##
## Vaex plotting:
#ds = vaex.open('big_file.hdf5')
#ds = vaex.from_arrays(x=x, y=y)     # load from arrays
#ds = vaex.from_csv('mydata.csv')

## Stats:
#ds.mean("x"), ds.std("x"), ds.correlation("vx**2+vy**2+vz**2", "E")
#ds.plot(....)
#http://vaex.astro.rug.nl/latest/tutorial_ipython_notebook.html

##--------------------------------------------------------------------------##
## Plot config:

# gridspec examples:
# https://matplotlib.org/users/gridspec.html

#gs1 = gridspec.GridSpec(4, 4)
#gs1.update(wspace=0.025, hspace=0.05)  # set axis spacing

#ax1 = plt.subplot2grid((3, 3), (0, 0), colspan=3) # top-left + center + right
#ax2 = plt.subplot2grid((3, 3), (1, 0), colspan=2) # mid-left + mid-center
#ax3 = plt.subplot2grid((3, 3), (1, 2), rowspan=2) # mid-right + bot-right
#ax4 = plt.subplot2grid((3, 3), (2, 0))            # bot-left
#ax5 = plt.subplot2grid((3, 3), (2, 1))            # bot-center


##--------------------------------------------------------------------------##
#plt.style.use('bmh')   # Bayesian Methods for Hackers style
fig_dims = (12, 10)
fig = plt.figure(1, figsize=fig_dims)
plt.gcf().clf()
#fig, axs = plt.subplots(2, 2, sharex=True, figsize=fig_dims, num=1)
# sharex='col' | sharex='row'
#fig.frameon = False # disable figure frame drawing
#fig.subplots_adjust(left=0.07, right=0.95)
#ax1 = plt.subplot(gs[0, 0])
ax1 = fig.add_subplot(111)
#ax1 = fig.add_axes([0, 0, 1, 1])
#ax1.patch.set_facecolor((0.8, 0.8, 0.8))
ax1.grid(True)
#ax1.axis('off')

## Disable axis offsets:
#ax1.xaxis.get_major_formatter().set_useOffset(False)
#ax1.yaxis.get_major_formatter().set_useOffset(False)
stride = slice(None, None, 10)
ss = (stride, stride)
#ax1.quiver(xx[ss], yy[ss], im_dx[ss], im_dy[ss])
ax1.quiver(xx[ss], yy[ss], (im_dx - ak_dx)[ss], (im_dy - ak_dy)[ss])

#ax1.plot(kde_pnts, kde_vals)

#blurb = "some text"
#ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes)
#ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes,
#      va='top', ha='left', bbox=dict(facecolor='white', pad=10.0))
#      fontdict={'family':'monospace'}) # fixed-width

#colors = cm.rainbow(np.linspace(0, 1, len(plot_list)))
#for camid, c in zip(plot_list, colors):
#    cam_data = subsets[camid]
#    xvalue = cam_data['CCDATEMP']
#    yvalue = cam_data['PIX_MED']
#    yvalue = cam_data['IMEAN']
#    ax1.scatter(xvalue, yvalue, color=c, lw=0, label=camid)

#mtickpos = [2,5,7]
#ndecades = 1.0   # for symlog, set width of linear portion in units of dex
#nonposx='mask' | nonposx='clip' | nonposy='mask' | nonposy='clip'
#ax1.set_xscale('log', basex=10, nonposx='mask', subsx=mtickpos)
#ax1.set_xscale('log', nonposx='clip', subsx=[3])
#ax1.set_yscale('symlog', basey=10, linthreshy=0.1, linscaley=ndecades)
#ax1.xaxis.set_major_formatter(formatter) # re-format x ticks
#ax1.set_ylim(ax1.get_ylim()[::-1])
#ax1.set_xlabel('whatever', labelpad=30)  # push X label down 

#ax1.set_xticks([1.0, 3.0, 10.0, 30.0, 100.0])
#ax1.set_xticks([1, 2, 3], ['Jan', 'Feb', 'Mar'])
#for label in ax1.get_xticklabels():
#    label.set_rotation(30)
#    label.set_fontsize(14) 

#ax1.xaxis.label.set_fontsize(18)
#ax1.yaxis.label.set_fontsize(18)

#ax1.set_xlim(nice_limits(xvec, pctiles=[1,99], pad=1.2))
#ax1.set_ylim(nice_limits(yvec, pctiles=[1,99], pad=1.2))

#spts = ax1.scatter(x, y, lw=0, s=5)
#cbar = fig.colorbar(spts, orientation='vertical')
#cbar.formatter.set_useOffset(False)
#cbar.update_ticks()

plot_name = 'im_minus_ak.png'
fig.tight_layout() # adjust boundaries sensibly, matplotlib v1.1+
plt.draw()
fig.savefig(plot_name, bbox_inches='tight')

# cyclical colormap ... cmocean.cm.phase
# cmocean: https://matplotlib.org/cmocean/




######################################################################
# CHANGELOG (compare_distortion.py):
#---------------------------------------------------------------------
#
#  2019-09-04:
#     -- Increased __version__ to 0.1.0.
#     -- First created compare_distortion.py.
#
#!/usr/bin/env python # vim: set fileencoding=utf-8 ts=4 sts=4 sw=4 et tw=80 : # # Compare <NAME>' 5th-order solution to what Spitzer currently provides. # # <NAME> # Created:       2019-09-04 # Last modified: 2019-11-11 #-------------------------------------------------------------------------- #************************************************************************** #-------------------------------------------------------------------------- ## Current version: ## Python version-agnostic module reloading: # Python 2.7 # Python 3.4+ # Python 3.0 - 3.3 ## Modules: #import shutil #import resource #import glob #import gc #from numpy.lib.recfunctions import append_fields #import datetime as dt #from dateutil import parser as dtp #import scipy.linalg as sla #import scipy.signal as ssig #import scipy.ndimage as ndi #import scipy.optimize as opti #import scipy.interpolate as stp #import scipy.spatial.distance as ssd #import matplotlib.cm as cm #import matplotlib.ticker as mt #import matplotlib._pylab_helpers as hlp #from matplotlib.colors import LogNorm #from matplotlib import colors #import matplotlib.colors as mplcolors #import matplotlib.gridspec as gridspec #from functools import partial #from collections import OrderedDict #import multiprocessing as mp #np.set_printoptions(suppress=True, linewidth=160) #import pandas as pd #import statsmodels.formula.api as smf #from statsmodels.regression.quantile_regression import QuantReg #import PIL.Image as pli #import seaborn as sns #import cmocean #import theil_sen as ts #import window_filter as wf #import itertools as itt ## Because obviously: #import warnings #if not sys.warnoptions: #    warnings.simplefilter("ignore", category=DeprecationWarning) #    warnings.simplefilter("ignore", category=UserWarning) #    warnings.simplefilter("ignore") #with warnings.catch_warnings(): #    some_risky_activity() #with warnings.catch_warnings(): #    warnings.filterwarnings("ignore", category=DeprecationWarning) #    import problem_child1, problem_child2 ## Attempt Kraus solution import: ##--------------------------------------------------------------------------## ## Projections with cartopy: #try: #    import cartopy.crs as ccrs #    from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter #    from cartopy.feature.nightshade import Nightshade #    #from cartopy import config as cartoconfig #except ImportError: #    sys.stderr.write("Error: cartopy module not found!\n") #    sys.exit(1) ##--------------------------------------------------------------------------## ## Disable buffering on stdout/stderr: ##--------------------------------------------------------------------------## ## Home-brew robust statistics: #try: #    import robust_stats #    reload(robust_stats) #    rs = robust_stats #except ImportError: #    sys.stderr.write("\nError!  robust_stats module not found!\n" #           "Please install and try again ...\n\n") #    sys.exit(1) ## Fast FITS I/O: #try: #    import fitsio #except ImportError: #    sys.stderr.write("\nError: fitsio module not found!\n") #    sys.exit(1) ## FITS I/O: ## Various from astropy: #    import astropy.io.ascii as aia #    import astropy.table as apt #    import astropy.time as astt #    from astropy import coordinates as coord #    from astropy import units as uu ## Star extraction: #try: #    import easy_sep #    reload(easy_sep) #except ImportError: #    sys.stderr.write("Error: easy_sep module not found!\n\n") #    sys.exit(1) #pse = easy_sep.EasySEP() ##--------------------------------------------------------------------------## ## Colors for fancy terminal output: ## Suppress colors in cron jobs: ## Fancy text: ## Dividers: ##--------------------------------------------------------------------------## ## Save FITS image with clobber (astropy / pyfits): # pad header ##--------------------------------------------------------------------------## ## Save FITS image with clobber (fitsio): #def qsave(iname, idata, header=None, **kwargs): #    this_func = sys._getframe().f_code.co_name #    parent_func = sys._getframe(1).f_code.co_name #    sys.stderr.write("Writing to '%s' ... " % iname) #    #if os.path.isfile(iname): #    #    os.remove(iname) #    fitsio.write(iname, idata, clobber=True, header=header, **kwargs) #    sys.stderr.write("done.\n") ##--------------------------------------------------------------------------## ##--------------------------------------------------------------------------## ## Rotation matrix (radians): ## Rotation matrix (degrees): ##--------------------------------------------------------------------------## ##------------------         Parse Command Line             ----------------## ##--------------------------------------------------------------------------## ## Parse arguments and run script: ## Enable raw text AND display of defaults: ## Parse the command line: # ------------------------------------------------------------------ Spitzer coordinate distortion comparison.          Version: %s # ------------------------------------------------------------------ #parser.set_defaults(thing1='value1', thing2='value2') # ------------------------------------------------------------------ #parser.add_argument('firstpos', help='first positional argument') #parser.add_argument('-w', '--whatever', required=False, default=5.0, #        help='some option with default [def: %(default)s]', type=float) #parser.add_argument('-s', '--site', #        help='Site to retrieve data for', required=True) #parser.add_argument('-n', '--number_of_days', default=1, #        help='Number of days of data to retrieve.') #parser.add_argument('-o', '--output_file', #        default='observations.csv', help='Output filename.') #parser.add_argument('remainder', help='other stuff', nargs='*') # ------------------------------------------------------------------ # ------------------------------------------------------------------ #iogroup.add_argument('-o', '--output_file', default=None, required=True, #        help='Output filename', type=str) #iogroup.add_argument('-R', '--ref_image', default=None, required=True, #        help='KELT image with WCS') # ------------------------------------------------------------------ # ------------------------------------------------------------------ # Miscellany: # ------------------------------------------------------------------ ##--------------------------------------------------------------------------## ## Load Spitzer image: ## Check channel: ##--------------------------------------------------------------------------## ## Load WCS from FITS header: ##--------------------------------------------------------------------------## ## Spitzer CCD coordinate grids: #x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered) #y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered) #xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered) #xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered) #xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute #yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute #yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute #yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array) # absolute (pixel) ## Relative coordinates: ##--------------------------------------------------------------------------## ##--------------------------------------------------------------------------## ## Intermediate coords with AK solution: #ak_xx, ak_yy = akp.xform_xy(xx.flatten(), yy.flatten(), irac_channel) #ak_xx, ak_yy = akp.xform_xy(xx, yy, irac_channel) ## Corresponding corrections relative to naive grid: #fiverows_xx = xx[:5] #fiverows_yy = yy[:5] #fr_xr = fiverows_xx - 128. #fr_yr = fiverows_yy - 128. #akx5r, aky5r = akp.xy2focal(fiverows_xx, fiverows_yy, irac_channel) #xpar = akp._a_pars[irac_channel] #ypar = akp._b_pars[irac_channel] # #xresult = np.zeros_like(fr_xr) #xresult += xpar[0] # #yresult = np.zeros_like(fr_xr) #yresult += ypar[0] ##--------------------------------------------------------------------------## ##--------------------------------------------------------------------------## ## Focal plane coordinates from WCS (de-distort using SIP terms): ## Corresponding deltas: ##--------------------------------------------------------------------------## ##--------------------------------------------------------------------------## ##--------------------------------------------------------------------------## ## How to compare dx and dy vs. rotation angle: ##--------------------------------------------------------------------------## ## Quick FITS I/O: #data_file = 'image.fits' #img_vals = pf.getdata(data_file) #hdr_keys = pf.getheader(data_file) #img_vals, hdr_keys = pf.getdata(data_file, header=True) #img_vals, hdr_keys = pf.getdata(data_file, header=True, uint=True) # USHORT #img_vals, hdr_keys = fitsio.read(data_file, header=True) #date_obs = hdr_keys['DATE-OBS'] #site_lat = hdr_keys['LATITUDE'] #site_lon = hdr_keys['LONGITUD'] ## Initialize time: ##img_time = astt.Time(hdr_keys['DATE-OBS'], scale='utc', format='isot') ##img_time += astt.TimeDelta(0.5 * hdr_keys['EXPTIME'], format='sec') ##jd_image = img_time.jd # ### Initialize location: ##observer = ephem.Observer() ##observer.lat = np.radians(site_lat) ##observer.lon = np.radians(site_lon) ##observer.date = img_time.datetime # ##pf.writeto('new.fits', img_vals) ##qsave('new.fits', img_vals) ##qsave('new.fits', img_vals, header=hdr_keys) ## Star extraction: #pse.set_image(img_vals, gain=3.6) #objlist = pse.analyze(sigthresh=5.0) ##--------------------------------------------------------------------------## ## Misc: #def log_10_product(x, pos): #   """The two args are the value and tick position. #   Label ticks with the product of the exponentiation.""" #   return '%.2f' % (x)  # floating-point # #formatter = plt.FuncFormatter(log_10_product) # wrap function for use ## Convenient, percentile-based plot limits: #def nice_limits(vec, pctiles=[1,99], pad=1.2): #    ends = np.percentile(vec[~np.isnan(vec)], pctiles) #    middle = np.average(ends) #    return (middle + pad * (ends - middle)) ## Convenient plot limits for datetime/astropy.Time content: #def nice_time_limits(tvec, buffer=0.05): #    lower = tvec.min() #    upper = tvec.max() #    ndays = upper - lower #    return ((lower - 0.05*ndays).datetime, (upper + 0.05*ndays).datetime) ## Convenient limits for datetime objects: #def dt_limits(vec, pad=0.1): #    tstart, tstop = vec.min(), vec.max() #    trange = (tstop - tstart).total_seconds() #    tpad = dt.timedelta(seconds=pad*trange) #    return (tstart - tpad, tstop + tpad) ##--------------------------------------------------------------------------## ## Solve prep: #ny, nx = img_vals.shape #x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered) #y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered) #xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered) #xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered) #xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute #yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute #yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute #yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array) #yy, xx = np.mgrid[1:ny+1, 1:nx+1].astype('uint16')   # absolute (pixel) ## 1-D vectors: #x_pix, y_pix, ivals = xx.flatten(), yy.flatten(), img_vals.flatten() #w_vec = np.ones_like(ivals)            # start with uniform weights #design_matrix = np.column_stack((np.ones(x_pix.size), x_pix, y_pix)) ## Image fitting (statsmodels etc.): #data = sm.datasets.stackloss.load() #ols_res = sm.OLS(ivals, design_matrix).fit() #rlm_res = sm.RLM(ivals, design_matrix).fit() #rlm_model = sm.RLM(ivals, design_matrix, M=sm.robust.norms.HuberT()) #rlm_res = rlm_model.fit() #data = pd.DataFrame({'xpix':x_pix, 'ypix':y_pix}) #rlm_model = sm.RLM.from_formula("ivals ~ xpix + ypix", data) ##--------------------------------------------------------------------------## ##------------------         Fit for im-ak Difference       ----------------## ##--------------------------------------------------------------------------## ### 1-D vectors: #xx_1d = xx.flatten() #yy_1d = yy.flatten() #xdiff = (im_dx - ak_dx).flatten() #ydiff = (im_dy - ak_dy).flatten() # #design_matrix = np.column_stack((np.ones_like(xdiff), xx_1d, yy_1d)) #dx_ols = sm.OLS(xdiff, design_matrix).fit() #dy_ols = sm.OLS(ydiff, design_matrix).fit() ##--------------------------------------------------------------------------## ## Theil-Sen line-fitting (linear): #model = ts.linefit(xvals, yvals) #icept, slope = ts.linefit(xvals, yvals) ## Theil-Sen line-fitting (loglog): #xvals, yvals = np.log10(original_xvals), np.log10(original_yvals) #xvals, yvals = np.log10(df['x'].values), np.log10(df['y'].values) #llmodel = ts.linefit(np.log10(xvals), np.log10(yvals)) #icept, slope = ts.linefit(xvals, yvals) #fit_exponent = slope #fit_multiplier = 10**icept #bestfit_x = np.arange(5000) #bestfit_y = fit_multiplier * bestfit_x**fit_exponent ## Log-log evaluator: #def loglog_eval(xvals, model): #    icept, slope = model #    return 10**icept * xvals**slope #def loglog_eval(xvals, icept, slope): #    return 10**icept * xvals**slope ##--------------------------------------------------------------------------## ## KDE: #kde_pnts, kde_vals = mk.go(data_vec) ##--------------------------------------------------------------------------## ## Vaex plotting: #ds = vaex.open('big_file.hdf5') #ds = vaex.from_arrays(x=x, y=y)     # load from arrays #ds = vaex.from_csv('mydata.csv') ## Stats: #ds.mean("x"), ds.std("x"), ds.correlation("vx**2+vy**2+vz**2", "E") #ds.plot(....) #http://vaex.astro.rug.nl/latest/tutorial_ipython_notebook.html ##--------------------------------------------------------------------------## ## Plot config: # gridspec examples: # https://matplotlib.org/users/gridspec.html #gs1 = gridspec.GridSpec(4, 4) #gs1.update(wspace=0.025, hspace=0.05)  # set axis spacing #ax1 = plt.subplot2grid((3, 3), (0, 0), colspan=3) # top-left + center + right #ax2 = plt.subplot2grid((3, 3), (1, 0), colspan=2) # mid-left + mid-center #ax3 = plt.subplot2grid((3, 3), (1, 2), rowspan=2) # mid-right + bot-right #ax4 = plt.subplot2grid((3, 3), (2, 0))            # bot-left #ax5 = plt.subplot2grid((3, 3), (2, 1))            # bot-center ##--------------------------------------------------------------------------## #plt.style.use('bmh')   # Bayesian Methods for Hackers style #fig, axs = plt.subplots(2, 2, sharex=True, figsize=fig_dims, num=1) # sharex='col' | sharex='row' #fig.frameon = False # disable figure frame drawing #fig.subplots_adjust(left=0.07, right=0.95) #ax1 = plt.subplot(gs[0, 0]) #ax1 = fig.add_axes([0, 0, 1, 1]) #ax1.patch.set_facecolor((0.8, 0.8, 0.8)) #ax1.axis('off') ## Disable axis offsets: #ax1.xaxis.get_major_formatter().set_useOffset(False) #ax1.yaxis.get_major_formatter().set_useOffset(False) #ax1.quiver(xx[ss], yy[ss], im_dx[ss], im_dy[ss]) #ax1.plot(kde_pnts, kde_vals) #blurb = "some text" #ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes) #ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes, #      va='top', ha='left', bbox=dict(facecolor='white', pad=10.0)) #      fontdict={'family':'monospace'}) # fixed-width #colors = cm.rainbow(np.linspace(0, 1, len(plot_list))) #for camid, c in zip(plot_list, colors): #    cam_data = subsets[camid] #    xvalue = cam_data['CCDATEMP'] #    yvalue = cam_data['PIX_MED'] #    yvalue = cam_data['IMEAN'] #    ax1.scatter(xvalue, yvalue, color=c, lw=0, label=camid) #mtickpos = [2,5,7] #ndecades = 1.0   # for symlog, set width of linear portion in units of dex #nonposx='mask' | nonposx='clip' | nonposy='mask' | nonposy='clip' #ax1.set_xscale('log', basex=10, nonposx='mask', subsx=mtickpos) #ax1.set_xscale('log', nonposx='clip', subsx=[3]) #ax1.set_yscale('symlog', basey=10, linthreshy=0.1, linscaley=ndecades) #ax1.xaxis.set_major_formatter(formatter) # re-format x ticks #ax1.set_ylim(ax1.get_ylim()[::-1]) #ax1.set_xlabel('whatever', labelpad=30)  # push X label down #ax1.set_xticks([1.0, 3.0, 10.0, 30.0, 100.0]) #ax1.set_xticks([1, 2, 3], ['Jan', 'Feb', 'Mar']) #for label in ax1.get_xticklabels(): #    label.set_rotation(30) #    label.set_fontsize(14) #ax1.xaxis.label.set_fontsize(18) #ax1.yaxis.label.set_fontsize(18) #ax1.set_xlim(nice_limits(xvec, pctiles=[1,99], pad=1.2)) #ax1.set_ylim(nice_limits(yvec, pctiles=[1,99], pad=1.2)) #spts = ax1.scatter(x, y, lw=0, s=5) #cbar = fig.colorbar(spts, orientation='vertical') #cbar.formatter.set_useOffset(False) #cbar.update_ticks() # adjust boundaries sensibly, matplotlib v1.1+ # cyclical colormap ... cmocean.cm.phase # cmocean: https://matplotlib.org/cmocean/ ###################################################################### # CHANGELOG (compare_distortion.py): #--------------------------------------------------------------------- # #  2019-09-04: #     -- Increased __version__ to 0.1.0. #     -- First created compare_distortion.py. #