# -*- coding: utf-8 -*-
"""
This program traces light rays of different colors through a series of 
prismatic glass objects.

For a number of optical configurations of an atmoshperic dispersion corrector,
the main program defines the parameters of the glass elements and their
orientations which produce the maximal and minimal effect, respectively.
For each configuration and orientation, the program plots the paths for red,
green and blue light through the glass objects.

Created on Mon Oct 14 06:48:38 2013

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

import numpy
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

def normalize (vec):
    """Scale a 2D vector to unit length and return the scaled vector.
    
    Arguments:
    vec        -- list with x and y coordinates of input vector
    
    """
    [dx,dy] = vec
    return [dx/numpy.sqrt(dx**2+dy**2), dy/numpy.sqrt(dx**2+dy**2)]
    
def ortho (vec):
    """Compute and return the orthogonal vector for a given 2D vector.
    
    Arguments:
    vec        -- list with x and y coordinates of input vector
    
    """
    [dx, dy] = vec
    # The vector should "face down", i.e., have a negative y component.
    if dx <= 0.:
        return [-dy, dx]
    else:
        return [dy, -dx]
        
def reverse (vec):
    """Compute and return the reversed vector for a given 2D vector.
    
    Arguments:
    vec        -- list with x and y coordinates of input vector
    
    """
    return [-1.*vec[0], -1.*vec[1]]
        
def angle (vec1, vec2):
    """Compute and return the angle between two 2D vectors.
    
    Arguments:
    vec1       -- list with x and y coordinates of first input vector
    vec2       -- list with x and y coordinates of second input vector
    
    """
    [x1, y1] = normalize (vec1)
    [x2, y2] = normalize (vec2)
    product = x1*x2+y1*y2
    # Because of rounding errors the scalar product might be out of range.
    product = numpy.min([product, 1.])
    product = numpy.max([product, -1.])
    return numpy.arccos(product)
    
def rotate (vec1, alpha):
    """Rotate a vector by a given angle and return the rotated vector.
    
    Arguments:
    vec1       -- list with x and y coordinates of input vector
    alpha      -- angle (in radian)
    
    """
    [x,y] = vec1
    return [x*numpy.cos(alpha)-y*numpy.sin(alpha),
            x*numpy.sin(alpha)+y*numpy.cos(alpha)]

def intersect (pos1, dir1, pos2, dir2):
    """Compute and return the intersection point between two straight lines.
    
    Arguments:
    pos1       -- list with x and y coordinates of a point on the first line
    dir1       -- list with x and y components of a vector parallel to the
                  first line
    pos2       -- list with x and y coordinates of a point on the second line
    dir2       -- list with x and y components of a vector parallel to the
                  second line
    
    Return value: [[x, y], [s, t]] with
    [x, y]:    -- x and y coordinates of the intersection point
    [s, t]:    -- scaling parameters, such that
                  [x, y] = pos1 + t*dir1 and [x, y] = pos2 + s*dir2    
    
    """
    [x1, y1] = pos1
    [dx1, dy1] = dir1   
    [x2, y2] = pos2
    [dx2, dy2] = dir2
    s=((x2-x1)*dy1 - (y2-y1)*dx1)/(dx1*dy2-dy1*dx2)
    # dx1 or dy1 may be zero, but not both of them.
    if abs(dx1) > 1.e-10:
        t=(x2-x1+dx2*s)/dx1
    else:
        t=(y2-y1+dy2*s)/dy1
    x3 = x1+dx1*t
    y3 = y2+dy2*s
    return [[x3, y3], [s, t]]
    
def refract (pos_in, dir_in, interfaces, i):
    """Compute and return refracted ray section behind an air/glas interface.
    
    Arguments:
    pos_in     -- list with x and y coordinates of a point on the
                  incoming light ray
    dir_in     -- list with x and y components of a vector parallel to the
                  incoming light ray
    interfaces -- list of air/glas interfaces. Each element is a list of the
                  form [pos, dir, n], where
                  pos = list with x and y coordinates of a point on the
                        interface
                  dir = list with x and y components of a vector parallel to
                        the interface
                  n   = refractive index of material behind the interface
    i          -- index of the interface for which the refraction is to be
                  computed
    
    Return value: [intersect_point, dir_out] with
    [x, y]:    -- x and y coordinates of the intersection point
    [s, t]:    -- scaling parameters, such that
                  [x, y] = pos1 + t*dir1 and [x, y] = pos2 + s*dir2    
    
    """
    if i>0:
        [pos1, dir1, n1] = interfaces[i-1]
    else:
        # Special case i=0: Upper border of plot window, no refraction.
        n1 = 1.
    
    [pos2, dir2, n2] = interfaces[i]
    # Compute the intersection point between incoming ray and interface.
    [intersect_point, [s, t]] = intersect (pos_in, dir_in, pos2, dir2)
    if i>=len(interfaces)-1:
        # Lower border of plot window, no refracting interface.
        dir_out = dir_in
    else:
        # Compute the vector perpendicular to refracting interface.
        normal_vec = ortho (dir2)
        # Compute the angle between incoming ray and normal vector.
        incidence_angle = angle (dir_in, normal_vec)
        # Compute the emergent angle according to Snellius' law.
        emergent_angle  = numpy.arcsin (numpy.sin(incidence_angle)*n1/n2)
        dir_out_1 = rotate (normal_vec, emergent_angle)
        dir_out_2 = rotate (normal_vec, -1.*emergent_angle)
        # Choose the correct direction of rotation for emergent ray.
        if angle (reverse(dir_in), dir_out_1) < angle (reverse(dir_in), \
                                                       dir_out_2):
            dir_out = dir_out_2
        else:
            dir_out = dir_out_1
    return [intersect_point, dir_out]
    
def trace (pos_in, dir_in, interfaces):
    """Compute and return the path of a ray through a series of interfaces.
    
    Arguments:
    pos_in     -- list with x and y coordinates of a point on the
                  incoming light ray
    dir_in     -- list with x and y components of a vector parallel to the
                  incoming light ray
    interfaces -- list of air/glas interfaces. Each element is a list of the
                  form [pos, dir, n], where
                  pos = list with x and y coordinates of a point on the
                        interface
                  dir = list with x and y components of a vector parallel to
                        the interface
                  n   = refractive index of material behind the interface
                  The first and last interfaces are the upper and lower
                  borders of the plot window, respectively.
        
    Return value: [[xi, yi]] with
    [xi, yi]:  -- x and y coordinates of the intersection point of the ray
                  with the i-th interface
        
    """
    path = []
    for i in range(len(interfaces)):
        [pos_in, dir_in] = refract (pos_in, dir_in, interfaces, i)
        path.append(pos_in)
    return path
    
def compute_rindex (wavelength):
    """Compute and return the refractive index of glas for a given wavelength.
    
    Arguments:
    wavelength -- wavelength in nanometers
    
    """
    # nr and nh are the refractive indices for the r and h lines, respectively.
    # for BK7 glass, nr = 1.51289 and nh = 1.53024
    # Extremely increase the dispersion to exaggerate the effect.
    nr = 1.4
    nh = 1.6
    # Compute the refractive index by linear interpolation.
    n_glass = nh+(wavelength-404.6561)/(706.5188-404.6561)*(nr-nh)
    return n_glass
           
def create_plot (plot, interfaces, polygons, paths, colors, title):
    """Draw a sub-plot for a single configuration.
    
    Arguments:
    plot       -- sub-plot object to draw on
    interfaces -- list of air/glass interfaces
    polygons   -- polygons around glas objects in the light path
    paths      -- paths of light rays
    colors     -- color of each light ray
    title      -- character string with sub-plot title
    
    """
    for poly in polygons:
        plot.add_patch(poly)
            
    for [path, color] in zip(paths, colors):
        # Extract x and y coordinates of light path points.
        x = [i[0] for i in path]
        y = [i[1] for i in path]
        # Plot a single light ray and set its color.
        lines = plot.plot(x, y)
        plt.setp(lines[0], color=color)
    # Draw the sub-plot title.
    plot.set_title(title)
    # Remove tick marks and numbers from x and y axes.
    nulllocator = plt.NullLocator()
    plot.xaxis.set_major_locator(nulllocator)
    plot.yaxis.set_major_locator(nulllocator)
 
##############################################################################   
# Main Program: Set the configuration and produce a figure.
#
# Two Prisms with non-zero dispersion are oriented such that the entrance
# and exit windows are perpendicular to the optical axis.
#
# Size of the 2D plot in inches
figsize = (10,6)
fig = plt.figure(figsize=figsize)
# The figure contains two plots: maximum dispersion correction (left)
#                                and minimum correction (right).
plot_left = fig.add_subplot(121)

# Create polygons around the glass objects.
polygons = []
verts = [(0., -3.), (10.,-3.), (10.,0.), (0.,-2.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))
verts = [(0., 3.), (10.,3.), (10.,0.), (0.,2.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))

# Three light rays of different color are traced.
wavelengths = [400., 550., 700.]
# Adjust the positions of incoming rays such that the outgoing rays converge
# in a single point.
pos_ins = [[5.011013998, 10.], [5.392, 10.], [5.765638871, 10.]]
# All incoming rays are in y-direction.
dir_ins = [[0., -1.], [0., -1.], [0., -1.]]
colors = ['b', 'g', 'r']

# Define the air/glass interfaces and use the "trace" function to compute the
# path of the light ray.
paths = []
for [wavelength, pos_in, dir_in, color] \
    in zip(wavelengths, pos_ins, dir_ins, colors):
    n_glass = compute_rindex(wavelength)
    interfaces = []
    interfaces.append ([[0., 10.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,  3.], normalize([ 1., 0.]), n_glass])
    interfaces.append ([[0.,  2.], normalize([10.,-2.]), 1. ])
    interfaces.append ([[0., -2.], normalize([10., 2.]), n_glass])
    interfaces.append ([[0., -3.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,-10.], normalize([ 1., 0.]), 1. ])
    paths.append(trace (pos_in, dir_in, interfaces))
#    print "wavelength=", wavelength, ", x(-10.)=", paths[-1][-1][0]

# Plot the glass objects and the paths of all light rays.
create_plot (plot_left, interfaces, polygons, paths, colors, \
             "Maximale Wirkung")

# Repeat the computations above for a configuration where the prisms are set
# to minimal effect, and draw the results in right-hand sub-plot.
plot_right = fig.add_subplot(122)
polygons = []
verts = [(0., -3.), (10.,-3.), (10.,-2.), (0.,0.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))
verts = [(0., 3.), (10.,3.), (10.,0.), (0.,2.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))

paths = []
for [wavelength, pos_in, dir_in, color] \
    in zip(wavelengths, pos_ins, dir_ins, colors):
    n_glass = compute_rindex(wavelength)
    interfaces = []
    interfaces.append ([[0., 10.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,  3.], normalize([ 1., 0.]), n_glass])
    interfaces.append ([[0.,  2.], normalize([10.,-2.]), 1. ])
    interfaces.append ([[0.,  0.], normalize([10.,-2.]), n_glass])
    interfaces.append ([[0., -3.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,-10.], normalize([ 1., 0.]), 1. ])
    paths.append(trace (pos_in, dir_in, interfaces))

create_plot (plot_right, interfaces, polygons, paths, colors, \
             "Minimale Wirkung")

# Save the complete figure into a file.            
plt.savefig("Dispersions_Korrektur_Prismen-getrennt.png", dpi=150)

##############################################################################
# Start a new figure for a different optical configuration.
#
# Two Prisms with non-zero dispersion are in cotact with each other. The
# separating plane is perpendicular to the optical axis. The entrance
# and exit windows are skewed with respect to the optical axis.
#
# Size of the 2D plot in inches
figsize = (10,6)
fig = plt.figure(figsize=figsize)
# The figure contains two plots: maximum dispersion correction (left)
#                                and minimum correction (right).
plot_left = fig.add_subplot(121)

# Create polygons around the glass objects.
polygons = []
verts = [(0., -1.), (10.,-3.), (10.,0.), (0.,0.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))
verts = [(0., 1.), (10.,3.), (10.,0.), (0.,0.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))

# Three light rays of different color are traced.
wavelengths = [400., 550., 700.]
# Adjust the positions of incoming rays such that the outgoing rays converge
# in a single point.
pos_ins = [[5., 10.], [5.392, 10.], [5.77331665, 10.]]
# All incoming rays are in y-direction.
dir_ins = [[0., -1.], [0., -1.], [0., -1.]]
colors = ['b', 'g', 'r']

# Define the air/glass interfaces and use the "trace" function to compute the
# path of the light ray.
paths = []
for [wavelength, pos_in, dir_in, color] \
    in zip(wavelengths, pos_ins, dir_ins, colors):
    n_glass = compute_rindex(wavelength)
    interfaces = []
    interfaces.append ([[0., 10.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,  1.], normalize([10., 2.]), n_glass])
    interfaces.append ([[0., -1.], normalize([10.,-2.]), 1. ])
    interfaces.append ([[0.,-10.], normalize([ 1., 0.]), 1. ])
    paths.append(trace (pos_in, dir_in, interfaces))
#    print "wavelength=", wavelength, ", x(-10.)=", paths[-1][-1][0]

# Plot the glass objects and the paths of all light rays.
create_plot (plot_left, interfaces, polygons, paths, colors, \
             "Maximale Wirkung")

# Repeat the computations above for a configuration where the prisms are set
# to minimal effect, and draw the results in right-hand sub-plot.
plot_right = fig.add_subplot(122)
polygons = []
verts = [(0., -3.), (10.,-1.), (10.,0.), (0.,0.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))
verts = [(0., 1.), (10.,3.), (10.,0.), (0.,0.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))

paths = []
for [wavelength, pos_in, dir_in, color] \
    in zip(wavelengths, pos_ins, dir_ins, colors):
    n_glass = compute_rindex(wavelength)
    interfaces = []
    interfaces.append ([[0., 10.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,  1.], normalize([10., 2.]), n_glass])
    interfaces.append ([[0., -3.], normalize([10., 2.]), 1. ])
    interfaces.append ([[0.,-10.], normalize([ 1., 0.]), 1. ])
    paths.append(trace (pos_in, dir_in, interfaces))

create_plot (plot_right, interfaces, polygons, paths, colors, \
             "Minimale Wirkung")

# Save the complete figure into a file.            
plt.savefig("Dispersions_Korrektur_Prismen-in-Kontakt.png", dpi=150)

##############################################################################
# Start a new figure for a different optical configuration.
#
# Two identical compound objects are located above and below the x axis,
# respectively. Their outer faces are perpendicular to the optical axis.
# Each of them is in turn made from two cemented prismatic slabs, the outer
# one having non-zero dispersion, while the inner one has zero dispersion.
#
# Size of the 2D plot in inches
figsize = (10,6)
fig = plt.figure(figsize=figsize)
# The figure contains two plots: maximum dispersion correction (left)
#                                and minimum correction (right).
plot_left = fig.add_subplot(121)

# Create polygons around the glass objects.
polygons = []
verts = [(0., 4.), (10., 4.), (10., 1.), (0., 3.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))
verts = [(0., 0.), (10., 0.), (10., 1.), (0., 3.)]
polygons.append(Polygon(verts, facecolor='#87CEFA', edgecolor='w'))
verts = [(0., 0.), (10., 0.), (10.,-1.), (0.,-3.)]
polygons.append(Polygon(verts, facecolor='#87CEFA', edgecolor='w'))
verts = [(0.,-4.), (10.,-4.), (10.,-1.), (0.,-3.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))

# Three light rays of different color are traced.
wavelengths = [400., 550., 700.]
# Adjust the positions of incoming rays such that the outgoing rays converge
# in a single point.
pos_ins = [[5.050918772, 10.], [5.392, 10.], [5.740686883, 10.]]
# All incoming rays are in y-direction.
dir_ins = [[0., -1.], [0., -1.], [0., -1.]]
colors = ['b', 'g', 'r']

# Define the air/glass interfaces and use the "trace" function to compute the
# path of the light ray.
paths = []
for [wavelength, pos_in, dir_in, color] \
    in zip(wavelengths, pos_ins, dir_ins, colors):
    n_glass = compute_rindex(wavelength)
    interfaces = []
    interfaces.append ([[0., 10.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,  4.], normalize([ 1., 0.]), n_glass])
    interfaces.append ([[0.,  3.], normalize([10.,-2.]), 1.5])
    interfaces.append ([[0., -3.], normalize([10., 2.]), n_glass])
    interfaces.append ([[0., -4.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,-10.], normalize([ 1., 0.]), 1. ])
    paths.append(trace (pos_in, dir_in, interfaces))
#    print "wavelength=", wavelength, ", x(-10.)=", paths[-1][-1][0]

# Plot the glass objects and the paths of all light rays.
create_plot (plot_left, interfaces, polygons, paths, colors, \
             "Maximale Wirkung")

# Repeat the computations above for a configuration where the prisms are set
# to minimal effect, and draw the results in right-hand sub-plot.
plot_right = fig.add_subplot(122)
polygons = []
verts = [(0., 4.), (10., 4.), (10., 1.), (0., 3.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))
verts = [(0., 0.), (10., 0.), (10., 1.), (0., 3.)]
polygons.append(Polygon(verts, facecolor='#87CEFA', edgecolor='w'))
verts = [(0., 0.), (10., 0.), (10.,-3.), (0.,-1.)]
polygons.append(Polygon(verts, facecolor='#87CEFA', edgecolor='w'))
verts = [(0.,-4.), (10.,-4.), (10.,-3.), (0.,-1.)]
polygons.append(Polygon(verts, facecolor='#B0E0E6', edgecolor='w'))

paths = []
for [wavelength, pos_in, dir_in, color] \
    in zip(wavelengths, pos_ins, dir_ins, colors):
    n_glass = compute_rindex(wavelength)
    interfaces = []
    interfaces.append ([[0., 10.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,  4.], normalize([ 1., 0.]), n_glass])
    interfaces.append ([[0.,  3.], normalize([10.,-2.]), 1.5])
    interfaces.append ([[0., -1.], normalize([10.,-2.]), n_glass])
    interfaces.append ([[0., -4.], normalize([ 1., 0.]), 1. ])
    interfaces.append ([[0.,-10.], normalize([ 1., 0.]), 1. ])
    paths.append(trace (pos_in, dir_in, interfaces))
    
create_plot (plot_right, interfaces, polygons, paths, colors, \
             "Minimale Wirkung")

# Save the complete figure into a file.            
plt.savefig("Dispersions_Korrektur_Achromatisch.png", dpi=150)