Created
March 23, 2023 22:30
-
-
Save raphaeldussin/ba99af6acfc70d66774697072cc80c40 to your computer and use it in GitHub Desktop.
rotating globe movie for surface speed
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import xarray as xr | |
import numpy as np | |
from matplotlib import pyplot as plt | |
from matplotlib import animation | |
import matplotlib.image as image | |
import cartopy.crs as ccrs | |
from owslib.wms import WebMapService | |
from matplotlib import cm | |
import os | |
from xgcm import Grid | |
import warnings | |
warnings.filterwarnings("ignore") | |
# output directory for pngs | |
outdir = '/local2/home/PLOTS_OM4p25/hist' | |
# start/end frame numbers | |
ntstart = 1 | |
ntend = 365 * 165 | |
# ---------------- read data ----------------------------------------------------------------------------- | |
ppdir = "/archive/Raphael.Dussin/FMS2019.01.03_devgfdl_20221223/CM4_historical_c192_OM4p25/gfdl.ncrc4-intel18-prod-openmp/pp/ocean_daily" | |
ds_grid = xr.open_dataset(f'{ppdir}/ocean_daily.static.nc') | |
dsu = xr.open_mfdataset(f'{ppdir}/ts/daily/5yr/ocean_daily.*.ssu.nc', chunks={'time':12}) | |
dsv = xr.open_mfdataset(f'{ppdir}/ts/daily/5yr/ocean_daily.*.ssv.nc', chunks={'time':12}) | |
# ---------------- compute speed from u and v ------------------------------------------------------------ | |
grid = Grid(ds_grid, coords={'X': {'center': 'xh', 'outer': 'xq'}, | |
'Y': {'center': 'yh', 'outer': 'yq'} }, periodic=[]) | |
speed2 = (grid.interp(dsu.ssu, 'X', boundary='fill') * grid.interp(dsu.ssu, 'X', boundary='fill') ) + \ | |
(grid.interp(dsv.ssv, 'Y', boundary='fill') * grid.interp(dsv.ssv, 'Y', boundary='fill') ) | |
speed2 = speed2.assign_coords({'geolon': ds_grid['geolon'], | |
'geolat': ds_grid['geolat']}) | |
speed = xr.apply_ufunc(np.sqrt, speed2, dask='parallelized', output_dtypes=[speed2.dtype]) | |
# ----------------- First set up the figure, the axis, and the plot element we want to animate ------------ | |
# blue marble for land | |
url = 'http://map1c.vis.earthdata.nasa.gov/wmts-geo/wmts.cgi' | |
# logo | |
im1 = image.imread('../../logos/noaaLogo_NoLetters_XL.png') | |
ct=0 # counter for image | |
for kt in range(ntstart, ntend): | |
lonstart=-50.0 | |
latstart=60.0 | |
lonframe=np.mod(lonstart-0.05*ct, 360) | |
latframe=latstart*np.cos(0.33*2*np.pi*ct/365/10) | |
subplot_kws={'projection': ccrs.NearsidePerspective(central_longitude=lonframe, | |
central_latitude=latframe), | |
'facecolor':'grey'} | |
fig = plt.figure() | |
ax = plt.axes(**subplot_kws) | |
frame = str(ct).zfill(6) | |
fileout = f'{outdir}/rotate_speed_frame_{frame}.png' | |
if not os.path.exists(fileout): | |
print(f"generating figure for frame {kt}") | |
data = speed.isel({'time': kt}) | |
p = data.plot(ax=ax,x='geolon', y='geolat', | |
vmin=0, vmax=1, | |
cmap=cm.Blues_r, extend='both', | |
#subplot_kws=subplot_kws, | |
transform=ccrs.PlateCarree(), | |
add_colorbar=False ) | |
cb = plt.colorbar(p, ticks=[0, 0.25, 0.5, 0.75, 1.0], | |
shrink=0.99) | |
cb.ax.tick_params(labelsize=18, labelcolor='w') | |
cb.set_label('Ocean Surface Speed (m/s)', color='w', fontsize=14) | |
run='NOAA-GFDL CM4 c192p25 historical' | |
date=str(dsu['time'].isel(time=kt).dt.strftime("%Y %m %d").values) | |
plt.title(run + '\n' + date, fontsize=16, color='w') | |
# background | |
p.axes.add_wmts(url, 'BlueMarble_NextGeneration') | |
newax = fig.add_axes([0.01, 0.01, 0.1, 0.1]) | |
newax.imshow(im1) | |
newax.axis('off') | |
plt.savefig(f'{fileout}', facecolor='black', edgecolor='black', | |
format='png', dpi=300) | |
plt.close() | |
p = None | |
cb = None | |
ct=ct+1 | |
# after all the images are generated, make movie with ffmpeg: | |
# ffmpeg -r 15 -f image2 -s 1920x1080 -i rotate_speed_frame_%06d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p movie.mp4 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Hi @raphaeldussin !
I made a bit more changes to ensure each frame being generated within 3 seconds on my Mac:
p.axes.add_wmts(url, 'BlueMarble_NextGeneration')
withp.axes.background_img(name='BM', resolution='high')
Thanks!