udg/ecoms/dataserver/interfaces/python: load_system4.py

File load_system4.py, 10.8 KB (added by maru, 8 years ago)
Line 
1#!/usr/bin/python
2
3import netCDF4 as ncdf
4import numpy as np
5import matplotlib.pyplot as plt
6from copy import deepcopy
7import pickle
8import os.path
9
10class UserData:
11   pass
12
13class Grid:
14   def __init__(self, lons=None, lats=None, type="regularll"):
15      if type == "regularll":
16         self.nx = len(lons)
17         self.ny = len(lats)
18         self.lon = np.transpose(np.reshape(np.repeat(lons, self.ny), (self.nx, self.ny)))
19         self.lat = np.reshape(np.repeat(lats, self.nx), (self.ny, self.nx))
20   def getCoordinates(self):
21      return np.concatenate([
22        np.reshape(self.lon, (1,self.ny*self.nx)),
23        np.reshape(self.lat, (1,self.ny*self.nx,))
24      ], axis=0)
25
26def deaccumulated_pr(arr):
27   """DEACCUMULATED PRECIPITATION"""
28   return arr[:,:,1:,:,:] - arr[:,:,:-1,:,:]
29
30
31def getLonFilter(lons,lonlim):
32   lons_180=np.where(lons <= 180, lons, lons-360)
33   indices=np.argsort(lons_180)
34   filter=(lons_180>=lonlim[0])*(lons_180<=lonlim[1])
35   true_indices=indices[filter]
36   true_indices_filter=np.empty(lons_180.shape,dtype=np.bool)
37   true_indices_filter[:]=False
38   true_indices_filter[true_indices]=True
39   lonfilter=indices[true_indices_filter]
40   return lonfilter
41
42
43def getLatFilter(lats,latlim):
44   indices=np.argsort(lats)
45   filter= (lats>=latlim[0])*(lats<=latlim[1])
46   true_indices=indices[filter]
47   true_indices_filter=np.empty(lats.shape,dtype=np.bool)
48   true_indices_filter[:]=False
49   true_indices_filter[true_indices]=True
50   latfilter=indices[true_indices_filter]
51   return latfilter
52
53
54def vecFilterRuns(rundates, period, leadMonth, year):
55  """Filter runs
56     year - list of years to be filtered
57            The year is defined by the first month of the period requested
58  """
59  rval = []
60  for rundate in rundates:
61    validYear = False
62    validMonth = False
63    inimonth = (period[0]-leadMonth-1)%12+1
64    if rundate.month == inimonth:
65      validMonth = True
66    if year:
67      if inimonth > period[0]: # .. then inimonth is in the prev. year
68        if rundate.year+1 in year:
69          validYear = True
70      else:
71        if rundate.year in year:
72          validYear = True
73    rval.append(validYear and validMonth)
74  return np.array(rval)
75
76
77def filterRuntimeTimeRanges(runvar, fctvar, period, leadMonth, year):
78   irun = vecFilterRuns(ncdf.num2date(runvar[:],runvar.units), period, leadMonth, year)
79   dates = ncdf.num2date(fctvar[:],fctvar.units)
80   dates = dates[irun]
81   match_month_and_years = np.vectorize(lambda x: x.month in period and (x.year in year or ((x.year)-1) in year))
82   ifct = match_month_and_years(dates)
83   runs = runvar[irun]
84   runs = ncdf.num2date(runs,runvar.units)
85   return irun, runs, ifct, dates[ifct]
86
87
88def loadSystem4(dataset, var, period, leadMonth, lonlim=None, latlim=None, year=[], members=None):
89   varname = {
90      "tasmax":"Maximum_temperature_at_2_metres_since_last_24_hours_surface",
91      "tasmin":"Minimum_temperature_at_2_metres_since_last_24_hours_surface",
92      "tas":"Mean_temperature_at_2_metres_since_last_24_hours_surface",
93      "pr":"Total_precipitation_surface",
94      "mslp":"Mean_sea_level_pressure_surface"
95   }
96   nc = ncdf.Dataset(dataset,"r")
97   #
98   #  Filter lons
99   #
100   lons=nc.variables["lon"][:]
101   if lonlim != None:
102     lonFilter = getLonFilter(lons,lonlim)
103   else:
104     lonFilter = slice(None)
105   lons=lons[lonFilter]
106   lons=np.where(lons <= 180, lons, lons-360)   
107   #
108   #  Filter lats
109   #
110   lats = nc.variables["lat"][:]
111   if latlim != None:
112      latFilter = getLatFilter(lats,latlim)
113   else:
114      latFilter = slice(None)
115   lats = lats[latFilter]
116   #
117   #  Filter runs and times
118   #
119   runFilter, runs, timeFilter, times = filterRuntimeTimeRanges(nc.variables["run"],nc.variables['time'], period, leadMonth, year)
120   #
121   #  Filter members
122   #
123   # TODO: Members should be retrieved by value within this variable,
124   #       not by position! (in this case it is the same, though)
125   memFilter = np.zeros(nc.variables["member"].shape[0], dtype=bool)
126   memFilter[members] = True
127   #
128   # Read variable
129   #
130   ncvar=nc.variables[varname[var]][memFilter,runFilter,:,latFilter,:]
131   if var=="pr":
132      ncvar=deaccumulated_pr(ncvar)
133   ncvar=ncvar[:,:,:,:,lonFilter]
134   ncvar=ncvar[:,timeFilter,:,:] 
135   ncvar.shape = ncvar.shape[:-2]+(ncvar.shape[-2]*ncvar.shape[-1],)
136   if not ncvar.any():
137      return "Conection time out. Sorry for the inconvenience but you are trying to load too much data, try again." 
138   ud = UserData() 
139   ud.grid = Grid(lons, lats)
140   ud.runtime = runs
141   ud.times = times
142   ud.LatLonCoords = ud.grid.getCoordinates()
143   ud.units = nc.variables[varname[var]].units
144   ud.short_name = var
145   rval = []
146   for imember in range(len(members)):
147      udc = deepcopy(ud)
148      udc.data = ncvar[imember]
149      if ud.units=="m":
150         udc.units = "mm/day"
151         udc.data = np.where(udc.data<=0,0,udc.data)
152         udc.data = udc.data*1000
153      udc.member = members[imember]
154      rval.append(udc)
155   return rval
156
157
158def l2string(list):
159   num2mon={1:"J", 2:"F", 3:"M", 4:"A", 5:"M", 6:"J", 7:"J", 8:"A", 9:"S", 10:"O", 11:"N", 12:"D"}
160   return "".join(map(lambda x: num2mon[x], list))
161
162
163def monthly_mean(meanthis,year,season,ud):
164   for y in year:
165      for m in season:
166         mask=[d.year==y and d.month==m for d in ud.times]
167         monthly_mean_value=np.mean(meanthis[np.array(mask)],axis=0)
168         if y==year[0] and m==season[0]:
169            monthly_mean_array=np.array([monthly_mean_value])
170         else:
171            monthly_mean_array=np.append(monthly_mean_array,monthly_mean_value)
172   return monthly_mean_array
173
174
175def plot_map(plotthis,ud,period,var,range="",cities_latlon=None,cities_name=None,method=None):
176   """Plot the map """
177   #import matplotlib.pyplot as plt
178   from matplotlib import cm
179   import matplotlib.dates as dates
180   from mpl_toolkits.basemap import Basemap
181   map = Basemap(projection='cyl',llcrnrlon=min(ud.grid.lon[0]), llcrnrlat=min(ud.grid.lat[:,0]),
182              urcrnrlon=max(ud.grid.lon[0]), urcrnrlat=max(ud.grid.lat[:,0]), resolution='i')
183   map.drawcoastlines()
184   parallels = np.arange(-90., 90, 5)
185   meridians = np.arange(-180, 180, 5)
186   map.drawparallels(parallels, labels=[1,1,1,1], fontsize=10)
187   map.drawmeridians(meridians,labels=[1,1,1,1], fontsize=10)
188   lonmat, latmat = map.makegrid(ud.grid.nx, ud.grid.ny)
189   plotthis=plotthis.reshape(ud.grid.ny, ud.grid.nx)
190   if method=="pixels":
191      cs = map.pcolormesh(lonmat, latmat, plotthis, cmap=cm.jet)
192   elif range!="":
193       cs = map.contourf(lonmat, latmat, plotthis, cmap=cm.jet,levels=range)
194   else:
195      cs = map.contourf(lonmat, latmat, plotthis, cmap=cm.jet)
196   if cities_latlon != None:
197      lat = cities_latlon[:,0]
198      lon = cities_latlon[:,1]
199      map.scatter(lon,lat,c="black")
200      for i in xrange(len(cities_name)):
201         plt.annotate(cities_name[i], xy = (lon[i],lat[i]))
202   cbar = map.colorbar(cs,location='bottom',pad="8%")
203   cbar.set_label('%s (%s)' %(var,ud.units))
204   plt.savefig("Map_%s_%s.png" %(var,l2string(period)))
205   plt.close()
206
207
208def plot_serie(spatial_mean,spatial_max,spatial_min, monthly_mean_array, ud, var, members,season,name_fig):
209   """Plot the time serie """
210   fig=plt.figure()
211   ax = fig.add_subplot(111)
212   x=np.arange(len(ud.times))
213   xticks_mask = [d.month in season and d.day == 1 for d in ud.times]
214   xticks = x[np.array(xticks_mask)]
215   ax.set_xticks(xticks)
216   ticklabels=[t.strftime("%Y-%m") for t in ud.times[np.array(xticks_mask)]]
217   ax.set_xticklabels(ticklabels)
218   plt.fill_between(x,spatial_min,spatial_max,color="grey")
219   plt.xticks(rotation=25)
220   plt.ylabel("%s" %var)
221   plt.plot([0,max(x)+10],[spatial_mean.mean(),spatial_mean.mean()],color="black",ls= '-.')
222   plt.plot(xticks,monthly_mean_array,'o',color="red")
223   for t in range(len(ud.times)):
224      plt.plot(t, spatial_mean[t], '+', color="blue")
225   plt.savefig("%s_%s.pdf" %(name_fig,var))
226   plt.close()
227
228
229def plot_serie_cities(cities_latlon,cities_name,ud):
230   pcolors = {
231     "0":"blue",
232     "1":"green",
233     "2":"red",
234     "3":"purple",
235     "4":"cyan",
236     "5":"yellow",
237     "6":"magenta",
238     "7":"pink",
239     "8":"orange",
240     "9":"brown",
241     "10":"grey",
242   }
243   fig=plt.figure()
244   ax = fig.add_subplot(111)
245   x=np.arange(len(ud.times))
246   xticks_mask = [d.month in season and d.day == 1 for d in ud.times]
247   xticks = x[np.array(xticks_mask)]
248   ax.set_xticks(xticks)
249   ticklabels=[t.strftime("%Y-%m") for t in ud.times[np.array(xticks_mask)]]
250   ax.set_xticklabels(ticklabels)
251   plt.xticks(rotation=25)
252   plt.ylabel("%s (%s)" %(var,ud.units))
253   lat = cities_latlon[:,0]
254   lon = cities_latlon[:,1]
255   for i in np.arange(len(cities_name)):
256      city_position=(((ud.LatLonCoords[0,:]-lon[i])**2)+((ud.LatLonCoords[1,:]-lat[i])**2)).argmin()
257      city = ud.data[:,city_position]
258      c=("%("+str(i)+")s") % pcolors
259      plt.plot(city,color=c, label=cities_name[i])
260   plt.legend(loc="lower left",prop={'size':'small'})
261   plt.savefig("map_serie_all_cities.png")
262
263
264#######################
265####  MAIN PROGRAM ####
266#######################
267
268if __name__ == "__main__":
269   World= {
270     "lonlim1":-180,
271     "lonlim2":180,
272     "latlim1":-90,
273     "latlim2":90
274   }
275   Europe= {
276     "lonlim1":-10,
277     "lonlim2":40,
278     "latlim1":35,
279     "latlim2":70
280   }
281   Spain= {
282     "lonlim1":-10,
283     "lonlim2":5,
284     "latlim1":35,
285     "latlim2":45
286   }
287   domains,domain=Europe,"Europe"
288   var = "tas" 
289   season = [4,5,6]
290   leadMonth = 4  # I want December initialization for April, May, June season. 
291   lonlim=[float("%(lonlim1)s" %domains),float("%(lonlim2)s" %domains)]
292   latlim=[float("%(latlim1)s" %domains),float("%(latlim2)s" %domains)]
293   year=[2005,2006,2007]
294   members=[0,1,2,3,4]
295   username="my_username"
296   password="my_password"
297   dataset="http://%s:%s@www.meteo.unican.es/tds5/dodsC/system4/System4_Seasonal_15Members.ncml" %(username,password)
298   uds=[]
299   if os.path.isfile("%s.pkl" %domain):
300      uds=pickle.load(open("%s.pkl"%domain))
301   else:
302      for i in members:
303         uds.extend(loadSystem4(dataset,var,season, leadMonth,lonlim, latlim, year, members=[i]))
304      pickle.dump(uds,open("%s.pkl" %domain,'w'),-1)
305
306   ud=deepcopy(uds[0])
307   ud.data = np.mean(map(lambda x:x.data, uds),axis=0)   # The ensemble mean
308   ud_max=deepcopy(uds[0])
309   ud_max.data=np.max(map(lambda x:x.data, uds),axis=0) # We take the ensemble maximum
310   ud_min=deepcopy(uds[0])
311   ud_min.data=np.min(map(lambda x:x.data, uds),axis=0) # We take the ensemble minimum
312
313   temporal_mean=ud.data.mean(axis=0)     
314   cities_name=np.array(["Paris","Madrid","Dublin","Oslo","Warsaw","Rome","Athens","Berlin","Helsinki"])
315   cities_latlon=np.array([[(48.83),(2.37)],[(40.40),(-3.72)],[(53.34),(-6.25)],[(59.92),(10.76)],[(52.24),(21.03)],[(41.89),(12.49)],[(37.98),(23.72)],[(52.52),(13.40)],[(60.17),(24.94)]])
316   #plot_map(temporal_mean,ud,season,var,cities_latlon,cities_name)  # If you want to mark certain locations
317   plot_map(temporal_mean,ud,season,var,cities_latlon,cities_name,method="pixels") # If you do not want to draw the map interpolated/with filled contours...
318
319