occurrence

#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import os as os

def listevar(fichier):
	datago = Dataset(fichier, 'r')
	x = datago.variables.keys()
	print x

def extraction(fichier):
	datago = Dataset(fichier, 'r')
	lg = datago.variables['longitude'][:]
	lt = datago.variables['latitude'][:]
	ti = datago.variables['time'][:]
	atb = datago.variables['ATB'][:,:]
	SR = datago.variables['instant_SR'][:,:]
	sol_alt = datago.variables['SE'][:]
	alt = datago.variables['alt_bound'][:,:]
	u, v = np.shape(alt)
	alt_bas = alt[0,:]
	alt_haut = alt[1,:]
	alt = np.zeros((v,1))
	for i in range(0,v):
		alt[i] = alt_bas[i]
	#alt[v] = alt_haut[v-1]
	alt = alt[:,0]
	return lg, lt, ti, atb,alt, sol_alt, SR


def affiche(n1,n2,lg,lt,ti,atb,alt,sol_alt):
	ti = ti[n1:n2]
	sol_alt = sol_alt[n1:n2]
	u,v = np.shape(atb[n1:n2,:])
	trait = ti[int (((n2-n1)/2))] 
	plt.plot ([trait,trait] , [alt[0],alt[-1]],'w--',linewidth=2)
	plt.plot(ti, sol_alt,'w',linewidth=2)
	plt.pcolormesh(ti,alt,atb[n1:n2,:].T)
	plt.ylim(-1, 16)
	plt.xlim(ti[0], ti[-1])
	plt.clim(0, 0.03)
	plt.xlabel('Time')
	plt.ylabel('Altitude [km]')
	plt.colorbar()	  
	plt.show() 

def afficheSR(lg,lt,ti,atb,alt,sol_alt):
	n1 = 0
	n2,x = np.shape(atb)
	lt = lt[n1:n2]
	sol_alt = sol_alt[n1:n2]
	u,v = np.shape(atb[n1:n2,:])	
 	sol_alt = sol_alt[n1:n2]
	fig = plt.figure(figsize=(15, 10), dpi=300)
	plt.plot(lt, sol_alt,'w',linewidth=2)
	plt.pcolormesh(lt,alt,atb[n1:n2,:].T)
	plt.ylim(-1, 16)
	plt.clim(1e-5, 1e-2)
	plt.xlabel('Latitude')
	plt.ylabel('Altitude [km]')
	plt.colorbar()	  
	plt.show() 

def lonlat(n1, n2,lg,lt,ti,atb):
	from mpl_toolkits.basemap import Basemap
	m = Basemap()
	m.plot(lg, lt)
	lg = lg[n1:n2]
	lt = lt[n1:n2]
	m.plot(lg, lt, 'r')
	m.drawcoastlines()
	m.fillcontinents()
	plt.show()


def profil(n,lg,lt,ti,atb,alt):
	plt.semilogx(atb[n,:], alt)
	#plt.xlim(1e-7, 1e-1)
	plt.ylim(-20,20)
	plt.xlabel('ATB [km-1 sr-1]')
	plt.ylabel('Altitude [km]')
	plt.title('Profil %d' % (n))
	plt.show()

def profil_SR(n,lg,lt,ti,atb,alt,fichier):
	plt.plot(atb[n,:], alt)
	plt.xlim(-0.1, 30)
	plt.ylim(-20,20)
	plt.xlabel('ATB [km-1 sr-1]')
	plt.ylabel('Altitude [km]')
	plt.title('Profil %d' % (n))
	plt.show()
	path = fichier
	path += '.png'
	#plt.savefig(path)
	#plt.close()
	

def coord_pole(lg,lt,ti,atb, sol_alt):
	idx = ( lt < -60 ) & ( lt >= -90 )
	lt6090 = lt[idx]
	lg6090 = lg[idx]
	ti6090 = ti[idx]
	sol_alt = sol_alt[idx]
	atb = atb [idx,:]
	n = len(lt6090)
	return n, lt6090, lg6090, ti6090,atb, sol_alt

def alt_max(SR, alt, n):
	u, v = np.shape(SR)
	j = - 1
	i = 0
	while SR[n,i] < 0.01 and i < v and SR[n,i] > -10 :
		if SR[n,i-1]  and SR[n,i-2] :
			j = i
		i = i + 1
	return j

