[Yt-svn] yt-commit r792 - in branches/yt-1.0/yt: lagos/hop raven

[mturk at wrangler.dreamhost.com]

Author: mturk
Date: Fri Sep 19 08:22:54 2008
New Revision: 792
URL: http://yt.spacepope.org/changeset/792

Log:
Backporting speed ups in the Hop outputting as well as the HOP Circle callback



Modified:
   branches/yt-1.0/yt/lagos/hop/HopOutput.py
   branches/yt-1.0/yt/raven/Callbacks.py

Modified: branches/yt-1.0/yt/lagos/hop/HopOutput.py
==============================================================================
--- branches/yt-1.0/yt/lagos/hop/HopOutput.py	(original)
+++ branches/yt-1.0/yt/lagos/hop/HopOutput.py	Fri Sep 19 08:22:54 2008
@@ -43,16 +43,17 @@
         self.__run_hop()
         mylog.info("Parsing outputs")
         self.__parse_output()
-        mylog.debug("Finished.")
+        mylog.debug("Finished. (%s)", len(self))
 
     def __obtain_particles(self):
         if self.dm_only: ii = self.__get_dm_indices()
         else: ii = slice(None)
-        self._base_indices = ii
         self.particle_fields = {}
         for field in ["particle_position_%s" % ax for ax in 'xyz'] + \
                      ["ParticleMassMsun"]:
+            tot_part = self.data_source[field].size
             self.particle_fields[field] = self.data_source[field][ii]
+        self._base_indices = na.arange(tot_part)[ii]
 
     def __run_hop(self):
         self.densities, self.tags = \
@@ -80,6 +81,7 @@
         counts = na.bincount(self.tags+1)
         sort_indices = na.argsort(self.tags)
         grab_indices = na.indices(self.tags.shape).ravel()[sort_indices]
+        dens = self.densities[sort_indices]
         cp = 0
         for i in unique_ids:
             cp_c = cp + counts[i+1]
@@ -88,10 +90,10 @@
                 continue
             group_indices = grab_indices[cp:cp_c]
             self._groups.append(HopGroup(self, i, group_indices))
-            md_i = na.argmax(self.densities[sort_indices][cp:cp_c])
+            md_i = na.argmax(dens[cp:cp_c])
             px, py, pz = [self.particle_fields['particle_position_%s'%ax][group_indices]
                                             for ax in 'xyz']
-            self._max_dens[i] = (self.densities[sort_indices][cp:cp_c][md_i],
+            self._max_dens[i] = (dens[cp:cp_c][md_i],
                                  px[md_i], py[md_i], pz[md_i])
             cp += counts[i+1]
 
@@ -147,8 +149,7 @@
         self.hop_output = hop_output
         self.id = id
         self.data = hop_output.data_source
-        self.indices = indices
-        self._base_indices = hop_output._base_indices
+        self.indices = hop_output._base_indices[indices]
         
     def center_of_mass(self):
         """
@@ -210,7 +211,7 @@
         return r.max()
 
     def __getitem__(self, key):
-        return self.data[key][self._base_indices][self.indices]
+        return self.data[key][self.indices]
 
     def get_sphere(self, center_of_mass=True):
         """

Modified: branches/yt-1.0/yt/raven/Callbacks.py
==============================================================================
--- branches/yt-1.0/yt/raven/Callbacks.py	(original)
+++ branches/yt-1.0/yt/raven/Callbacks.py	Fri Sep 19 08:22:54 2008
@@ -359,3 +359,32 @@
         self.rv = plot._axes.contour(buff, len(self.clumps)+1,
                                      **self.plot_args)
         plot._axes.hold(False)
+
+
+class HopCircleCallback(PlotCallback):
+    def __init__(self, hop_output, axis, max_number=None,
+                 annotate=False):
+        self.axis = axis
+        self.hop_output = hop_output
+        self.max_number = max_number
+        self.annotate = annotate
+
+    def __call__(self, plot):
+        from matplotlib.patches import Circle
+        x0, x1 = plot.xlim
+        y0, y1 = plot.ylim
+        l, b, width, height = plot._axes.bbox.get_bounds()
+        xi = lagos.x_dict[plot.data.axis]
+        yi = lagos.y_dict[plot.data.axis]
+        dx = plot.image._A.shape[0] / (x1-x0)
+        dy = plot.image._A.shape[1] / (y1-y0)
+        for halo in self.hop_output[:self.max_number]:
+            radius = halo.maximum_radius() * dx
+            center = halo.center_of_mass()
+            center_x = (center[xi] - x0)*dx
+            center_y = (center[yi] - y0)*dy
+            cir = Circle((center_x, center_y), radius, fill=False)
+            plot._axes.add_patch(cir)
+            if self.annotate:
+                plot._axes.text(center_x, center_y, "%s" % halo.id)
+