Commits

Matthew Turk committed 6283ac4

Adding ability to sample AMRSurface objects at vertices. This enables .PLY
files to be directly exported to SketchFab.

Comments (0)

Files changed (3)

yt/data_objects/data_containers.py

         self.data_source = data_source
         self.surface_field = surface_field
         self.field_value = field_value
+        self.vertex_samples = YTFieldData()
         center = data_source.get_field_parameter("center")
         AMRData.__init__(self, center = center, fields = None, pf =
                          data_source.pf)
         self._grids = self.data_source._grids.copy()
 
-    def get_data(self, fields = None):
+    def get_data(self, fields = None, sample_type = "face"):
         if isinstance(fields, list) and len(fields) > 1:
             for field in fields: self.get_data(field)
             return
             pb.update(i)
             my_verts = self._extract_isocontours_from_grid(
                             g, self.surface_field, self.field_value,
-                            fields)
+                            fields, sample_type)
             if fields is not None:
                 my_verts, svals = my_verts
                 samples.append(svals)
             samples = np.concatenate(samples)
             samples = self.comm.par_combine_object(samples, op='cat',
                                 datatype='array')
-            self[fields] = samples
+            if sample_type == "face":
+                self[fields] = samples
+            elif sample_type == "vertex":
+                self.vertex_samples[fields] = samples
+        
 
     @restore_grid_state
     def _extract_isocontours_from_grid(self, grid, field, value,
-                                       sample_values = None):
+                                       sample_values = None,
+                                       sample_type = "face"):
         mask = self.data_source._get_cut_mask(grid) * grid.child_mask
         vals = grid.get_vertex_centered_data(field, no_ghost = False)
         if sample_values is not None:
             svals = grid.get_vertex_centered_data(sample_values)
         else:
             svals = None
+        sample_type = {"face":1, "vertex":2}[sample_type]
         my_verts = march_cubes_grid(value, vals, mask, grid.LeftEdge,
-                                    grid.dds, svals)
+                                    grid.dds, svals, sample_type)
         return my_verts
 
     def calculate_flux(self, field_x, field_y, field_z, fluxing_field = None):
                     ff, mask, grid.LeftEdge, grid.dds)
 
     def export_ply(self, filename, bounds = None, color_field = None,
-                   color_map = "algae", color_log = True):
+                   color_map = "algae", color_log = True, sample_type = "face"):
         r"""This exports the surface to the PLY format, suitable for visualization
         in many different programs (e.g., MeshLab).
 
         >>> surf.export_ply("my_galaxy.ply", bounds = bounds)
         """
         if self.vertices is None:
-            self.get_data(color_field)
-        elif color_field is not None and color_field not in self.field_data:
-            self[color_field]
-        self._export_ply(filename, bounds, color_field, color_map, color_log)
+            self.get_data(color_field, sample_type)
+        elif color_field is not None:
+            if sample_type == "face" and \
+                color_field not in self.field_data:
+                self[color_field]
+            elif sample_type == "vertex" and \
+                color_field not in self.vertex_data:
+                self.get_data(color_field, sample_type)
+        self._export_ply(filename, bounds, color_field, color_map, color_log,
+                         sample_type)
+
+    def _color_samples(self, cs, color_log, color_map, arr):
+            if color_log: cs = np.log10(cs)
+            mi, ma = cs.min(), cs.max()
+            cs = (cs - mi) / (ma - mi)
+            from yt.visualization.image_writer import map_to_colors
+            cs = map_to_colors(cs, color_map)
+            arr["red"][:] = cs[0,:,0]
+            arr["green"][:] = cs[0,:,1]
+            arr["blue"][:] = cs[0,:,2]
 
     @parallel_root_only
     def _export_ply(self, filename, bounds = None, color_field = None,
-                   color_map = "algae", color_log = True):
+                   color_map = "algae", color_log = True, sample_type = "face"):
         f = open(filename, "wb")
         if bounds is None:
             DLE = self.pf.domain_left_edge
             DRE = self.pf.domain_right_edge
             bounds = [(DLE[i], DRE[i]) for i in range(3)]
