In [ ]:
import os, sys
sys.path.append('/nfs/see-fs-01_users/eepdw/python_scripts/modules')

from update_pp_cube_coords import update_coordsimport iris

import iris.analysis.cartography

import h5py

import numpy as np

import pdb

import scipy

#Load specific humidity and wind

#/nfs/a90/eepdw/Data/EMBRACE/Pressure_level\_means/sp_hum_pressure_levels_interp_djzns_mean_masked/

#experiment_ids = ['djznw', 'djznq', 'djzny', 'djzns', 'dkmbq', 'dklyu', 'dklwu', 'dklzq']
experiment_ids = ['dkmbq', 'dklyu']
p_levels = [1000, 950, 925, 850, 700, 500, 400, 300, 250, 200, 150, 100, 70, 50, 30, 20, 10]

for experiment_id in experiment_ids:
    
   expmin1 = experiment_id[:-1]

   fname_h = '/nfs/a90/eepdw/Data/EMBRACE/On_Heights_Interpolation_Data/sp_hum_pressure_levels_interp_%s' % (experiment_id)

   with h5py.File(fname_h, 'r') as i:
        
#  q = i['%s' % 'mean'][. . .]
    q = i['%s' % 'sh_on_p'][. . .]
   pdb.set_trace()


   f_oro =  '/nfs/a90/eepdw/Data/EMBRACE/Mean_State/pp_files/%s/%s/33.pp' % (expmin1, experiment_id)
   oro = iris.load_cube(f_oro)
   oro,lats,lons = update_coords(oro)

   fu = '/nfs/a90/eepdw/Data/EMBRACE/Mean_State/pp_files/%s/%s/30201.pp' % (expmin1, experiment_id)
    
   u_wind,v_wind = iris.load(fu)

   u_wind,lats_w,lons_w = update_coords(u_wind)
   v_wind,lats_w,lons_w = update_coords(v_wind)

   print u_wind
   print u_wind.coord('pressure')

   qu_div = np.empty((u_wind.shape[1], u_wind.shape[2], u_wind.shape[0]))
   qv_div = np.empty((u_wind.shape[1], u_wind.shape[2], u_wind.shape[0]))
 
   fl_la_lo = (lats.flatten(),lons.flatten())

   p_lev_delta = np.diff( np.append(u_wind.coord('pressure').points, u_wind.coord('pressure').points[-1]))
   
   for p, pressure_cube in enumerate(u_wind.slices(['grid_latitude', 'grid_longitude'])):
    
    print p
    s = np.searchsorted(p_levels[::-1], p)
    #sc =  np.searchsorted(p_levs, p)

    q_slice = q[:,:,-(s+1)]
    q_interp = scipy.interpolate.griddata(fl_la_lo, q_slice.flatten(), (lats_w, lons_w), method='linear')

    #pdb.set_trace()
    qu_div[:,:,pn] = (u_wind.coord('pressure').points[pn]*q_interp)/9.81
    qv_div[:,:,pn] = (v_wind.coord('pressure').points[pn]*q_interp)/9.81

   Qu_div = np.sum(qu_div*p_lev_delta, axis=-1)
   Qv_div = np.sum(qv_div*p_lev_delta, axis=-1)

In [2]:
q.shape


Out[2]:
(504, 600, 600, 17)

In [4]:
print q[p,:,:,-(s+1)].shape


(600, 600)

In [5]:
pressure_cube.shape


Out[5]:
(599, 600)

In [7]:
q_interp = scipy.interpolate.griddata(fl_la_lo, q[p,:,:,-(s+1)].flatten(), (lats_w, lons_w), method='linear')

In [8]:
print q_interp.shape


(599, 600)

In [9]:
qu_div[:,:,pn] = (*q_interp)/9.81
qv_div[:,:,pn] = (v_wind.coord('pressure').points[pn]*q_interp)/9.81


---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-9-cd1c87289d7b> in <module>()
----> 1 qu_div[:,:,pn] = (u_wind.coord('pressure').points[pn]*q_interp)/9.81
      2 qv_div[:,:,pn] = (v_wind.coord('pressure').points[pn]*q_interp)/9.81

NameError: name 'pn' is not defined

In [10]:
print u_wind


