Finally fixed PBC

Signed-off-by: Rafael Ravedutti Lucio Machado <rafael.r.ravedutti@fau.de>

ast/loops.py

@@ -21,11 +21,17 @@ class Iter():
         return Type_Int
 
     def is_mutable(self):
-        return False
+        # FIXME: This should be set to False, but currently the flattening transformation is reusing
+        # expressions that are not alive anymore (used first within if, then outside it), causing
+        # the generated code to be uncompilable
+        return True
 
     def scope(self):
         return self.loop.block
 
+    def __add__(self, other):
+        return Expr(self.sim, self, other, '+')
+
     def __sub__(self, other):
         return Expr(self.sim, self, other, '-')
 

sim/particle_simulation.py

@@ -175,22 +175,17 @@ class ParticleSimulation:
     def generate(self):
         timestep = Timestep(self, self.ntimesteps, [
             (EnforcePBC(self.pbc).lower(), 20),
-            (SetupPBC(self.pbc).lower(), 20),
-            UpdatePBC(self.pbc).lower(),
+            (SetupPBC(self.pbc).lower(), UpdatePBC(self.pbc).lower(), 20),
             (CellListsBuild(self.cell_lists).lower(), 20),
             PropertiesResetVolatile(self).lower(),
             self.kernels.lower()
         ])
 
-        timestep.add(Block(self, VTKWrite(self, self.vtk_file, timestep.timestep())))
+        timestep.add(Block(self, VTKWrite(self, self.vtk_file, timestep.timestep() + 1)))
 
         body = Block.from_list(self, [
             CellListsStencilBuild(self.cell_lists).lower(),
             self.setups.lower(),
-            EnforcePBC(self.pbc).lower(),
-            SetupPBC(self.pbc).lower(),
-            UpdatePBC(self.pbc).lower(),
-            CellListsBuild(self.cell_lists).lower(),
             Block(self, VTKWrite(self, self.vtk_file, 0)),
             timestep.as_block()
         ])

sim/pbc.py

@@ -95,20 +95,13 @@ class SetupPBC:
                             if capacity_exceeded:
                                 resize.set(Select(sim, resize > npbc, resize + 1, npbc))
                             else:
+                                pbc_map[npbc].set(i)
                                 pbc_mult[npbc][d].set(1)
+                                positions[nlocal + npbc][d].set(positions[i][d] + grid.length(d))
 
-                                for is_local in Branch(sim, i < nlocal):
-                                    # TODO: VecFilter.others generator?
-                                    if is_local:
-                                        pbc_map[npbc].set(i)
-                                    else:
-                                        pbc_map[npbc].set(pbc_map[i - nlocal])
-
-                                    for d_ in [x for x in range(0, ndims) if x != d]:
-                                        if is_local:
-                                            pbc_mult[npbc][d_].set(0)
-                                        else:
-                                            pbc_mult[npbc][d_].set(pbc_mult[i - nlocal][d_])
+                                for d_ in [x for x in range(0, ndims) if x != d]:
+                                    pbc_mult[npbc][d_].set(0)
+                                    positions[nlocal + npbc][d_].set(positions[i][d_])
 
                                 npbc.add(1)
 
@@ -117,18 +110,13 @@ class SetupPBC:
                             if capacity_exceeded:
                                 resize.set(Select(sim, resize > npbc, resize + 1, npbc))
                             else:
+                                pbc_map[npbc].set(i)
                                 pbc_mult[npbc][d].set(-1)
-                                for is_local in Branch(sim, i < nlocal):
-                                    if is_local:
-                                        pbc_map[npbc].set(i)
-                                    else:
-                                        pbc_map[npbc].set(pbc_map[i - nlocal])
-
-                                    for d_ in [x for x in range(0, ndims) if x != d]:
-                                        if is_local:
-                                            pbc_mult[npbc][d_].set(0)
-                                        else:
-                                            pbc_mult[npbc][d_].set(pbc_mult[i - nlocal][d_])
+                                positions[nlocal + npbc][d].set(positions[i][d] - grid.length(d))
+
+                                for d_ in [x for x in range(0, ndims) if x != d]:
+                                    pbc_mult[npbc][d_].set(0)
+                                    positions[nlocal + npbc][d_].set(positions[i][d_])
 
                                 npbc.add(1)
 

sim/timestep.py

@@ -8,29 +8,34 @@ class Timestep:
     def __init__(self, sim, nsteps, item_list=None):
         self.sim = sim
         self.block = Block(sim, [])
-        self.timestep_loop = For(sim, 1, nsteps + 1, self.block)
+        self.timestep_loop = For(sim, 0, nsteps + 1, self.block)
 
         if item_list is not None:
             for item in item_list:
                 if isinstance(item, tuple):
-                    self.add(item[0], item[1])
+                    if len(item) >= 3:
+                        self.add(item[0], item[2], item[1])
+                    else:
+                        self.add(item[0], item[1])
                 else:
                     self.add(item)
 
     def timestep(self):
         return self.timestep_loop.iter()
 
-    def add(self, item, exec_every=0):
-        assert exec_every >= 0, \
-            "exec_every parameter must be higher or equal than zero!"
-
+    def add(self, item, exec_every=0, item_else=None):
+        assert exec_every >= 0, "exec_every parameter must be higher or equal than zero!"
         stmts = item if not isinstance(item, Block) else item.statements()
+        stmts_else = None
         ts = self.timestep_loop.iter()
+
+        if item_else is not None:
+            stmts_else = item_else if not isinstance(item_else, Block) else item_else.statements()
+
         if exec_every > 0:
             self.block.add_statement(
-                Branch(self.sim,
-                       Expr.cmp(ts % exec_every, 0),
-                       True, Block(self.sim, stmts), None))
+                Branch(self.sim, Expr.cmp(ts % exec_every, 0), True if stmts_else is None else False,
+                Block(self.sim, stmts), None if stmts_else is None else Block(self.sim, stmts_else)))
         else:
             self.block.add_statement(stmts)