def plot_cs(layer_data,hrz_grid=None,ax=None,crs_data=None,crs_plot=None,show_colorbar=True):
n_cs, is_gmao = guess_cs_grid(layer_data.shape)
if is_gmao:
# Use data as-is
layer_data_3D = layer_data
else:
# Reshape data to be "GMAO-style" (3D)
layer_data_3D = np.reshape(layer_data,(6,n_cs,n_cs))
if hrz_grid is None:
# Try to figure out the grid from the layer data
hrz_grid = cubedsphere.csgrid_GMAO(n_cs)
cs_threshold = 5.0
masked_data = np.ma.masked_where(np.abs(hrz_grid['lon'] - 180.0) < cs_threshold, layer_data_3D)
im_vec = []
for i_face in range(6):
im = ax.pcolormesh(hrz_grid['lon_b'][i_face,:,:],hrz_grid['lat_b'][i_face,:,:],masked_data[i_face,:,:],transform=crs_data)
im_vec.append(im)
c_lim = [np.min(layer_data),np.max(layer_data)]
if (c_lim[0] == c_lim[1]):
c_lim = [-1.0,1.0]
for im in im_vec:
im.set_clim(c_lim)
return im_vec
def gen_grid(n_cs):
cs_temp = cubedsphere.csgrid_GMAO(n_cs)
return xr.Dataset({'nf': (['nf'],np.array(range(6))),
'Ydim': (['Ydim'],np.array(range(n_cs))),
'Xdim': (['Xdim'],np.array(range(n_cs))),
'Ydim_b': (['Ydim_b'],np.array(range(n_cs+1))),
'Xdim_b': (['Xdim_b'],np.array(range(n_cs+1))),
'lat': (['nf','Ydim','Xdim'], cs_temp['lat']),
'lon': (['nf','Ydim','Xdim'], cs_temp['lon']),
'lat_b': (['nf','Ydim_b','Xdim_b'], cs_temp['lat_b']),
'lon_b': (['nf','Ydim_b','Xdim_b'], cs_temp['lon_b']),
'area': (['nf','Ydim','Xdim'], grid_area(cs_temp))})
def regrid_cs(cs_data,new_grid,cs_grid=None,regridder_list=None):
'''
# regrid cs to ll
ds_dev_cmp = np.zeros([nlev, cmpgrid['lat'].size, cmpgrid['lon'].size])
for i in range(6):
regridder = devregridder_list[i]
ds_dev_cmp += regridder(ds_dev_reshaped[i])
'''
full_data = cs_data.copy()
# Assume the CS data is 3D
single_layer = len(full_data.shape) == 3
if single_layer:
layer_shape = full_data.shape
full_data = np.reshape(full_data,layer_shape.insert(0,1))
else:
layer_shape = full_data.shape[1:]
full_shape = full_data.shape
n_lev = full_shape[0]
# Figure out the grid size
n_cs, is_gmao = guess_cs_grid(layer_shape)
if cs_grid is None:
# Try to figure out the grid from the layer data
cs_grid = cubedsphere.csgrid_GMAO(n_cs)
n_lon = len(new_grid['lon'])
n_lat = len(new_grid['lat'])
if regridder_list is None:
regridder_list = []
for i_face in range(6):
sub_grid = {'lat': cs_grid['lat'][i_face],
'lon': cs_grid['lon'][i_face],
'lat_b': cs_grid['lat_b'][i_face],
'lon_b': cs_grid['lon_b'][i_face]}
regridder_list.append(xesmf.Regridder(sub_grid,new_grid,method='conservative',reuse_weights=True,filename='conservative_c{:d}f{:d}_{:d}x{:d}'.format(n_cs,i_face,n_lat,n_lon)))
ll_data = np.zeros((n_lev,n_lat,n_lon))
for i_lev in range(n_lev):
for i_face in range(6):
ll_data[i_lev,:,:] += regridder_list[i_face](full_data[i_lev,i_face,:,:])
if single_layer:
ll_data = np.squeeze(ll_data)
return ll_data
def grid_area(cs_grid=None,cs_res=None):
"""Return area in m2 for each cell in a cubed-sphere grid
Uses GMAO indexing convention (6xNxN)
"""
# Calculate area on a cubed sphere
if cs_res is None:
cs_res = cs_grid['lon_b'].shape[-1] - 1
elif cs_grid is None:
cs_grid = cubedsphere.csgrid_GMAO(cs_res)
elif cs_grid is not None and cs_res is not None:
assert cs_res == cs_grid['lon_b'].shape[-1], 'Routine grid_area received inconsistent inputs'
cs_area = np.zeros((6,cs_res,cs_res))
cs_area[0,:,:] = face_area(cs_grid['lon_b'][0,:,:],cs_grid['lat_b'][0,:,:])
for i_face in range(1,6):
cs_area[i_face,:,:] = cs_area[0,:,:].copy()
return cs_area
def test_csgrid_gmap():
# read reference data
ds_FV3 = cs.open_FV3data("./example_notebooks/sample_data/FV3diag_C48/",
'atmos_daily')
# calculate grid online
grid = cs.csgrid_GMAO(48)
# fix a single polar point before testing. should not have real effect
index = (5, 24, 24) # index of the pole point
assert grid['lat_b'][index] == -90.0
assert grid['lon_b'][index] == 35.0
assert ds_FV3['lon_b'][index] == 350.0
grid['lon_b'][index] = 350.0 # change to FV3 value
for varname in ['lon', 'lon_b', 'lat', 'lat_b']:
assert_almost_equal(grid[varname], ds_FV3[varname], decimal=4)
import numpy as np
import gcgridobj
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cubedsphere
# Generate test grids
# Cubed sphere grids (C24, C90)
c24 = cubedsphere.csgrid_GMAO(24)
c90 = cubedsphere.csgrid_GMAO(90)
# Lat-lon grids (Generic 2x2, standard 4x5)
ll2x2 = gcgridobj.latlontools.gen_hrz_grid(lon_stride=2,
lat_stride=2,
half_polar=False,
center_180=False)
ll4x5 = gcgridobj.gc_horizontal.gmao_4x5_global
# Lat-lon grid data comes pre-packaged with areas in m2 ("ll2x2['area']), but for cubed sphere we must generate it ourselves
c24_area = gcgridobj.cstools.calc_cs_area(c24)
c90_area = gcgridobj.cstools.calc_cs_area(c90)
# Regridder objects, all generated through a single helper function
# Lat-lon (2x2) to cubed-sphere (c90)
ro_2x2_c90 = gcgridobj.regrid.gen_regridder(ll2x2, c90)
# Cubed-sphere (c90) to cubed-sphere (c24)
ro_c90_c24 = gcgridobj.regrid.gen_regridder(c90, c24)
# Cubed-sphere (c24) to lat-lon (2x2)
ro_c24_2x2 = gcgridobj.regrid.gen_regridder(c24, ll2x2)
# Lat-lon (2x2) to lat-lon (4x5)
ro_2x2_4x5 = gcgridobj.regrid.gen_regridder(ll2x2, ll4x5)