x_wind / (m s-1)                    (forecast_period: 6; forecast_reference_time: 84; pressure: 12; grid_latitude: 599; grid_longitude: 600)
     Dimension coordinates:
          forecast_period                           x                           -             -                  -                    -
          forecast_reference_time                   -                           x             -                  -                    -
          pressure                                  -                           -             x                  -                    -
          grid_latitude                             -                           -             -                  x                    -
          grid_longitude                            -                           -             -                  -                    x
     Auxiliary coordinates:
          time                                      x                           x             -                  -                    -
     Attributes:
          STASH: m01s30i201
          source: Data from Met Office Unified Model 8.02

In [14]:
print time_cube


x_wind / (m s-1)                    (pressure: 12; grid_latitude: 599; grid_longitude: 600)
     Dimension coordinates:
          pressure                           x                  -                    -
          grid_latitude                      -                  x                    -
          grid_longitude                     -                  -                    x
     Scalar coordinates:
          forecast_period: 6.0 hours
          forecast_reference_time: 2011-09-07 18:00:00
          time: 2011-09-08 00:00:00
     Attributes:
          STASH: m01s30i201
          source: Data from Met Office Unified Model 8.02

In [15]:
print q.shape


(504, 600, 600, 17)

In [22]:
print pressure_cube


x_wind / (m s-1)                    (grid_latitude: 599; grid_longitude: 600)
     Dimension coordinates:
          grid_latitude                           x                    -
          grid_longitude                          -                    x
     Scalar coordinates:
          forecast_period: 6.0 hours
          forecast_reference_time: 2011-09-07 18:00:00
          pressure: 1000.0 hPa
          time: 2011-09-08 00:00:00
     Attributes:
          STASH: m01s30i201
          source: Data from Met Office Unified Model 8.02

In [15]:
pressure_cube.coord('pressure').points[0]


Out[15]:
100.0

In [14]:
pq= iris.Constraint(pressure=pressure_cube.coord('pressure').points[0])
v_slice = time_cube.extract(pq)
print v_slice


x_wind / (m s-1)                    (grid_latitude: 599; grid_longitude: 600)
     Dimension coordinates:
          grid_latitude                           x                    -
          grid_longitude                          -                    x
     Scalar coordinates:
          forecast_period: 0.999999996275 hours
          forecast_reference_time: 2011-08-18 00:00:00
          pressure: 100.0 hPa
          time: 2011-08-18 01:00:00
     Attributes:
          STASH: m01s30i201
          source: Data from Met Office Unified Model 8.02

In [20]:
print qu_div


