[Yt-svn] yt-commit r416 - trunk/yt/lagos

mturk at wrangler.dreamhost.com mturk at wrangler.dreamhost.com
Tue Apr 29 11:47:16 PDT 2008


Author: mturk
Date: Tue Apr 29 11:47:14 2008
New Revision: 416
URL: http://yt.spacepope.org/changeset/416

Log:
Some leftover issues with the profiles have been fixed, including storing
information about the log spacing of bins, and getting rid of some weave cruft.

EnzoFortranRoutines is no longer even attempted to be imported, which is as it
should be until I either remove it or get it to match the fortran API of the
public release of enzo 1.5.

Derived Quantities now work, and are accessible via Enzo3DData.quantities.  I
have not *extensively* tested them, but so far they have passed all I've thrown
at them.  Unit tests are forthcoming, because I need to do them correctly. (
Inverse TDD! )  I'm almost ready to close #82, once the unit tests are in place
and I have both tested spin parameter and added clumping factor.

Additionally, I have tested the average quantities and the center of mass, but
not yet the spin parameter, as it is not quite ready yet.

Making progress toward 0.3...



Modified:
   trunk/yt/lagos/BaseDataTypes.py
   trunk/yt/lagos/DerivedQuantities.py
   trunk/yt/lagos/Profiles.py
   trunk/yt/lagos/__init__.py

Modified: trunk/yt/lagos/BaseDataTypes.py
==============================================================================
--- trunk/yt/lagos/BaseDataTypes.py	(original)
+++ trunk/yt/lagos/BaseDataTypes.py	Tue Apr 29 11:47:14 2008
@@ -1118,6 +1118,13 @@
     gridLevels = property(__get_gridLevels, __set_gridLevels,
                              __del_gridLevels)
 
