Screeps map processor utility and max-flow-based wall planner

flow_walls.py

## We want to compute the minimum number of squares to make impassable (by
## adding walls/ramparts) such that all of a certain set of locations (the
## "protected" locations) are unreachable from all room exits. We model this as
## a max-flow/min-cut problem on a graph.
##
## Each passable square (swamp, plain, or exit) becomes two vertices, a "top"
## and a "bottom", with an edge of capacity 1 from the top to the bottom of each
## square. Additionally, there is an edge of infinite capacity from the bottom
## of each square to the top of each adjacent square.
##
## There are edges of infinite capacity from the source to each protected
## square, and from each exit square to the sink.
##
## Then the only edges with finite capacity are the ones from the top to the
## bottom of a node, so those are the only edges that can be in a min-cut of the
## graph. Making a square impassable corresponds to cutting one of these edges,
## so the edges in a min-cut exactly correspond to a minimum-size set of walls
## needed to protect the area.
##
## One last thing in the setup: certain squares can't be blocked (the exits and
## the squares next to them); we represent this by making the top-to-bottom
## edges for those squares have infinite capacity, so they can't be in any
## min-cuts.

import sys

import Image
import numpy as np
import maxflow
import screeps_map
np.set_printoptions(threshold=np.nan, linewidth=np.nan)

room, x0, y0, x1, y1 = sys.argv[1:6]
x0 = int(x0)
y0 = int(y0)
x1 = int(x1)
y1 = int(y1)

types = screeps_map.get_room_types(*screeps_map.room_coords(room))
grid = (types == screeps_map.Terrain.plain) | (types == screeps_map.Terrain.swamp) | (types == screeps_map.Terrain.port)
H, W = grid.shape

graph = maxflow.Graph[int](2 * H * W, 0)
top_nodes = graph.add_grid_nodes((H, W))
bottom_nodes = graph.add_grid_nodes((H, W))

## the "infinite" capacity value
M = 10 ** 4

## set up edges between normal nodes
for a in xrange(H):
    for b in xrange(W):
        ## ignore this square if it's not passable
        if not grid[a, b]: continue

        ## add top-to-bottom edges
        graph.add_edge(top_nodes[a, b], bottom_nodes[a, b], 1, 0)

        ## add bottom-to-top edges to the neighbors
        for da, db in [[0, -1], [-1, -1], [-1, 0], [-1, 1], [0, 1], [1, 1], [1, 0], [1, -1]]:
            a2 = a + da
            b2 = b + db

            if 0 <= a2 < H and 0 <= b2 < W and grid[a2, b2]:
                graph.add_edge(bottom_nodes[a, b], top_nodes[a2, b2], M, 0)


## mark a point as being unavailable for removal
def disable(a, b):
    graph.add_edge(top_nodes[a, b], bottom_nodes[a, b], M, 0)

## mark a point as a potential entrance
def add_sink(a, b):
    if grid[a, b]:
        graph.add_tedge(bottom_nodes[a, b], 0, M)
        disable(a, b)

## mark a point as needing to be protected
def add_source(a, b):
    if grid[a, b]:
        graph.add_tedge(top_nodes[a, b], M, 0)
        #disable(a, b)

## mark all the protected points
for b in xrange(x0, x1 + 1):
    for a in xrange(y0, y1 + 1):
        add_source(a, b)

## mark dangerous points at all exits, and block off those points and the
## adjacent locations where walls can't be built
for a in xrange(H):
    add_sink(a, 0)
    add_sink(a, -1)
    if (a > 0 and grid[a-1, 0]) or (grid[a, 0]) or (a < H - 1 and grid[a+1, 0]): disable(a, 1)
    if (a > 0 and grid[a-1, -1]) or (grid[a, -1]) or (a < H - 1 and grid[a+1, -1]): disable(a, -2)
for b in xrange(W):
    add_sink(0, b)
    add_sink(-1, b)
    if (b > 0 and grid[0, b-1]) or (grid[0, b]) or (b < W - 1 and grid[0, b+1]): disable(1, b)
    if (b > 0 and grid[-1, b-1]) or (grid[-1, b]) or (b < W - 1 and grid[-1, b+1]): disable(-2, b)


## finally, do the flow
flow = graph.maxflow()
print 'need', flow, 'walls'

## the squares in the min-cut are the ones whose top nodes are cut away from
## their bottom nodes
top_seg = graph.get_grid_segments(top_nodes)
bottom_seg = graph.get_grid_segments(bottom_nodes)
walls = grid & (top_seg != bottom_seg)

# def pr(a): print a.astype(np.int8)
# pr(grid)
# print
# pr(walls)
# print
# pr(top_seg)
# print
# pr(bottom_seg)

img = np.zeros((H, W), dtype=np.uint8)
img[grid] = 43
img[walls] = 255

Image.fromarray(img).save('walls.png')
print 'saved to walls.png'

screeps_map.py

import Image
import numpy as np

color_map = {
    'wall': np.array([0, 0, 0], dtype=np.uint8),
    'plain': np.array([43, 43, 43], dtype=np.uint8),
    'swamp': np.array([35, 37, 19], dtype=np.uint8),
    'source': np.array([255, 242, 70], dtype=np.uint8),
    'controller': np.array([80, 80, 80], dtype=np.uint8),
    'port': np.array([50, 50, 50], dtype=np.uint8),
    'keeper': np.array([100, 0, 0], dtype=np.uint8),
}

color_items = color_map.items()
color_names = [i[0] for i in color_items]
color_vals = [i[1] for i in color_items]
color_packed = np.hstack(color_vals).view(np.dtype((np.void, 3)))
color_inds = np.argsort(color_packed)
color_names = [color_names[i] for i in color_inds]
color_vals = [color_vals[i] for i in color_inds]

class C(object): pass
Terrain = C()
# for n, i in color_items:
#     setattr(Terrain, n, i)
for i, n in enumerate(color_names):
    setattr(Terrain, n, i)

try:
    img = np.array(Image.open('map.png'))
except IOError:
    print '\nmap.png missing; please download http://static.screeps.com/map.png into this directory\n'
    exit()

img_types = np.searchsorted(color_packed, img.view(np.dtype((np.void, 3)))[:,:,0], sorter=color_inds)

## returns an array representing the given room; each element is one of
## Terrain.wall, Terrain.plain, etc.
def get_room_types(x, y):
    return img_types[1550 + 50 * y : 1600 + 50 * y,
                     1550 + 50 * x : 1600 + 50 * x].copy()

def get_room_image(x, y):
    return img[1550 + 50 * y : 1600 + 50 * y,
               1550 + 50 * x : 1600 + 50 * x].copy()

__all__ = ['get_room', 'get_room_image', 'Terrain']