Paste #89

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
from yt.mods import *

#field = "density"
field = "Density"

smoothed = False

# load the data
#data = OrionStaticOutput('/ebisu/jsoishi/orion/StellaTurbTest/pltgmlcs5600/')
data = load("galaxy1200.dir/galaxy1200")
v,center = data.h.find_max(field)
r_max = 10.0/data['kpc'] # or whatever

# construct the r, theta, phi grids
n_r = 200
n_theta = 100
n_phi = 100

## syntatic sugar to quickly create a grid, evenly spaced in r, theta, phi
r,theta,phi = na.mgrid[0:r_max:n_r*1j,0:na.pi:n_theta*1j,0:2*na.pi:n_phi*1j]

# convert to x,y,z arrays for interpolation
sph_grid = {}
sph_grid['x'] = r*na.sin(theta)*na.cos(phi) + center[0]
sph_grid['y'] = r*na.sin(theta)*na.sin(phi) + center[1]
sph_grid['z'] = r*na.cos(theta)             + center[2]
sph_grid['handled'] = na.zeros(sph_grid['x'].shape, dtype='bool')
sph_grid[field] = na.zeros(sph_grid['x'].shape, dtype='float64')

def regrid_points(grid, new_grid):
    """
    This routine takes a grid and a dict of desired locations
    and then re-grids from one onto the other
    """
    # get one ghost zone
    cg = grid.retrieve_ghost_zones(1, [field], smoothed=smoothed)

    # makes x0,x1,y0,y1,z0,z1
    bounds = na.concatenate(zip(cg.left_edge, cg.right_edge)) 

    # Set up our interpolator, with its bounds set with the ghost zones
    interpolator = lagos.TrilinearFieldInterpolator(
        cg[field],bounds,['x','y','z'])
    
    # Now we figure out which of our points are inside this grid
    # Note that we're only looking at the grid, not the grid-with-ghost-zones
    point_ind = na.ones(new_grid['handled'].shape, dtype='bool') # everything at first
    for i,ax in enumerate('xyz'): # i = 0,1,2 ; ax = x, y, z
        # &= does a logical_and on the array
        point_ind &= ( ( grid.LeftEdge[i] <= new_grid[ax]      )
                     & ( new_grid[ax]     <= grid.RightEdge[i] ) )
    point_ind &= (new_grid['handled'] == False) # only want unhandled points

    # If we don't have any, we can just leave
    if point_ind.sum() == 0: return

    # because of the funky way the interpolator takes points, we have to make a
    # new dict of just the points inside this grid
    point_grid = {'x' : new_grid['x'][point_ind],
                  'y' : new_grid['y'][point_ind],
                  'z' : new_grid['z'][point_ind]}

    # Now we know which of the points in new_grid are inside this grid
    new_grid[field][point_ind] = interpolator(point_grid)
    new_grid['handled'][point_ind] = True

# Just to figure out which grids are in our sphere...
sphere = data.h.sphere(center, r_max) 

# sphere._grids needs to be sorted by levels
# argsort returns the indices so that they are sorted
grid_order = na.argsort(sphere.gridLevels)

ng = len(sphere._grids)

for i,g in enumerate(sphere._grids[grid_order][::-1]):
    print "Doing grid % 4i / % 4i (%s - %s)" % (i, ng, g.id, g.Level)
    regrid_points(g, sph_grid)

print "Finished with %s dangling points" % \
    (sph_grid['handled'].size - sph_grid['handled'].sum())