+    def __get_quantities(self):
+        if self.__quantities is None:
+            self.__quantities = DerivedQuantityCollection(self)
+        return self.__quantities
+    __quantities = None
+    quantities = property(__get_quantities)
+
     def extract_connected_sets(self, field, num_levels, min_val, max_val,
                                 log_space=True, cumulative=True):
         """

Modified: trunk/yt/lagos/DerivedQuantities.py
==============================================================================
--- trunk/yt/lagos/DerivedQuantities.py	(original)
+++ trunk/yt/lagos/DerivedQuantities.py	Tue Apr 29 11:47:14 2008
@@ -29,65 +29,112 @@
 
 from yt.lagos import *
 
-quantityInfo = {}
-
-def addQuantity(func):
-    n = func.func_name
-    quantityInfo[n] = tuple([a() for a in func()])
-
-def Phi():
-    def GetLog():
-        return False
-    def GetUnits():
-        return "radians"
-    def GetFunc():
-        def getPhi(data, vector, r_vec = None):
-            # Get the declination here
-            if r_vec == None:
-                r_vec = (data.coords - na.reshape(data.center,(3,1)))
-            vec = vector.reshape((3,1))
-            nv = vec[0,0] * vec[0,0] + \
-                 vec[1,0] * vec[1,0] + \
-                 vec[2,0] * vec[2,0]
-            nv = na.sqrt(nv)
-            nr = r_vec[0,:] * r_vec[0,:] + \
-                 r_vec[1,:] * r_vec[1,:] + \
-                 r_vec[2,:] * r_vec[2,:]
-            nr = na.sqrt(nr)
-            dp = r_vec[0,:] * vec[0,0] + \
-                 r_vec[1,:] * vec[1,0] + \
-                 r_vec[2,:] * vec[2,0]
-            phi = na.arccos(dp / (nr * nv))
-            return phi
-        return getPhi
-    def GetHelp():
-        return "This returns an array of angles-of-declination."
-    return GetLog, GetUnits, GetFunc, GetHelp
-addQuantity(Phi)
-
-def AngularMomentumVector():
-    def GetLog():
-        return False
-    def GetUnits():
-        return ""
-    def GetFunc():
-        def getVector(data, weight="CellMassCode"):
-            # We want the mass-weighted angular momentum vector for the entire
-            # system
-            # So we first want to get the angular momentum vectors for every
-            # cell:  L = mass * r x v
-            #          = cellmass * angular_velocity
-            # Then we mass-weight a sum
-            #        L' = sum(mass * L) / sum(mass)
-            v = na.sum( \
-                data[weight]* data["CellMassCode"] * \
-                    data["SpecificAngularMomentum"], axis=1) / \
-                na.sum(data[weight])
-            vec = v/na.sqrt(na.sum(v*v))
-            return v, vec
-        return getVector
-    def GetHelp():
-        return "Returns a single mass-weighted angular momentum vector"
-    return GetLog, GetUnits, GetFunc, GetHelp
-addQuantity(AngularMomentumVector)
+quantity_info = {}
 
+class GridChildMaskWrapper:
+    def __init__(self, grid):
+        self.grid = grid
+    def __getattr__(self, attr):
+        return getattr(self.grid, attr)
+    def __getitem__(self, item):
+        return self.grid[item]*self.grid.child_mask
+
+def func_wrapper(quantities_collection, quantities_object):
+    func = quantities_object.function
+    c_func = quantities_object.combine_function
+    def call_func_unlazy(args, kwargs):
+        retval = func(quantities_collection.data_object, *args, **kwargs)
+        return c_func(quantities_collection.data_object, *retval)
+    def call_func_lazy(args, kwargs):
+        n_ret = quantities_object.n_ret
+        retvals = [ [] for i in range(n_ret)]
+        for g in quantities_collection.data_object._grids:
+            rv = func(GridChildMaskWrapper(g), *args, **kwargs)
+            for i in range(n_ret): retvals[i].append(rv[i])
+            g.clear_data()
+        retvals = [na.array(retvals[i]) for i in range(n_ret)]
+        return c_func(quantities_collection.data_object, *retvals)
+    def call_func(*args, **kwargs):
+        lazy_reader = kwargs.pop('lazy_reader', False)
+        if lazy_reader:
+            return call_func_lazy(args, kwargs)
+        else:
+            return call_func_unlazy(args, kwargs)
+    return call_func
+
+class DerivedQuantity(object):
+    def __init__(self, name, function,
+                 combine_function, units = "",
+                 n_ret = 0):
+        self.name = name
+        self.function = function
+        self.combine_function = combine_function
+        self.n_ret = n_ret
+
+def add_quantity(name, **kwargs):
+    if 'function' not in kwargs or 'combine_function' not in kwargs:
+        mylog.error("Not adding field %s because both function and combine_function must be provided" % name)
+        return
+    f = kwargs.pop('function')
+    c = kwargs.pop('combine_function')
+    quantity_info[name] = DerivedQuantity(name, f, c, **kwargs)
+
+class DerivedQuantityCollection(object):
+    functions = quantity_info
+    def __init__(self, data_object):
+        self.data_object = data_object
+
+    def __getitem__(self, key):
+        if key not in self.functions:
+            raise KeyError(key)
+        return func_wrapper(self, self.functions[key])
+
+def _CenterOfMass(data):
+    x = (data["x"] * data["CellMassMsun"]).sum()
+    y = (data["y"] * data["CellMassMsun"]).sum()
+    z = (data["z"] * data["CellMassMsun"]).sum()
+    den = data["CellMassMsun"].sum()
+    return x,y,z, den
+def _combCenterOfMass(data, x,y,z, den):
+    return na.array([x.sum(), y.sum(), z.sum()])/den.sum()
+add_quantity("CenterOfMass", function=_CenterOfMass,
+             combine_function=_combCenterOfMass, n_ret = 4)
+
+def _WeightedAverageQuantity(data, field, weight):
+    num = (data[field] * data[weight]).sum()
+    den = data[weight].sum()
+    return num, den
+def _combWeightedAverageQuantity(data, field, weight):
+    return field.sum()/weight.sum()
+add_quantity("WeightedAverageQuantity", function=_WeightedAverageQuantity,
+             combine_function=_combWeightedAverageQuantity, n_ret = 2)
+
+"""
+# SpecificAngularMomentum is in cm^2 / s
+j_vec = na.sum(sp["SpecificAngularMomentum"]*sp["CellMassMsun"],axis=1)/m_weight
+j_mag = na.sqrt(na.sum(j_vec**2.0)) # cm^2 / s
+
+eterm = na.sqrt(0.5*na.sum(sp["CellMassMsun"]*sp["VelocityMagnitude"]**2.0)/m_weight)
+
+G = 6.67e-8 # cm^3 g^-1 s^-2
+m_vir = na.sum(sp["CellMassMsun"]) + na.sum(sp["ParticleMassMsun"])# g
+print "MVIR: %0.5e m_sun" % m_vir
+m_vir *= 1.989e33
+spin = j_mag * eterm / (m_vir * G)
+"""
+
+def _SpinParameter(data):
+    weight = data["CellMassMsun"].sum()
+    j_vec = (data["SpecificAngularMomentum"]*data["CellMassMsun"]).sum(axis=1)
+    m_enc = data["CellMassMsun"].sum() + data["ParticleMassMsun"].sum()
+    j_mag_pre = na.sum(j_vec**2.0)
+    e_term_pre = 0.5*na.sum(data["CellMassMsun"]*data["VelocityMagnitude"]**2.0)
+    return j_mag_pre, weight, m_enc, e_term_pre
+def _combSpinParameter(data, j_mag_pre, weight, m_enc, e_term_pre):
+    e_term = na.sqrt(e_term_pre.sum())/weight.sum()
+    j_mag = na.sqrt(j_mag_pre.sum())/weight.sum()
+    G = 6.67e-8 # cm^3 g^-1 s^-2
+    spin = j_mag * e_term / (m_enc.sum() * G * 1.989e33)
+    return spin
+add_quantity("SpinParameter", function=_SpinParameter,
+             combine_parameter=_combSpinParameter, n_ret=4)

Modified: trunk/yt/lagos/Profiles.py
==============================================================================
--- trunk/yt/lagos/Profiles.py	(original)
+++ trunk/yt/lagos/Profiles.py	Tue Apr 29 11:47:14 2008
@@ -1,5 +1,5 @@
 """
