1
\$\begingroup\$

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()
\$\endgroup\$

1 Answer 1

0
\$\begingroup\$

Found a better solution using numpy's roll function:

shift = center - snake_head
snake_view.grid = np.roll(field.grid, shift=(shift[0], shift[1]), axis=(0,1))
\$\endgroup\$
1
  • 2
    \$\begingroup\$ Even for self-answers, we're looking for more than just code-only answers - among other reasons, for the benefit of other readers on the internet. Can you add some context? \$\endgroup\$
    – Reinderien
    Mar 1 at 20:25

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.