Commits

Sam Skillman committed fc92564

Working on annotated volume renders, and the ability to show them interactively when using ipython with a compatible matplotlib backend. This needs cleaning up.

  • Participants
  • Parent commits e3f254a
  • Branches yt

Comments (0)

Files changed (2)

File yt/visualization/volume_rendering/camera.py

     pixelize_healpix, arr_fisheye_vectors
 
 class Camera(ParallelAnalysisInterface):
+    _pylab = None
+    _tf_figure = None
+    _render_figure = None
+
     def __init__(self, center, normal_vector, width,
                  resolution, transfer_function,
                  north_vector = None, steady_north=False,
             transfer_function = ProjectionTransferFunction()
         self.transfer_function = transfer_function
         self.log_fields = log_fields
+        if self.log_fields is None:
+            self.log_fields = [self.pf.field_info[f].take_log for f in self.fields]
         self.use_kd = use_kd
         self.l_max = l_max
         self.no_ghost = no_ghost
                                            log_fields = log_fields)
         else:
             self.use_kd = isinstance(volume, AMRKDTree)
-        self.volume = volume
+        self.volume = volume        
 
     def _setup_box_properties(self, width, center, unit_vectors):
         self.width = width
         self.finalize_image(image)
         return image
 
+    def show_tf(self):
+        if self._pylab is None: 
+            import pylab
+            self._pylab = pylab
+        if self._tf_figure is None:
+            self._tf_figure = self._pylab.figure(2)
+            self.transfer_function.show(ax=self._tf_figure.axes)
+        self._pylab.draw()
+
+    def annotate(self, ax, enhance=True):
+        ax.get_xaxis().set_visible(False)
+        ax.get_xaxis().set_ticks([])
+        ax.get_yaxis().set_visible(False)
+        ax.get_yaxis().set_ticks([])
+        cb = self._pylab.colorbar(ax.images[0], pad=0.0, fraction=0.05, drawedges=True, shrink=0.9)
+        label = self.pf.field_info[self.fields[0]].get_label()
+        if self.log_fields[0]:
+            label = '$\\rm{log}\\/ $' + label
+        self.transfer_function.vert_cbar(ax=cb.ax, label=label)
+
+
+
+    def show(self, im, enhance=True):
+        if self._pylab is None:
+            import pylab
+            self._pylab = pylab
+        if self._render_figure is None:
+            self._render_figure = self._pylab.figure(1)
+        self._render_figure.clf()
+
+        if enhance:
+            nz = im[im > 0.0]
+            nim = im / (nz.mean() + 6.0 * na.std(nz))
+            nim[nim > 1.0] = 1.0
+            nim[nim < 0.0] = 0.0
+            del nz
+        else:
+            nim = im
+        ax = self._pylab.imshow(nim/nim.max(), origin='lower')
+        return ax
+
+    def draw(self):
+        self._pylab.draw()
+    
+    def save_annotated(self, fn, im=None, enhance=True, dpi=100):
+        ax = self.show(im, enhance=enhance)
+        self.annotate(ax.axes, enhance)
+        self._pylab.savefig(fn, bbox_inches='tight', facecolor='black', dpi=dpi)
+        
     def save_image(self, fn, clip_ratio, image):
         if self.comm.rank is 0 and fn is not None:
             if clip_ratio is not None:
             else:
                 write_bitmap(image, fn)
 
-
     def initialize_source(self):
         return self.volume.initialize_source()
 
 
 class InteractiveCamera(Camera):
     frames = []
-    _pylab = None
-    def snapshot(self, fn = None, clip_ratio = None):
-        if self._pylab is None: 
-            import pylab
-            self._pylab = pylab
-        self._pylab.figure(2)
-        self.transfer_function.show()
-        self._pylab.draw()
+
+    def snapshot(self, fn = None, clip_ratio = None, enhance=True):
         im = Camera.snapshot(self, fn, clip_ratio)
-        self._pylab.figure(1)
-        self._pylab.imshow(im/im.max())
-        self._pylab.draw()
+        self.show(im, enhance)
         self.frames.append(im)
         
     def rotation(self, theta, n_steps, rot_vector=None):
         del self.frames
         self.frames = []
         
+    def save(self,fn):
+        self._pylab.savefig(fn, bbox_inches='tight', facecolor='black')
+        
     def save_frames(self, basename, clip_ratio=None):
         for i, frame in enumerate(self.frames):
             fn = basename + '_%04i.png'%i

File yt/visualization/volume_rendering/transfer_functions.py

         ax.set_xlabel("Value")
         pyplot.savefig(filename)
 
-    def show(self):
+    def show(self, ax=None):
         r"""Display an image of the transfer function
 
         This function loads up matplotlib and displays the current transfer function.
         """
         from matplotlib import pyplot
         from matplotlib.ticker import FuncFormatter