-Profile class, to deal with generating and obtaining profiles
+Profile classes, to deal with generating and obtaining profiles
 
 @author: U{Matthew Turk<http://www.stanford.edu/~mturk/>}
 @organization: U{KIPAC<http://www-group.slac.stanford.edu/KIPAC/>}
@@ -47,48 +47,56 @@
         self._data = {}
         self._lazy_reader = lazy_reader
 
+    def _lazy_add_fields(self, fields, weight, accumulation):
+        data = {}
+        weight_data = {}
+        for field in fields:
+            data[field] = self._get_empty_field()
+            weight_data[field] = self._get_empty_field()
+        used = self._get_empty_field().astype('bool')
+        for grid in self._data_source._grids:
+            mylog.debug("Binner adding fields from grid %s", grid)
+            args = self._get_bins(grid, check_cut=True)
+            if not args:
+                continue
+            for field in fields:
+                f, w, u = self._bin_field(grid, field, weight, accumulation,
+                                          args=args, check_cut=True)
+                data[field] += f
+                weight_data[field] += w
+                used = (used | u)
+            grid.clear_data()
+        ub = na.where(used)
+        for field in fields:
+            if weight:
+                data[field][ub] /= weight_data[field][ub]
+            self[field] = data[field]
+        self["UsedBins"] = used
+
+    def _unlazy_add_fields(self, fields, weight, accumulation):
+        for field in fields:
+            f, w, u = self._bin_field(self._data_source, field, weight,
+                                      accumulation, self._args, check_cut = False)
+            ub = na.where(u)
+            if weight:
+                f[ub] /= w[ub]
+            self[field] = f
+        self["UsedBins"] = u
+
     def add_fields(self, fields, weight = "CellMassMsun", accumulation = False):
         # Note that the specification has to be the same for all of these
         fields = ensure_list(fields)
-        if not self._lazy_reader:
-            for field in fields:
-                f, w, u = self._bin_field(self._data_source, field, weight,
-                                          accumulation, self._args, check_cut = False)
-                ub = na.where(u)
-                if weight:
-                    f[ub] /= w[ub]
-                self[field] = f
-            self["UsedBins"] = u
+        if self._lazy_reader:
+            self._lazy_add_fields(fields, weight, accumulation)
         else:
-            data = {}
-            weight_data = {}
-            for field in fields:
-                data[field] = self._get_empty_field()
-                weight_data[field] = self._get_empty_field()
-            used = self._get_empty_field().astype('bool')
-            for grid in self._data_source._grids:
-                mylog.debug("Binner adding fields from grid %s", grid)
-                args = self._get_bins(grid, check_cut=True)
-                if not args:
-                    continue
-                for field in fields:
-                    f, w, u = self._bin_field(grid, field, weight, accumulation,
-                                              args=args, check_cut=True)
-                    data[field] += f
-                    weight_data[field] += w
-                    used = (used | u)
-                grid.clear_data()
-            ub = na.where(used)
-            for field in fields:
-                if weight:
-                    data[field][ub] /= weight_data[field][ub]
-                self[field] = data[field]
-            self["UsedBins"] = used
+            self._unlazy_add_fields(fields, weight, accumulation)
 
+    def keys(self):
+        return self._data.keys()
 
     def __getitem__(self, key):
         # This raises a KeyError if it doesn't exist
-        # This is because we explicitly want to
+        # This is because we explicitly want to add all fields
         return self._data[key]
 
     def __setitem__(self, key, value):
@@ -102,6 +110,7 @@
                  log_space = True, lazy_reader=False):
         BinnedProfile.__init__(self, data_source, lazy_reader)
         self.bin_field = bin_field
+        self._x_log = log_space
         if log_space:
             func = na.logspace
             lower_bound, upper_bound = na.log10(lower_bound), na.log10(upper_bound)
@@ -160,22 +169,6 @@
             inv_bin_indices[bin] = na.where(bin_indices == bin)
         return inv_bin_indices
 
-    @preserve_source_parameters
-    def old_get_bins(self, source, check_cut=False):
-        if check_cut:
-            cm = self._data_source._get_point_indices(source)
-            source_data = source[self.bin_field][cm]
-        else:
-            source_data = source[self.bin_field]
-        bin_order = na.argsort(source_data)
-        bin_indices = na.searchsorted(self[self.bin_field],
-                                      source_data)
-        # Now we set up our inverse bin indices
-        inv_bin_indices = {}
-        for bin in range(self[self.bin_field].size):
-            inv_bin_indices[bin] = na.where(bin_indices == bin)
-        return inv_bin_indices
-
 
 class BinnedProfile2D(BinnedProfile):
     def __init__(self, data_source,
@@ -187,6 +180,8 @@
         BinnedProfile.__init__(self, data_source, lazy_reader)
         self.x_bin_field = x_bin_field
         self.y_bin_field = y_bin_field
+        self._x_log = x_log
+        self._y_log = y_log
         if x_log:
             self[x_bin_field] = na.logspace(na.log10(x_lower_bound*0.99),
                                             na.log10(x_upper_bound*1.01),
@@ -203,6 +198,7 @@
                 y_lower_bound*0.99, y_upper_bound*1.01, y_n_bins)
         if not lazy_reader:
             self._args = self._get_bins(data_source)
+
     def _get_empty_field(self):
         return na.zeros((self[self.x_bin_field].size,
                          self[self.y_bin_field].size), dtype='float64')
@@ -223,8 +219,8 @@
         binned_field = self._get_empty_field()
         weight_field = self._get_empty_field()
         used_field = self._get_empty_field()
-        bin_indices_x = args[0].ravel()
-        bin_indices_y = args[1].ravel()
+        bin_indices_x = args[0].ravel().astype('int64')
+        bin_indices_y = args[1].ravel().astype('int64')
         self.total_cells += bin_indices_x.size
         nx = bin_indices_x.size
         mylog.debug("Binning %s / %s times", source_data.size, nx)
@@ -262,17 +258,3 @@
         bin_indices_y = na.digitize(sd_y, self[self.y_bin_field])
         # Now we set up our inverse bin indices
         return (bin_indices_x, bin_indices_y)
-
-
-_2d_profile_code = r"""
-       int i,j;
-        using namespace std;
-       for(int n = 0; n < nx ; n++) {
-         //cerr << bin_indices_x(n) << "\t" << bin_indices_y(n) << endl;
-         i = bin_indices_x(n)-1;
-         j = bin_indices_y(n)-1;
-         weight_field(i,j) += weight_data(n);
-         binned_field(i,j) += source_data(n) * weight_data(n);
-         used_field(i,j) = 1.0;
-       }
-       """

Modified: trunk/yt/lagos/__init__.py
==============================================================================
--- trunk/yt/lagos/__init__.py	(original)
+++ trunk/yt/lagos/__init__.py	Tue Apr 29 11:47:14 2008
@@ -57,10 +57,8 @@
         pass
 
 if ytcfg.getboolean("lagos","usefortran"):
-    try:
-        import EnzoFortranRoutines
-    except ImportError:
-        pass
+    pass
+    #import EnzoFortranRoutines
 
 # Now we import all the subfiles
 
@@ -69,7 +67,7 @@
 from WeaveStrings import *
 from EnzoDefs import *
 from DerivedFields import *
-from DerivedQuantities import quantityInfo
+from DerivedQuantities import DerivedQuantityCollection
 from DataReadingFuncs import *
 from ClusterFiles import *
 from ContourFinder import *



More information about the yt-svn mailing list