Evaluate positions for sprinklers in rim factory

sprinkler-test.py

from typing import *
import numpy as np
import itertools
import functools
import math
import statistics
import matplotlib
import matplotlib.pyplot as plt
import random
from enum import IntEnum, auto


RADIUS = 12
MAP_SIZE = 24

def cartesian(start: int, end: int) -> Iterator[Tuple[int,int]]:
    """
    Turn a single range iterator as in range(start,end+1) into a cartesian product
    """
    return itertools.product(range(start,end+1),repeat=2)

def map_tiles():
    r = RADIUS+1
    z = itertools.product(range(-MAP_SIZE,MAP_SIZE), repeat=2)
    z = list(z)
    random.shuffle(z)
    for x in z:
        yield x

def sprinkler(position:Tuple[int,int]) -> Iterator[Tuple[int,int]]:
    for i, j in cartesian(0, RADIUS+1):
        if i==0 and j == 0:
            continue
        s = i**2+j**2
        if math.sqrt(s) <= RADIUS:
            yield (i+position[0], j+position[1])
            yield (position[0]+i,position[1]-j)
            yield (position[0]-i,j+position[1])
            yield (position[0]-i,position[1]-j)



field = {}
class Fitness(IntEnum):
    MEAN = auto()
    PENALTY = auto()
    #: Count the cells that are covered
    COVERAGE = auto()
    MEDIAN = auto()
    #: Attempt to minimize the overall impact of the sprinklers
    MINIMIZE = auto()
    TOTAL = auto()
    SPECIAL_REGION = auto()

fitness_metric = Fitness.SPECIAL_REGION
special_region = set(cartesian(-5,5))

def evaluate_fitness(field: Dict[Tuple[int,int], int]) -> float:
    # return functools.reduce(lambda x,y:x*y,field.values())**(1.0/len(field))
    r=0
    if fitness_metric == Fitness.MEAN:
        r = sum(field.values())/(MAP_SIZE*2)**2
    elif fitness_metric == Fitness.PENALTY:
        average = sum(field.values())/len(field)
        r = map(lambda x: -x if x < average*0.8 else x, field.values())
        r = sum(r)
    elif fitness_metric == Fitness.COVERAGE:
        r = len(list(filter(lambda x:x>0, field.values())))
    elif fitness_metric == Fitness.MEDIAN:
        r = statistics.median(field.values())
    elif fitness_metric == Fitness.MINIMIZE:
        r = MAP_SIZE**2 - len(list(filter(lambda x:x>0, field.values())))
    elif fitness_metric == Fitness.TOTAL:
        r = sum(field.values())
    elif fitness_metric == Fitness.SPECIAL_REGION:
        r = sum(map(lambda x:x[1], filter(lambda x: x[0] in special_region,field.items())))
        penalty_squares = len([(x,y) for x,y in special_region
                              if (x,y) not in field.keys()
                              if (x,y) not in excluded_positions])
        
        r/= penalty_squares+1
    return r
    return sum(field.values())/len(field)

def simulate(positions: Iterable[Tuple[int,int]], rounds: int)\
             -> Dict[Tuple[int,int],int]:
    field = {(x,y):0 for x,y in cartesian(-MAP_SIZE, MAP_SIZE)
             if (x,y) not in positions}
    sprinklers = list(map(lambda x: itertools.cycle(sprinkler(x)), positions))
    for _ in range(rounds):
        for s in sprinklers:
            pos = next(s)
            if pos in field:
                field[pos] += 1
    return field

excluded_positions = frozenset([(0,0)])

def evaluate_alternatives(positions: Iterable[Tuple[int,int]],
                          to_try: int,
                          rounds: int) -> List[Tuple[int,int]]:
    # Eliminate already extant positions
    options = set(map_tiles()) - set(positions) - excluded_positions
    options = list(options)
    random.shuffle(options)
    to_try = min(to_try, len(options)-1)
    options = options[0:to_try]
    best_fitness = -2**48
    best_found = None
    for i in options:
        f = simulate(positions + [i], rounds)
        fitness = evaluate_fitness(f)
        if fitness > best_fitness:
            best_found = i
            best_fitness = fitness
    print(f"best fitness found this round {best_fitness}")
    return positions + [best_found]
    
def optimize(n: int, candidates:int, rounds: int):
    positions = []
    while len(positions) < n:
        positions = evaluate_alternatives(positions, candidates, rounds)
    fields = simulate(positions, rounds)
    heatmap = np.array([[(fields[x,y] if (x,y) in fields and (x,y) not in excluded_positions else -5)
                             for x in range(-MAP_SIZE, MAP_SIZE)]
                        for y in range(-MAP_SIZE,MAP_SIZE)])
    fig, ax = plt.subplots()
    im = ax.imshow(heatmap)
    for x,y in positions:
        ax.text(x+MAP_SIZE,y+MAP_SIZE, "S" + ('e' if (x,y) in special_region else ''), ha='center', va='center', color='w')
    for x,y in excluded_positions:
        ax.text(x+MAP_SIZE,y+MAP_SIZE, "E", ha='center', va='center', color='w')
##    for (x,y), v in fields.items():
##        ax.text(x+MAP_SIZE,y+MAP_SIZE, math.floor(fields[x,y]), ha='center', va='center', color='w')
    fig.tight_layout()
    plt.show()

optimize(10, 80, 1500)