#! /usr/bin/env python
#
# Python script to plot stuff from GS2 NetCDF file
#
# To execute from the shell, do:
#
# ./gs2plot dimits00_Lti_hv_S_4.out.nc
#
# Or you can execute from inside python by doing:
#
# python
# >>> nc_filename = 'dimits00_Lti_hv_S_4.out.nc'
# >>> execfile('gs2plot.py')
#

import sys
from Numeric import *
from Scientific.IO.NetCDF import *

if (len(sys.argv) > 1):
	nc_filename=sys.argv[1]

file = NetCDFFile(nc_filename, 'r') 


# Note, the name of the integer variable that indexes into
# the time dimension is "time".  The time array is variable "t".
#
t = file.variables['t'][:]
flux = file.variables['es_heat_flux'][:]
tprim = file.variables['tprim'][:]
fprim = file.variables['fprim'][:]

print "chi_i calculated assuming R/a_norm = 1, etc."
chi_i = flux[:,0]/tprim[0]/fprim[0]

# Arrays in NetCDF are accessed by indexing, [:] gets the whole array.
# Read non-indexed scalar variables a slightly different way:
nkx = file.variables['nkx'].getValue()
nky = file.variables['nky'].getValue()

phi2_by_ky = file.variables['phi2_by_ky'][:]

import Gnuplot
g = Gnuplot.Gnuplot()
g('set data style lines')
g.xlabel('t')
g.ylabel('chi_i')
d = Gnuplot.Data(t, chi_i, title='chi_i')
g.plot(d)
raw_input('Please press return to continue...\n')

dphi = Gnuplot.Data(t, phi2_by_ky[:,0]*2./100, title='phi2_00')
g.title('predator/prey oscillations between instabilities and zonal flows')
g.plot(d,dphi)
raw_input('Please press return to continue...\n')

ky = file.variables['ky'][:]
kx = file.variables['kx'][:]
phi2_by_kx = file.variables['phi2_by_kx'][:]
phi2_by_ky = file.variables['phi2_by_ky'][:]

g.xlabel('k_x or k_y')
g.ylabel('<|Phi|^2>')
g.title('at t=end (looks funny because k_x wraps around')

dy = Gnuplot.Data(ky, phi2_by_ky[-1,:],title='<|Phi|^2> vs. k_y')
dx = Gnuplot.Data(kx, phi2_by_kx[-1,:],title='<|Phi|^2> vs. k_x')
g.plot(dx,dy)
raw_input('Please press return to continue...\n')

# To close netCDF file:
# file.close()