Animate cloud water isosurfaces using matplotlib

plot_3d_cloud_field.py

"""
Make movie of 3D cloud fields
"""
import xarray as xr
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LightSource
from skimage import measure
from scipy.interpolate import interp1d
from scipy.ndimage import gaussian_filter
import numpy as np
import os

rgas = 287.0
sun = LightSource(azdeg = -10, altdeg = 90)
sim = '/nfs/twcroninlab002/thabbott/ConvectiveInvigoration/SAM_RUNS/' + \
    'OUT_3D/RCE_SST300_NC50_RESTART_MOVIE'
dset = xr.open_mfdataset(sim + '/*.nc',
        decode_cf = False, combine = 'by_coords')

for itime in range(0, dset['time'].size):
    
    print('%d/%d' % (itime + 1, dset['time'].size))
    
    fig = plt.figure(figsize = (9.6, 2.7), dpi = 400)
    ax = fig.add_subplot(111, projection = '3d')
    ax.view_init(elev = 45.0, azim = 10.0)
    ax.set_proj_type('ortho')
    ax.margins(0)
    ax.grid(False)
    ax.set_zticks([0, 5, 10, 15])
    ax.set_zlabel('z (km)')
    ax.set_xticks([0, 50, 100])
    ax.set_xlabel('x (km)')
    ax.set_yticks([0, 50, 100])
    ax.set_ylabel('y (km)')
    ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    ax.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    ax.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    
    # Set up plot
    snap = dset.isel(time = itime)
    
    # Plot temperature on bottom surface
    x = snap['x'].values
    y = snap['y'].values
    tabs = snap['TABS'].isel(z = 0).values
    cset = ax.contourf(x/1e3, y/1e3, tabs, 100, 
        zdir = 'z', offset = 0, cmap = 'magma',
        vmin = 296, vmax = 299)
    plt.colorbar(cset, ax = ax, label = 'Surface air T (K)', shrink = 0.6,
        ticks = np.arange(290, 310, 1), pad = -0.05)
    
    ax.axes.set_xlim3d(left = x[0]/1e3 - 1, right = x[-1]/1e3 - 1)
    ax.axes.set_ylim3d(bottom = y[0]/1e3, top = y[-1]/1e3)
    ax.axes.set_zlim3d(bottom = 0, top = 16)
                 
    # Plot black rectanges on other surfaces
    z0 = [0, 16e3]
    Y, Z = np.meshgrid(y, z0)
    ax.contourf(np.zeros(Y.shape), Y/1e3, Z/1e3, 
        zdir = 'x', offset = 0, cmap = 'Greys_r', vmin = 0, vmax = 1)
    X, Z = np.meshgrid(x, z0)
    ax.contourf(X/1e3, np.zeros(X.shape), Z/1e3, 
        zdir = 'y', offset = 0, cmap = 'Greys_r', vmin = 0, vmax = 1)
                 
    # Create triangulation of cloud water isosurfaces
    rho = snap['p']*1e2/(rgas*snap['TABS'])
    qcloud = ((snap['QN'] + snap['QP'])*rho).values
    z = 1e3*np.linspace(0.05, 16.0, 128)
    dz = z[1] - z[0]
    z0 = snap['z'].values
    f = interp1d(z0, qcloud, axis = 0)
    qcloud = f(z)
    qcloud = gaussian_filter(qcloud, sigma = 0.8, mode = 'wrap')
    verts, faces, _, _ = measure.marching_cubes(qcloud, 0.1, 
        spacing = (dz/1e3, 1, 1))
    
    # Plot triangulation
    ax.plot_trisurf(verts[:,2], verts[:,1], faces[:], verts[:,0],
                linewidth = 0.0, antialiased = True, color = 'white', alpha = 0.8, 
                shade = True, lightsource = sun)
    ax.dist = 10
    ax.patch.set_alpha(0)
    
    # Add text
    ax.text2D(0.85, 0.9, 
            'Isosurfaces: 0.1 g m$^{-3}$ total condensate\n%d minutes' % 
            (1*(itime + 1)), transform = fig.transFigure,
            horizontalalignment = 'right', verticalalignment = 'top')
    
    # Save
    plt.savefig('3d_cloud_field/frame_%010d.png' % itime, bbox_inches = 'tight',
               pad_inches = 0.5, dpi = fig.dpi)
    ax.cla()
    fig.clf()
    plt.close()