def alt_max_orbite(alt, sol_alt, SR,lt):
	p,v = np.shape(SR)
	tab1 = np.zeros((p))
	tab2 = np.zeros((p))
	for o in range (0,p):
		j = alt_max(SR,alt,o)
		if j != -1 :
			tab1[o] = alt[j]
			#if tab1[o] > 5.: print o
		else : tab1[o] = -10
		tab2[o] = o

	tab1 = np.ma.masked_where(tab1 <= -1, tab1)
	fig = plt.figure(figsize=(15, 10), dpi=300)
	plt.plot(lt,tab1)
	plt.plot(lt, sol_alt,'g',linewidth=2)
	plt.xlim(np.min(lt), np.max(lt))
	plt.ylabel('Altitude')
	plt.xlabel('indice')
	plt.show()

def alt_max_profil(alt,SR,n,fichier):
	j = alt_max(SR,alt,n)
	s = np.shape(alt)
	if j != -1 :alt_ext= alt[j]
	else : alt_ext= -10
	print alt_ext
	profil_SR(n,lg,lt,ti,SR,alt,fichier)
	
def affiche_flag(lg,lt,atb,alt,sol_alt):
	n1 = 0
	n2, x = np.shape(atb)
	lt = lt[n1:n2]
	sol_alt = sol_alt[n1:n2]
	u,v = np.shape(atb[n1:n2,:])
	fig = plt.figure(figsize=(15, 10), dpi=300)
	plt.plot(lt, sol_alt,'w',linewidth=2)
	plt.pcolormesh(lt,alt,atb.T)
	plt.ylim(-1, 20)
	plt.xlim(lt[0], lt[-1])
	plt.clim(0, 4)
	plt.xlabel('Latitude')
	plt.ylabel('Altitude')
	plt.colorbar()	  
	plt.show() 

def flag_orbite(alt,SR):	
	p,t = np.shape(SR)
	u,v = np.shape(SR)
	flag = np.zeros((u,v))
	for n in range (0,p):
		j = alt_max(SR,alt,n)
		if j != -1 :
			i = 0
			while i < v :
				if SR[n,i] < 0.01 and  SR[n,i] > -10 and alt[i] < alt [j]:
					flag[n,i] = 4 #vrai eteint
				elif SR[n,i] < 0.01 and  SR[n,i] > -10 and alt[i] > alt [j] :
					if SR[n, i-1] > 0.01 and SR[n, i-2] > 0.01 :
						flag[n,i] = 2 # faux eteint
					elif alt[i] < alt[j]:
						flag[n,i] = 4 # vrai eteint
				elif SR[n,i] < 0.01 and  SR[n,i] > -10 and alt[i] == alt [j]:
					flag[n,i] = 3 # nuage opaque
				elif SR[n,i] > 5 and alt[i] > alt [j]:
					flag[n,i] = 1 # nuage
				i += 1
	return flag
	
def lecture_dossier(repertoire):
	u = np.shape(os.listdir(repertoire))
	
	return u, os.listdir(repertoire)

def occur(flag,alt):
	clouds = (flag == 1).sum(axis=0)
	faux_eteint = (flag == 2).sum(axis=0) 
	cloud_opa = (flag == 3).sum(axis=0)
	vrai_eteint = (flag == 4).sum(axis=0)
	return clouds, faux_eteint, cloud_opa, vrai_eteint

def freq_occur(flag,alt):
	u,v = np.shape(flag)
	clouds, faux_eteint, cloud_opa, vrai_eteint = occur(flag,alt)
	clouds, faux_eteint, cloud_opa, vrai_eteint = (1.)*clouds / u, (1.)*faux_eteint / u, (1.)*cloud_opa / u, (1.)*vrai_eteint / u
	fig = plt.figure(figsize=(9, 15), dpi=100)
	p1, = plt.plot(vrai_eteint,alt)
	sum1 = vrai_eteint + cloud_opa
	p2, = plt.plot(sum1,alt)
	sum2 = vrai_eteint + cloud_opa + clouds
	p3, = plt.plot(sum2,alt)
	sum3 = vrai_eteint + cloud_opa + clouds + faux_eteint
	p4, = plt.plot(sum3,alt)
	#plt.xlim(0.0001, 0.1)
	plt.legend([p1,p2,p3,p4], ["Nuages","Faux eteint","Nuages opaques","Vrai eteint"], loc=1)
	plt.xlabel("Frequence d'occurence")
	plt.ylabel("Altitude en  km")
	plt.show()
	
def occur_in_region(flag,alt,lt,latrange):
	idx = (lt < latrange[1]) & (lt > latrange[0])
	flag = flag[idx]
	clouds = (flag == 1).sum(axis=0)
	faux_eteint = (flag == 2).sum(axis=0) 
	cloud_opa = (flag == 3).sum(axis=0)
	vrai_eteint = (flag == 4).sum(axis=0)
	return clouds, faux_eteint, cloud_opa, vrai_eteint