[[[  2.35882378e-310   2.35885344e-310   2.35885344e-310   2.35885344e-310
     0.00000000e+000   7.95445690e-322]
  [  2.35885346e-310   2.35885346e-310   3.64908000e+005   6.37344683e-322
     2.35885346e-310   2.35882378e-310]
  [  3.64912000e+005   4.79243676e-322   2.35885346e-310   2.35885345e-310
     3.64912000e+005   3.21142670e-322]
  ..., 
  [  2.35885345e-310   2.35885345e-310   2.37151510e-322   3.16696079e-321
     2.35885345e-310   2.35882378e-310]
  [  0.00000000e+000   0.00000000e+000   2.35885346e-310   2.35882795e-310
     5.65176949e-315   4.79243676e-322]
  [  2.35885346e-310   2.35885345e-310   1.58101007e-322   1.77863633e-322
     2.35885346e-310   2.35882795e-310]]

 [[  3.16202013e-322   1.77863633e-322   2.35885345e-310   2.35882795e-310
     7.90505033e-322   1.77863633e-322]
  [  2.35885346e-310   2.35882795e-310   9.48606040e-322   1.77863633e-322
     2.35885346e-310   2.35882795e-310]
  [  3.64908000e+005   7.95445690e-322   2.35885346e-310   2.35885346e-310
     3.64908000e+005   6.37344683e-322]
  ..., 
  [  2.35885346e-310   2.35882378e-310   3.64912000e+005   1.63041663e-322
     2.35885346e-310   2.35885346e-310]
  [  6.32404027e-322   1.77863633e-322   2.35885346e-310   2.35882795e-310
     5.45448473e-321   2.56914136e-322]
  [  2.35885346e-310   2.35885345e-310   2.35885345e-310   2.35885345e-310
     2.37151510e-322   2.77170827e-321]]

 [[  2.35885345e-310   2.35882378e-310   0.00000000e+000   0.00000000e+000
     2.35885346e-310   2.35882795e-310]
  [  5.61841903e-315   1.63041663e-322   2.35885346e-310   2.35885346e-310
     4.74303020e-322   2.56914136e-322]
  [  2.35885346e-310   2.35885345e-310   2.35885345e-310   2.35885345e-310
     2.37151510e-322   2.06025374e-321]
  ..., 
  [  2.35885346e-310   2.35885346e-310   4.00000000e+000   4.79243676e-322
     2.35885346e-310   2.35885346e-310]
  [  3.64908000e+005   3.21142670e-322   2.35885346e-310   2.35882378e-310
     3.64913000e+005   1.63041663e-322]
  [  2.35885346e-310   2.35885346e-310   8.45840386e-321   2.56914136e-322
     2.35885346e-310   2.35885345e-310]]

 ..., 
 [[  3.64926000e+005   1.63041663e-322   2.35885346e-310   2.35885346e-310
     3.25688074e-320   2.56914136e-322]
  [  2.35885346e-310   2.35885346e-310   2.35885346e-310   2.35885346e-310
     2.35885346e-310   6.24993042e-321]
  [  2.35885346e-310   2.35885348e-310   0.00000000e+000   0.00000000e+000
     2.37151510e-322   1.77863633e-322]
  ..., 
  [  7.11454530e-322   2.56914136e-322   2.35885346e-310   2.35885346e-310
     2.35885346e-310   2.35885346e-310]
  [  1.97626258e-321   1.77863633e-322   2.35885346e-310   2.35882795e-310
     6.00000000e+000   4.79243676e-322]
  [  2.35885348e-310   2.35885346e-310   1.58101007e-322   1.77863633e-322
     2.35885346e-310   2.35882795e-310]]

 [[  3.64926000e+005   1.63041663e-322   2.35885346e-310   2.35885346e-310
     2.60866661e-321   1.77863633e-322]
  [  2.35885346e-310   2.35882795e-310   6.00000000e+000   2.53455676e-321
     2.35885346e-310   2.35885346e-310]
  [  0.00000000e+000   0.00000000e+000   2.35885346e-310   2.35882795e-310
     3.64926000e+005   1.63041663e-322]
  ..., 
  [  2.35885346e-310   2.35885346e-310   6.32404027e-322   2.56914136e-322
     2.35885346e-310   2.35885346e-310]
  [  2.35885346e-310   2.35885346e-310   8.69555537e-322   2.56914136e-322
     2.35885346e-310   2.35885346e-310]
  [  2.35885346e-310   2.35885346e-310   2.37151510e-322   1.63041663e-322
     2.35885346e-310   2.35885346e-310]]

 [[  5.29638372e-321   1.77863633e-322   2.35885348e-310   2.35882795e-310
     5.45448473e-321   1.77863633e-322]
  [  2.35885348e-310   2.35882795e-310   3.64926000e+005   6.37344683e-322
     2.35885348e-310   2.35885348e-310]
  [  3.64927000e+005   4.79243676e-322   2.35885348e-310   2.35885348e-310
     1.58101007e-322   1.77863633e-322]
  ..., 
  [  1.02765654e-321   2.56914136e-322   2.35885346e-310   2.35885346e-310
     2.35885346e-310   2.35885346e-310]
  [  2.37151510e-322   1.42784972e-321   2.35885346e-310   2.35885346e-310
     1.58101007e-322   1.77863633e-322]
  [  2.35885346e-310   2.35882795e-310   3.16202013e-322   1.77863633e-322
     2.35885346e-310   2.35882795e-310]]]

In [21]:
u_wind.shape


Out[21]:
(6, 84, 12, 599, 600)

In [22]:
u_wind


Out[22]:
<iris 'Cube' of x_wind / (m s-1) (forecast_period: 6; forecast_reference_time: 84; pressure: 12; grid_latitude: 599; grid_longitude: 600)>

In [ ]: