 
#######################################################################
#
# Program:      Plotting for trap simulations
#
# Author:       Nathanael Lampe, School of Physics, Monash University
#
# Created:      December 4, 2012
#
# Changelog:    December 4, 2012: Created from twodtrap.py
#           
#               
#
# Purpose:      Contain functions that will plot trap data from HDF5 files
#               
########################################################################

#Imports
from __future__ import division

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.cm as cm

from scripts.trapsave import *


def nicecmap():
  '''
  Nice colormap that looks good in greyscale
  From: Rappaport 2002, doi:10.1109/MAP.2002.1028735
  '''
  cdict = {'red' : ((0.000, 0.00, 0.00),
                    (0.125, 0.15, 0.15),
                    (0.250, 0.30, 0.30),
                    (0.375, 0.60, 0.60),
                    (0.500, 1.00, 1.00),
                    (0.625, 0.90, 0.90),
                    (0.750, 0.90, 0.90),
                    (0.875, 0.90, 0.90),
                    (1.000, 1.00, 1.00)),
         
         'green' : ((0.000, 0.00, 0.00),
                    (0.125, 0.15, 0.15),
                    (0.250, 0.15, 0.15),
                    (0.375, 0.20, 0.20),
                    (0.500, 0.25, 0.25),
                    (0.625, 0.50, 0.50),
                    (0.750, 0.75, 0.75),
                    (0.875, 0.90, 0.90),
                    (1.000, 1.00, 1.00)),
         
         'blue'  : ((0.000, 0.00, 0.00),
                    (0.125, 0.50, 0.50),
                    (0.250, 0.75, 0.75),
                    (0.375, 0.50, 0.50),
                    (0.500, 0.15, 0.15),
                    (0.625, 0.00, 0.00),
                    (0.750, 0.10, 0.10),
                    (0.875, 0.50, 0.50),
                    (1.000, 1.00, 1.00))
         }
  return matplotlib.colors.LinearSegmentedColormap('rgb-grey_cmap',cdict,256)
  
def twodani(filename):
  '''
  Plot the simulation generated in the file filename.
  Filename is a binary hd5 file generated by twodsave.
  '''
  
  #Setup------------------------------------------------------------------------
  infile = h5file( filename, erase = False, read = True )
  fig = plt.figure()
  
  steps = infile.head[ 'steps' ]
  npt   = infile.head[ 'npt' ]
  
  xvals, yvals = infile.readxy() #read x-y array
  
  p = infile.readpsi('0.0') #read initial state
  
  #Plot probability density-----------------------------------------------------
  
  ax = fig.add_subplot(221, aspect = 'equal')
  
  plt.xlim( [ -5, 5 ] )
  plt.ylim( [ -5, 5 ] )
  
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  im = plt.imshow( abs( p ) ** 2 / ( abs( p ) ** 2 ).max(),
                  extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                  vmin = 0., vmax = 1., cmap = nicecmap() ) #(abs(p)**2).max())
 
  cbar = plt.colorbar(im, ticks = np.linspace( 0, 1, 5 ) )
  
  #if not autoscaling density:
  #v = np.linspace(0, ( abs( p ) ** 2. ).max() , 10)
  #for jj in range(0,len(v)): v[jj] = round(v[jj],3)
  #cbar = plt.colorbar(ticks = v)
  
  #Phase------------------------------------------------------------------------
  
  ax2 = fig.add_subplot( 222, aspect = 'equal' )
  plt.xlim( [ -5, 5 ] )
  plt.ylim( [ -5, 5 ] )
 
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  pi = np.pi
  im2 = plt.imshow( np.angle( p ),
                    extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                    vmin = -np.pi, vmax = np.pi, cmap=cm.gist_rainbow ) 
  cbar2 = plt.colorbar(im2,
                       ticks = [round(j,2) for j in np.linspace(-pi,pi,5)]
                      )
  
  #--------------------------------------------------------
  
  #x = 0.
  ax3 = fig.add_subplot( 223 )
  plt.xlim( [-10,10] )
  plt.ylim( [0,1] )
 
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'$|\psi|^2$ (atoms $a_{0}^{-1}$)')
  
  line1, = ax3.plot( xvals, ( abs( p ) ** 2. )[ npt / 2., : ] / 
                          ( abs( p ) ** 2. )[ npt / 2., : ].max() )
    
  #--------------------------------------------------------
  
  #y = 0.
  
  ax4 = fig.add_subplot(224)
  plt.xlim( [-10,10] )
  plt.ylim( [0, 1] )
  
  plt.xlabel( r'y ($a_0$)' )
  plt.ylabel( r'$|\psi|^2$ (atoms $a_{0}^{-1}$)' )
  
  line2, = ax4.plot( yvals, ( abs( p ) **2. )[ : , npt / 2. ] / 
                          ( abs( p ) **2. )[ : , npt / 2.].max() )
  
  #Update the plot--------------------------------------------------------------
  def update_fig(i):
    p = infile.readpsi( str(i) )
    
    im.set_array( abs( p ) ** 2 / (abs(p)**2).max() )
    im2.set_array( np.angle( p ) )
    line1.set_ydata( ( abs( p ) ** 2. )[ npt / 2., : ] / 
                     ( abs( p ) ** 2. )[ npt / 2., : ].max() )
    line2.set_ydata( ( abs( p ) ** 2. )[ :, npt / 2. ] / 
                     ( abs( p ) ** 2. )[ :, npt / 2. ].max() )
    
    return im, im2, line1, line2

  ani = animation.FuncAnimation( fig, update_fig,
                                 np.arange( 1, steps + 1., dtype=float ),
                                 interval=20., repeat_delay = 3000., 
                                 blit = False )
  plt.show()
  return ani
  infile.f.close()
  

