Last active
March 8, 2024 02:52
-
-
Save ashao/d9ef630f6567a4c44a50dfcdcf8a933e to your computer and use it in GitHub Desktop.
Make animations of OM4_025 in parallel
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 cartopy.crs as ccrs | |
import cartopy.feature as cfeature | |
import xarray as xr | |
import matplotlib.pyplot as plt | |
import matplotlib | |
import cmocean | |
import os | |
import json | |
import numpy as np | |
import pathlib | |
import urllib | |
import zipfile | |
import multiprocessing | |
import tqdm | |
import argparse | |
# TODO: Make these runtime parameters | |
data_path = "./" | |
image_path = "./images" | |
force_write = False | |
plot_options = { | |
"global": { | |
"proj":ccrs.Robinson(central_longitude=200), | |
"plot_kwargs":{ | |
'tos':{ | |
"vmin": -2, | |
"vmax": 34, | |
"cmap": cmocean.cm.thermal, | |
}, | |
'zos': { | |
"vmin": -2, | |
"vmax": 2, | |
"cmap": cmocean.cm.tarn, | |
} | |
} | |
}, | |
"southern_ocean": { | |
"proj":ccrs.Orthographic(central_latitude=-90), | |
"plot_kwargs":{ | |
'tos':{ | |
"vmin": -2, | |
"vmax": 20, | |
"cmap": cmocean.cm.thermal, | |
}, | |
'zos': { | |
"vmin": -2, | |
"vmax": 0, | |
"cmap": cmocean.cm.tarn, | |
} | |
} | |
}, | |
"north_atlantic": { | |
"proj":ccrs.NearsidePerspective(central_longitude=-50, central_latitude=30, satellite_height=35785831*0.25), | |
"plot_kwargs":{ | |
'tos':{ | |
"vmin": -2, | |
"vmax": 28, | |
"cmap": cmocean.cm.thermal, | |
}, | |
'zos': { | |
"vmin": -1, | |
"vmax": 1, | |
"cmap": cmocean.cm.tarn, | |
} | |
} | |
}, | |
} | |
def plot_image(plot_type, grid, time_level): | |
proj = plot_options[plot_type]["proj"] | |
plot_kwargs = plot_options[plot_type]["plot_kwargs"] | |
replace_deg = lambda a: a.replace('deg',u'\N{DEGREE SIGN}') | |
for var in plot_kwargs: | |
outfile = f"{var}_{time_level:08d}.png" | |
outpath = pathlib.Path(image_path, plot_type, var, outfile) | |
ds = xr.open_mfdataset(f'{data_path}/output/*.ocean_daily.nc').isel(time=time_level).load() | |
if not outpath.exists() or force_write: | |
fig = plt.figure(figsize=(12,8)) | |
ax = plt.axes(projection=proj) | |
im = ax.pcolormesh( | |
grid.geolon, | |
grid.geolat, | |
ds[var], | |
transform=ccrs.PlateCarree(), | |
**plot_kwargs[var] | |
) | |
ax.set_title(ds.time.values.item().strftime('%b %d')) | |
fig.colorbar( | |
im, | |
ax = ax, | |
label = replace_deg(f'{ds[var].attrs["long_name"]} [{ds[var].attrs["units"]}]'), | |
shrink = 0.75, | |
orientation = "horizontal", | |
pad = 0.05, | |
) | |
ax.background_img(name='NaturalEarthRelief', resolution='high') | |
ax.coastlines(linewidth=0.5) | |
fig.savefig(outpath, dpi=300) | |
plt.close(fig) | |
ds.close() | |
def setup_background(): | |
# Download background image if necessary | |
background_path = pathlib.Path(f"{data_path}/Backgrounds") | |
background_path.mkdir(parents=True, exist_ok=True) | |
os.environ['CARTOPY_USER_BACKGROUNDS'] = str(background_path) | |
backgrounds = { | |
"__comment__": "JSON file specifying background images. env CARTOPY_USER_BACKGROUNDS, ax.background_img()", | |
"NaturalEarthRelief": { | |
"__comment__": "Natural Earth I with shaded Relief, water, and drainage", | |
"__source__": "https://naciscdn.org/naturalearth/50m/raster/NE1_50M_SR_W.zip", | |
"__projection__": "PlateCarree", | |
"high": "NE1_50M_SR_W/NE1_50M_SR_W.tif" | |
}, | |
} | |
with open(f"{background_path}/images.json", "w") as outfile: | |
json.dump(backgrounds, outfile) | |
for key, background in backgrounds.items(): | |
if key != "__comment__": | |
source = background["__source__"] | |
basename = pathlib.Path(source).name | |
destination = pathlib.Path(background_path, basename) | |
if not destination.exists(): | |
urllib.request.urlretrieve(source, str(destination)) | |
with zipfile.ZipFile(destination, 'r') as zip_ref: | |
zip_ref.extractall(background_path) | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
"plot_type", | |
default = "global", | |
choices = ["global", "north_atlantic", "southern_ocean"], | |
help = "The type of plot to make" | |
) | |
parser.add_argument( | |
"num_cpus", | |
default = 128, | |
type=int, | |
help = "Number of images to make in parallel", | |
) | |
args = parser.parse_args() | |
setup_background() | |
ds = xr.open_mfdataset(f"{data_path}/output/*.nc") | |
num_times = len(ds.time) | |
ds.close() | |
grid = xr.open_dataset('ocean_static.nc').load() | |
iterables = ( (args.plot_type, grid, timelevel) for timelevel in range(num_times) ) | |
with multiprocessing.Pool(args.num_cpus) as pool: | |
pool.starmap(plot_image, tqdm.tqdm(iterables, total=num_times)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
To make an animation, go into the images directory and do: