%matplotlib inline
from yt.mods import *
from IPython.core.display import Image
from yt.visualization.volume_rendering.transfer_function_helper import TransferFunctionHelper

def showme(np_im):
    np_im[np_im != np_im] = 0.0
    imb = write_bitmap(np_im, None)
    return Image(imb)

pf = load('Enzo_64/DD0043/data0043')

tfh = TransferFunctionHelper(pf)

# Build a transfer function that is a multivariate gaussian in Density
tfh = TransferFunctionHelper(pf)
tfh.set_field('Temperature')
tfh.set_log(True)
tfh.set_bounds()
tfh.build_transfer_function()
tfh.tf.add_layers(5)
tfh.plot()

tfh.plot(profile_field='CellMass')

tfh = TransferFunctionHelper(pf)
tfh.set_field('Temperature')
tfh.set_bounds()
tfh.set_log(True)
tfh.build_transfer_function()
tfh.tf.add_layers(8, w=0.01, mi=4.0, ma=8.0, col_bounds=[4.,8.], alpha=np.logspace(-1,2,7), colormap='RdBu_r')
tfh.tf.map_to_colormap(6.0, 8.0, colormap='Reds', scale=10.0)
tfh.tf.map_to_colormap(-1.0, 6.0, colormap='Blues_r', scale=1.)

tfh.plot(profile_field='CellMass')

L = [-0.1, -1.0, -0.1]
c = pf.domain_center
W = 1.5*pf.domain_width
Npixels = 512 
cam = pf.h.camera(c, L, W, Npixels, tfh.tf, fields=['Temperature'],
                  north_vector=[1.,0.,0.], steady_north=True, 
                  sub_samples=5, no_ghost=False, l_max=0)

# Here we substitute the TransferFunction we constructed earlier.
cam.transfer_function = tfh.tf


im = cam.snapshot()
showme(im[:,:,:3])

im2 = cam.snapshot(clip_ratio=2.0)
showme(im2[:,:,:3])

im3 = cam.snapshot(clip_ratio=20.0)
showme(im3[:,:,:3])

print (im == im2).all()
print (im == im3).all()