Merge pull request #12 from jameslawlor/feature/add-frame-num-to-anim…

…ation add support for drawing frame number in animation
jameslawlor · Mar 30, 2024 · bc6c502 · bc6c502
2 parents 7989fe7 + 8f561cb
commit bc6c502
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 9 deletions.
diff --git a/src/pyturmite/constants.py b/src/pyturmite/constants.py
@@ -2,8 +2,8 @@
 N_STEPS = 6000
 PLOT_MODE = "animate"
 CMAP = "Oranges"
-RULESET = "RL"
+RULESET = "RLLR"
 ANIMATION_INTERVAL = 1
-PADDING_SIZE = 2 # expand in each direction by this amount of cells
+PADDING_SIZE = 2  # expand in each direction by this amount of cells
 SAVE_ANIMATION = False
-FRAME_SKIP = 2
+FRAME_SKIP = 5
diff --git a/src/pyturmite/turmites.py b/src/pyturmite/turmites.py
@@ -104,14 +104,14 @@ def update(self):
 
         # expand canvas/grid if we moved outside it
         if (
-            (self.x >= self.grid.shape[0]) or
-            (self.x < 0) or
-            (self.y >= self.grid.shape[1]) or
-            (self.y < 0)        
+            (self.x >= self.grid.shape[0])
+            or (self.x < 0)
+            or (self.y >= self.grid.shape[1])
+            or (self.y < 0)
         ):
             self.expand_grid()
 
     def expand_grid(self):
         self.grid = np.pad(self.grid, PADDING_SIZE)
         self.x += PADDING_SIZE
-        self.y += PADDING_SIZE
+        self.y += PADDING_SIZE
diff --git a/src/pyturmite/utils/plotters.py b/src/pyturmite/utils/plotters.py
@@ -46,6 +46,17 @@ def animate(self, turmite, n_steps):
             vmax=turmite.n_colours + 1,
             cmap=turmite.cmap,
         )
+        ax = plt.gca()
+        frame_num = ax.text(
+            0.99, 
+            0.01,
+            '',
+            fontsize=5,
+            horizontalalignment='right',
+            verticalalignment='bottom',
+            transform=ax.transAxes,
+            )
+        frame_num.set_bbox(dict(facecolor='white', alpha=0.3, edgecolor='black', lw=.1))
         plt.axis("off")
 
         def update(step, turmite, skip=FRAME_SKIP):
@@ -54,7 +65,8 @@ def update(step, turmite, skip=FRAME_SKIP):
             data = np.array(turmite.grid)
             data[turmite.x][turmite.y] = turmite.n_colours + 1
             im.set_data(data)
-            return [im]
+            frame_num.set_text(f"Step: {step*skip:,}")
+            return [im, frame_num]
 
         anim = FuncAnimation(  # noqa: F841
             fig,