barakuda_ncio.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565
  1. import numpy as nmp
  2. import sys
  3. from netCDF4 import Dataset
  4. from os import path
  5. from os import stat
  6. def wrt_1d_series(vt, vd, cvar, cinfo,
  7. cu_t='unknown', cu_d='unknown', cln_d='unknown', nsmooth=0,
  8. vd2=[], vd3=[], vd4=[], vd5=[],
  9. cvar2='', cvar3='', cvar4='', cvar5='',
  10. cln_d2='', cln_d3='', cln_d4='', cln_d5=''):
  11. cf_o = cvar+'_'+cinfo+'.nc'
  12. lsmooth = False
  13. if nsmooth > 0:
  14. import barakuda_stat as bs
  15. lsmooth = True
  16. if nsmooth == 11:
  17. vd_sm = bs.running_mean_11(vd, l_fill_bounds=False)
  18. elif nsmooth == 5:
  19. vd_sm = bs.running_mean_5(vd, l_fill_bounds=False)
  20. else:
  21. print 'ERROR: wrt_1d_series.barakuda_ncio => smoothing with nsmooth='+str(nsmooth)+' not supported!'; sys.exit(0)
  22. f_o = Dataset(cf_o, 'w', format='NETCDF4')
  23. nt = len(vt)
  24. if len(vd) != nt: print 'ERROR: wrt_1d_series.barakuda_ncio => data & time have different lengths!'; sys.exit(0)
  25. l_do_v2=False ; l_do_v3=False ; l_do_v4=False ; l_do_v5=False
  26. if len(vd2) == nt: l_do_v2=True
  27. if len(vd3) == nt: l_do_v3=True
  28. if len(vd4) == nt: l_do_v4=True
  29. if len(vd5) == nt: l_do_v5=True
  30. f_o.createDimension('time', None)
  31. id_t = f_o.createVariable('time','f4',('time',), zlib=True, complevel=5) ; id_t.units = cu_t
  32. id_d = f_o.createVariable(cvar,'f4',('time',), zlib=True, complevel=5)
  33. id_d.units = cu_d ; id_d.long_name = cln_d
  34. if l_do_v2: id_d2 = f_o.createVariable(cvar2,'f4',('time',), zlib=True, complevel=5); id_d2.units = cu_d; id_d2.long_name = cln_d2
  35. if l_do_v3: id_d3 = f_o.createVariable(cvar3,'f4',('time',), zlib=True, complevel=5); id_d3.units = cu_d; id_d3.long_name = cln_d3
  36. if l_do_v4: id_d4 = f_o.createVariable(cvar4,'f4',('time',), zlib=True, complevel=5); id_d4.units = cu_d; id_d4.long_name = cln_d4
  37. if l_do_v5: id_d5 = f_o.createVariable(cvar5,'f4',('time',), zlib=True, complevel=5); id_d5.units = cu_d; id_d5.long_name = cln_d5
  38. if lsmooth:
  39. id_sm = f_o.createVariable(cvar+'_'+str(nsmooth)+'yrm','f4',('time',), zlib=True, complevel=5)
  40. id_sm.units = cu_d ; id_sm.long_name = str(nsmooth)+'-year running mean of '+cln_d
  41. for jt in range(nt):
  42. id_t[jt] = vt[jt]
  43. id_d[jt] = vd[jt]
  44. if lsmooth: id_sm[jt] = vd_sm[jt]
  45. if l_do_v2: id_d2[jt] = vd2[jt]
  46. if l_do_v3: id_d3[jt] = vd3[jt]
  47. if l_do_v4: id_d4[jt] = vd4[jt]
  48. if l_do_v5: id_d5[jt] = vd5[jt]
  49. f_o.Author = 'L. Brodeau (barakuda_ncio.py of Barakuda)'
  50. f_o.close()
  51. print ' * wrt_1d_series => '+cf_o+' written!\n'
  52. return 0
  53. def read_1d_series(cf_i, cv_i, cv_t='time', l_return_time=True):
  54. if not path.exists(cf_i): print 'ERROR: read_1d_series.barakuda_ncio => '+cf_i+' not found!'; sys.exit(0)
  55. id_i = Dataset(cf_i)
  56. list_variables = id_i.variables.keys()
  57. if l_return_time: vt = id_i.variables[cv_t][:]
  58. if cv_i in list_variables[:]:
  59. vd = id_i.variables[cv_i][:]
  60. else:
  61. print 'WARNING: read_1d_series.barakuda_ncio => '+cv_i+' not found into '+cf_i+' !'
  62. if l_return_time:
  63. vd = nmp.zeros(len(vt))
  64. else:
  65. vd = [ 0. ]
  66. id_i.close()
  67. if l_return_time:
  68. print ' * read_1d_series => just read '+cv_t+' and '+cv_i+' into '+cf_i+'\n'
  69. return vt, vd
  70. else:
  71. print ' * read_1d_series => just read '+cv_i+' into '+cf_i+'\n'
  72. return vd
  73. def read_1d_series_2(cf_i, cv_i, cv_t='time'):
  74. if not path.exists(cf_i): print 'ERROR: read_1d_series_2.barakuda_ncio => '+cf_i+' not found!'; sys.exit(0)
  75. id_i = Dataset(cf_i)
  76. vt = id_i.variables[cv_t][:]
  77. nt = len(vt)
  78. xout = nmp.zeros((2,nt))
  79. xout[0,:] = vt[:]
  80. xout[1,:] = id_i.variables[cv_i][:]
  81. id_i.close()
  82. print ' * read_1d_series_2 => just read '+cv_t+' and '+cv_i+' into '+cf_i+'\n'
  83. return xout
  84. def wrt_appnd_1d_series(vt, vd, cf, cvar1, cu_t='year', cu_d='unknown', cln_d='unknown',
  85. vd2=[], cvar2='', cln_d2='', cun2='unknown',
  86. vd3=[], cvar3='', cln_d3='', cun3='unknown',
  87. vd4=[], cvar4='', cln_d4='', cun4='unknown',
  88. vd5=[], cvar5='', cln_d5='', cun5='unknown',
  89. vd6=[], cvar6='', cln_d6='', cun6='unknown',
  90. vd7=[], cvar7='', cln_d7='', cun7='unknown',
  91. vd8=[], cvar8='', cln_d8='', cun8='unknown',
  92. vd9=[], cvar9='', cln_d9='', cun9='unknown',
  93. vd10=[],cvar10='',cln_d10='',cun10='unknown' ):
  94. l_nc_is_new = not path.exists(cf)
  95. Nt = len(vt)
  96. if len(vd) != Nt:
  97. print 'ERROR: wrt_appnd_1d_series.barakuda_ncio => data & time have different lengths!'
  98. sys.exit(0)
  99. l_do_v2=False ; l_do_v3=False ; l_do_v4=False ; l_do_v5=False
  100. l_do_v6=False ; l_do_v7=False ; l_do_v8=False ; l_do_v9=False ; l_do_v10=False
  101. if len(vd2) == Nt: l_do_v2=True
  102. if len(vd3) == Nt: l_do_v3=True
  103. if len(vd4) == Nt: l_do_v4=True
  104. if len(vd5) == Nt: l_do_v5=True
  105. if len(vd6) == Nt: l_do_v6=True
  106. if len(vd7) == Nt: l_do_v7=True
  107. if len(vd8) == Nt: l_do_v8=True
  108. if len(vd9) == Nt: l_do_v9=True
  109. if len(vd10)== Nt: l_do_v10=True
  110. if l_nc_is_new:
  111. f_out = Dataset(cf, 'w', format='NETCDF4')
  112. else:
  113. f_out = Dataset(cf, 'a', format='NETCDF4')
  114. if l_nc_is_new:
  115. jrec2write = 0
  116. f_out.createDimension('time', None)
  117. id_t = f_out.createVariable('time','f4',('time',), zlib=True, complevel=5) ; id_t.units = 'year'
  118. id_v01 = f_out.createVariable(cvar1 ,'f4',('time',), zlib=True, complevel=5)
  119. id_v01.units = cu_d
  120. id_v01.long_name = cln_d
  121. if l_do_v2:
  122. id_v02 = f_out.createVariable(cvar2 ,'f4',('time',), zlib=True, complevel=5)
  123. id_v02.units = cu_d
  124. if not cun2=='unknown': id_v02.units = cun2
  125. id_v02.long_name = cln_d2
  126. if l_do_v3:
  127. id_v03 = f_out.createVariable(cvar3 ,'f4',('time',), zlib=True, complevel=5)
  128. id_v03.units = cu_d
  129. if not cun3=='unknown': id_v03.units = cun3
  130. id_v03.long_name = cln_d3
  131. if l_do_v4:
  132. id_v04 = f_out.createVariable(cvar4 ,'f4',('time',), zlib=True, complevel=5)
  133. id_v04.units = cu_d
  134. if not cun4=='unknown': id_v04.units = cun4
  135. id_v04.long_name = cln_d4
  136. if l_do_v5:
  137. id_v05 = f_out.createVariable(cvar5 ,'f4',('time',), zlib=True, complevel=5)
  138. id_v05.units = cu_d
  139. if not cun5=='unknown': id_v05.units = cun5
  140. id_v05.long_name = cln_d5
  141. if l_do_v6:
  142. id_v06 = f_out.createVariable(cvar6 ,'f4',('time',), zlib=True, complevel=5)
  143. id_v06.units = cu_d
  144. if not cun6=='unknown': id_v06.units = cun6
  145. id_v06.long_name = cln_d6
  146. if l_do_v7:
  147. id_v07 = f_out.createVariable(cvar7 ,'f4',('time',), zlib=True, complevel=5)
  148. id_v07.units = cu_d
  149. if not cun7=='unknown': id_v07.units = cun7
  150. id_v07.long_name = cln_d7
  151. if l_do_v8:
  152. id_v08 = f_out.createVariable(cvar8 ,'f4',('time',), zlib=True, complevel=5)
  153. id_v08.units = cu_d
  154. if not cun8=='unknown': id_v08.units = cun8
  155. id_v08.long_name = cln_d8
  156. if l_do_v9:
  157. id_v09 = f_out.createVariable(cvar9 ,'f4',('time',), zlib=True, complevel=5)
  158. id_v09.units = cu_d
  159. if not cun9=='unknown': id_v09.units = cun9
  160. id_v09.long_name = cln_d9
  161. if l_do_v10:
  162. id_v10 = f_out.createVariable(cvar10 ,'f4',('time',), zlib=True, complevel=5)
  163. id_v10.units = cu_d
  164. if not cun7=='unknown': id_v10.units = cun10
  165. id_v10.long_name = cln_d10
  166. jrw = 0
  167. for jt in range(Nt):
  168. jrw = jrec2write + jt
  169. id_t[jrw] = vt[jt]
  170. id_v01[jrw] = vd[jt]
  171. if l_do_v2: id_v02[jrw] = vd2[jt]
  172. if l_do_v3: id_v03[jrw] = vd3[jt]
  173. if l_do_v4: id_v04[jrw] = vd4[jt]
  174. if l_do_v5: id_v05[jrw] = vd5[jt]
  175. if l_do_v6: id_v06[jrw] = vd6[jt]
  176. if l_do_v7: id_v07[jrw] = vd7[jt]
  177. if l_do_v8: id_v08[jrw] = vd8[jt]
  178. if l_do_v9: id_v09[jrw] = vd9[jt]
  179. if l_do_v10:id_v10[jrw] = vd10[jt]
  180. f_out.about = 'Diagnostics created with BaraKuda (https://github.com/brodeau/barakuda)'
  181. else:
  182. vtr = f_out.variables['time']
  183. jrec2write = len(vtr)
  184. v01 = f_out.variables[cvar1]
  185. if l_do_v2: v02 = f_out.variables[cvar2]
  186. if l_do_v3: v03 = f_out.variables[cvar3]
  187. if l_do_v4: v04 = f_out.variables[cvar4]
  188. if l_do_v5: v05 = f_out.variables[cvar5]
  189. if l_do_v6: v06 = f_out.variables[cvar6]
  190. if l_do_v7: v07 = f_out.variables[cvar7]
  191. if l_do_v8: v08 = f_out.variables[cvar8]
  192. if l_do_v9: v09 = f_out.variables[cvar9]
  193. if l_do_v10:v10 = f_out.variables[cvar10]
  194. jrw = 0
  195. for jt in range(Nt):
  196. jrw = jrec2write + jt
  197. vtr[jrw] = vt[jt]
  198. v01[jrw] = vd[jt]
  199. if l_do_v2: v02[jrw] = vd2[jt]
  200. if l_do_v3: v03[jrw] = vd3[jt]
  201. if l_do_v4: v04[jrw] = vd4[jt]
  202. if l_do_v5: v05[jrw] = vd5[jt]
  203. if l_do_v6: v06[jrw] = vd6[jt]
  204. if l_do_v7: v07[jrw] = vd7[jt]
  205. if l_do_v8: v08[jrw] = vd8[jt]
  206. if l_do_v9: v09[jrw] = vd9[jt]
  207. if l_do_v10:v10[jrw] = vd10[jt]
  208. f_out.close()
  209. print ' *** '+str(Nt)+' snapshots of fields saved into '+cf+' !\n'
  210. # Write a Array(x,y,t) to netcdf:
  211. def wrt_appnd_2dt_series(vx, vy, vt, xd, cf, cvar1, missing_val=-9999.,
  212. cxdnm='x', cydnm='y', cxvnm='lon', cyvnm='lat',
  213. cu_t='year', cu_d='unknown', cln_d='unknown',
  214. xd2=[], cvar2='', cln_d2='', cun2='unknown',
  215. xd3=[], cvar3='', cln_d3='', cun3='unknown',
  216. xd4=[], cvar4='', cln_d4='', cun4='unknown',
  217. xd5=[], cvar5='', cln_d5='', cun5='unknown',
  218. xd6=[], cvar6='', cln_d6='', cun6='unknown',
  219. xd7=[], cvar7='', cln_d7='', cun7='unknown',
  220. xd8=[], cvar8='', cln_d8='', cun8='unknown',
  221. xd9=[], cvar9='', cln_d9='', cun9='unknown',
  222. xd10=[],cvar10='',cln_d10='',cun10='unknown' ):
  223. l_nc_is_new = not path.exists(cf)
  224. [ Nt, ny, nx ] = nmp.shape(xd)
  225. if Nt != len(vt):
  226. print 'ERROR: wrt_appnd_2dt_series.barakuda_ncio => data & time have different lengths!'
  227. sys.exit(0)
  228. l_do_v2=False ; l_do_v3=False ; l_do_v4=False ; l_do_v5=False
  229. l_do_v6=False ; l_do_v7=False ; l_do_v8=False ; l_do_v9=False ; l_do_v10=False
  230. if nmp.shape(xd2) == ( Nt, ny, nx ): l_do_v2=True
  231. if nmp.shape(xd3) == ( Nt, ny, nx ): l_do_v3=True
  232. if nmp.shape(xd4) == ( Nt, ny, nx ): l_do_v4=True
  233. if nmp.shape(xd5) == ( Nt, ny, nx ): l_do_v5=True
  234. if nmp.shape(xd6) == ( Nt, ny, nx ): l_do_v6=True
  235. if nmp.shape(xd7) == ( Nt, ny, nx ): l_do_v7=True
  236. if nmp.shape(xd8) == ( Nt, ny, nx ): l_do_v8=True
  237. if nmp.shape(xd9) == ( Nt, ny, nx ): l_do_v9=True
  238. if nmp.shape(xd10)== ( Nt, ny, nx ): l_do_v10=True
  239. if l_nc_is_new:
  240. f_out = Dataset(cf, 'w', format='NETCDF4')
  241. else:
  242. f_out = Dataset(cf, 'a', format='NETCDF4')
  243. if l_nc_is_new:
  244. jrec2write = 0
  245. f_out.createDimension('time', None)
  246. f_out.createDimension(cydnm , ny)
  247. f_out.createDimension(cxdnm , nx)
  248. id_t = f_out.createVariable('time','f4',('time',), zlib=True, complevel=5) ; id_t.units = 'year'
  249. id_lat = f_out.createVariable(cyvnm ,'f4',(cydnm,), zlib=True, complevel=5)
  250. id_lon = f_out.createVariable(cxvnm ,'f4',(cxdnm,), zlib=True, complevel=5)
  251. id_lat[:] = vy[:]
  252. id_lon[:] = vx[:]
  253. id_x01 = f_out.createVariable(cvar1 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  254. id_x01.units = cu_d
  255. id_x01.long_name = cln_d
  256. if l_do_v2:
  257. id_x02 = f_out.createVariable(cvar2 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  258. id_x02.units = cu_d
  259. if not cun2=='unknown': id_x02.units = cun2
  260. id_x02.long_name = cln_d2
  261. if l_do_v3:
  262. id_x03 = f_out.createVariable(cvar3 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  263. id_x03.units = cu_d
  264. if not cun3=='unknown': id_x03.units = cun3
  265. id_x03.long_name = cln_d3
  266. if l_do_v4:
  267. id_x04 = f_out.createVariable(cvar4 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  268. id_x04.units = cu_d
  269. if not cun4=='unknown': id_x04.units = cun4
  270. id_x04.long_name = cln_d4
  271. if l_do_v5:
  272. id_x05 = f_out.createVariable(cvar5 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  273. id_x05.units = cu_d
  274. if not cun5=='unknown': id_x05.units = cun5
  275. id_x05.long_name = cln_d5
  276. if l_do_v6:
  277. id_x06 = f_out.createVariable(cvar6 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  278. id_x06.units = cu_d
  279. if not cun6=='unknown': id_x06.units = cun6
  280. id_x06.long_name = cln_d6
  281. if l_do_v7:
  282. id_x07 = f_out.createVariable(cvar7 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  283. id_x07.units = cu_d
  284. if not cun7=='unknown': id_x07.units = cun7
  285. id_x07.long_name = cln_d7
  286. if l_do_v8:
  287. id_x08 = f_out.createVariable(cvar8 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  288. id_x08.units = cu_d
  289. if not cun8=='unknown': id_x08.units = cun8
  290. id_x08.long_name = cln_d8
  291. if l_do_v9:
  292. id_x09 = f_out.createVariable(cvar9 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  293. id_x09.units = cu_d
  294. if not cun9=='unknown': id_x09.units = cun9
  295. id_x09.long_name = cln_d9
  296. if l_do_v10:
  297. id_x10 = f_out.createVariable(cvar10 ,'f4',('time',cydnm,cxdnm,), fill_value=missing_val, zlib=True, complevel=5)
  298. id_x10.units = cu_d
  299. if not cun7=='unknown': id_x10.units = cun10
  300. id_x10.long_name = cln_d10
  301. jrw = 0
  302. for jt in range(Nt):
  303. jrw = jrec2write + jt
  304. id_t[jrw] = vt[jt]
  305. id_x01[jrw,:,:] = xd[jt,:,:]
  306. if l_do_v2: id_x02[jrw,:,:] = xd2[jt,:,:]
  307. if l_do_v3: id_x03[jrw,:,:] = xd3[jt,:,:]
  308. if l_do_v4: id_x04[jrw,:,:] = xd4[jt,:,:]
  309. if l_do_v5: id_x05[jrw,:,:] = xd5[jt,:,:]
  310. if l_do_v6: id_x06[jrw,:,:] = xd6[jt,:,:]
  311. if l_do_v7: id_x07[jrw,:,:] = xd7[jt,:,:]
  312. if l_do_v8: id_x08[jrw,:,:] = xd8[jt,:,:]
  313. if l_do_v9: id_x09[jrw,:,:] = xd9[jt,:,:]
  314. if l_do_v10:id_x10[jrw,:,:] = xd10[jt,:,:]
  315. f_out.about = 'Diagnostics created with BaraKuda (https://github.com/brodeau/barakuda)'
  316. #
  317. #
  318. else:
  319. vtr = f_out.variables['time']
  320. jrec2write = len(vtr)
  321. x01 = f_out.variables[cvar1]
  322. if l_do_v2: x02 = f_out.variables[cvar2]
  323. if l_do_v3: x03 = f_out.variables[cvar3]
  324. if l_do_v4: x04 = f_out.variables[cvar4]
  325. if l_do_v5: x05 = f_out.variables[cvar5]
  326. if l_do_v6: x06 = f_out.variables[cvar6]
  327. if l_do_v7: x07 = f_out.variables[cvar7]
  328. if l_do_v8: x08 = f_out.variables[cvar8]
  329. if l_do_v9: x09 = f_out.variables[cvar9]
  330. if l_do_v10:x10 = f_out.variables[cvar10]
  331. jrw = 0
  332. for jt in range(Nt):
  333. jrw = jrec2write + jt
  334. vtr[jrw] = vt[jt]
  335. x01[jrw,:,:] = xd[jt,:,:]
  336. if l_do_v2: x02[jrw,:,:] = xd2[jt,:,:]
  337. if l_do_v3: x03[jrw,:,:] = xd3[jt,:,:]
  338. if l_do_v4: x04[jrw,:,:] = xd4[jt,:,:]
  339. if l_do_v5: x05[jrw,:,:] = xd5[jt,:,:]
  340. if l_do_v6: x06[jrw,:,:] = xd6[jt,:,:]
  341. if l_do_v7: x07[jrw,:,:] = xd7[jt,:,:]
  342. if l_do_v8: x08[jrw,:,:] = xd8[jt,:,:]
  343. if l_do_v9: x09[jrw,:,:] = xd9[jt,:,:]
  344. if l_do_v10:x10[jrw,:,:] = xd10[jt,:,:]
  345. f_out.close()
  346. print ' *** '+str(Nt)+' snapshots of fields saved into '+cf+' !\n'
  347. def write_2d_mask(cf_out, MSK, xlon=[], xlat=[], name='mask'):
  348. (nj,ni) = nmp.shape(MSK)
  349. f_out = Dataset(cf_out, 'w', format='NETCDF4')
  350. # Dimensions:
  351. f_out.createDimension('x', ni)
  352. f_out.createDimension('y', nj)
  353. if (xlon != []) and (xlat != []):
  354. if (xlon.shape == (nj,ni)) and (xlon.shape == xlat.shape):
  355. id_lon = f_out.createVariable('nav_lon' ,'f4',('y','x',), zlib=True, complevel=5)
  356. id_lat = f_out.createVariable('nav_lat' ,'f4',('y','x',), zlib=True, complevel=5)
  357. id_lon[:,:] = xlon[:,:]
  358. id_lat[:,:] = xlat[:,:]
  359. id_msk = f_out.createVariable(name ,'i1',('y','x',), zlib=True, complevel=5)
  360. id_msk[:,:] = MSK[:,:]
  361. f_out.about = 'Diagnostics created with BaraKuda (https://github.com/brodeau/barakuda)'
  362. f_out.close()
  363. return
  364. def dump_2d_field( cf_out, XFLD, xlon=[], xlat=[], name='field', unit='', long_name='', mask=[] ):
  365. (nj,ni) = nmp.shape(XFLD)
  366. f_out = Dataset(cf_out, 'w', format='NETCDF4')
  367. # Dimensions:
  368. f_out.createDimension('y', nj)
  369. f_out.createDimension('x', ni)
  370. if (xlon != []) and (xlat != []):
  371. if (xlon.shape == (nj,ni)) and (xlon.shape == xlat.shape):
  372. id_lon = f_out.createVariable('nav_lon' ,'f4',('y','x',), zlib=True, complevel=5)
  373. id_lat = f_out.createVariable('nav_lat' ,'f4',('y','x',), zlib=True, complevel=5)
  374. id_lon[:,:] = xlon[:,:]
  375. id_lat[:,:] = xlat[:,:]
  376. id_fld = f_out.createVariable(name ,'f4',('y','x',), zlib=True, complevel=5)
  377. if long_name != '': id_fld.long_name = long_name
  378. if unit != '': id_fld.units = unit
  379. if mask != []:
  380. idx_land = nmp.where( mask < 0.5)
  381. XFLD[idx_land] = nmp.nan
  382. id_fld[:,:] = XFLD[:,:]
  383. f_out.about = 'Diagnostics created with BaraKuda (https://github.com/brodeau/barakuda)'
  384. f_out.close()
  385. return
  386. def dump_2d_multi_field(cf_out, XFLD, vnames, vndim=[], xlon=[], xlat=[], vtime=[]):
  387. if len(vtime)>0:
  388. l_add_time = True
  389. (nbfld, Nbt, nj, ni) = nmp.shape(XFLD)
  390. if Nbt != len(vtime): print 'ERROR (dump_2d_multi_field): array and time vector disagree!'; sys.exit(0)
  391. else:
  392. l_add_time = False
  393. (nbfld, nj, ni) = nmp.shape(XFLD)
  394. vnbdim = nmp.zeros(nbfld)
  395. if vndim == []:
  396. vnbdim[:] = 3 ; # default dim is 3 (time_counter,y,x)
  397. else:
  398. nn0 = len(vndim)
  399. if nbfld != nn0: print 'ERROR (dump_2d_multi_field): vndim and main array dont agree in shape!'; sys.exit(0)
  400. vnbdim[:] = vndim[:]
  401. nf = len(vnames)
  402. if nbfld != nf: print 'ERROR (dump_2d_multi_field): list of names and main array dont agree in shape!'; sys.exit(0)
  403. f_out = Dataset(cf_out, 'w', format='NETCDF4')
  404. l_coord_2d = False
  405. cnm_dim_x = 'lon'
  406. cnm_dim_y = 'lat'
  407. if (xlon != []) and (xlat != []):
  408. if (xlon.shape == (nj,ni)) and (xlon.shape == xlat.shape):
  409. l_coord_2d = True
  410. cnm_dim_x = 'x'
  411. cnm_dim_y = 'y'
  412. # Dimensions:
  413. f_out.createDimension(cnm_dim_x, ni)
  414. f_out.createDimension(cnm_dim_y, nj)
  415. if l_add_time: f_out.createDimension('time_counter', None)
  416. if (xlon != []) and (xlat != []):
  417. if l_coord_2d:
  418. id_lon = f_out.createVariable('nav_lon' ,'f4',(cnm_dim_y,cnm_dim_x,), zlib=True, complevel=5)
  419. id_lat = f_out.createVariable('nav_lat' ,'f4',(cnm_dim_y,cnm_dim_x,), zlib=True, complevel=5)
  420. id_lon[:,:] = xlon[:,:]
  421. id_lat[:,:] = xlat[:,:]
  422. else:
  423. id_lon = f_out.createVariable(cnm_dim_x ,'f4',(cnm_dim_x,), zlib=True, complevel=5)
  424. id_lat = f_out.createVariable(cnm_dim_y ,'f4',(cnm_dim_y,), zlib=True, complevel=5)
  425. id_lon[:] = xlon[:]
  426. id_lat[:] = xlat[:]
  427. if l_add_time:
  428. id_tim = f_out.createVariable('time_counter' ,'f8',('time_counter',))
  429. id_tim[:] = vtime[:]
  430. #id_fld = nmp.zeros(nbfld, dtype=int)
  431. for jv in range(nbfld):
  432. if (not l_add_time) or (vnbdim[jv]==2):
  433. id_fld = f_out.createVariable(vnames[jv] ,'f8',(cnm_dim_y,cnm_dim_x,), zlib=True, complevel=5)
  434. if l_add_time:
  435. id_fld[:,:] = XFLD[jv,0,:,:]
  436. else:
  437. id_fld[:,:] = XFLD[jv,:,:]
  438. else:
  439. id_fld = f_out.createVariable(vnames[jv] ,'f8',('time_counter',cnm_dim_y,cnm_dim_x,), zlib=True, complevel=5)
  440. id_fld[:,:,:] = XFLD[jv,:,:,:]
  441. f_out.about = 'Diagnostics created with BaraKuda (https://github.com/brodeau/barakuda)'
  442. f_out.close()
  443. return