-        pyplot.clf()
-        ax = pyplot.axes()
+        #pyplot.clf()
+        if ax is None:
+            ax = pyplot.axes()
         alpha = self.red.y + self.green.y + self.blue.y
         max_alpha = alpha.max()
         norm = max_alpha
         i_data[:,:,1] = na.outer(na.ones(self.alpha.x.size), self.funcs[1].y/norm)
         i_data[:,:,2] = na.outer(na.ones(self.alpha.x.size), self.funcs[2].y/norm)
         ax.imshow(i_data, origin='lower', aspect='auto')
-        print max_alpha
-        #ax.fill_between(na.arange(self.alpha.y.size), alpha, y2=max_alpha, color='black')
+
+        #ax.fill_between(na.arange(self.alpha.y.size), alpha, y2=max_alpha, color='white')
         ax.plot(na.arange(self.alpha.y.size), alpha, 'w')
         ax.set_xlim(0, self.alpha.x.size)
         xticks = na.arange(na.ceil(self.alpha.x[0]), na.floor(self.alpha.x[-1]) + 1, 1) - self.alpha.x[0]
         ax.set_ylim(0., max_alpha)
         ax.set_ylabel("Opacity")
         ax.set_xlabel("Value")
+
+    def vert_cbar(self, ax=None, label=None):
+        r"""Display an image of the transfer function
+
+        This function loads up matplotlib and displays the current transfer function.
+
+        Parameters
+        ----------
+
+        Examples
+        --------
+
+        >>> tf = TransferFunction( (-10.0, -5.0) )
+        >>> tf.add_gaussian(-9.0, 0.01, 1.0)
+        >>> tf.show()
+        """
+        from matplotlib import pyplot
+        from matplotlib.ticker import FuncFormatter
+        #pyplot.clf()
+        if ax is None:
+            ax = pyplot.axes()
+        if label is None:
+            label = ''
+        alpha = self.red.y + self.green.y + self.blue.y
+        max_alpha = alpha.max()
+        norm = max_alpha
+        i_data = na.zeros((self.alpha.x.size, self.funcs[0].y.size, 3))
+        i_data[:,:,0] = na.outer(self.funcs[0].y/norm, na.ones(self.alpha.x.size))
+        i_data[:,:,1] = na.outer(self.funcs[1].y/norm, na.ones(self.alpha.x.size))
+        i_data[:,:,2] = na.outer(self.funcs[2].y/norm, na.ones(self.alpha.x.size))
+        ax.imshow(i_data, origin='lower', aspect='auto')
+
+        #ax.fill_between(na.arange(self.alpha.y.size), alpha, y2=max_alpha, color='white')
+        #ax.plot(alpha, na.arange(self.alpha.y.size), 'w')
+        ax.set_ylim(0, self.alpha.x.size)
+        xticks = na.arange(na.ceil(self.alpha.x[0]), na.floor(self.alpha.x[-1]) + 1, 1) - self.alpha.x[0]
+        xticks *= self.alpha.x.size / (self.alpha.x[-1] - self.alpha.x[0])
+        if len(xticks) > 5:
+            xticks = xticks[::len(xticks)/5]
+        ax.yaxis.set_ticks(xticks)
+        def x_format(x, pos):
+            return "%.1f" % (x * (self.alpha.x[-1] - self.alpha.x[0]) / (self.alpha.x.size) + self.alpha.x[0])
+        ax.yaxis.set_major_formatter(FuncFormatter(x_format))
+        yticks = na.linspace(0,1,2,endpoint=True) * max_alpha
+        ax.xaxis.set_ticks(yticks)
+        #balls
+        def y_format(y, pos):
+            s = '%0.2f' % ( y )
+            return s
+        ax.xaxis.set_major_formatter(FuncFormatter(y_format))
+        ax.set_xlim(0., max_alpha)
+        #ax.set_xlabel("Opacity")
+        ax.get_xaxis().set_visible(False)
+        ax.get_xaxis().set_ticks([])
+        ax.tick_params(axis='y', colors='white')
+        ax.set_ylabel(label, color='white')
+
         
     def sample_colormap(self, v, w, alpha=None, colormap="gist_stern", col_bounds=None):
         r"""Add a Gaussian based on an existing colormap.
         cmap = get_cmap(colormap)
         r,g,b,a = cmap(rel)
         if alpha is None: alpha = a
-        r *= alpha
-        g *= alpha
-        b *= alpha    
+        r *= alpha/3.
+        g *= alpha/3.
+        b *= alpha/3. 
         self.add_gaussian(v, w, [r, g, b, alpha])
         mylog.debug("Adding gaussian at %s with width %s and colors %s" % (
                 v, w, (r,g,b,alpha)))