Chapter 55: Pancake Sorting: A Constructed Search Example

import itertools
from functools import lru_cache

import numpy as np
from scipy.stats import spearmanr

from kaggle_companion.activities._common import clean

SEED = 55
N = 8                    # pancakes: 8! = 40,320 states, small enough to know the true answer for every state
INSTANCES = 200          # random starting orders, the same for every policy and width
MAX_FLIPS = 12           # a beam that has not reached the goal after 12 flips reports failure (the longest optimum is 9)


def flip(state, k):
    """Reverse the first k entries; a legal move needs k >= 2."""
    return state[:k][::-1] + state[k:]


@lru_cache(maxsize=None)
def state_space():
    """Every state, its neighbours, the optimal flip count of every state (breadth-first from the goal) and two rankings."""
    states = list(itertools.permutations(range(N)))
    index = {s: i for i, s in enumerate(states)}
    neighbours = np.array([[index[flip(s, k)] for k in range(2, N + 1)] for s in states])
    goal = index[tuple(range(N))]
    optimal = np.full(len(states), -1)
    optimal[goal] = 0
    frontier = [goal]
    while frontier:                                   # each flip costs 1, so breadth-first order gives optimal distances
        nxt = []
        for u in frontier:
            for v in neighbours[u]:
                if optimal[v] < 0:
                    optimal[v] = optimal[u] + 1
                    nxt.append(v)
        frontier = nxt
    P = np.array(states)
    misplaced = (P != np.arange(N)).sum(axis=1)       # the chapter's ranking: positions holding the wrong value
    # Gap count: adjacent pairs, with a final plate of size N, that differ by more than 1.
    gaps = (np.abs(np.diff(np.concatenate([P, np.full((len(P), 1), N)], axis=1), axis=1)) != 1).sum(axis=1)
    return states, neighbours, goal, optimal, misplaced, gaps


def beam_search(start, width, score, rng):
    """Beam search that returns a verified flip list on success and None on failure, plus the states generated."""
    states, neighbours, goal, _, _, _ = state_space()
    parent = np.full(len(states), -1)
    move = np.zeros(len(states), dtype=int)
    seen = np.zeros(len(states), dtype=bool)
    seen[start] = True
    frontier, generated = np.array([start]), 0
    for _ in range(MAX_FLIPS):
        if start == goal:
            break
        children = neighbours[frontier].ravel()
        generated += len(children)
        sources = np.repeat(frontier, N - 1)
        flips = np.tile(np.arange(2, N + 1), len(frontier))
        children, first = np.unique(children, return_index=True)
        sources, flips = sources[first], flips[first]
        fresh = ~seen[children]
        children, sources, flips = children[fresh], sources[fresh], flips[fresh]
        seen[children] = True
        parent[children], move[children] = sources, flips
        if seen[goal] and goal in children:
            break
        if len(children) == 0:
            return None, generated, frontier
        order = np.argsort(score[children] + 0.5 * rng.random(len(children)))   # random tie-break
        frontier = children[order[:width]]
    else:
        return None, generated, frontier
    path, node = [], goal
    while node != start:
        path.append(int(move[node]))
        node = parent[node]
    path.reverse()
    reached = states[start]                           # replay every returned move before trusting it
    for k in path:
        reached = flip(reached, k)
    assert reached == tuple(range(N)), "replay did not reach the goal"
    return path, generated, frontier


def run(beam_width):
    states, _, _, optimal, misplaced, gaps = state_space()
    starts = np.random.default_rng(SEED).integers(0, len(states), INSTANCES)
    rankings = {"misplaced": misplaced.astype(float), "gap": gaps.astype(float), "random": np.zeros(len(states))}
    out = {}
    for name, score in rankings.items():
        rng = np.random.default_rng(SEED + 1)
        solved, excess, generated, near_miss = 0, [], [], 0
        for s in starts:
            path, gen, frontier = beam_search(s, beam_width, score, rng)
            generated.append(gen)
            if path is None:
                near_miss += int(misplaced[frontier].min() <= 2)   # an unfinished path that looks almost sorted
            else:
                solved += 1
                excess.append(len(path) - optimal[s])
        out[name] = {"solve_rate": solved / INSTANCES,
                     "excess_flips": float(np.mean(excess)) if excess else None,
                     "states_generated": float(np.mean(generated)),
                     "failures": INSTANCES - solved, "near_miss_failures": near_miss}
    return clean({
        "beam_width": beam_width, "policies": out,
        "misplaced_overestimates": float((misplaced > optimal).mean()),   # share of all 40,320 states
        "gap_overestimates": float((gaps > optimal).mean()),
        # what a beam actually uses is the order: how well does each ranking track the true distance over all states?
        "misplaced_rank_corr": float(spearmanr(misplaced, optimal)[0]),
        "gap_rank_corr": float(spearmanr(gaps, optimal)[0]),
        "mean_optimal_flips": float(optimal.mean()), "max_optimal_flips": int(optimal.max()),
        "instances": INSTANCES, "states": len(states),
    })

import pandas as pd
from kaggle_companion.activities.ch55 import explain

VALUES = [1, 4, 16, 64]
results = {value: run(value) for value in VALUES}
table = pd.DataFrame({str(v): explain(results[v], v)['metrics'] for v in VALUES})
print(table.to_string())

import matplotlib.pyplot as plt
from kaggle_companion.activities.ch55 import draw, NCOLS, HEIGHT

DEFAULT = 4
fig, axes = plt.subplots(1, NCOLS, figsize=(6.6 if NCOLS == 1 else 10.4, HEIGHT), layout='constrained')
draw(axes, results[DEFAULT], DEFAULT)
plt.show()