Fix usage of MomentsTexture in volumetric examples

5 files changed, 57 insertions, 19 deletions

diff --git a/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py b/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py
index a51eaa1..41d2958 100644
--- a/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py
+++ b/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py
@@ -313,7 +313,7 @@ lattice = Lattice(
 
 lattice.setup_channel_with_sdf_obstacle(grid_fin)
 
-moments_texture = MomentsTexture(lattice)
+moments_texture = MomentsTexture(lattice, include_materials = False)
 
 projection = Projection(distance = 2*lattice_x)
 rotation = Rotation([-0.5*lattice_x, -0.5*lattice_y, -0.5*lattice_z])
diff --git a/channel_3d_volumetric_rendering_gl_interop.py b/channel_3d_volumetric_rendering_gl_interop.py
index b97c29b..28ed003 100644
--- a/channel_3d_volumetric_rendering_gl_interop.py
+++ b/channel_3d_volumetric_rendering_gl_interop.py
@@ -301,7 +301,7 @@ lattice = Lattice(
 
 lattice.setup_channel_with_sdf_obstacle(channel)
 
-moments_texture = MomentsTexture(lattice)
+moments_texture = MomentsTexture(lattice, include_materials = False)
 
 projection = Projection(distance = 2*lattice_x)
 rotation = Rotation([-0.5*lattice_x, -0.5*lattice_y, -0.5*lattice_z])
diff --git a/ldc_3d_gl_interop.py b/ldc_3d_gl_interop.py
index a8f6f8d..130d8ec 100644
--- a/ldc_3d_gl_interop.py
+++ b/ldc_3d_gl_interop.py
@@ -142,11 +142,8 @@ lattice.apply_material_map(
     get_cavity_material_map(lattice.geometry))
 lattice.sync_material()
 
-moments_vbo = MomentsVertexBuffer(lattice)
-
 particles = Particles(
     lattice,
-    moments_vbo,
     numpy.mgrid[
         8*lattice.geometry.size_x//10:9*lattice.geometry.size_x//10:10j,
         lattice.geometry.size_y//10:9*lattice.geometry.size_y//10:particle_count/100j,
@@ -162,7 +159,7 @@ def on_display():
     for i in range(0,updates_per_frame):
         lattice.evolve()
 
-    moments_vbo.collect()
+    lattice.update_moments()
 
     for i in range(0,updates_per_frame):
         particles.update(aging = True)
diff --git a/template/opengl.mako b/template/opengl.mako
index 2029623..6958db8 100644
--- a/template/opengl.mako
+++ b/template/opengl.mako
@@ -64,12 +64,12 @@ def neighbor_offset(c_i):
 
 %>
 
-__kernel void collect_gl_moments_to_texture(__global ${float_type}* f,
-                                            __global int* material,
+__kernel void collect_gl_moments_and_materials_to_texture(__global ${float_type}* f,
+                                                          __global int* material,
 % if descriptor.d == 2:
-                                            __write_only image2d_t moments)
+                                                          __write_only image2d_t moments)
 % elif descriptor.d == 3:
-                                            __write_only image3d_t moments)
+                                                          __write_only image3d_t moments)
 % endif
 {
     const unsigned int gid = ${gid()};
@@ -107,3 +107,39 @@ __kernel void collect_gl_moments_to_texture(__global ${float_type}* f,
 
     write_imagef(moments, ${moments_cell()}, data);
 }
+
+__kernel void collect_gl_moments_to_texture(__global ${float_type}* f,
+% if descriptor.d == 2:
+                                            __write_only image2d_t moments)
+% elif descriptor.d == 3:
+                                            __write_only image3d_t moments)
+% endif
+{
+    const unsigned int gid = ${gid()};
+
+    __global ${float_type}* preshifted_f = f + gid;
+
+% for i in range(0,descriptor.q):
+    const ${float_type} f_curr_${i} = preshifted_f[${pop_offset(i)}];
+% endfor
+
+% for i, expr in enumerate(moments_subexpr):
+    const ${float_type} ${expr[0]} = ${ccode(expr[1])};
+% endfor
+
+    float4 data;
+
+% if descriptor.d == 2:
+    data.x = ${ccode(moments_assignment[0].rhs)};
+    data.y = ${ccode(moments_assignment[1].rhs)};
+    data.z = ${ccode(moments_assignment[2].rhs)};
+    data.w = sqrt(data.y*data.y + data.z*data.z);
+% elif descriptor.d == 3:
+    data.x = ${ccode(moments_assignment[0].rhs)};
+    data.y = ${ccode(moments_assignment[1].rhs)};
+    data.z = ${ccode(moments_assignment[2].rhs)};
+    data.w = ${ccode(moments_assignment[3].rhs)};
+% endif
+
+    write_imagef(moments, ${moments_cell()}, data);
+}
diff --git a/utility/opengl.py b/utility/opengl.py
index 1037cbd..f305eae 100644
--- a/utility/opengl.py
+++ b/utility/opengl.py
@@ -65,8 +65,9 @@ class MomentsVertexBuffer:
 
 
 class MomentsTexture:
-    def __init__(self, lattice):
+    def __init__(self, lattice, include_materials = True):
         self.lattice = lattice
+        self.include_materials = include_materials
         self.gl_texture_buffer = numpy.ndarray(shape=(self.lattice.memory.volume, 4), dtype=self.lattice.memory.float_type)
         self.gl_texture_buffer[:,:] = 0.0
 
@@ -114,15 +115,13 @@ class MomentsTexture:
         glActiveTexture(location);
         glBindTexture(self.gl_texture_type, self.gl_moments)
 
-    def collect(self):
-        cl.enqueue_acquire_gl_objects(self.lattice.queue, [self.cl_gl_moments])
-
-        if self.lattice.tick:
-            self.program.collect_gl_moments_to_texture(
+    def collect_moments_from_pop_to_texture(self, population):
+        if self.include_materials:
+            self.program.collect_gl_moments_and_materials_to_texture(
                 self.lattice.queue,
                 self.lattice.grid.size(),
                 self.lattice.layout,
-                self.lattice.memory.cl_pop_b,
+                population,
                 self.lattice.memory.cl_material,
                 self.cl_gl_moments)
         else:
@@ -130,7 +129,13 @@ class MomentsTexture:
                 self.lattice.queue,
                 self.lattice.grid.size(),
                 self.lattice.layout,
-                self.lattice.memory.cl_pop_a,
-                self.lattice.memory.cl_material,
+                population,
                 self.cl_gl_moments)
 
+    def collect(self):
+        cl.enqueue_acquire_gl_objects(self.lattice.queue, [self.cl_gl_moments])
+
+        if self.lattice.tick:
+            self.collect_moments_from_pop_to_texture(self.lattice.memory.cl_pop_b)
+        else:
+            self.collect_moments_from_pop_to_texture(self.lattice.memory.cl_pop_a)