Matthew Turk avatar Matthew Turk committed 7c6f769

Initial import of Axisem-reading code.

Comments (0)

Files changed (2)

+from yt.mods import *
+import os
+import numpy as np
+import pyximport; pyximport.install()
+import grid_interp
+
+coords = np.concatenate([np.array([[float(i) for i in line.split()]
+                        for line in open("glob_grid_%04i.dat" % g)]) for g in range(4)])
+
+QT = grid_interp.QTree(coords.min(axis=0), coords.max(axis=0))
+QT.add(coords)
+level_counts = QT.stats()
+for i, c in enumerate(level_counts):
+    if c == 0: break
+    print "% 3i % 8i" % (i, c)
+
+# Now we have our x, y coords in the plane of the simulation.  We need to
+# interpolate into 3D.
+
+# It's supposed to be a sphere, so these will define our left edge and right
+# edge for our grids.
+mi = coords.min()
+ma = coords.max()
+
+# Now we set up our grids.  But only if it's not already on-disk.
+
+d = np.zeros((512,512,512))
+arr = np.array((512,512,512), dtype="int64")
+pf = load_uniform_grid({"Perturbation": d}, arr, 100.0,
+                       bbox = np.array( [[mi, ma], [mi, ma], [mi, ma]] ),
+                       nprocs = 32)
+#pf.h.print_stats()
+values = (coords**2.0).sum(axis=1)**0.5
+
+#pb = get_pbar("Interpolating ", len(pf.h.grids))
+for i, g in enumerate(pf.h.grids):
+    print i
+    #pb.update(i+1)
+    #print g.LeftEdge, g.RightEdge
+    QT.interp(g.LeftEdge, g.RightEdge, g["Perturbation"], values)
+#pb.finish()
+p = SlicePlot(pf, "x", "Perturbation")
+p.save()
+cimport numpy as np
+import numpy as np
+from libc.stdlib cimport malloc, free
+
+DEF NREF=128
+
+cdef struct qnode:
+    # Note we do not store any center or width info
+    # That is only held during the traversal
+    qnode **children
+    np.int64_t *indices
+    np.float64_t *pos[2]
+    np.int32_t n
+
+cdef class QTree:
+    cdef np.float64_t left_edge[2]
+    cdef np.float64_t right_edge[2]
+    cdef qnode* base
+
+    def __init__(self, left_edge, right_edge):
+        cdef int i
+        for i in range(2):
+            self.left_edge[i] = left_edge[i]
+            self.right_edge[i] = right_edge[i]
+        self.base = create_node()
+
+    cdef np.int64_t find_nearest(self, np.float64_t pos[2]):
+        cdef qnode* node = self.find(pos)
+        # We have our node, now we need the nearest neighbor
+        cdef int i, mi
+        cdef np.float64_t min_r2 = 1e300
+        cdef np.float64_t radius2 = 0.0
+        mi = -1
+        for i in range(node.n):
+            radius2  = (node.pos[0][i] - pos[0])**2.0
+            radius2 += (node.pos[1][i] - pos[1])**2.0
+            if radius2 < min_r2:
+                mi = i
+                min_r2 = radius2
+        if mi == -1:
+            print "COULD NOT FIND ANYTHING", min_r2, radius2, node.n
+            raise RuntimeError
+        return node.indices[mi]
+
+    cdef qnode *find(self, np.float64_t pos[2]):
+        cdef np.float64_t c[2], w[2]
+        cdef int i, level, ci
+        cdef qnode *node = self.base
+        for i in range(2):
+            c[i] = (self.right_edge[i] + self.left_edge[i])/2.0
+            w[i] = (self.right_edge[i] - self.left_edge[i])
+        level = 0
+        check_refine_node(node, c, w)
+        while node.children != NULL:
+            level += 1
+            if level > 40: raise RuntimeError
+            ci = 0
+            if pos[0] > c[0]:
+                ci += 2
+                c[0] += w[0] / 4.0
+            else:
+                c[0] -= w[0] / 4.0
+            if pos[1] > c[1]:
+                ci += 1
+                c[1] += w[1] / 4.0
+            else:
+                c[1] -= w[1] / 4.0
+            w[0] /= 2.0
+            w[1] /= 2.0
+            node = node.children[ci]
+            check_refine_node(node, c, w)
+        return node
+
+    def add(self, np.ndarray[np.float64_t, ndim=2] positions):
+        cdef int n, i
+        cdef np.float64_t pos[2]
+        cdef qnode *node
+        for i in range(positions.shape[0]):
+            pos[0] = positions[i, 0]
+            pos[1] = positions[i, 1]
+            node = self.find(pos)
+            node.pos[0][node.n] = pos[0]
+            node.pos[1][node.n] = pos[1]
+            node.indices[node.n] = i
+            node.n += 1
+
+    def __dealloc__(self):
+        clear_node(self.base)
+
+    def stats(self):
+        cdef np.ndarray[np.int64_t, ndim=1] level_counts
+        level_counts = np.zeros(40, dtype="int64")
+        cdef int level = 0
+        self.recursively_count(self.base, level_counts, 0)
+        return level_counts
+
+    cdef void recursively_count(self, qnode *node,
+                                np.ndarray[np.int64_t, ndim=1] counts,
+                                int level):
+        cdef int i
+        if level > 40: return
+        counts[level] += 1
+        if node.children == NULL:
+            return
+        for i in range(4):
+            self.recursively_count(node.children[i], counts, level + 1)
+
+    def interp(self,
+               np.ndarray[np.float64_t, ndim=1] LE,
+               np.ndarray[np.float64_t, ndim=1] RE,
+               np.ndarray[np.float64_t, ndim=3] dest,
+               np.ndarray[np.float64_t, ndim=1] values):
+        # First we get the bounds of our grid in x, y plane.
+        cdef np.float64_t pos[3], cpos[2]
+        cdef qnode *node
+        cdef int i, j, k
+        cdef np.float64_t dds[3]
+        cdef np.int64_t ind
+        for i in range(3):
+            dds[i] = (RE[i]-LE[i])/dest.shape[i]
+        for i in range(dest.shape[0]):
+            pos[0] = LE[0] + (i + 0.5) * dds[0]
+            for j in range(dest.shape[1]):
+                pos[1] = LE[1] + (j + 0.5) * dds[1]
+                for k in range(dest.shape[2]):
+                    pos[2] = LE[2] + (k + 0.5) * dds[2]
+                    cpos[0] = (pos[0]*pos[0] + pos[1]*pos[1])**0.5
+                    cpos[1] = pos[2]
+                    node = self.find(cpos)
+                    if node.n == 0:
+                        # We have nothing to interpolate
+                        continue
+                    #print pos[0], pos[1], pos[2], cpos[0], cpos[1],
+                    #print node.n
+                    #continue
+                    ind = self.find_nearest(cpos)
+                    dest[i,j,k] = values[ind]
+        #raise RuntimeError
+
+cdef qnode* create_node():
+    cdef int i
+    cdef qnode* nn = <qnode*> malloc(sizeof(qnode))
+    nn.children = NULL
+    nn.n = 0
+    nn.indices = <np.int64_t*> malloc(sizeof(np.int64_t) * NREF)
+    nn.pos[0] = <np.float64_t*> malloc(sizeof(np.float64_t) * NREF)
+    nn.pos[1] = <np.float64_t*> malloc(sizeof(np.float64_t) * NREF)
+    for i in range(NREF):
+        nn.indices[i] = -1
+    return nn
+
+cdef void clear_node(qnode *tc):
+    cdef int i
+    if tc.children != NULL:
+        for i in range(4):
+            clear_node(tc.children[i])
+        free(tc.children) # Free the pointer
+    else:
+        free(tc.pos[0])
+        free(tc.pos[1])
+        free(tc.indices)
+    free(tc)
+
+cdef void check_refine_node(qnode *node, np.float64_t c[2],
+                            np.float64_t w[2]):
+    cdef int i, ci
+    cdef np.float64_t w2[2], c2[2]
+    cdef qnode *child
+    if node.n < NREF: return
+    # We need to refine this node
+    node.children = <qnode **> malloc(sizeof(qnode*) * 4)
+    for i in range(4): # Four new ones
+        node.children[i] = create_node()
+    for i in range(NREF):
+        # ci == 0: left of both
+        # ci == 1: left of c[0], right of c[0]
+        # ci == 2: right of c[0], left of c[0]
+        # ci == 3: right of both
+        ci = 0
+        if node.pos[0][i] > c[0]:
+            ci += 2
+        if node.pos[1][i] > c[1]:
+            ci += 1
+        child = node.children[ci]
+        child.pos[0][child.n] = node.pos[0][i]
+        child.pos[1][child.n] = node.pos[1][i]
+        child.indices[child.n] = node.indices[i]
+        child.n += 1
+    free(node.indices)
+    free(node.pos[0])
+    free(node.pos[1])
+    node.indices = node.pos[0] = node.pos[1] = NULL
+    node.n = 0
+    w2[0] = w[0]/2.0
+    w2[1] = w[1]/2.0
+    for i in range(4):
+        if i < 2:
+            c2[0] = c[0] - w[0]/4.0
+        else:
+            c2[0] = c[0] + w[0]/4.0
+        if i == 0 or i == 2:
+            c2[1] = c[1] - w[1]/4.0
+        else:
+            c2[1] = c[1] + w[1]/4.0
+        check_refine_node(node.children[i], c2, w2)
+
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.