def aniprob(filename):
  '''
  Plot the simulation generated in the file filename, plotting only 
  the probability density
  Filename is a binary hd5 file generated by twodsave.
  '''
  
  #Setup------------------------------------------------------------------------
  infile = h5file( filename, erase = False, read = True )
  fig = plt.figure()
  matplotlib.rcParams.update({'font.size':16}) #set inital font size
  
  steps = infile.head[ 'steps' ]
  npt   = infile.head[ 'npt' ]
  
  xvals, yvals = infile.readxy() #read x-y array
  
  p = infile.readpsi('0.0') #read initial state
  
  #Plot probability density-----------------------------------------------------
  
  ax = fig.add_subplot(111, aspect = 'equal')
  
  #plt.xlim( [ -5, 5 ] )
  #plt.ylim( [ -5, 5 ] )
  
  plt.xlim( [ -10, 10 ] )
  plt.ylim( [ -10, 10 ] )
  
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  im = plt.imshow( abs( p ) ** 2 / ( abs( p ) ** 2 ).max(),
                  extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                  vmin = 0., vmax = 1., cmap = nicecmap() ) #(abs(p)**2).max())
 
  cbar = plt.colorbar(im, ticks = np.linspace( 0, 1, 5 ) )
  
  #if not autoscaling density:
  #v = np.linspace(0, ( abs( p ) ** 2. ).max() , 10)
  #for jj in range(0,len(v)): v[jj] = round(v[jj],3)
  #cbar = plt.colorbar(ticks = v)
  
  
  #Update the plot--------------------------------------------------------------
  def update_fig(i):
    p = infile.readpsi( str(i) )
    
    im.set_array( abs( p ) ** 2 / (abs(p)**2).max() )
        
    return im,

  ani = animation.FuncAnimation( fig, update_fig,
                                 np.arange( 1, steps + 1., dtype=float ),
                                 interval=20., repeat_delay = 3000., 
                                 blit = False )
  plt.show()
  return ani
  infile.f.close()

def aniphase(filename):
  '''
  Plot the simulation generated in the file filename, plotting only 
  the phase
  Filename is a binary hd5 file generated by twodsave.
  '''
  
  #Setup------------------------------------------------------------------------
  infile = h5file( filename, erase = False, read = True )
  fig = plt.figure()
  matplotlib.rcParams.update({'font.size':16}) #set inital font size
  
  steps = infile.head[ 'steps' ]
  npt   = infile.head[ 'npt' ]
  
  xvals, yvals = infile.readxy() #read x-y array
  
  p = infile.readpsi('0.0') #read initial state
  
  #Plot phase-------------------------------------------------------------------
  
  ax = fig.add_subplot(111, aspect = 'equal')
  
  #plt.xlim( [ -5, 5 ] )
  #plt.ylim( [ -5, 5 ] )
  
  plt.xlim( [ -10, 10 ] )
  plt.ylim( [ -10, 10 ] )
  
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  pi = np.pi
  im = plt.imshow( np.angle( p ),
                    extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                    vmin = -np.pi, vmax = np.pi, cmap=cm.gist_rainbow ) 
  cbar = plt.colorbar(im,
                       ticks = [round(j,2) for j in np.linspace(-pi,pi,5)]
                      )
  
  
  #Update the plot--------------------------------------------------------------
  def update_fig(i):
    p = infile.readpsi( str(i) )
    
    im.set_array( np.angle( p ) )
        
    return im,

  ani = animation.FuncAnimation( fig, update_fig,
                                 np.arange( 1, steps + 1., dtype=float ),
                                 interval=20., repeat_delay = 3000., 
                                 blit = False )
  plt.show()
  return ani
  infile.f.close() 


  

  