langtons_ant.py

"""
Langton's ant

@ https://en.wikipedia.org/wiki/Langton%27s_ant
@ https://upload.wikimedia.org/wikipedia/commons/0/09/LangtonsAntAnimated.gif
"""

fromfunctoolsimportpartial

frommatplotlibimportpyplotasplt
frommatplotlib.animationimportFuncAnimation

WIDTH=80
HEIGHT=80


classLangtonsAnt:
"""
 Represents the main LangonsAnt algorithm.

 >>> la = LangtonsAnt(2, 2)
 >>> la.board
 [[True, True], [True, True]]
 >>> la.ant_position
 (1, 1)
 """

def__init__(self, width: int, height: int) ->None:
# Each square is either True or False where True is white and False is black
self.board= [[True] *widthfor_inrange(height)]
self.ant_position: tuple[int, int] = (width//2, height//2)

# Initially pointing left (similar to the wikipedia image)
# (0 = 0° | 1 = 90° | 2 = 180 ° | 3 = 270°)
self.ant_direction: int=3

defmove_ant(self, axes: plt.Axes|None, display: bool, _frame: int) ->None:
"""
 Performs three tasks:
 1. The ant turns either clockwise or anti-clockwise according to the colour
 of the square that it is currently on. If the square is white, the ant
 turns clockwise, and if the square is black the ant turns anti-clockwise
 2. The ant moves one square in the direction that it is currently facing
 3. The square the ant was previously on is inverted (White -> Black and
 Black -> White)

 If display is True, the board will also be displayed on the axes

 >>> la = LangtonsAnt(2, 2)
 >>> la.move_ant(None, True, 0)
 >>> la.board
 [[True, True], [True, False]]
 >>> la.move_ant(None, True, 0)
 >>> la.board
 [[True, False], [True, False]]
 """
directions= {
0: (-1, 0), # 0°
1: (0, 1), # 90°
2: (1, 0), # 180°
3: (0, -1), # 270°
 }
x, y=self.ant_position

# Turn clockwise or anti-clockwise according to colour of square
ifself.board[x][y] isTrue:
# The square is white so turn 90° clockwise
self.ant_direction= (self.ant_direction+1) %4
else:
# The square is black so turn 90° anti-clockwise
self.ant_direction= (self.ant_direction-1) %4

# Move ant
move_x, move_y=directions[self.ant_direction]
self.ant_position= (x+move_x, y+move_y)

# Flip colour of square
self.board[x][y] =notself.board[x][y]

ifdisplayandaxes:
# Display the board on the axes
axes.get_xaxis().set_ticks([])
axes.get_yaxis().set_ticks([])
axes.imshow(self.board, cmap="gray", interpolation="nearest")

defdisplay(self, frames: int=100_000) ->None:
"""
 Displays the board without delay in a matplotlib plot
 to visually understand and track the ant.

 >>> _ = LangtonsAnt(WIDTH, HEIGHT)
 """
fig, ax=plt.subplots()
# Assign animation to a variable to prevent it from getting garbage collected
self.animation=FuncAnimation(
fig, partial(self.move_ant, ax, True), frames=frames, interval=1
 )
plt.show()


if__name__=="__main__":
importdoctest

doctest.testmod()

LangtonsAnt(WIDTH, HEIGHT).display()