graduate_code/solutions/data.py

"""G15: group-disjoint splits, training-only scaling, and bounded Lloyd clustering."""
import math
import random


def split_groups(groups, seed=21):
    if not groups or any(type(g) is not str or not g for g in groups):
        raise ValueError("nonempty string group labels required")
    unique = sorted(set(groups))
    if len(unique) < 5:
        raise ValueError("at least five groups required for this split policy")
    random.Random(seed).shuffle(unique)
    a,b = max(1,int(.6*len(unique))), max(1,int(.2*len(unique)))
    partitions = [set(unique[:a]),set(unique[a:a+b]),set(unique[a+b:])]
    return {name:[i for i,g in enumerate(groups) if g in part]
            for name,part in zip(("train","validation","test"),partitions)}


def matrix(X):
    if not X or not X[0] or len(X) > 10000 or len(X[0]) > 64:
        raise ValueError("bounded nonempty matrix required")
    d = len(X[0])
    if any(len(r) != d or any(not math.isfinite(v) for v in r) for r in X):
        raise ValueError("finite rectangular features required")
    return [[float(v) for v in r] for r in X]


def fit_scaler(X):
    X = matrix(X)
    n,d = len(X),len(X[0])
    mean = [math.fsum(r[j] for r in X)/n for j in range(d)]
    scale = [math.sqrt(math.fsum((r[j]-mean[j])**2 for r in X)/n) for j in range(d)]
    if not all(math.isfinite(x) for x in [*mean,*scale]):
        raise ValueError("unsupported feature scale")
    return {"mean":mean,"scale":[s if s > 0 else 1. for s in scale]}


def transform(X, scaler):
    X = matrix(X)
    mu,s = scaler["mean"],scaler["scale"]
    if len(mu) != len(X[0]) or len(s) != len(mu) or any(v <= 0 for v in s):
        raise ValueError("invalid scaler dimensions or scales")
    if not all(math.isfinite(v) for v in [*mu,*s]):
        raise ValueError("nonfinite scaler")
    return [[(x-m)/sd for x,m,sd in zip(row,mu,s)] for row in X]


def kmeans(X, initial, steps=50):
    X,C = matrix(X),matrix(initial)
    if len(C) > len(X) or len(C[0]) != len(X[0]) or type(steps) is not int or not 1 <= steps <= 1000:
        raise ValueError("matching cluster dimensions and bounded step count required")
    d = len(X[0])
    def assign(centers):
        return [min(range(len(centers)),
                    key=lambda j:math.fsum((a-b)**2 for a,b in zip(x,centers[j]))) for x in X]
    history=[]
    for _ in range(steps):
        labels = assign(C)
        new = []
        for j in range(len(C)):
            members = [x for x,l in zip(X,labels) if l == j]
            new.append([math.fsum(x[k] for x in members)/len(members) for k in range(d)]
                       if members else list(C[j]))
        labels = assign(new)
        sse = math.fsum((v-c)**2 for x,l in zip(X,labels) for v,c in zip(x,new[l]))
        if not math.isfinite(sse):
            raise ValueError("nonfinite clustering objective")
        history.append(sse)
        if new == C:
            C = new
            break
        C = new
    return {"centers":C,"labels":labels,"sse":history[-1],"history":history}


def fixture(seed=21, n_groups=30):
    rng = random.Random(seed)
    rows=[]
    for g in range(n_groups):
        group_effect = rng.uniform(-.3,.3)
        for r in range(4):
            x = [rng.uniform(-2,2),rng.uniform(-2,2)]
            noise = rng.uniform(-.4,.4)
            rows.append({"group":f"g{g:03}","x":x,
                         "y":int(x[0]+.6*x[1]+group_effect+noise > 0)})
    return rows


if __name__ == "__main__":
    rows=fixture()
    print({k:len(v) for k,v in split_groups([r["group"] for r in rows]).items()})
    print(kmeans([[0.],[2.],[8.],[10.]],[[0.],[10.]]))