Pygame Astar Algorithm

Pygame Astar Algorithm

from pygame import Vector2

# Make pygame.Vector2 hashable
class Vector(Vector2):
    def __hash__(self):
        return int(self.x + self.y)

class StepPathing:
    def __init__(self, start, goal, map_level, map_size, pattern, block_value):
        start = start
        self.goal = goal
        self.break_out = 0
        self.came_from = {}
        self.pattern = pattern
        self.open_set = [start]
        self.gscore = {start: 0}
        self.fscore = {start: start.distance_to(self.goal)}
        self.map_size = map_size
        self.map_level = map_level
        self.block_value = block_value
        self.current = start

    def get_low_score(self):
        low_score = -1
        gvector = None
        for vector in self.open_set:
            if low_score == -1:
                gvector = vector
                low_score = self.fscore[vector]
            elif self.fscore[vector] < low_score:
                gvector = vector
                low_score = self.fscore[vector]

        return gvector

    def get_path(self):
        path = None
        while path is None:
            path = self.step()

        return path

    def get_neighbors(self):
        straight = list(map(Vector, ((0, 1), (0, -1), (1, 0), (-1, 0))))
        neighbors = []

        if self.current is None:
            return []

        for n in straight:
            new_neighbor = self.current + n
            # Is neighbor with in grid
            if self.map_size.x > int(new_neighbor.x) >= 0 and \
               self.map_size.y > int(new_neighbor.y) >= 0:

                if self.map_level[int(new_neighbor.x)][int(new_neighbor.y)] \
                   < self.block_value:
                    if self.map_size.x > new_neighbor.x >= 0 and self.map_size.y \
                       > new_neighbor.y >= 0:
                        neighbors.append(Vector(new_neighbor))

        # moves on angles
        if self.pattern == '*':
            angles = list(map(Vector, ((1, 1), (-1, -1), (1, -1), (-1, 1))))
            for n in angles:
                new_neighbor = self.current + n
                x = int(new_neighbor.x)
                y = int(new_neighbor.y)
                if self.map_size.x > x >= 0 and self.map_size.y > y >= 0:
                    wall = self.map_level[x][y] > self.block_value - 1
                    block = (self.map_level[x][int(self.current.y)] >
                             self.block_value - 1 or
                             self.map_level[int(self.current.x)][y] >
                             self.block_value - 1)

                    # Don't allow angle movement next to wall.
                    if not wall and not block:
                        neighbors.append(Vector(new_neighbor))

        return neighbors

    def step(self):
        if len(self.open_set) > 0:
            self.current = self.get_low_score()
            if self.current == self.goal:
                return self.reconstructed_path()

            self.open_set.remove(self.current)
            neighbors = self.get_neighbors()

            for neighbor in neighbors:
                gscore = self.gscore[self.current] + \
                         neighbor.distance_to(self.current)

                if neighbor not in self.gscore.keys():
                    self.gscore[neighbor] = gscore
                    self.fscore[neighbor] = gscore + \
                        neighbor.distance_to(self.goal)
                    self.came_from[neighbor] = self.current
                    self.open_set.append(neighbor)
                elif gscore < self.gscore[neighbor]:
                    self.gscore[neighbor] = gscore
                    self.fscore[neighbor] = gscore + \
                        neighbor.distance_to(self.goal)
                    self.came_from[neighbor] = self.current
                    self.open_set.append(neighbor)
        else:
            return []

    def reconstructed_path(self):
        current = self.current
        path = [current]
        while current in self.came_from.keys():
            current = self.came_from[current]
            path.append(current)

        return path

Example How Algorithm Works.

Left mouse button create walls. Right mouse remove walls. Space key show algorithm steps. M key will quickly show path. S key will set left mouse button to place start position. E Key will set left mouse button to place goal position. C key will stop algorithm and clear data.

import pygame
import random
from statemachine import State
from astar import StepPathing, Vector

class DeltaTimer:
    def __init__(self, delay):
        self.delta = 0
        self.delay = delay
        self.stop = False

    def add(self, delta):
        if not self.stop:
            self.delta += delta

    def is_tripped(self):
        if self.delta > self.delay:
            self.delta = 0
            return True

        return False

