[Yt-svn] yt: Fixing a bug in the merger tree and adding another vector fu...

[hg at spacepope.org]

hg Repository: yt
details:   yt/rev/73751c26c64a
changeset: 1975:73751c26c64a
user:      Stephen Skory <stephenskory at yahoo.com>
date:
Fri Aug 27 16:19:48 2010 -0700
description:
Fixing a bug in the merger tree and adding another vector function.

diffstat:

 yt/extensions/merger_tree.py |   4 ++--
 yt/math_utils.py             |  48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diffs (69 lines):

diff -r 4674c79377d6 -r 73751c26c64a yt/extensions/merger_tree.py
--- a/yt/extensions/merger_tree.py	Fri Aug 27 12:55:38 2010 -0600
+++ b/yt/extensions/merger_tree.py	Fri Aug 27 16:19:48 2010 -0700
@@ -777,8 +777,8 @@
         self.cursor.execute(string)
         results = self.cursor.fetchone()
         while results:
-            if results[1] < minz:
-                minz = results[1]
+            if abs(results[1] - z) < minz:
+                minz = abs(results[1] - z)
                 this_halo = results[0]
             results = self.cursor.fetchone()
         return this_halo
diff -r 4674c79377d6 -r 73751c26c64a yt/math_utils.py
--- a/yt/math_utils.py	Fri Aug 27 12:55:38 2010 -0600
+++ b/yt/math_utils.py	Fri Aug 27 16:19:48 2010 -0700
@@ -361,3 +361,51 @@
         res[i] = na.dot(temp, temp)**0.5
     return res
 
+def compute_cylindrical_radius(CoM, L, P, V):
+    r"""Compute the radius for some data around an axis in cylindrical
+    coordinates.
+    
+    This is primarily for halo computations.
+    Given some data, this computes the cylindrical radius for each point.
+    This is accomplished by converting the reference frame of the center of
+    mass of the halo.
+    
+    Parameters
+    ----------
+    CoM : array
+        The center of mass in 3D.
+    
+    L : array
+        The angular momentum vector.
+    
+    P : array
+        The positions of the data to be modified (i.e. particle or grid cell
+        postions). The array should be Nx3.
+    
+    V : array
+        The velocities of the data to be modified (i.e. particle or grid cell
+        velocities). The array should be Nx3.
+    
+    Returns
+    -------
+    cyl_r : array
+        An array N elements long that gives the radial velocity for
+        each datum (particle).
+    
+    Examples
+    --------
+    >>> CoM = na.array([0, 0, 0])
+    >>> L = na.array([0, 0, 1])
+    >>> P = na.array([[1, 0, 0], [1, 1, 1], [0, 0, 1], [1, 1, 0]])
+    >>> V = na.array([[0, 1, 10], [-1, -1, -1], [1, 1, 1], [1, -1, -1]])
+    >>> cyl_r = compute_cylindrical_radius(CoM, L, P, V)
+    >>> cyl_r
+    array([ 1.        ,  1.41421356,  0.        ,  1.41421356])
+    """
+    # First we translate into the simple coordinates.
+    L, P, V = modify_reference_frame(CoM, L, P, V)
+    # Demote all the positions to the z=0 plane, which makes the distance
+    # calculation very easy.
+    P[:,2] = 0
+    return na.sqrt((P * P).sum(axis=1))
+