snapshot...

modules/amv/intercompare.py

@@ -13,6 +13,7 @@ import h5py
 import time
 from scipy.interpolate import interp1d
 from util.util import minimize_quadratic
+from netCDF4 import Dataset
 
 #-- AMV intercompare stuff ------------------------------------------
 
@@ -106,6 +107,116 @@ def filter_amvs(amvs, qitype=amv_cqi_idx, qival=50, lat_range=None):
     else:
         return amvs[keep, :]
 
+def get_raob_dict_cdf(filename):
+    rtgrp = Dataset(filename, 'r', format='NETCDF3')
+
+    sta_lats_var = rtgrp['staLat']
+    sta_lats_var.set_always_mask(False)
+    sta_lats_var.set_auto_mask(False)
+
+    sta_lons_var = rtgrp['staLon']
+    sta_lons_var.set_always_mask(False)
+    sta_lons_var.set_auto_mask(False)
+
+    sta_lats = sta_lats_var[:]
+    sta_lons = sta_lons_var[:]
+
+    sta_msk = np.logical_and(sta_lats > -90.0, sta_lats < 90.0)
+    # sta_msk = np.invert(sta_lats._mask)
+    num_sta = np.sum(sta_msk)
+    sta_lats = sta_lats[sta_msk]
+    sta_lons = sta_lons[sta_msk]
+
+    # mandatory levels
+    man_levs = rtgrp['prMan']
+    man_levs = man_levs[:, :]
+    num_man_levs = rtgrp['numMand']
+    num_man_levs = num_man_levs[:]._data
+    man_temp = rtgrp['tpMan']
+    man_temp = man_temp[:, :]
+    man_spd = rtgrp['wsMan']
+    man_spd = man_spd[:, :]
+    man_dir = rtgrp['wdMan']
+    man_dir = man_dir[:, :]
+
+    # significant levels
+    sig_levs = rtgrp['prSigW']
+    sig_levs = sig_levs[:, :]
+    num_sig_levs = rtgrp['numSigW']
+    num_sig_levs = num_sig_levs[:]._data
+    sig_temp = rtgrp['tpSigW']
+    sig_temp = sig_temp[:, :]
+    sig_spd = rtgrp['wsSigW']
+    sig_spd = sig_spd[:, :]
+    sig_dir = rtgrp['wdSigW']
+    sig_dir = sig_dir[:, :]
+
+    man_levs = man_levs[sta_msk,]
+    num_man_levs = num_man_levs[sta_msk]
+
+    sig_levs = sig_levs[sta_msk,]
+    num_sig_levs = num_sig_levs[sta_msk]
+
+    man_temp = man_temp[sta_msk,]
+    man_spd = man_spd[sta_msk,]
+    man_dir = man_dir[sta_msk,]
+
+    sig_temp = sig_temp[sta_msk,]
+    sig_spd = sig_spd[sta_msk,]
+    sig_dir = sig_dir[sta_msk,]
+
+    raob_dct = {}
+    for k in range(num_sta):
+        lon = sta_lons[k]
+        lat = sta_lats[k]
+
+        vld_man = np.invert(man_levs[k]._mask)
+        vld_man = np.logical_and(vld_man, np.invert(man_spd[k]._mask))
+        vld_man = np.logical_and(vld_man, np.invert(man_dir[k]._mask))
+
+        vld_sig = np.invert(sig_levs[k]._mask)
+        vld_sig = np.logical_and(vld_sig, np.invert(sig_spd[k]._mask))
+        vld_sig = np.logical_and(vld_sig, np.invert(sig_dir[k]._mask))
+
+        vld_man_levs = man_levs[k, vld_man]
+        vld_man_spd = man_spd[k, vld_man]
+        vld_man_dir = man_dir[k, vld_man]
+        vld_man_temp = man_temp[k, vld_man]
+
+        s_idxs = np.argsort(vld_man_levs)
+        s_man_levs = vld_man_levs[s_idxs]
+        s_man_spd = vld_man_spd[s_idxs]
+        s_man_dir = vld_man_dir[s_idxs]
+        s_man_temp = vld_man_temp[s_idxs]
+
+        vld_sig_levs = sig_levs[k, vld_sig]
+        vld_sig_spd = sig_spd[k, vld_sig]
+        vld_sig_dir = sig_dir[k, vld_sig]
+        vld_sig_temp = sig_temp[k, vld_sig]
+
+        all_levs = np.append(s_man_levs, vld_sig_levs)
+        all_temp = np.append(s_man_temp, vld_sig_temp)
+        all_spd = np.append(s_man_spd, vld_sig_spd)
+        all_dir = np.append(s_man_dir, vld_sig_dir)
+
+        if len(all_levs) == 0:
+            continue
+
+        s_idxs = np.argsort(all_levs)
+        all_levs = all_levs[s_idxs]
+        all_temp = all_temp[s_idxs]
+        all_spd = all_spd[s_idxs]
+
+        all_levs = all_levs[::-1]
+        all_temp = all_temp[::-1]
+        all_spd = all_spd[::-1]
+        all_dir = all_dir[::-1]
+
+        raob = np.stack([all_levs._data, all_temp._data, all_dir._data, all_spd._data], axis=1)
+
+        raob_dct[(lat, lon)] = raob
+
+    return raob_dct
 
 def get_raob_dict(filename):
     header = None