Small fixes

Signed-off-by: Rafael Ravedutti <rafaelravedutti@gmail.com>

src/pairs/mapping/keywords.py

@@ -79,6 +79,7 @@ class Keywords:
         length = self.keyword_length([vector])
         inv_length = Lit(self.sim, 1.0) / length
         return Select(self.sim, length > Lit(self.sim, 0.0), vector * inv_length, ZeroVector(self.sim))
+        #return vector * inv_length
 
     def keyword_squared_length(self, args):
         assert len(args) == 1, "length() keyword requires one parameter!"

src/pairs/sim/cell_lists.py

@@ -122,18 +122,20 @@ class PartitionCellLists(Lowerable):
             end = self.sim.add_temp_var(0)
 
             for shape in For(self.sim, 0, self.sim.max_shapes()):
+                shape_start = self.sim.add_temp_var(start)
                 Assign(self.sim, end, self.cell_lists.cell_sizes[cell] - 1)
 
-                for _ in While(self.sim, start < end):
+                for _ in While(self.sim, start <= end):
                     particle = cell_particles[cell][start]
 
                     for unmatch in Branch(self.sim, ScalarOp.neq(self.sim.particle_shape[particle], shape)):
                         if unmatch:
-                            Assign(self.sim, cell_particles[cell][start], cell_particles[cell][end])
-                            Assign(self.sim, cell_particles[cell][end], particle)
+                            for _ in Filter(self.sim, ScalarOp.neq(start, end)):
+                                Assign(self.sim, cell_particles[cell][start], cell_particles[cell][end])
+                                Assign(self.sim, cell_particles[cell][end], particle)
+
                             Assign(self.sim, end, end - 1)
 
                         else:
                             Assign(self.sim, start, start + 1)
-
-                    Assign(self.sim, self.cell_lists.nshapes[cell][shape], start)
+                            Assign(self.sim, self.cell_lists.nshapes[cell][shape], start - shape_start)

src/pairs/sim/pbc.py

@@ -20,12 +20,16 @@ class EnforcePBC(Lowerable):
         positions = sim.position()
         sim.module_name("enforce_pbc")
 
+        # Particles with one of the following flags are ignored
+        flags_to_exclude = (Flags.Infinite | Flags.Fixed | Flags.Global)
+
         for i in ParticleFor(sim):
-            # TODO: VecFilter?
-            for d in range(0, ndims):
-                if sim._pbc[d] is True:
-                    for _ in Filter(sim, positions[i][d] < grid.min(d)):
-                        Assign(sim, positions[i][d], positions[i][d] + grid.length(d)) 
+            for _ in Filter(self.sim, ScalarOp.cmp(self.sim.particle_flags[i] & flags_to_exclude, 0)):
+                # TODO: VecFilter?
+                for d in range(0, ndims):
+                    if sim._pbc[d] is True:
+                        for _ in Filter(sim, positions[i][d] < grid.min(d)):
+                            Assign(sim, positions[i][d], positions[i][d] + grid.length(d))
 
-                    for _ in Filter(sim, positions[i][d] > grid.max(d)):
-                        Assign(sim, positions[i][d], positions[i][d] - grid.length(d)) 
+                        for _ in Filter(sim, positions[i][d] > grid.max(d)):
+                            Assign(sim, positions[i][d], positions[i][d] - grid.length(d))