yt_renders.py

import sys
import yt
import spectral_cube
from spectral_cube import SpectralCube
import numpy as np
import pylab as pl
from astropy.utils.console import ProgressBar
import paths
import subprocess
from itertools import izip

red = pl.mpl.colors.LinearSegmentedColormap('red',
                                            {'red':[(0,1,1),(1,1,1)],
                                             'green':[(0,0,0),(1,1,1)],
                                             'blue':[(0,0,0),(1,1,1)]})
green = pl.mpl.colors.LinearSegmentedColormap('green',
                                              {'red':[(0,0,0),(1,1,1)],
                                               'green':[(0,1,1),(1,1,1)],
                                               'blue':[(0,0,0),(1,1,1)]})
blue = pl.mpl.colors.LinearSegmentedColormap('blue', {'red':[(0,0,0),(1,1,1)],
                                                      'green':[(0,0,0),(1,1,1)],
                                                      'blue':[(0,1,1),(1,1,1)]})

cubes = [
    'Brick_C2H.fits',
    'Brick_H13CN.fits',
    'Brick_H13COP.fits',
    'Brick_H2CS.fits',
    'Brick_HCN.fits',
    'Brick_HCOP.fits',
    'Brick_HNCO.fits',
    'Brick_SO.fits',
    'Brick_SiO.fits',
]

def mpath(x, dirname='/Volumes/passport/brick/'):
    return os.path.join(dirname, x)

def load_cube(cubefilename):

    cube = spectral_cube.SpectralCube.read(cubefilename)
    #noise = cube.spectral_slab(-125*u.km/u.s, -37*u.km/u.s).std(axis=0)
    noise = cube.spectral_slab(-125*u.km/u.s, -37*u.km/u.s)._apply_numpy_function(np.std, axis=0)
    sn = cube.filled_data[:] / noise
    mcube = cube.with_mask(sn>1)
    ytcube = mcube.to_yt()

    return mcube,ytcube

def render_cube(cubefilename, outdir='yt_renders',
                size=512, scale=700., nframes=120,
                movie=True, camera_angle=[0, 0, 1],
                north_vector=[1, 0, 0],
                rot_vector1=[1, 0, 0],
                rot_vector2=[0.5, 0.5, 0],
               ):

    cube,pf = load_cube(cubefilename)
    prefix = os.path.splitext(os.path.split(cubefilename)[1])[0]

    mi = cube.min()
    ma = cube.max()

    if not os.path.exists(mpath(outdir)):
        os.makedirs(mpath(outdir))

    tf = yt.ColorTransferFunction([0,ma], grey_opacity=True)
    tf.map_to_colormap(mi, ma, colormap='GREEN-PINK', scale=1)

    center = pf.domain_dimensions /2.
    cam = pf.h.camera(center, camera_angle, scale, size, tf,
                      north_vector=north_vector, fields='flux')

    im  = cam.snapshot()

    images = [im]

    if movie:
        pb = ProgressBar(nframes)
        for ii,im in enumerate(cam.rotation(2 * np.pi, nframes/2, rot_vector=rot_vector1)):
            images.append(im)
            im.write_png(mpath(os.path.join(outdir,"%s_%04i.png" % (prefix,ii))),
                                        rescale=False)
            pb.update(ii)
        for jj,im in enumerate(cam.rotation(2 * np.pi, nframes/2, rot_vector=rot_vector2)):
            images.append(im)
            im.write_png(mpath(os.path.join(outdir,"%s_%04i.png" % (prefix, ii+jj))),
                                        rescale=False)
            pb.update(ii+jj)
        pb.next()

        save_images(images, prefix, mpath(outdir))

        pipe = make_movie(mpath(outdir), outname=prefix+".mp4")
        pipe.terminate()
        
        return images
    else:
        return im

def save_images(images, prefix, outdir):
    """Save a sequence of images, at a common scaling
    Reduces flickering
    """
    if not os.path.exists(outdir):
        os.makedirs(outdir)
 
    cmax = max(np.percentile(i[:, :, :3].sum(axis=2), 99.5) for i in images)
    amax = max(np.percentile(i[:, :, 3], 95) for i in images)
    print cmax, amax
 
    for i, image in enumerate(images):
        image = image.rescale(cmax=cmax, amax=amax).swapaxes(0,1)
        image.write_png(os.path.join(outdir, "%s_%04i.png" % (prefix, i)), rescale=False)

def make_movie(moviepath, overwrite=True, outname='out.mp4'):
    """
    Use ffmpeg to generate a movie from the image series
    """

    outpath = os.path.join(moviepath, outname)

    if os.path.exists(outpath) and overwrite:
        command = ['ffmpeg', '-y', '-r','1','-i',
                   os.path.join(moviepath,'%04d.png'),
                   '-c:v','libx264','-r','30','-pix_fmt', 'yuv420p',
                   outpath]
    elif os.path.exists(outpath):
        log.info("File {0} exists - skipping".format(outpath))
    else:
        command = ['ffmpeg', '-r', '1', '-i',
                   os.path.join(moviepath,'%04d.png'),
                   '-c:v','libx264','-r','30','-pix_fmt', 'yuv420p',
                   outpath]

    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, close_fds=True)

    pipe.wait()

    return pipe

if __name__ == '__main__':

    for cubename in cubes:
        render_cube(mpath(cubename))