I wrote a little game of snake that where you can see the field in which the snake moves fixed and you can also see the "viewpoint" of the snake, which is basically calculating the positions of all the things in the board taking the snake's head as the center.
Right now I'm doing it through two for loops in the following way:
for x in range(ROWS):
n_x = (x - snake_head[0] - center) % ROWS
for y in range(ROWS):
n_y = (y - snake_head[1] - center) % ROWS
snake_view.grid[n_x, n_y] = field.grid[x,y]
ROWS is the number of rows the field's grid has (currently 15) and it has to be an odd number so there's a center. snake_head is a numpy array of shape (2,) which has the coordinates of head in the field grid. field.grid is a numpy array of shape (ROWS, ROWS) that holds the value of what is in each square of the grid, -1 for snake body, -2 for snake head, 0 for nothing or empty, 1 for a fruit or snack.
snake_view.grid is a numpy array of the same dimensions as field.grid but has the viewpoint as the snake sees it.
This works fine, as you can see in the code that I have pasted below (it is in two different .py files but I have pasted it all into one so it can be copy/pasted easily). But since I want to use this game to train an AI, I need help optimizing it.
Is there a way to vectorize this?
import pygame
import time
import random
import numpy as np
WIN_WIDTH, WIN_HEIGHT = 800, 720
GRID_WIDTH = 400
ROWS = 15
GRID_SPACE = GRID_WIDTH // ROWS
class Snake:
"""
Snake Object
"""
def __init__(self, x, y, body_color = (255,0,0), head_color = (255,0,255)):
"""
Initializes the Snake Object on a given coordinate
"""
# Position is a list of tuples where each tuple
# is a coordinate of the position of a certain
# block in the field
self.positions = np.array([np.array([x, y])])
self.direction = np.array([0, 1])
# dir is an array of [x,y] values for the direction
# of the snake's movement
self.vel = 1
self.body_color = body_color
self.head_color = head_color
self.eaten = False
def head(self):
# Returns the position of the head
return self.positions[-1]
def move(self):
# Funcion that moves the snake in a particular direction
new_head = np.array([self.positions[-1] + self.direction*self.vel])
if not self.eaten:
self.positions[:-1] = self.positions[1:]
self.positions[-1] = new_head
else:
self.eaten = False
self.positions = np.concatenate((self.positions, new_head))
self.positions = np.mod(self.positions, ROWS)
def bite(self):
return np.unique(self.positions, axis=0).shape != self.positions.shape
def eat(self):
self.eaten = True
class Fruit:
"""
Fruits the snake can eat
"""
def __init__(self, grid):
x, y = random.sample(grid, 1)[0]
grid.remove((x,y))
self.position = np.array([x, y])
class Field():
COLORS = {
-1: (255, 0, 0),
-2: (255, 0, 255),
0: (50,50,50),
1: (0,255,0)
}
available_grid = set(
(i,j)
for i in range(ROWS)
for j in range(ROWS)
)
available_grid.remove((5,5))
def __init__(self):
self.grid = np.zeros((ROWS, ROWS), dtype=np.int8)
def empty_grid(self):
self.grid = np.zeros((ROWS, ROWS), dtype=np.int8)
def update_grid(self, snake=None, fruits=None):
self.empty_grid()
self.available_grid = set(
(i,j)
for i in range(ROWS)
for j in range(ROWS)
)
for pos in snake.positions[:-1]:
self.grid[pos[0], pos[1]] = -1
if (pos[0], pos[1]) in self.available_grid:
self.available_grid.remove((pos[0], pos[1]))
pos = snake.head()
self.grid[pos[0], pos[1]] = -2
if (pos[0], pos[1]) in self.available_grid:
self.available_grid.remove((pos[0], pos[1]))
for fruit in fruits:
self.grid[fruit.position[0], fruit.position[1]] = 1
if (fruit.position[0], fruit.position[1]) in self.available_grid:
self.available_grid.remove((fruit.position[0], fruit.position[1]))
def draw(self, win, offset_x = 0, offset_y = 0):
x, y = 0, 0
for row in range(ROWS):
pos_x = x * GRID_SPACE + offset_x
for col in range(ROWS):
pos_y = y * GRID_SPACE + offset_y
pygame.draw.rect(win, self.COLORS[self.grid[x,y]], (
pos_x + 2,
pos_y + 2,
GRID_SPACE - 2,
GRID_SPACE -2
))
y+=1
x+=1
y=0
pygame.font.init()
STAT_FONT = pygame.font.SysFont("arial", 50)
SCORE = 0
try:
with open("HIGH_SCORE.txt", "r") as f:
HIGH_SCORE = f.read()
if len(HIGH_SCORE) > 0:
HIGH_SCORE = int(HIGH_SCORE)
else:
HIGH_SCORE = 0
except FileNotFoundError:
HIGH_SCORE = 0
def record_hs():
if SCORE > HIGH_SCORE:
with open("HIGH_SCORE.txt", "w") as f:
f.write(str(SCORE))
def draw_window(win, field, snake_view=None):
win.fill((0,0,0))
pygame.draw.line(win, (150, 150, 150), (GRID_WIDTH, 0), (GRID_WIDTH, GRID_WIDTH))
field.draw(win)
snake_view.draw(win, offset_x=GRID_WIDTH)
text = STAT_FONT.render(f"HIGH SCORE: {HIGH_SCORE}", 1,(255,255,255))
win.blit(text, (10,10+GRID_WIDTH))
text = STAT_FONT.render(f"Score: {SCORE}", 1,(255,255,255))
win.blit(text, (10,10+GRID_WIDTH+text.get_height()))
pygame.display.update()
def main():
# Initializes the screen
center = ROWS // 2 + 1 if ROWS % 2 == 1 else ROWS // 2
n_fruits = 5
global HIGH_SCORE, SCORE
win = pygame.display.set_mode((WIN_WIDTH, WIN_HEIGHT))
field = Field()
snake = Snake(5,5)
snake_view = Field()
fruits = [Fruit(field.available_grid) for i in range(n_fruits)]
clock = pygame.time.Clock()
run = True
while run:
clock.tick(8)
for event in pygame.event.get():
if event.type == pygame.QUIT:
run = False
pygame.quit()
quit()
# Get the direction the snake should turn from keyboard
keys = pygame.key.get_pressed()
if keys[pygame.K_DOWN]:
if snake.direction[1] != -1:
snake.direction[1] = 1
snake.direction[0] = 0
elif keys[pygame.K_LEFT]:
if snake.direction[0] != 1:
snake.direction[0] = -1
snake.direction[1] = 0
elif keys[pygame.K_UP]:
if snake.direction[1] != 1:
snake.direction[1] = -1
snake.direction[0] = 0
elif keys[pygame.K_RIGHT]:
if snake.direction[0] != -1:
snake.direction[0] = 1
snake.direction[1] = 0
snake.move()
if snake.bite():
record_hs()
run = False
snake_head = snake.head()
rem = None
a = np.abs(snake_head- fruits[0].position)
d = np.sqrt(np.square(a[0]) + np.square(a[1]))
for fruit in fruits:
if abs(snake_head[0] - fruit.position[0]) < 1 and abs(snake_head[1] - fruit.position[1]) < 1:
rem = fruit
if rem:
fruits.remove(rem)
fruits.append(Fruit(field.available_grid))
snake.eat()
SCORE += 1
field.update_grid(snake, fruits)
for x in range(ROWS):
n_x = (x - snake_head[0] - center) %ROWS
for y in range(ROWS):
n_y = (y - snake_head[1] - center) % ROWS
snake_view.grid[n_x, n_y] = field.grid[x,y]
draw_window(win, field, snake_view)
if __name__ == "__main__":
main()
