Text rule trends

#!/usr/bin/env python3
"""Group items by meaning: k-means on one sample, nearest centroid for any other.

Fit (the reference sample):
  python3 cluster.py fit SAMPLE_DIR --k 40 [--n-init 6] [--seed 0] [--show 7] [--centroids-out FILE]
                     [--match-to OLD_centroids.json]
    -> SAMPLE_DIR/labels.json {"k", "labels", "centroids_id"} and SAMPLE_DIR/centroids.json
    Prints each cluster with its size, reach and the items nearest its centre, so a person can
    name the clusters and group them into themes (see examples/*/themes.json).
    --match-to renumbers the new clusters to line up with older centroids (Hungarian matching on
    cosine similarity) and prints how close each pair is. Even then the new clusters are not the
    old ones: a themes file names one set of centroids, and analyze.py refuses a themes file
    whose centroids_id differs from the labels'. To reuse old themes on new data, assign.

Assign (any other sample, to saved centroids):
  python3 cluster.py assign SAMPLE_DIR --centroids centroids.json
    -> SAMPLE_DIR/labels.json

Convert saved centroids between .npy and JSON:
  python3 cluster.py convert IN.npy OUT.json      (or IN.json OUT.npy)

k-means is sklearn's KMeans with k-means++ init. The saved centroids are the mean embedding of
each cluster's items, unit-normalised. Assignment picks the highest cosine similarity.
centroids.json is plain text: each value is the shortest decimal that reads back as the same
float32, so a JSON round trip gives exactly the same labels as the .npy.
"""
import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import centroids_id, load_centroids, save_centroids  # noqa: E402


def show(life, E, labels, centres, n):
    for c in np.argsort(-np.bincount(labels, minlength=len(centres))):
        idx = np.where(labels == c)[0]
        if not len(idx):
            continue
        files = len({life[i]["repo"] for i in idx})
        alive = np.mean([life[i]["died"] is None for i in idx])
        print(f"\n#{c}  {len(idx)} items in {files} files, still there {alive:.0%}")
        for i in idx[np.argsort(-(E[idx] @ centres[c]))[:n]]:
            print("   -", life[i]["text"][:110])


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("mode", choices=["fit", "assign", "convert"])
    ap.add_argument("sample", help="SAMPLE_DIR (convert: the input file)")
    ap.add_argument("target", nargs="?", help="convert: the output file")
    ap.add_argument("--k", type=int, default=40)
    ap.add_argument("--n-init", type=int, default=6)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--centroids", help="assign: centroids.json or .npy")
    ap.add_argument("--centroids-out", help="fit: where to write the centroids (default SAMPLE_DIR/centroids.json)")
    ap.add_argument("--show", type=int, default=7, help="fit: items to print per cluster (0 = none)")
    ap.add_argument("--match-to", help="fit: renumber clusters to match these centroids (Hungarian on cosine)")
    a = ap.parse_args()

    if a.mode == "convert":
        C = load_centroids(a.sample)
        if a.target.endswith(".npy"):
            np.save(a.target, C)
        else:
            save_centroids(a.target, C)
        back = load_centroids(a.target)
        print(f"wrote {a.target}: {C.shape}, round trip exact: {bool(np.array_equal(back, C))}")
        return

    life = json.load(open(os.path.join(a.sample, "rules.json")))["life"]
    E = np.load(os.path.join(a.sample, "emb.npy"))
    if len(E) != len(life):
        raise SystemExit("emb.npy does not match rules.json; run embed.py again")

    if a.mode == "fit":
        from sklearn.cluster import KMeans
        if len(E) < a.k:
            raise SystemExit(f"{len(E)} items cannot make {a.k} clusters; lower --k")
        km = KMeans(a.k, n_init=a.n_init, random_state=a.seed).fit(E)
        labels = km.labels_
        C = np.stack([E[labels == k].mean(0) for k in range(a.k)])
        C /= np.linalg.norm(C, axis=1, keepdims=True)
        centres = km.cluster_centers_
        if a.match_to:
            from scipy.optimize import linear_sum_assignment
            old = load_centroids(a.match_to)
            if old.shape != C.shape:
                raise SystemExit(f"--match-to has shape {old.shape}, the fit {C.shape}")
            sim = old @ C.T                                    # [old, new]
            rows, cols = linear_sum_assignment(-sim)
            perm = np.empty(a.k, dtype=int)
            perm[cols] = rows                                  # new cluster cols[i] becomes old id rows[i]
            labels = perm[labels]
            order = np.argsort(perm)
            C, centres = C[order], centres[order]
            s = sim[rows, cols]
            same = int(np.sum(np.abs(old - C).max(1) == 0))
            print(f"matched to {a.match_to}: cosine per pair min {s.min():.3f}, median {np.median(s):.3f}; "
                  f"{same} of {a.k} centroids identical", flush=True)
            if same < a.k:
                print("warning: the new clusters are not the old ones. A themes file written for the old "
                      "centroids will be refused by analyze.py; assign to the old centroids instead, or "
                      "rename the clusters for this fit.", flush=True)
        out = a.centroids_out or os.path.join(a.sample, "centroids.json")
        if out.endswith(".npy"):
            np.save(out, C.astype(np.float32))
        else:
            save_centroids(out, C, n_init=a.n_init, seed=a.seed)
        k = a.k
        cid = centroids_id(C.astype(np.float32))
    else:
        if not a.centroids:
            raise SystemExit("assign needs --centroids")
        C = load_centroids(a.centroids)
        labels = (E / np.linalg.norm(E, axis=1, keepdims=True) @ C.T).argmax(1)
        k = len(C)
        cid = centroids_id(C)
    json.dump({"k": k, "labels": [int(x) for x in labels], "centroids_id": cid}, open(os.path.join(a.sample, "labels.json"), "w"))
    print(f"{a.mode}: {len(labels)} items in {k} clusters -> {a.sample}/labels.json", flush=True)
    if a.mode == "fit" and a.show:
        show(life, E, labels, centres, a.show)


if __name__ == "__main__":
    main()