Fix PBC and VTK output

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

code_gen/cgen.py

@@ -126,6 +126,7 @@ class CGen:
                  "ASCII\n" \
                  "DATASET UNSTRUCTURED_GRID\n"
 
+        end = start + n
         filename_var = f"filename{id}"
         filehandle_var = f"vtk{id}"
         printer.print(f"char {filename_var}[128];")
@@ -137,7 +138,7 @@ class CGen:
 
         # Write positions
         printer.print(f"fprintf({filehandle_var}, \"POINTS %d double\\n\", {n.generate()});")
-        CGen.generate_for_preamble("i", start, n.generate())
+        CGen.generate_for_preamble("i", start.generate(), end.generate())
         printer.add_ind(4)
         printer.print(f"fprintf({filehandle_var}, \"%.4f %.4f %.4f\\n\", position[i * 3], position[i * 3 + 1], position[i * 3 + 2]);")
         printer.add_ind(-4)
@@ -146,16 +147,16 @@ class CGen:
 
         # Write cells
         printer.print(f"fprintf({filehandle_var}, \"CELLS %d %d\\n\", {n.generate()}, {n.generate()} * 2);")
-        CGen.generate_for_preamble("i", start, n.generate())
+        CGen.generate_for_preamble("i", start.generate(), end.generate())
         printer.add_ind(4)
-        printer.print(f"fprintf({filehandle_var}, \"1 %d\\n\", i);")
+        printer.print(f"fprintf({filehandle_var}, \"1 %d\\n\", i - {start.generate()});")
         printer.add_ind(-4)
         CGen.generate_for_epilogue()
         printer.print(f"fwrite(\"\\n\\n\", 1, 2, {filehandle_var});")
 
         # Write cell types
         printer.print(f"fprintf({filehandle_var}, \"CELL_TYPES %d\\n\", {n.generate()});")
-        CGen.generate_for_preamble("i", start, n.generate())
+        CGen.generate_for_preamble("i", start.generate(), end.generate())
         printer.add_ind(4)
         printer.print(f"fwrite(\"1\\n\", 1, 2, {filehandle_var});")
         printer.add_ind(-4)
@@ -166,9 +167,10 @@ class CGen:
         printer.print(f"fprintf({filehandle_var}, \"POINT_DATA %d\\n\", {n.generate()});")
         printer.print(f"fprintf({filehandle_var}, \"SCALARS mass double\\n\");")
         printer.print(f"fprintf({filehandle_var}, \"LOOKUP_TABLE default\\n\");")
-        CGen.generate_for_preamble("i", start, n.generate())
+        CGen.generate_for_preamble("i", start.generate(), end.generate())
         printer.add_ind(4)
-        printer.print(f"fprintf({filehandle_var}, \"%4.f\\n\", mass[i]);")
+        #printer.print(f"fprintf({filehandle_var}, \"%4.f\\n\", mass[i]);")
+        printer.print(f"fprintf({filehandle_var}, \"1.0\\n\");")
         printer.add_ind(-4)
         CGen.generate_for_epilogue()
         printer.print(f"fwrite(\"\\n\\n\", 1, 2, {filehandle_var});")

sim/pbc.py

@@ -95,11 +95,15 @@ 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):
                                     # 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)
@@ -113,10 +117,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)

sim/vtk.py

@@ -12,7 +12,14 @@ class VTKWrite:
         return []
 
     def generate(self):
-        self.sim.code_gen.generate_vtk_writing(VTKWrite.vtk_id, self.filename, 0, self.sim.nlocal, self.timestep)
+        self.sim.code_gen.generate_vtk_writing(
+            VTKWrite.vtk_id * 2, f"{self.filename}_local",
+            as_lit_ast(self.sim, 0), self.sim.nlocal, self.timestep)
+
+        self.sim.code_gen.generate_vtk_writing(
+            VTKWrite.vtk_id * 2 + 1, f"{self.filename}_pbc",
+            self.sim.nlocal, self.sim.pbc.npbc, self.timestep)
+
         VTKWrite.vtk_id += 1
 
     def transform(self, fn):