Information gain — which question to ask first

QUESTIONS = {
    "star":   "Is there a star the audience turns up for?",
    "summer": "Did it open in summer?",
    "sequel": "Is it a sequel?",
    "wide":   "Did it open on 3,000 screens or more?",
    "budget": "Did it cost over $100 million?",
}
KEYS = list(QUESTIONS)

#                          star summer sequel wide budget  hit
FILMS = {
    "Harbour Lights":       (1,   1,     0,     0,   0,     1),
    "Ironwake II":          (1,   1,     1,     1,   0,     1),
    "The Quiet Ledger II":  (1,   0,     1,     0,   0,     1),
    "Nightfall Divide":     (1,   1,     0,     1,   1,     1),
    "Ironwake III":         (1,   0,     1,     1,   1,     1),
    "Cinder Coast II":      (1,   0,     1,     1,   1,     1),
    "Redline Returns":      (1,   0,     1,     1,   1,     1),
    "Paper Kingdoms":       (1,   1,     0,     1,   0,     1),
    "Glass Monsoon":        (1,   1,     0,     0,   1,     0),
    "Cinder Coast III":     (1,   0,     1,     1,   0,     0),
    "Vermilion Rising II":  (0,   0,     1,     1,   1,     1),
    "Saltwater Sunday":     (0,   1,     0,     1,   1,     1),
    "Field of Static II":   (0,   0,     1,     0,   0,     1),
    "Wildflower County":    (0,   1,     0,     1,   0,     1),
    "Tin Sky":              (0,   1,     0,     0,   0,     0),
    "The Cartographer":     (0,   0,     0,     0,   0,     0),
    "Neon Bazaar":          (0,   1,     0,     0,   1,     0),
    "Ash & Ivory II":       (0,   0,     1,     1,   0,     0),
    "Slow River":           (0,   0,     0,     0,   0,     0),
    "The Understudy":       (0,   1,     0,     0,   0,     0),
    "Meridian":             (0,   0,     0,     1,   1,     0),
    "Copper Harbour":       (0,   1,     0,     1,   1,     0),
    "The Winter Post":      (0,   0,     0,     0,   0,     0),
    "Little Eden":          (0,   1,     0,     0,   1,     0),
}

def answers(name):
    return FILMS[name][:len(KEYS)]

def hit(name):
    return FILMS[name][-1]

ALL = list(FILMS)
record("n_films", len(ALL))
record("n_hits", sum(hit(f) for f in ALL))

# The pair the whole lesson turns on: same answers, opposite outcomes.
same = {}
for f in ALL:
    same.setdefault(answers(f), []).append(f)
twins = [g for g in same.values() if len(g) > 1 and len({hit(f) for f in g}) > 1][0]
def row(name):
    return [name] + ["yes" if a else "no" for a in answers(name)] + ["hit" if hit(name) else "flop"]

record("twin_hit", next(f for f in twins if hit(f)))
record("twin_flop", next(f for f in twins if not hit(f)))
# Transposed — questions down the side, the two films across. Seven columns of
# yes/no is unreadable projected (tests/present-fits.test.ts caps a slide at
# six), and this way the two answer columns sit next to each other, which is
# the entire thing the table is there to show.
pair = sorted(twins, key=lambda f: -hit(f))
record("twin_columns", [""] + pair)
record("twin_rows",
       [[QUESTIONS[q]] + ["yes" if answers(f)[i] else "no" for f in pair]
        for i, q in enumerate(KEYS)]
       + [["Result"] + ["hit" if hit(f) else "flop" for f in pair]])
print(twins, "answer identically:", dict(zip(KEYS, answers(twins[0]))))

def split(films, q):
    i = KEYS.index(q)
    return ([f for f in films if answers(f)[i]], [f for f in films if not answers(f)[i]])

print(f"{'question':<9} {'yes':>3} {'no':>3}   {'hits on the yes side':>22}   {'hits on the no side':>21}")
for q in KEYS:
    yes, no = split(ALL, q)
    print(f"{q:<9} {len(yes):>3} {len(no):>3}   "
          f"{sum(hit(f) for f in yes):>10} of {len(yes):<9}   {sum(hit(f) for f in no):>9} of {len(no):<9}")

evenest = min(KEYS, key=lambda q: abs(len(split(ALL, q)[0]) - len(split(ALL, q)[1])))
ey, en = split(ALL, evenest)
record("even_question", QUESTIONS[evenest])
record("even_yes_hits", sum(hit(f) for f in ey))
record("even_yes_n", len(ey))

