[Yt-svn] yt-commit r1628 - trunk/yt/extensions/volume_rendering

[mturk at wrangler.dreamhost.com]

Author: mturk
Date: Thu Feb 11 15:44:41 2010
New Revision: 1628
URL: http://yt.enzotools.org/changeset/1628

Log:
Adding sample_colormap and add_layers to the ColorTransferFunction object



Modified:
   trunk/yt/extensions/volume_rendering/TransferFunction.py

Modified: trunk/yt/extensions/volume_rendering/TransferFunction.py
==============================================================================
--- trunk/yt/extensions/volume_rendering/TransferFunction.py	(original)
+++ trunk/yt/extensions/volume_rendering/TransferFunction.py	Thu Feb 11 15:44:41 2010
@@ -24,6 +24,7 @@
 """
 
 import numpy as na
+from matplotlib.cm import get_cmap
 
 class TransferFunction(object):
     def __init__(self, x_bounds, nbins=256):
@@ -84,6 +85,23 @@
         pylab.ylabel("Transmission")
         pylab.savefig(filename)
 
+    def sample_colormap(self, v, w, alpha=None, colormap="gist_stern"):
+        rel = (v - self.x_bounds[0])/(self.x_bounds[1] - self.x_bounds[0])
+        cmap = get_cmap(colormap)
+        r,g,b,a = cmap(rel)
+        if alpha is None: alpha = a
+        self.add_gaussian(v, w, [r,g,b,alpha])
+        print "Adding gaussian at %s with width %s and colors %s" % (
+                v, w, (r,g,b,alpha))
+
+    def add_layers(self, N, w=None, mi=None, ma=None, colormap="gist_stern"):
+        dist = (self.x_bounds[1] - self.x_bounds[0])
+        if mi is None: mi = self.x_bounds[0] + dist/(10.0*N)
+        if ma is None: ma = self.x_bounds[1] - dist/(10.0*N)
+        if w is None: w = 0.001 * (ma-mi)/N
+        for v, a in zip(na.mgrid[mi:ma:N*1j], na.logspace(-2.0, 0.0,N)):
+            self.sample_colormap(v, w, a, colormap=colormap)
+
 if __name__ == "__main__":
     tf = ColorTransferFunction((-20, -5))
     tf.add_gaussian(-16.0, 0.4, [0.2, 0.3, 0.1])