jweyn / DLWP-CS

Deep learning models for global weather prediction on a cubed sphere
GNU General Public License v3.0
92 stars 40 forks source link

How to visualize the Figure 1b as in the paper #3

Closed manmeet3591 closed 4 years ago

manmeet3591 commented 4 years ago

I am unable to plot the cubed sphere projection as shown in Figure 1b of the paper. Can you please help me with the same.

jweyn commented 4 years ago

I remember not successfully being able to use cartopy for this particular projection, so I just plotted by lat/lon coordinates and used a land-sea mask contour.

Modify paths etc. per your needs. You'll need to have some data and the land-sea mask in cubed sphere format.

import matplotlib.pyplot as plt
import xarray as xr
import numpy as np

root_directory = '/home/disk/wave2/jweyn/Data/DLWP'
predictor_file = '/home/gold/jweyn/Data/era5_2deg_3h_CS2_1979-2018_z-tau-t2_500-1000_tcwv_psi850.nc'
validation_file = '%s/era5_2deg_3h_validation_z500_t2m_ILL.nc' % root_directory
lsm_file = '/home/gold/jweyn/Data/era5_2deg_3h_CS2_land_sea_mask.nc'

data = xr.open_dataset(predictor_file)
scale = xr.open_dataset(validation_file)
lsm = xr.open_dataset(lsm_file)

field = data.predictors.sel(sample='2018-01-05 00:00', varlev='t2m/0') * scale['std'].sel(varlev='t2m/0') \
    + scale['mean'].sel(varlev='t2m/0')

square_lon = np.zeros(data.lon.shape)
square_lat = np.zeros(data.lat.shape)
n_side = data.lon.shape[1]
for f in range(4):
    square_lon[f], square_lat[f] = np.meshgrid(np.linspace(90. * f, 90. * (f + 1), n_side),
                                               np.linspace(-45, 45, n_side))
square_lon[4], square_lat[4] = np.meshgrid(np.linspace(0., 90., n_side),
                                           np.linspace(-135., -45., n_side))
square_lon[5], square_lat[5] = np.meshgrid(np.linspace(0., 90., n_side),
                                           np.linspace(45., 135., n_side))

fig = plt.figure(figsize=(10, 6))
ax = plt.gca()
for face in range(6):
    cf = ax.contourf(square_lon[face], square_lat[face], field.isel(face=face).values,
                     np.arange(240, 311, 5), cmap='Spectral_r', extend='both')
    ax.contour(square_lon[face], square_lat[face], lsm['lsm'].isel(face=face).values,
               [0.5], colors='k')

plt.colorbar(cf, label='$T_2$ (K)')
ax.set_aspect('equal', 'box')
plt.savefig('cs-example-t2m.pdf', dpi=200, bbox_inches='tight')
plt.show()
manmeet3591 commented 4 years ago

@jweyn Thanks. This worked.