def gini(films):
    if not films:
        return 0.0
    p = sum(hit(f) for f in films) / len(films)
    return 1 - p * p - (1 - p) * (1 - p)

root = record("root_gini", round(gini(ALL), 3))
print(f"all 24 films          {gini(ALL):.3f}   (12 hits, 12 flops — as bad as it gets)")
yes, no = split(ALL, "star")
print(f"the 10 with a star    {gini(yes):.3f}   ({sum(hit(f) for f in yes)} hits, {len(yes) - sum(hit(f) for f in yes)} flops)")
print(f"the 14 without        {gini(no):.3f}   ({sum(hit(f) for f in no)} hits, {len(no) - sum(hit(f) for f in no)} flops)")

def gain(films, q):
    yes, no = split(films, q)
    after = (len(yes) * gini(yes) + len(no) * gini(no)) / len(films)
    return gini(films) - after

scores = sorted(((gain(ALL, q), q) for q in KEYS), reverse=True)
for g, q in scores:
    yes, no = split(ALL, q)
    print(f"{QUESTIONS[q]:<42} {gini(ALL):.3f} -> "
          f"{(len(yes) * gini(yes) + len(no) * gini(no)) / len(ALL):.3f}   gain {g:.3f}")

best_gain, best_q = scores[0]
record("best_question", QUESTIONS[best_q])
record("best_gain", round(best_gain, 3))
record("best_yes_hits", sum(hit(f) for f in split(ALL, best_q)[0]))
record("best_yes_n", len(split(ALL, best_q)[0]))
record("best_no_hits", sum(hit(f) for f in split(ALL, best_q)[1]))
record("best_no_n", len(split(ALL, best_q)[1]))
record("even_gain", round(gain(ALL, evenest), 3))
record("budget_gain", round(gain(ALL, "budget"), 3))

# The table the lesson prints, in the order it prints it. Recorded rather than
# retyped into the lesson: five rows of counts is exactly the sort of thing
# that goes stale silently when a film is edited.
record("question_columns", ["Question", "Splits into", "Hits on the yes side", "Hits on the no side"])
record("question_rows", [
    [QUESTIONS[q],
     f"{len(split(ALL, q)[0])} / {len(split(ALL, q)[1])}",
     f"{sum(hit(f) for f in split(ALL, q)[0])} of {len(split(ALL, q)[0])}",
     f"{sum(hit(f) for f in split(ALL, q)[1])} of {len(split(ALL, q)[1])}"]
    for _, q in scores
])


def plot2(ax):
    labels = [QUESTIONS[q].rstrip("?") for _, q in scores][::-1]
    vals = [g for g, _ in scores][::-1]
    colors = ["#e2574c" if v == max(vals) else "#9aa0aa" for v in vals]
    ax.barh(labels, vals, color=colors)
    for y, (v, q) in enumerate(zip(vals, [q for _, q in scores][::-1])):
        yes, no = split(ALL, q)
        ax.text(v + 0.003, y, f"{len(yes)} / {len(no)} films", va="center", fontsize=9, color="#77777f")
    ax.set_xlabel("Impurity removed (gain)")
    ax.set_xlim(0, max(vals) * 1.55)
    ax.grid(axis="y", visible=False)

save_fig("gain-by-question", plot2, figsize=(7, 3.9))

def plot3(ax):
    pairs = [(evenest, "the evenest split"), (best_q, "the most useful split")]
    x = 0
    for q, label in pairs:
        yes, no = split(ALL, q)
        for side, name in ((yes, "yes"), (no, "no")):
            h = sum(hit(f) for f in side)
            ax.bar([x], [h], color="#e2574c", width=0.7)
            ax.bar([x], [len(side) - h], bottom=[h], color="#9aa0aa", width=0.7)
            ax.text(x, len(side) + 0.4, f"{h}/{len(side)}", ha="center", fontsize=9)
            ax.text(x, -1.4, name, ha="center", fontsize=9, color="#77777f")
            x += 1
        ax.text(x - 1.5, -2.9, f"{label}\ngain {gain(ALL, q):.3f}", ha="center", fontsize=9.5)
        x += 0.8
    ax.set_ylabel("Films")
    ax.set_ylim(-3.6, 17)
    ax.set_xticks([])
    ax.grid(axis="x", visible=False)
    ax.spines["bottom"].set_visible(False)

save_fig("even-against-useful", plot3, figsize=(7, 4.0))