Make animations of OM4_025 in parallel

make_om4_025_animations.py

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)) 

To make an animation, go into the images directory and do:

ffmpeg -framerate 60 -pattern_type glob -i '*.png'  -c:a copy -shortest -c:v libx264 -pix_fmt yuv420p out.mp4