class Example(State):
    def __init__(self):
        self.block_size = 40
        self.colors = [
            pygame.Color("lawngreen"),
            pygame.Color("firebrick"),
            pygame.Color("dodgerblue"),
            pygame.Color("purple"),
            pygame.Color("Gray20")]

        self.fill_color = 4
        self.fill = self.fill_color
        self.pattern = "*"
        self.rect = pygame.Rect(0, 0, self.block_size -1, self.block_size -1)
        self.grid_reset()
        self.grid_start_positions()
        self.grid_update()
        self.create_background()
        self.grid_timer = DeltaTimer(100)
        self.grid_timer.stop = True

    def create_background(self):
        self.background = pygame.Surface(State.machine.rect.size)
        self.background.fill(pygame.Color('black'))
        color = pygame.Color("gray80")
        # Grid lines
        for x in range(1, self.grid_size[0]):
            s_pos = x * self.block_size, 0
            e_pos = x * self.block_size, State.machine.rect.bottom
            pygame.draw.line(self.background, color, s_pos, e_pos)

        for y in range(1, self.grid_size[1]):
            s_pos = 0, y * self.block_size
            e_pos = State.machine.rect.right, y * self.block_size
            pygame.draw.line(self.background, color, s_pos, e_pos)

    def grid_find_values(self, pos):
        size = State.machine.rect.size
        grid_values = [[Vector(pos).distance_to(Vector(x, y)) for y in
            range(self.grid_size[1])] for x in range(self.grid_size[0])]

        grid_image = pygame.Surface(size, pygame.SRCALPHA)
        grid_image.fill((0, 0, 0, 0))
        font = pygame.font.Font(None, 12)
        rect = pygame.Rect(0, 0, self.block_size, self.block_size)
        for x, row in enumerate(grid_values):
            for y, value in enumerate(row):
                rect.topleft = x * self.block_size, y * self.block_size
                text = font.render("{:.1f}".format(value), 1,
                                   pygame.Color("white"))
                t_rect = text.get_rect(center=rect.center)
                grid_image.blit(text, t_rect)

        return grid_image

    def grid_reset(self):
        size = State.machine.rect.size
        self.grid_size = size[0] // self.block_size, size[1] // self.block_size
        self.grid = [[-1 for y in range(self.grid_size[1])]
                     for x in range(self.grid_size[0])]

    def grid_start_positions(self):
        self.start_pos = 0, 0
        self.end_pos = self.grid_size[0] - 1, self.grid_size[1] - 1

    def grid_clear_trace(self):
        for x in range(self.grid_size[0]):
            for y in range(self.grid_size[1]):
                if self.grid[x][y] == 3 or self.grid[x][y] == 2:
                    self.grid[x][y] = -1

    def grid_change_trace(self):
        for x in range(self.grid_size[0]):
            for y in range(self.grid_size[1]):
                if self.grid[x][y] == 2:
                    self.grid[x][y] = 3

    def grid_display_path(self):
        path = self.path_step.get_path()[1:-1]
        if path != []:
            for vector in path:
                x, y = int(vector.x), int(vector.y)
                self.grid[x][y] = 2

    def grid_update(self):
        self.grid_clear_trace()
        self.path_step = StepPathing(Vector(self.start_pos), Vector(self.end_pos),
            self.grid, Vector(self.grid_size), self.pattern, 4)

        self.grid[self.start_pos[0]][self.start_pos[1]] = 0
        self.grid[self.end_pos[0]][self.end_pos[1]] = 1

    def grid_step(self):
        result = self.path_step.step()
        if result is None:
            self.grid_change_trace()
            vector = self.path_step.current
            x, y = int(vector.x), int(vector.y)
            if self.grid[x][y] != 0:
                self.grid[x][y] = 2
        else:
            self.grid_change_trace()
            self.grid_timer.stop = True
            if result != []:
                path = self.path_step.reconstructed_path()[1:-1]
                for vector in path:
                    x, y = int(vector.x), int(vector.y)
                    self.grid[x][y] = 2

    def on_draw(self, surface):
        surface.blit(self.background, (0, 0))
        for x in range(self.grid_size[0]):
            for y in range(self.grid_size[1]):
                if self.grid[x][y] >= 0:
                    self.rect.topleft = (x * self.block_size + 1,
                                         y * self.block_size + 1)
                    surface.fill(self.colors[self.grid[x][y]], self.rect)

        for x in range(self.grid_size[0]):
            pygame.draw

    def on_event(self, event):
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                self.grid_update()
                self.grid_timer.stop = False
            elif event.key == pygame.K_c:
                if self.grid_timer.stop:
                    self.grid_clear_trace()
            elif event.key == pygame.K_s:
                self.fill = 0
            elif event.key == pygame.K_e:
                self.fill = 1
            elif event.key == pygame.K_m:
                self.grid_timer.stop = True
                self.grid_update()
                self.grid_display_path()
            elif event.key == pygame.K_ESCAPE:
                State.machine.running = False
        elif event.type == pygame.QUIT:
            State.machine.running = False

    def on_update(self, delta):
        mouse = pygame.mouse.get_pressed()
        self.grid_timer.add(delta)
        if self.grid_timer.is_tripped():
            self.grid_step()

        if mouse[0]:
            x, y = pygame.mouse.get_pos()
            x //= self.block_size
            y //= self.block_size
            if 0 <= x < self.grid_size[0]:
                if 0 <= y < self.grid_size[0]:
                    if self.grid[x][y] == -1:
                        self.grid[x][y] = self.fill

                    if self.fill == 0:
                        self.clear_pos(self.start_pos)
                        self.start_pos = x, y
                    elif self.fill == 1:
                        self.clear_pos(self.end_pos)
                        self.end_pos = x, y
        elif mouse[2]:
            x, y = pygame.mouse.get_pos()
            x //= self.block_size
            y //= self.block_size
            if 0 <= x < self.grid_size[0]:
                if 0 <= y < self.grid_size[0]:
                    if self.grid[x][y] == self.fill_color:
                        self.grid[x][y] = -1

    def clear_pos(self, pos):
        x, y = pos
        self.grid[x][y] = -1
        self.fill = self.fill_color

def main():
    pygame.init()
    # caption, width, height, center
    State.machine_setup("Trace Example", 800, 600, True)
    State.machine.flip(Example())
    State.machine.mainloop()

if __name__ == '__main__':
    main()