+        nv = self.vertices.shape[1]
+        vs = [("x", "<f"), ("y", "<f"), ("z", "<f"),
+              ("red", "uint8"), ("green", "uint8"), ("blue", "uint8") ]
         fs = [("ni", "uint8"), ("v1", "<i4"), ("v2", "<i4"), ("v3", "<i4"),
               ("red", "uint8"), ("green", "uint8"), ("blue", "uint8") ]
-        v = np.empty((self.vertices.shape[1], 3), "<f")
-        for i in range(3):
-            v[:,i] = self.vertices[i,:]
-            np.subtract(v[:,i], bounds[i][0], v[:,i])
-            w = bounds[i][1] - bounds[i][0]
-            np.divide(v[:,i], w, v[:,i])
-            np.subtract(v[:,i], 0.5, v[:,i]) # Center at origin.
         f.write("ply\n")
         f.write("format binary_little_endian 1.0\n")
-        f.write("element vertex %s\n" % (v.shape[0]))
+        f.write("element vertex %s\n" % (nv))
         f.write("property float x\n")
         f.write("property float y\n")
         f.write("property float z\n")
-        f.write("element face %s\n" % (v.shape[0]/3))
+        if color_field is not None and sample_type == "vertex":
+            f.write("property uchar red\n")
+            f.write("property uchar green\n")
+            f.write("property uchar blue\n")
+            v = np.empty(self.vertices.shape[1], dtype=vs)
+            cs = self.vertex_samples[color_field]
+            self._color_samples(cs, color_log, color_map, v)
+        else:
+            v = np.empty(self.vertices.shape[1], dtype=vs[:3])
+        f.write("element face %s\n" % (nv/3))
         f.write("property list uchar int vertex_indices\n")
-        if color_field is not None:
+        if color_field is not None and sample_type == "face":
             f.write("property uchar red\n")
             f.write("property uchar green\n")
             f.write("property uchar blue\n")
             # Now we get our samples
             cs = self[color_field]
-            if color_log: cs = np.log10(cs)
-            mi, ma = cs.min(), cs.max()
-            cs = (cs - mi) / (ma - mi)
-            from yt.visualization.image_writer import map_to_colors
-            cs = map_to_colors(cs, color_map)
             arr = np.empty(cs.shape[1], dtype=np.dtype(fs))
-            arr["red"][:] = cs[0,:,0]
-            arr["green"][:] = cs[0,:,1]
-            arr["blue"][:] = cs[0,:,2]
+            self._color_samples(cs, color_log, color_map, arr)
         else:
-            arr = np.empty(v.shape[0]/3, np.dtype(fs[:-3]))
+            arr = np.empty(nv/3, np.dtype(fs[:-3]))
+        for i, ax in enumerate("xyz"):
+            v[ax][:] = self.vertices[i,:]
+            np.subtract(v[ax][:], bounds[i][0], v[ax][:])
+            w = bounds[i][1] - bounds[i][0]
+            np.divide(v[ax][:], w, v[ax][:])
+            np.subtract(v[ax][:], 0.5, v[ax][:]) # Center at origin.
         f.write("end_header\n")
         v.tofile(f)
         arr["ni"][:] = 3
-        vi = np.arange(v.shape[0], dtype="<i")
-        vi.shape = (v.shape[0]/3, 3)
+        vi = np.arange(nv, dtype="<i")
+        vi.shape = (nv/3, 3)
         arr["v1"][:] = vi[:,0]
         arr["v2"][:] = vi[:,1]
         arr["v3"][:] = vi[:,2]

yt/data_objects/tests/test_fluxes.py

     flux = surf.calculate_flux("Ones", "Zeros", "Zeros", "Ones")
     yield assert_almost_equal, flux, 1.0, 12
 
+def test_sampling():
+    pf = fake_random_pf(64, nprocs = 4)
+    dd = pf.h.all_data()
+    for i, ax in enumerate('xyz'):
+        surf = pf.h.surface(dd, ax, 0.51)
+        surf.get_data(ax, "vertex")
+        yield assert_equal, surf.vertex_samples[ax], surf.vertices[i,:]

