wrf-python/test/ipynb/loop_and_fill.ipynb

In [24]:

from __future__ import print_function, division

import os
import numpy as np
from netCDF4 import Dataset
from wrf import getvar, ALL_TIMES, to_np
import xarray

path = '/scratch/mawagner/2015_GRELL3D/TESTFILES/TestTime'
filename_list = os.listdir(path)
filename_list.sort()

#filename_list = ["/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_00:00:00",
#                "/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_03:00:00",
#                "/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_06:00:00",
#                "/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_09:00:00"]

# Result shape 
result_shape = (6, 1, 290, 265)
#result_shape = (4, 1, 96, 96)

# Let's get the first time so we can copy the metadata later
f = Dataset(filename_list[0])
# By setting squeeze to False, you'll get all the dimension names.
z1 = getvar(f, "T2", 0, squeeze=False)
xlat = getvar(f, "XLAT", 0)
xlong = getvar(f, "XLONG", 0)


z_final = np.empty(result_shape, np.float32)

# Modify this number if using more than 1 time per file
times_per_file = 1
#times_per_file = 4

data_times = []
xtimes = []
for timeidx in range(result_shape[0]):
    # Compute the file index and the time index inside the file
    fileidx = timeidx // times_per_file
    file_timeidx = timeidx % times_per_file

    f = Dataset(filename_list[fileidx])
    z = getvar(f, "T2", file_timeidx)
    t = getvar(f, "Times", file_timeidx)
    xt = getvar(f, "xtimes", file_timeidx)
    data_times.append(to_np(t))
    xtimes.append(to_np(xt))
    z_final[timeidx,:] = z[:]
    f.close()
    
# Let's make the metadata. Dimension names should copy easily if you set sqeeze to False, 
# otherwise you can just set them yourself is a tuple of dimension names. Since you wanted
# to keep the bottom_top dimension for this 2D variable (which is normally removed), 
# I'm doing this manually.
z_dims = ["Time", "bottom_top", "south_north", "west_east"]

# Xarray doesn't copy coordinates easily (it always complains about shape mismatches), so do this
# manually
z_coords = {}
z_coords["Time"] = data_times
z_coords["XTIME"] = ("Time",), xtimes
z_coords["XLAT"] = ("south_north", "west_east"), xlat
z_coords["XLONG"] = ("south_north", "west_east"), xlong
z_name = "T2"

# Attributes copy nicely
z_attrs = {}
z_attrs.update(z1.attrs)

z_with_meta = xarray.DataArray(z_final, coords=z_coords, dims=z_dims, attrs=z_attrs, name=z_name)

print(z_with_meta)
    
    

<xarray.DataArray 'T2' (Time: 4, bottom_top: 1, south_north: 96, west_east: 96)>
array([[[[303.6255 , ..., 303.81546],
         ...,
         [304.7401 , ..., 300.15717]]],


       ...,


       [[[301.89346, ..., 302.56534],
         ...,
         [302.61246, ..., 298.01028]]]], dtype=float32)
Coordinates:
  * Time     (Time) datetime64[ns] 2005-08-28 2005-08-28T03:00:00 ...
    XTIME    (Time) float64 0.0 180.0 360.0 540.0
    XLAT     (south_north, west_east) float32 21.302004 21.302004 21.302004 ...
    XLONG    (south_north, west_east) float32 -90.57406 -90.484116 ...
Dimensions without coordinates: bottom_top, south_north, west_east
Attributes:
    FieldType:    104
    MemoryOrder:  XY 
    description:  TEMP at 2 M
    units:        K
    stagger:      
    coordinates:  XLONG XLAT XTIME
    projection:   Mercator(stand_lon=-89.0, moad_cen_lat=27.99999237060547, t...