Source

yt.lagrangian_volume / find_particles.py

Full commit
"""
Particle operations for Lagrangian Volume

Author: Matthew Turk <matthewturk@gmail.com>
Affiliation: Columbia University
Homepage: http://yt.enzotools.org/
License:
  Copyright (C) 2011 Matthew Turk.  All Rights Reserved.

  This file is part of yt.

  yt is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

from yt.mods import *
from os import environ
environ['CFLAGS'] = "-I"+na.get_include()

import pyximport; pyximport.install()
import particle_ops
import argparse

def correlate_particles(particle_ids, pf):
    # First we identify all the particles and find their maximum extent
    min_pos = pf.domain_right_edge.copy().astype(float)
    max_pos = pf.domain_left_edge.copy().astype(float)
    mask_to_get = na.zeros(particle_ids.size, dtype='int32')
    for g in pf.h.grids:
        if g.NumberOfParticles == 0: continue
        found_any, mask = particle_ops.mask_particles(
            particle_ids, g["particle_index"], mask_to_get)
        if found_any == 0: continue
        mask = mask.astype("bool")
        for i, ax in enumerate('xyz'):
            min_pos[i] = min(g["particle_position_%s" % ax][mask].min(),
                             min_pos[i])
            max_pos[i] = max(g["particle_position_%s" % ax][mask].max(),
                             max_pos[i])
    return min_pos, max_pos

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("later", type=str)
    parser.add_argument("earlier", type=str)
    args = parser.parse_args()
    pf1 = load(args.later)
    pf2 = load(args.earlier)
    if pf1 is None or pf2 is None:
        print "Sorry, couldn't load!"
        sys.exit(1)
    halos = HaloFinder(pf1)
    ind = halos[0]["particle_index"]
    min_pos, max_pos = correlate_particles(ind, pf2)
    print
    print "The bounding box:"
    print "    %0.8e %0.8e %0.8e" % tuple(min_pos)
    print "    %0.8e %0.8e %0.8e" % tuple(max_pos)
    print