Implement LeapFrog integrator

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

examples/lj.py

@@ -7,11 +7,16 @@ def lj(i, j):
     sr6 = sr2 * sr2 * sr2 * sigma6[i, j]
     apply(force, delta(i, j) * (48.0 * sr6 * (sr6 - 0.5) * sr2 * epsilon[i, j]))
 
-def euler(i):
-    linear_velocity[i] += dt * force[i] / mass[i]
+
+def initial_integrate(i):
+    linear_velocity[i] += (dt * 0.5) * force[i] / mass[i]
     position[i] += dt * linear_velocity[i]
 
 
+def final_integrate(i):
+    linear_velocity[i] += (dt * 0.5) * force[i] / mass[i]
+
+
 cmd = sys.argv[0]
 target = sys.argv[1] if len(sys.argv[1]) > 1 else "none"
 if target != 'cpu' and target != 'gpu':
@@ -25,7 +30,7 @@ sigma = 1.0
 epsilon = 1.0
 sigma6 = sigma ** 6
 
-psim = pairs.simulation("lj", [pairs.point_mass()], timesteps=200, double_prec=True, debug=True)
+psim = pairs.simulation("lj", [pairs.point_mass()], timesteps=200, double_prec=True)
 psim.add_position('position')
 psim.add_property('mass', pairs.real(), 1.0)
 psim.add_property('linear_velocity', pairs.vector())
@@ -39,8 +44,9 @@ psim.reneighbor_every(20)
 #psim.compute_half()
 psim.build_neighbor_lists(cutoff_radius + skin)
 psim.vtk_output(f"output/lj_{target}")
+psim.compute(initial_integrate, symbols={'dt': dt}, pre_step=True, skip_first=True)
 psim.compute(lj, cutoff_radius)
-psim.compute(euler, symbols={'dt': dt})
+psim.compute(final_integrate, symbols={'dt': dt}, skip_first=True)
 
 if target == 'gpu':
     psim.target(pairs.target_gpu())

src/pairs/mapping/funcs.py

@@ -273,7 +273,7 @@ class BuildParticleIR(ast.NodeVisitor):
         return op_class(self.sim, operand, None, BuildParticleIR.get_unary_op(node.op))
 
 
-def compute(sim, func, cutoff_radius=None, symbols={}):
+def compute(sim, func, cutoff_radius=None, symbols={}, pre_step=False, skip_first=False):
     src = inspect.getsource(func)
     tree = ast.parse(src, mode='exec')
     #print(ast.dump(ast.parse(src, mode='exec')))
@@ -318,7 +318,11 @@ def compute(sim, func, cutoff_radius=None, symbols={}):
 
             ir.visit(tree)
 
-    sim.build_module_with_statements()
+    if pre_step:
+        sim.build_pre_step_module_with_statements(skip_first=skip_first)
+
+    else:
+        sim.build_module_with_statements(skip_first=skip_first)
 
 
 def setup(sim, func, symbols={}):

src/pairs/sim/simulation.py

@@ -10,7 +10,7 @@ from pairs.ir.properties import Properties, ContactProperties
 from pairs.ir.symbols import Symbol
 from pairs.ir.types import Types
 from pairs.ir.variables import Variables
-from pairs.graph.graphviz import ASTGraph
+#from pairs.graph.graphviz import ASTGraph
 from pairs.mapping.funcs import compute, setup
 from pairs.sim.arrays import DeclareArrays
 from pairs.sim.cell_lists import CellLists, BuildCellLists, BuildCellListsStencil, PartitionCellLists
@@ -59,8 +59,9 @@ class Simulation:
         self._capture_statements = True
         self._block = Block(self, [])
         self.setups = Block(self, [])
-        self.setup_functions = Block(self, [])
-        self.functions = Block(self, [])
+        self.setup_functions = []
+        self.pre_step_functions = []
+        self.functions = []
         self.module_list = []
         self.kernel_list = []
         self._check_properties_resize = False
@@ -215,8 +216,8 @@ class Simulation:
         self.neighbor_lists = NeighborLists(self.cell_lists)
         return self.neighbor_lists
 
-    def compute(self, func, cutoff_radius=None, symbols={}):
-        return compute(self, func, cutoff_radius, symbols)
+    def compute(self, func, cutoff_radius=None, symbols={}, pre_step=False, skip_first=False):
+        return compute(self, func, cutoff_radius, symbols, pre_step, skip_first)
 
     def setup(self, func, symbols={}):
         return setup(self, func, symbols)
@@ -240,7 +241,7 @@ class Simulation:
             raise Exception("Two sizes assigned to same capacity!")
 
     def build_setup_module_with_statements(self):
