Rendering Perseus in yt

perseus_yt.py

import numpy as np
import os
from yt.frontends.stream.api import load_uniform_grid
from yt.mods import ColorTransferFunction
from astropy.io import fits
from scipy.ndimage.measurements import label
from scipy.ndimage.morphology import binary_opening


def thresh(image, lo, hi):
    """ Fancy-pants image thresholding.

    Parameters
    -----------
    image : array to threshold
    lo : Hard-cutoff. All pixels in thresholded image must be > lo
    hi : Second-cutoff. All pixels in output are part of a region with
         at least one pixel > hi

    Returns
    -------
    A thresholded image
    """
    hi = image > hi
    lo = image > lo
    lo = binary_opening(lo, np.ones((4, 4, 4), dtype=np.bool))
    hi = binary_opening(hi, np.ones((4, 4, 4), dtype=np.bool))

    lbl, nf = label(lo)
    mask = np.zeros(image.shape, dtype=np.bool)
    print nf

    for i in range(nf):
        m = lbl == i
        if hi[m].any():
            mask[m] = 1
    return image * mask


def perseus_data():

    #this data also available at http://bit.ly/15kFwUw
    path = '/users/beaumont/perseus/PerA_12co_FCRAO_xyv_smooth_v2.fits'
    cube = fits.getdata(path)
    cube /= 0.49  # convert to Tmb
    cube = cube[58:201, :, :]  # trim velocity channels
    cube[~np.isfinite(cube)] = 0
    cube = thresh(cube, .8, 2)  # filter out noisy pixels
    return cube


def save_images(images, prefix):
    """Save a sequence of images, at a common scaling
    Reduces flickering
    """
    if not os.path.exists(prefix):
        os.makedirs(prefix)

    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)
        image.write_png("%s/%4.4i.png" % (prefix, i), rescale=False)

def make_image(cube, prefix, movie=False):
    """Write 1 or more volume renderings to a directory

    Parameters
    ----------
    cube : Data cube to render
    prefix : Folder to write to.
    movie: If True, write many images that revolve around image
    """

    # flip cube up-down, to match fits convention (low pixels at bottom)
    data = dict(Density=cube[:, ::-1, :])
    pf = load_uniform_grid(data, cube.shape, 3.08e19,  nprocs=4)

    #construct the transfer function
    #isosurfaces at 3, 8, 15K
    def c(r, g, b, a):
        return r / 255., g / 255., b / 255., a

    tf = ColorTransferFunction([0, 25], grey_opacity=False)
    tf.add_gaussian(3, .5, c(239, 237, 245, .3))
    tf.add_gaussian(8, .5, c(188, 189, 220, 1))
    tf.add_gaussian(15, .5, c(230, 85, 13, 2))

    #look at cloud face-on
    c = [.5, .5, .5]  # point that camera centers on
    L = [-1.0, 0, 0]  # direction vector for camera
    W = 1.5           # fractional field of view
    Nvec = 512
    north_vector = [0, 1, 0]
    cam = pf.h.camera(c, L, W, Nvec, tf,
                      north_vector=north_vector, no_ghost=False)

    im = cam.snapshot()
    images = [im]

    if movie:
        images.extend(cam.rotation(2 * np.pi, 60))

    save_images(images, prefix)

    return cam

if __name__ == "__main__":
    make_image(perseus_data(), 'perseus_movie', movie=True)

Output available at
https://vimeo.com/67421373