snapshot...

modules/util/gfs_reader.py

@@ -20,6 +20,7 @@ lon_coords = np.linspace(0, 359.5, NX)
 plevs = np.array([10.0, 20.0, 30.0, 50.0, 70.0, 100.0, 150.0, 200.0, 250.0, 300.0,
                   350.0, 400.0, 450.0, 500.0, 550.0, 600.0, 650.0, 700.0, 750.0, 800.0,
                   850.0, 900.0, 925.0, 950.0, 975.0, 1000.0])
+NZ = plevs.shape[0]
 
 
 class MyGenericException(Exception):
@@ -310,11 +311,23 @@ def get_voxel(xr_dataset, fld_name, lon, lat, press, x_width=5, y_width=5, z_wid
     x_h = int(x_width / 2)
     p_h = int(z_width / 2)
 
-    sub_fld = fld[y_c-y_h:y_c+y_h+1, x_c-x_h:x_c+x_h+1, p_c-p_h:p_c+p_h+1]
+    y_start = y_c - y_h
+    x_start = x_c - x_h
+    z_start = p_c - p_h
+    if y_start < 0 or x_start < 0 or z_start < 0:
+        return None
+
+    y_stop = y_c + y_h + 1
+    x_stop = x_c + x_h + 1
+    z_stop = p_c + p_h + 1
+    if y_stop > NY-1 or x_stop > NX-1 or z_stop > NZ-1:
+        return None
+
+    sub_fld = fld[y_start:y_stop, x_start:x_stop, z_start:z_stop]
 
     sub_fld = sub_fld.expand_dims('channel')
-    sub_fld = sub_fld.assign_coords(channel=[fld_name], fakeDim2=lon_coords[x_c-x_h:x_c+x_h+1],
-                                    fakeDim1=lat_coords[y_c-y_h:y_c+y_h+1], fakeDim0=plevs[p_c-p_h:p_c+p_h+1])
+    sub_fld = sub_fld.assign_coords(channel=[fld_name], fakeDim2=lon_coords[x_start:x_stop],
+                                    fakeDim1=lat_coords[y_start:y_stop], fakeDim0=plevs[z_start:z_stop])
 
     return sub_fld
 
@@ -345,14 +358,26 @@ def get_voxel_s(xr_dataset, fld_name_s, lon, lat, press, x_width=5, y_width=5, z
     x_h = int(x_width / 2)
     p_h = int(z_width / 2)
 
+    y_start = y_c - y_h
+    x_start = x_c - x_h
+    z_start = p_c - p_h
+    if y_start < 0 or x_start < 0 or z_start < 0:
+        return None
+
+    y_stop = y_c + y_h + 1
+    x_stop = x_c + x_h + 1
+    z_stop = p_c + p_h + 1
+    if y_stop > NY-1 or x_stop > NX-1 or z_stop > NZ-1:
+        return None
+
     sub_fld_s = []
     for name in fld_name_s:
         fld = xr_dataset[name]
-        sub_fld = fld[y_c-y_h:y_c+y_h+1, x_c-x_h:x_c+x_h+1, p_c-p_h:p_c+p_h+1]
+        sub_fld = fld[y_start:y_stop, x_start:x_stop, z_start:z_stop]
         sub_fld_s.append(sub_fld)
     sub_fld = xr.concat(sub_fld_s, 'channel')
 
-    sub_fld = sub_fld.assign_coords(channel=fld_name_s, fakeDim2=lon_coords[x_c-x_h:x_c+x_h+1],
-                                    fakeDim1=lat_coords[y_c-y_h:y_c+y_h+1], fakeDim0=plevs[p_c-p_h:p_c+p_h+1])
+    sub_fld = sub_fld.assign_coords(channel=fld_name_s, fakeDim2=lon_coords[x_start:x_stop],
+                                    fakeDim1=lat_coords[y_start:y_stop], fakeDim0=plevs[z_start:z_stop])
 
     return sub_fld
\ No newline at end of file