# -*- coding: utf-8 -*-
"""
This program computes the effect of atmospheric dispersion as a function of
zenith distance for various wavelength intervals of the visible spectrum.
For the wavelength-dependent computation of atmospheric refraction it uses
the high-precision function sla_refro from the SLALIB package, which is based
on numerical integration along the line of sight through a model atmosphere.

For each wavelength interval a matplotlib graph plots the atmospheric
dispersion as a function of zenith distance in comparison to the angular
resolution of a perfect telescope of a given aperture. The graph is written
to a file in the working directory.

Created on Tue Oct  8 06:27:43 2013

@author: Rolf Hempel, German Aerospace Center (DLR)
"""

import numpy
import matplotlib.pyplot as plt
# pySLALIB is a Python-wrapped version of P.T. Wallace's SLALIB positional
# astronomy library, provided by Scott M. Ransom. Version v1.0.2  (Dec 2010)
# was usded for all computaions.
from pyslalib import slalib
# Conversion factor between degrees and radian
conv = 180./numpy.pi

def refraction (ZD, wavelength):
    """Set model atmosphere parameters and return refraction in arcseconds.
    
    Arguments:
    ZD         -- zenith distance in degrees
    wavelength -- wavelength in nanometers
    
    """
    # Conversion of zenith distance to radian
    ZOBS = ZD/conv
    # Height above sea level in meters   
    HM = 50.
    # Temperature in Kelvin: Celsius + 273.15 (standard: 288.15)
    TDK = 288.15
    # Atmospheric pressure in millibar (standard: 1013.25)   
    PMB = 1013.25
    # Relative humidity (between 0. and 1.)    
    RH = 0.5
    # The input wavelength is given in nanometers (standard: 575.3).
    # Conversion to microns.
    WL = wavelength / 1000.
    # Geographical location of site (59.771667 for Pulkovo)
    PHI = 51./conv
    # Temperature lapse rate in K/meter (suggested: 0.0065)    
    TLR = 0.0065
    # Targeted precision for iteration (in radian)    
    EPS = 1.E-8
    # sla_refro computes the refraction for the APPARENT zenith distance ZOBS.
    return slalib.sla_refro(ZOBS, HM, TDK, PMB, RH, WL, PHI, TLR, EPS) \
            *conv*3600.

def create_plot (filename):
    """Compute dispersion vs. zenith distance and write plot to file.
    
    Arguments:
    filename       -- name of file to be written to the working directory
    
    """
    # For the luminance channel: Set the range and step size of
    #                            zenith distance interval.
    x_l = numpy.arange(0., 84., 0.5)
    # Set the upper and lower bounds of the luminance wavelength interval.
    wavelength_ub_luminance = 685.
    wavelength_lb_luminance = 390.
    # Compute the atmospheric dispersion as the difference between refraction
    # at upper and lower limits of the luminance wavelength interval.
    y_l = [(refraction(ZD, wavelength_lb_luminance) \
          - refraction(ZD, wavelength_ub_luminance)) for ZD in x_l]
    # Repeat the same for all other wavelength intervals.
    x_b = numpy.arange(0., 86.5, 0.5)
    wavelength_ub_blue = 510.
    wavelength_lb_blue = 390.
    y_b = [(refraction(ZD, wavelength_lb_blue) \
          - refraction(ZD, wavelength_ub_blue)) for ZD in x_b]
    x_g = numpy.arange(0., 89., 0.5)
    wavelength_ub_green = 575.
    wavelength_lb_green = 495.
    y_g = [(refraction(ZD, wavelength_lb_green) \
          - refraction(ZD, wavelength_ub_green)) for ZD in x_g]
    x_r = numpy.arange(0., 89., 0.5)
    wavelength_ub_red = 685.
    wavelength_lb_red = 580.
    y_r = [(refraction(ZD, wavelength_lb_red) \
          - refraction(ZD, wavelength_ub_red)) for ZD in x_r]
    # Set the aperture of the model telescope in nanometers.
    aperture = 130.E6
    # Compute the wavelength at the middle of the luminance interval.
    wavelength = (wavelength_ub_luminance+wavelength_lb_luminance)/2.
    # Resolution (arcsec) of the model telescope using the Rayleigh criterion.
    resolution = 1.22*wavelength/aperture*conv*3600.
    x_res = numpy.arange(0., 90., 0.5)
    res = [resolution for ZD in x_res]   
    # Size of the 2D plot in inches
    figsize = (8,5)
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111)
    # Plot dispersion curves and straight line with telescope resolution.
    lines = ax.plot(x_l, y_l,
                    x_b, y_b,
                    x_g, y_g,
                    x_r, y_r,
                    x_res, res)
    # Select the luminance dispersion curve and set its color.
    disp_line_l = lines[0]
    plt.setp(disp_line_l, color='k')
    # Select the blue dispersion curve and set its color.
    disp_line_b = lines[1]
    plt.setp(disp_line_b, color='b')
    # Select the green dispersion curve and set its color.
    disp_line_g = lines[2]
    plt.setp(disp_line_g, color='g')
    # Select the red dispersion curve and set its color.
    disp_line_r = lines[3]
    plt.setp(disp_line_r, color='r')
    # Plot the telescope resolution as a black dashed line.
    res_line = lines[4]
    plt.setp(res_line, color='k', linestyle='--')
    # Add coordinate axis annotations.
    ax.set_xlabel('Zenitdistanz (Grad)')
    ax.set_ylabel('Dispersion (")')
    # Add a legend
    ax.legend((u"Luminanz-Kanal (390-685nm)",
               u"Blau-Kanal (390-510nm)",
               u"Grün-Kanal (495-575nm)",
               u"Rot-Kanal (580-685nm)",
               u"Auflösung eines 130mm-Refraktors"), loc='upper left')
    ax.grid(True)
    # Save the figure to a local file.
    plt.savefig(filename, dpi=150)


# Main program: For four wavelength intervals, plot the atmospheric
# dispersion as a function of zenith distance.
create_plot ('dispersion_combined.png')