def freq_occur_region(dossier):
	nbr_fichier,fichier = lecture_dossier(dossier)	

	dic_opa_cloud = dict()
	dic_cloud = dict()
	dic_vrai = dict()
	dic_faux = dict()
	nbr_fichier=15
	latbands = {'tropique' :[-20,20],'pole_nord' :[60,90],'pole_sud' : [-90,-60],'mid_sud' :[-50,-30],'mid_nord' :[30, 50 ]}
	for i in range(0,nbr_fichier):	
		fichier3 = dossier + fichier[i]
		lg, lt, ti, atb,alt, sol_alt, SR = extraction(fichier3)
		flag = flag_orbite(alt,SR)
		u,v = np.shape(flag)
		for region, tab_reg  in latbands.items():
			print region, tab_reg
			#flag = flag[tab_reg,:]
			clouds, faux_eteint, cloud_opa, vrai_eteint = occur_in_region(flag, alt, lt, latbands[region])
			if region in dic_opa_cloud:
				dic_opa_cloud[region] += 1.* cloud_opa/ (u * nbr_fichier)
				dic_cloud[region] += 1.* clouds/ (u * nbr_fichier)
				dic_vrai[region] += 1.* vrai_eteint/ (u * nbr_fichier)
				dic_faux[region] += 1.* faux_eteint/ (u * nbr_fichier)
			else:
				dic_opa_cloud[region] = 1.* cloud_opa/ (u * nbr_fichier)
				dic_cloud[region] = 1.* clouds/ (u * nbr_fichier)
				dic_vrai[region] = 1.* vrai_eteint/ (u * nbr_fichier)
				dic_faux[region] = 1.* faux_eteint/ (u * nbr_fichier)
	for	region in latbands.keys():
		
		affiche_freq(alt,dic_opa_cloud[region],dic_cloud[region],dic_vrai[region],dic_faux[region],region)

def affiche_freq(alt,op,cl,vr,fa,region):
	fig = plt.figure(figsize=(9, 15), dpi=100)
	p1, = plt.plot(vr,alt)
	sum1 = vr + op
	p2, = plt.plot(sum1,alt)
	sum2 = sum1 + cl
	p3, = plt.plot(sum2,alt)
	sum3 = sum2 + fa
	p4, = plt.plot(sum3,alt)
	plt.legend([p1,p2,p3,p4], ["Nuages","Faux eteint","Nuages opaques","Vrai eteint"], loc=1)
	plt.xlabel("Frequence d'occurence")
	plt.ylabel("Altitude en  km")
	path = region
	path += '.png'
	plt.savefig(path)
	plt.close()

if __name__=='__main__':

	fichier3="/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP/instant_SR_CR_DR/grid_L40/2013/201308/night/instant_SR_CR_DR_2013-08-24T23-19-30ZN_night_CFMIP2_2.70.nc"#38418 mid
	#fichier3="/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP/instant_SR_CR_DR/grid_L40/2013/201308/night/instant_SR_CR_DR_2013-08-23T19-18-36ZN_night_CFMIP2_2.70.nc" #20632 trop
	listevar(fichier3)
	#n1 , n2 = 0  , 20632
	lg, lt, ti, atb,alt, sol_alt, SR = extraction(fichier3)
 	#alt_max(SR, alt, n1,2*n2)
	#profil_SR(n1,2*n2,lg,lt,ti,SR,alt)
	#affiche(n1,n2,lg, lt, ti, SR,alt,sol_alt)
	'''
	n1 = 0
	n2, lt, lg, ti, atb, sol_alt = coord_pole(lg, lt, ti, atb,sol_alt)
	alt_max(SR, alt, n1,n2)
	profil(n1,n2,lg,lt,ti,atb,alt)	
	affiche(n1,n2,lg, lt, ti, atb,alt,sol_alt)
	lonlat(n1, n2,lg, lt, ti, atb)
	'''

	#lonlat(n1, n2,lg, lt, ti, atb)
	#affiche(n1,n2,lg, lt, ti, SR,alt,sol_alt)
	#alt_max_orbite(alt, sol_alt, SR,lt)
	'''
	# Tableau de flag
	flag = flag_orbite(alt,SR)
	#affiche_flag(lg, lt, flag,alt,sol_alt)
	freq_occur(flag,alt)
	'''
	dossier = '/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP/instant_SR_CR_DR/grid_L40/2013/201308/night/'
	freq_occur_region(dossier)