magvort.py

#
# render vorticity
#

from mpi4py import MPI
from yt.pmods import *

import os
import yt

#
# config
#

comm = MPI.COMM_WORLD
rank, size = comm.rank, comm.size

src = '/scratch2/scratchdirs/wqz/share/'
dst = './'

c = [ 0.0, 0.0, 0.1 ]
L = [ 1, 0, 0 ]
U = [ 0, 0, -1 ]
W = 0.17
N = 512

tf = ColorTransferFunction((5.0, 7.5))
tf.map_to_colormap(5.0, 7.5, colormap='jet')

#
# build list of plot files that still need to be rendered
#

if rank == 0:
    import glob
    plts = glob.glob(src + 'plt*')
    pngs = glob.glob(dst + 'plt*.png')

    plts = set([ os.path.basename(x) for x in plts ])
    pngs = set([ os.path.splitext(os.path.basename(x))[0] for x in pngs ])

    plts = list(plts.difference(pngs))

else:
    plts = None

plts = comm.bcast(plts, root=0)

#
# parallel render!
#

for plt in yt.parallel_objects(plts, size):
  print "==> rank %d rendering: %s" % (rank, plt)

  pf   = yt.frontends.boxlib.BoxlibDataset(src + plt)
  cam  = pf.camera(c, L, W, N, tf, U, fields=['magvort'], log_fields=[True])
  img  = cam.snapshot(dst + plt + '.png')