4

This is a very bare bones program where a population of blocks learn to reach a target on a 2D screen (my attempt at making a smart rockets program).

I've only used Python for 2 weeks, so I'm positive certain aspects of the program have extraneous lines of code and unnecessary junk. What are some ways I can optimize the program/algorithm and make it more readable?

import random
import math
import pygame
import time

WINDOW_SIZE = (600, 600)
COLOR_KEY = {"RED": [255, 0, 0, 10],
             "GREEN": [0, 255, 0, 128],
             "BLUE": [0, 0, 255, 128],
             "WHITE": [255, 255, 255],
             "BLACK": [0, 0, 0]}

target = [500, 500]
pop_size = 20
max_moves = 1500
mutation_rate = 30


class Block:
    def __init__(self):
        self.color = COLOR_KEY["GREEN"]
        self.size = (40, 40)
        self.move_set = []
        self.position = [0, 0]
        self.fitness = -1
        self.target_reached = False

    def evaluate(self):
        dx = target[0] - self.position[0]
        dy = target[1] - self.position[1]
        if dx == 0 and dy == 0:
            self.fitness = 1.0
        else:
            self.fitness = 1 / math.sqrt((dx * dx) + (dy * dy))

    def move(self, frame_count):
        if not self.target_reached:
            self.position[0] += self.move_set[frame_count][0]
            self.position[1] += self.move_set[frame_count][1]

            if self.position[0] == target[0] and self.position[1] == target[1]:
                self.target_reached = True
                self.color = COLOR_KEY["BLUE"]


def create_pop():
    pop = []
    for block in range(pop_size):
        pop.append(Block())

    return pop


def generate_moves(population):
    for block in population:
        for _ in range(max_moves+1):
            rand_x = random.randint(-1, 1)
            rand_y = random.randint(-1, 1)
            block.move_set.append([rand_x, rand_y])

    return population


def fitness(population):
    for block in population:
        block.evaluate()


def selection(population):
    population = sorted(population, key=lambda block: block.fitness, reverse=True)
    best_fit = round(0.1 * len(population))
    population = population[:best_fit]
    return population


def cross_over(population):
    offspring = []

    for _ in range(int(pop_size/2)):
        parents = random.sample(population, 2)
        child1 = Block()
        child2 = Block()

        split = random.randint(0, max_moves)
        child1.move_set = parents[0].move_set[0:split] + parents[1].move_set[split:max_moves]
        child2.move_set = parents[1].move_set[0:split] + parents[0].move_set[split:max_moves]

        offspring.append(child1)
        offspring.append(child2)

    return offspring


def mutation(population):
    chance = random.randint(0, 100)
    num_mutated = random.randint(0, pop_size)

    if chance >= 100 - mutation_rate:
        for _ in range(num_mutated):
            mutated_block = population[random.randint(0, len(population) - 1)]
            for _ in range(50):
                if chance >= 100 - mutation_rate/2:
                    rand_x = random.randint(0, 1)
                    rand_y = random.randint(0, 1)
                else:
                    rand_x = random.randint(-1, 1)
                    rand_y = random.randint(-1, 1)

                mutated_block.move_set[random.randint(0, max_moves - 1)] = [rand_x, rand_y]

    return population


def calc_avg_fit(population):
    avg_sum = sum(block.fitness for block in population)
    return avg_sum/pop_size


def ga(population):
    fitness(population)
    avg_fit = calc_avg_fit(population)
    population = selection(population)
    population = cross_over(population)
    population = mutation(population)

    returning = (avg_fit, population)
    return returning