-        self.setup_functions.add_statement(
+        self.setup_functions.append(
             Module(self,
                 name=self._module_name,
                 block=Block(self, self._block),
@@ -248,14 +249,31 @@ class Simulation:
                 check_properties_resize=self._check_properties_resize,
                 run_on_device=False))
 
-    def build_module_with_statements(self, run_on_device=True):
-        self.functions.add_statement(
-            Module(self,
-                name=self._module_name,
-                block=Block(self, self._block),
-                resizes_to_check=self._resizes_to_check,
-                check_properties_resize=self._check_properties_resize,
-                run_on_device=run_on_device))
+    def build_pre_step_module_with_statements(self, run_on_device=True, skip_first=False):
+        module = Module(self, name=self._module_name,
+                              block=Block(self, self._block),
+                              resizes_to_check=self._resizes_to_check,
+                              check_properties_resize=self._check_properties_resize,
+                              run_on_device=run_on_device)
+
+        if skip_first:
+            self.pre_step_functions.append((module, {'skip_first': True}))
+
+        else:
+            self.pre_step_functions.append(module)
+
+    def build_module_with_statements(self, run_on_device=True, skip_first=False):
+        module = Module(self, name=self._module_name,
+                              block=Block(self, self._block),
+                              resizes_to_check=self._resizes_to_check,
+                              check_properties_resize=self._check_properties_resize,
+                              run_on_device=run_on_device)
+
+        if skip_first:
+            self.functions.append((module, {'skip_first': True}))
+
+        else:
+            self.functions.append(module)
 
     def capture_statements(self, capture=True):
         self._capture_statements = capture
@@ -312,41 +330,45 @@ class Simulation:
     def generate(self):
         assert self._target is not None, "Target not specified!"
         comm = Comm(self, self._dom_part)
+        every_reneighbor_params = {'every': self.reneighbor_frequency}
 
-        timestep_procedures = [
-            (comm.exchange(), self.reneighbor_frequency),
-            (comm.borders(), comm.synchronize(), self.reneighbor_frequency),
-            (BuildCellLists(self, self.cell_lists), self.reneighbor_frequency),
-            (PartitionCellLists(self, self.cell_lists), self.reneighbor_frequency)
+        timestep_procedures = self.pre_step_functions + [
+            (comm.exchange(), every_reneighbor_params),
+            (comm.borders(), comm.synchronize(), every_reneighbor_params),
+            (BuildCellLists(self, self.cell_lists), every_reneighbor_params),
+            (PartitionCellLists(self, self.cell_lists), every_reneighbor_params)
         ]
 
         if self.neighbor_lists is not None:
-            timestep_procedures.append((BuildNeighborLists(self, self.neighbor_lists), self.reneighbor_frequency))
+            timestep_procedures.append(
+                (BuildNeighborLists(self, self.neighbor_lists), every_reneighbor_params))
 
         if self._use_contact_history:
             if self.neighbor_lists is not None:
-                timestep_procedures.append((BuildContactHistory(self, self._contact_history, self.cell_lists), self.reneighbor_frequency))
+                timestep_procedures.append(
+                    (BuildContactHistory(self, self._contact_history, self.cell_lists), every_reneighbor_params))
 
             timestep_procedures.append(ResetContactHistoryUsageStatus(self, self._contact_history))
 
-        timestep_procedures += [
-            ResetVolatileProperties(self),
-            self.functions
-        ]
+        timestep_procedures += [ResetVolatileProperties(self)] + self.functions
 
         if self._use_contact_history:
             timestep_procedures.append(ClearUnusedContactHistory(self, self._contact_history))
 
         timestep = Timestep(self, self.ntimesteps, timestep_procedures)
         self.enter(timestep.block)
-        timestep.add(VTKWrite(self, self.vtk_file, timestep.timestep() + 1, self.vtk_frequency))
+
+        if self.vtk_file is not None:
+            timestep.add(VTKWrite(self, self.vtk_file, timestep.timestep() + 1, self.vtk_frequency))
+
         self.leave()
 
+        vtk_item = [] if self.vtk_file is None else VTKWrite(self, self.vtk_file, 0, self.vtk_frequency)
         body = Block.from_list(self, [
             self.setups,
             self.setup_functions,
             BuildCellListsStencil(self, self.cell_lists),
-            VTKWrite(self, self.vtk_file, 0, self.vtk_frequency),
+            vtk_item,
             timestep.as_block()
         ])
 

src/pairs/sim/timestep.py

 from pairs.ir.scalars import ScalarOp
 from pairs.ir.block import Block
-from pairs.ir.branches import Branch
+from pairs.ir.branches import Branch, Filter
 from pairs.ir.loops import For
 
 
@@ -8,22 +8,30 @@ class Timestep:
     def __init__(self, sim, nsteps, item_list=None):
         self.sim = sim
         self.block = Block(sim, [])
-        self.timestep_loop = For(sim, 0, nsteps + 1, self.block)
+        self.timestep_loop = For(sim, 0, nsteps, self.block)
 
         if item_list is not None:
             for item in item_list:
                 if isinstance(item, tuple):
-                    if len(item) >= 3:
-                        self.add(item[0], item[2], item[1])
-                    else:
-                        self.add(item[0], item[1])
+                    stmt_else = None
+
+                    if len(item) == 2:
+                        stmt, params = item
+
+                    if len(item) == 3:
+                        stmt, stmt_else, params = item
+
+                    exec_every = 0 if 'every' not in params else params['every']
+                    skip_first = False if 'skip_first' not in params else params['skip_first']
+                    self.add(stmt, exec_every, stmt_else, skip_first)
+
                 else:
                     self.add(item)
 
     def timestep(self):
         return self.timestep_loop.iter()
 
-    def add(self, item, exec_every=0, item_else=None):
+    def add(self, item, exec_every=0, item_else=None, skip_first=False):
         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
@@ -37,6 +45,10 @@ class Timestep:
             self.block.add_statement(
                 Branch(self.sim, ScalarOp.inline(ScalarOp.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)))
+
+        elif skip_first:
+            self.block.add_statement(Filter(self.sim, ScalarOp.inline(ts > 0), Block(self.sim, stmts)))
+
         else:
             self.block.add_statement(stmts)