Gaussian mixture models

Like k-means, but each group is an ellipse, and each point belongs partly to every group. It repeats:

1. Color each point by how much it belongs to each group.
2. Fit each ellipse to the points, weighted by how much they belong.

Run this first. Rerun it for new points.

pie.clear()
pie.scalar(0, out='gridOpacity')
import numpy as np
k = 3
colors = [(225, 29, 72), (37, 99, 235), (217, 119, 6), (124, 58, 237), (13, 148, 136), (101, 163, 13)]
blobs = [((-2, 2), [[2, 1.2], [1.2, 0.9]]), ((2.5, 1), [[0.2, 0], [0, 1.5]]), ((0, -2.5), [[2.5, 0], [0, 0.15]])]
rng = np.random.default_rng()
coords = np.concatenate([rng.multivariate_normal(m, c, 15) for m, c in blobs]).clip(-4, 4).round(2)
n = len(coords)
ps = [pie.point(x, y, fill='#a8a29e') for x, y in coords]

# Each group: a weight (its share of the points), a center and a covariance (its shape).
weights = np.full(k, 1 / k)
means = coords[rng.choice(n, k, replace=False)]
covs = np.array([np.eye(2)] * k)

def ellipse(cov, last_angle=0):
    """The radii (2 standard deviations, longest first) and the angle in degrees of the ellipse for `cov`.
    Turning an ellipse by 180° gives the same ellipse, so the angle is the one nearest `last_angle`."""
    spreads, dirs = np.linalg.eigh(cov)
    angle = np.degrees(np.arctan2(dirs[1, 1], dirs[0, 1]))
    angle += 180 * round((last_angle - angle) / 180)
    return 2 * np.sqrt(spreads[1]), 2 * np.sqrt(spreads[0]), angle

# Each ellipse is built from scalars, so animating the scalars moves, stretches and turns it.
shapes, angles = [], []
for (x, y), cov, (r, g, b) in zip(means, covs, colors):
    rx, ry, angle = ellipse(cov)
    s = [pie.scalar(v) for v in (x, y, rx, ry, angle)]
    pie.ellipse(pie.point(s[0], s[1], hidden=True), s[2], s[3], s[4], fill=f'rgba({r}, {g}, {b}, 0.1)', stroke=f'rgb({r}, {g}, {b})')
    shapes.append(s)
    angles.append(angle)

def density(mean, cov):
    """How dense the group with this center and covariance is at each point."""
    d = coords - mean
    inside = np.einsum('ij,jk,ik->i', d, np.linalg.inv(cov), d)
    return np.exp(-inside / 2) / (2 * np.pi * np.sqrt(np.linalg.det(cov)))

def show(r, fit, done):
    pie.text(f'round {r}: fit {fit:.1f}' + (', done' if done else ''), 0, 4.8, 14, stroke='#1c1917', out='status')

Then run this. The fit at the top is the log-likelihood: higher is better.

last = -np.inf
for r in range(1, 200):
    # Step 1: how much each point belongs to each group.
    dens = np.array([w * density(m, c) for w, m, c in zip(weights, means, covs)]).T
    belong = dens / dens.sum(axis=1, keepdims=True)
    for p, mix in zip(ps, belong @ np.array(colors[:k])):
        pie.set_fill(p, 'rgb({:.0f}, {:.0f}, {:.0f})'.format(*mix))
    fit = np.log(dens.sum(axis=1)).sum()
    done = fit - last < 0.01
    show(r, fit, done)
    if done:
        break
    last = fit
    # Step 2: each group fits its ellipse to the points, each counted by how much it belongs.
    share = belong.sum(axis=0) + 1e-9  # + 1e-9: no dividing by 0 if a group has no points
    weights = share / n
    means = belong.T @ coords / share[:, None]
    # The small circle added to each shape keeps an ellipse from shrinking to nothing around one point.
    covs = np.array([(belong[:, j, None] * (coords - means[j])).T @ (coords - means[j]) / share[j] + 0.01 * np.eye(2) for j in range(k)])
    targets = []
    for j, (x, y) in enumerate(means):
        rx, ry, angles[j] = ellipse(covs[j], angles[j])
        targets += [float(v) for v in (x, y, rx, ry, angles[j])]
    pie.translate_scalar_to([v for s in shapes for v in s], targets, animate=True)

A bad start can end in a worse fit. Set k up to 6.