def main():
    pygame.init()
    window = pygame.display.set_mode(WINDOW_SIZE)
    pygame.display.set_caption("AI Algorithm")
    population = create_pop()
    population = generate_moves(population)

    my_font = pygame.font.SysFont("Arial", 16)
    frame_count = 0
    frame_rate = 0
    t0 = time.process_time()
    gen = 0
    avg_fit = 0

    while True:
        event = pygame.event.poll()
        if event.type == pygame.QUIT:
            break

        frame_count += 1
        if frame_count % max_moves == 0:
            t1 = time.process_time()
            frame_rate = 500 / (t1-t0)
            t0 = t1
            frame_count = 0
            data = ga(population)
            avg_fit = data[0]
            population = data[1]
            gen += 1

        window.fill(COLOR_KEY["BLACK"], (0, 0, WINDOW_SIZE[0], WINDOW_SIZE[1]))

        for block in population:
            block.move(frame_count)

        for block in population:
            window.fill(block.color, (block.position[0], block.position[1], block.size[0], block.size[1]))

        window.fill(COLOR_KEY["RED"], (target[0] + 10, target[1] + 10, 20, 20), 1)

        frame_rate_text = my_font.render("Frame = {0} rate = {1:.2f} fps Generation: {2}"
                                         .format(frame_count, frame_rate, gen), True, COLOR_KEY["WHITE"])

        fitness_text = my_font.render("Average Fitness: {0}".format(avg_fit), True, COLOR_KEY["WHITE"])

        window.blit(fitness_text, (WINDOW_SIZE[0] - 300, 40))
        window.blit(frame_rate_text, (WINDOW_SIZE[0] - 300, 10))

        pygame.display.flip()

    pygame.quit()


main()

1 Answer 1

Reset to default
2

UX

When I run the code, I don't always see the value for "Generation" printed in the GUI. It is sometimes clipped at the right edge of the screen. It would be better to move it to its own line.

Also, "Average Fitness" shows a lot more digits than are needed. It would be better to show fewer digits.

Here are suggested changes:

frame_rate_text = my_font.render("Frame = {0} rate = {1:.2f} fps"
                                 .format(frame_count, frame_rate), True, COLOR_KEY["WHITE"])

gen_text = my_font.render(f"Generation: {gen}", True, COLOR_KEY["WHITE"])

fitness_text = my_font.render(f"Average Fitness: {avg_fit:.9f}", True, COLOR_KEY["WHITE"])

window.blit(gen_text, (WINDOW_SIZE[0] - 300, 70))

This code makes use of f-strings to simplify the formatting code:

f"Generation: {gen}"

Documentation

The PEP 8 style guide recommends adding docstrings for classes and functions.

For example, Block is a generic name for a class. The docstring should describe what you mean by a block:

class Block:
    """
    Rectangle on the 2D screen.
    Represents a member of the population.
    """

    def evaluate(self):
        """Evaluate fitness of population member"""

Naming

PEP 8 recommends all caps for constants.

target = [500, 500]
pop_size = 20
max_moves = 1500
mutation_rate = 30

Would become:

TARGET = [500, 500]
POP_SIZE = 20
MAX_MOVES = 1500
MUTATION_RATE = 30

The word "pop" has many meanings. It would be better to spell out "population". Instead of:

def create_pop():
    pop = []
    for block in range(pop_size):
        pop.append(Block())

    return pop

consider:

def create_population():
    population = []
    for _ in range(POPULATION_SIZE):
        population.append(Block())
    return population

Since the block variable was unused, it was replaced by _.

ga is not a very descriptive name for a function. Again, it would be better to spell it out: genetic_algo.

Simpler

Instead of sqrt:

self.fitness = 1 / math.sqrt((dx * dx) + (dy * dy))

consider hypot:

self.fitness = 1 / math.hypot(dx, dy)

In the ga function, you can replace these 2 lines:

returning = (avg_fit, population)
return returning

with 1 line:

return (avg_fit, population)

This eliminates the generically named intermediate variable: returning

Similarly, in the main function, replace:

population = create_pop()
population = generate_moves(population)

with:

population = generate_moves(create_pop())

Main

It is customary to add a main guard:

if __name__ == "__main__":
    main()

You must log in to answer this question.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.