Genetic algorithm - basic example

ga.py

from __future__ import annotations
from enum import Enum, unique
from typing import Iterable, List, NewType, Text

from numpy.random import Generator, default_rng


@unique
class ItemType(Enum):
    POTATOES = 0
    WHEAT_FLOUR = 1
    RICE = 2
    BEANS = 3
    TOMATOES = 4
    STRAWBERRY_JAM = 5
    PEANUT_BUTTER = 6

    def __repr__(self) -> Text:
        return f'ItemType({self})'

    def __str__(self) -> Text:
        return self.name.title().replace('_', '')

    @classmethod
    def from_index(cls, index: int) -> ItemType:
        if index == ItemType.POTATOES.value:
            return ItemType.POTATOES

        if index == ItemType.WHEAT_FLOUR.value:
            return ItemType.WHEAT_FLOUR

        if index == ItemType.RICE.value:
            return ItemType.RICE

        if index == ItemType.BEANS.value:
            return ItemType.BEANS

        if index == ItemType.TOMATOES.value:
            return ItemType.TOMATOES

        if index == ItemType.STRAWBERRY_JAM.value:
            return ItemType.STRAWBERRY_JAM

        if index == ItemType.PEANUT_BUTTER.value:
            return ItemType.PEANUT_BUTTER

        raise ValueError(f'Invalid Item Type Index: {index}')


class Item:

    ITEM_STATISTICS = {
        ItemType.POTATOES: {
            'weight': 800,
            'calories': 1_502_000
        },
        ItemType.WHEAT_FLOUR: {
            'weight': 400,
            'calories': 1_444_000
        },
        ItemType.RICE: {
            'weight': 300,
            'calories': 1_122_000
        },
        ItemType.BEANS: {
            'weight': 300,
            'calories': 690_000
        },
        ItemType.TOMATOES: {
            'weight': 300,
            'calories': 237_000
        },
        ItemType.STRAWBERRY_JAM: {
            'weight': 50,
            'calories': 130_000
        },
        ItemType.PEANUT_BUTTER: {
            'weight': 20,
            'calories': 117_800
        }
    }

    def __init__(self, type: ItemType, quantity: int = 1) -> None:
        self.type = type
        self.quantity = quantity
        self._weight = Item._get_weight(self.type)
        self._calories = Item._get_calories(self.type)

    def __repr__(self) -> Text:
        return f'{self.type}(Quantity={self.quantity}, Weight={self.weight:,}, Calories={self.calories:,})'

    @property
    def weight(self) -> int:
        return self._weight * self.quantity

    @weight.setter
    def weight(self, _weight: int) -> None:
        self._weight = _weight

    @property
    def calories(self) -> int:
        return self._calories * self.quantity

    @calories.setter
    def calories(self, _calories: int) -> None:
        self._calories = _calories

    @classmethod
    def _get_weight(cls, type: ItemType) -> int:
        return cls.ITEM_STATISTICS[type]['weight']

    @classmethod
    def _get_calories(cls, type: ItemType) -> int:
        return cls.ITEM_STATISTICS[type]['calories']


Chromosome = NewType('Chromosome', Text)


def generate_chromosome(length: int = 7, rng: Generator = default_rng()) -> Chromosome:
    return ''.join(rng.choice(['0', '1'], size=length).tolist()) # type: ignore


def get_items_from_chromosome(chromosome: Chromosome) -> List[Item]:
    # Initialize the item list.
    items: List[Item] = []

    for index, character in enumerate(chromosome):
        # Use the current index as the determinant of which item type we are dealing with
        # at the moment.
        item_type = ItemType.from_index(index)

        # Use the item type we got above to create the item of the given type, although
        # we only do this if this chromosome actually calls for this item by indicating a
        # 1 on this index.
        item = Item(item_type) if character == '1' else None

        # Once the item has been instantiated, add it to the list of items. If the item
        # was not actually instantiated, there is no need to add it to the list.
        if item:
            items.append(item)

    # Return the final item list once we have finished putting it together.
    return items


def calculate_fitness(items: List[Item], limit: int = 1000) -> int:
    if sum(item.weight for item in items) > limit:
        return 0

    return sum(item.calories for item in items)


def generate_chromosomes(number: int = 100) -> Iterable[Chromosome]:
    for _ in range(number):
        yield generate_chromosome()


def find_optimum(population_size: int) -> List[Item]:
    optimal_fitness = 0
    optimal_item_list = []

    for chromosome in generate_chromosomes(population_size):
        item_list = get_items_from_chromosome(chromosome)
        fitness = calculate_fitness(item_list)

        if fitness > optimal_fitness:
            optimal_fitness = fitness
            optimal_item_list = item_list

    return optimal_item_list