diff --git a/src/pairs/sim/cell_lists.py b/src/pairs/sim/cell_lists.py
index 4055313c637d70c3954806572153c8dbce0bf045..4d3fb98d582416eae1c65b00bf86468665521947 100644
--- a/src/pairs/sim/cell_lists.py
+++ b/src/pairs/sim/cell_lists.py
@@ -28,8 +28,8 @@ class CellLists:
         self.ncells             =   self.sim.add_var('ncells', Types.Int32, 1)
         self.ncells_capacity    =   self.sim.add_var('ncells_capacity', Types.Int32, 100)
         self.cell_capacity      =   self.sim.add_var('cell_capacity', Types.Int32, 20)
-        self.dim_ncells         =   self.sim.add_static_array('dim_cells', self.sim.ndims(), Types.Int32)
-        self.shapes_buffer      =   self.sim.add_static_array('shapes_buffer', self.sim.max_shapes(), Types.Int32)
+        self.dim_ncells         =   self.sim.add_array('dim_cells', self.sim.ndims(), Types.Int32)
+        self.shapes_buffer      =   self.sim.add_array('shapes_buffer', self.sim.max_shapes(), Types.Int32)
         self.cell_particles     =   self.sim.add_array('cell_particles', [self.ncells_capacity, self.cell_capacity], Types.Int32)
         self.cell_sizes         =   self.sim.add_array('cell_sizes', self.ncells_capacity, Types.Int32)
         self.nshapes            =   self.sim.add_array('nshapes', [self.ncells_capacity, self.sim.max_shapes()], Types.Int32)
@@ -51,12 +51,16 @@ class BuildCellListsStencil(Lowerable):
         dim_ncells = self.cell_lists.dim_ncells
         ncells = self.cell_lists.ncells
         ncells_capacity = self.cell_lists.ncells_capacity
+        shapes_buffer = self.cell_lists.shapes_buffer
         index = None
         ntotal_cells = 1
 
         self.sim.module_name("build_cell_lists_stencil")
         self.sim.check_resize(ncells_capacity, ncells)
 
+        for s in range(self.sim.max_shapes()):
+            Assign(self.sim, shapes_buffer[s], self.sim.get_shape_id(s))
+
         for dim in range(self.sim.ndims()):
             dim_min = self.sim.grid.min(dim) - spacing[dim]
             dim_max = self.sim.grid.max(dim) + spacing[dim]
@@ -133,9 +137,6 @@ class PartitionCellLists(Lowerable):
         cell_particles = self.cell_lists.cell_particles
         shapes_buffer = self.cell_lists.shapes_buffer
 
-        for s in range(self.sim.max_shapes()):
-            Assign(self.sim, shapes_buffer[s], self.sim.get_shape_id(s))
-
         for cell in For(self.sim, 0, self.cell_lists.ncells):
             start = self.sim.add_temp_var(0)
             end = self.sim.add_temp_var(0)
@@ -144,6 +145,7 @@ class PartitionCellLists(Lowerable):
                 shape_id = shapes_buffer[shape]
                 shape_start = self.sim.add_temp_var(start)
                 Assign(self.sim, end, self.cell_lists.cell_sizes[cell] - 1)
+                Assign(self.sim, self.cell_lists.nshapes[cell][shape], 0)
 
                 for _ in While(self.sim, start <= end):
                     particle = cell_particles[cell][start]