yt/utilities/lib/marching_cubes.pyx

 cdef struct Triangle:
     Triangle *next
     np.float64_t p[3][3]
-    np.float64_t val
+    np.float64_t val[3] # Usually only use one value
 
 cdef struct TriangleCollection:
     int count
     return count
 
 cdef void FillTriangleValues(np.ndarray[np.float64_t, ndim=1] values,
-                             Triangle *first):
+                             Triangle *first, int nskip = 1):
     cdef Triangle *this = first
     cdef Triangle *last
     cdef int i = 0
+    cdef int j
     while this != NULL:
-        values[i] = this.val
+        for j in range(nskip):
+            values[i*nskip + j] = this.val[j]
         i += 1
         last = this
         this = this.next
                      np.ndarray[np.int32_t, ndim=3] mask,
                      np.ndarray[np.float64_t, ndim=1] left_edge,
                      np.ndarray[np.float64_t, ndim=1] dxs,
-                     obj_sample = None):
+                     obj_sample = None, int sample_type = 1):
     cdef int dims[3]
     cdef int i, j, k, n, m, nt
     cdef int offset
     if obj_sample is not None:
         sample = obj_sample
         sdata = <np.float64_t *> sample.data
-        do_sample = 1
+        do_sample = sample_type # 1 for face, 2 for vertex
     else:
         do_sample = 0
     for i in range(3):
                     offset_fill(dims, intdata, gv)
                     nt = march_cubes(gv, isovalue, dds, pos[0], pos[1], pos[2],
                                 &triangles)
-                    if do_sample == 1 and nt > 0:
+                    if nt == 0 or do_sample == 0:
+                        pos[2] += dds[2]
+                        continue
+                    if last == NULL and triangles.first != NULL:
+                        current = triangles.first
+                        last = NULL
+                    elif last != NULL:
+                        current = last.next
+                    if do_sample == 1:
                         # At each triangle's center, sample our secondary field
-                        if last == NULL and triangles.first != NULL:
-                            current = triangles.first
-                            last = NULL
-                        elif last != NULL:
-                            current = last.next
                         while current != NULL:
                             for n in range(3):
                                 point[n] = 0.0
                                     point[m] += (current.p[n][m]-pos[m])*idds[m]
                             for n in range(3):
                                 point[n] /= 3.0
-                            current.val = offset_interpolate(dims, point,
+                            current.val[0] = offset_interpolate(dims, point,
                                                              sdata + offset)
                             last = current
                             if current.next == NULL: break
                             current = current.next
+                    elif do_sample == 2:
+                        while current != NULL:
+                            for n in range(3):
+                                for m in range(3):
+                                    point[m] = (current.p[n][m]-pos[m])*idds[m]
+                                current.val[n] = offset_interpolate(dims,
+                                                    point, sdata + offset)
+                            last = current
+                            if current.next == NULL: break
+                            current = current.next
                 pos[2] += dds[2]
             pos[1] += dds[1]
         pos[0] += dds[0]
     # Hallo, we are all done.
     cdef np.ndarray[np.float64_t, ndim=2] vertices 
     vertices = np.zeros((triangles.count*3,3), dtype='float64')
+    if do_sample == 0:
+        FillAndWipeTriangles(vertices, triangles.first)
+    cdef int nskip
     if do_sample == 1:
-        sampled = np.zeros(triangles.count, dtype='float64')
-        FillTriangleValues(sampled, triangles.first)
-        FillAndWipeTriangles(vertices, triangles.first)
-        return vertices, sampled
+        nskip = 1
+    elif do_sample == 2:
+        nskip = 3
+    sampled = np.zeros(triangles.count * nskip, dtype='float64')
+    FillTriangleValues(sampled, triangles.first, nskip)
     FillAndWipeTriangles(vertices, triangles.first)
-    return vertices
+    return vertices, sampled
 
 @cython.boundscheck(False)
 @cython.wraparound(False)