Exercise 4 — VAE¶

X_train, y_train, X_val, y_val, X_test, y_test = load_mnist()
print(f"A) treino {tuple(X_train.shape)}, validação {tuple(X_val.shape)}, teste {tuple(X_test.shape)}; "
      f"pixels em [{X_train.min():.0f}, {X_train.max():.0f}]")

D = 16
print(f"B) {VAE(D)}")
print(f"   parâmetros: {sum(p.numel() for p in VAE(D).parameters())}")

t0 = time.time()
model, hist = train_vae(D, X_train, X_val)
print(f"C) treino de {EPOCHS} épocas em {time.time() - t0:.0f} s")
torch.save(model.state_dict(), DATA / "vae_d16.pt")

# O teste só é usado aqui, depois do treino, para os números finais
test_recon, test_kl = evaluate(model, X_test)
print(f"C) teste: reconstrução {test_recon:.2f}, KL {test_kl:.2f}, -ELBO {test_recon + test_kl:.2f} nats por imagem")
print(f"   validação: KL na época 1 = {hist['val_kl'][0]:.2f}, na época {EPOCHS} = {hist['val_kl'][-1]:.2f}")

epochs = np.arange(1, EPOCHS + 1)
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
for ax, term, label in [(axes[0], "recon", "reconstrução (BCE)"), (axes[1], "kl", "KL")]:
    ax.plot(epochs, hist[f"train_{term}"], "o-", ms=3, label="treino (média da época)")
    ax.plot(epochs, hist[f"val_{term}"], "s-", ms=3, label="validação (fim da época)")
    ax.set_xlabel("época"); ax.set_ylabel("nats por imagem")
    ax.set_title(label); ax.set_xticks(epochs[::2]); ax.legend()
fig.suptitle(f"Figure 1 — Exercise 1: os dois termos da perda por época (d = {D})")
fig.tight_layout()
fig.savefig(FIG / "fig1.png", bbox_inches="tight"); plt.show()

model.eval()
first_idx = [int((y_test == k).nonzero()[0]) for k in range(10)]    # primeira imagem de teste de cada dígito
x10 = X_test[first_idx]
with torch.no_grad():
    mu10, _ = model.encode(x10)
    rec10 = model.decode(mu10)                                       # reconstrução a partir de mu, sem sortear z
print(f"D) índices de teste da Figure 2 (dígitos 0 a 9): {first_idx}")

fig, axes = plt.subplots(2, 10, figsize=(12, 2.9))
for k in range(10):
    for row, imgs in enumerate([x10, rec10]):
        axes[row, k].imshow(imgs[k].reshape(28, 28), cmap="gray", vmin=0, vmax=1)
        axes[row, k].set_xticks([]); axes[row, k].set_yticks([])
    axes[0, k].set_title(f"dígito {k}", fontsize=9)
axes[0, 0].set_ylabel("original"); axes[1, 0].set_ylabel("reconstrução\n(decodifica μ)")
fig.supxlabel("dígito (primeira imagem de teste de cada um)", y=-0.02)
fig.suptitle("Figure 2 — Exercise 1: primeira imagem de teste de cada dígito e sua reconstrução (d = 16)")
fig.savefig(FIG / "fig2.png", bbox_inches="tight"); plt.show()

torch.manual_seed(0)
samples = model.sample(64)                                           # z ~ N(0, I), depois o decoder
grid = samples.reshape(8, 8, 28, 28).permute(0, 2, 1, 3).reshape(8 * 28, 8 * 28)
fig, ax = plt.subplots(figsize=(6, 6.3))
ax.imshow(grid, cmap="gray", vmin=0, vmax=1)
ax.set_xticks(np.arange(8) * 28 + 13.5); ax.set_xticklabels(range(1, 9))
ax.set_yticks(np.arange(8) * 28 + 13.5); ax.set_yticklabels(range(1, 9))
ax.set_xlabel("coluna da grade"); ax.set_ylabel("linha da grade")
ax.set_title("Figure 3 — Exercise 1: 64 amostras de z ~ N(0, I)\ndecodificadas (d = 16, torch.manual_seed(0))")
fig.savefig(FIG / "fig3.png", bbox_inches="tight"); plt.show()

# 1) KL subindo no começo: quanto a reconstrução caiu enquanto o KL subia
print(f"E1) validação antes do treino: reconstrução {hist['init_recon']:.2f}, KL {hist['init_kl']:.2f}")
print(f"    KL de validação por época: {np.round(hist['val_kl'], 2).tolist()}")
print(f"    reconstrução de validação por época: {np.round(hist['val_recon'], 2).tolist()}")
k = int(np.argmax(np.diff(hist["val_kl"]) < 0))          # última época da subida inicial do KL (índice)
r0, r1 = hist["init_recon"], hist["val_recon"][k]
k0, k1 = hist["init_kl"], hist["val_kl"][k]
print(f"    do início à época {k + 1} (fim da subida contínua): KL {k0:.2f} -> {k1:.2f} ({k1 - k0:+.2f}), "
      f"reconstrução {r0:.2f} -> {r1:.2f} ({r1 - r0:+.2f}), total {r0 + k0:.2f} -> {r1 + k1:.2f}")
print(f"    da época 1 à época {k + 1}: KL {hist['val_kl'][0]:.2f} -> {k1:.2f} ({k1 - hist['val_kl'][0]:+.2f}), "
      f"reconstrução {hist['val_recon'][0]:.2f} -> {r1:.2f} ({r1 - hist['val_recon'][0]:+.2f})")

with torch.no_grad():
    mu_v, logvar_v = model.encode(X_val)
kl_dim = 0.5 * (mu_v.pow(2) + logvar_v.exp() - 1 - logvar_v).mean(0)  # KL médio por dimensão latente
print(f"E2) KL médio por dimensão (validação): {np.round(kl_dim.numpy(), 2).tolist()}")
active = kl_dim >= 0.1
print(f"    dimensões ativas (KL >= 0.1 nat): {active.sum().item()} de {D}")

# 2) Onde as amostras da prior caem: distância até o código mu de imagem real mais próximo, só nas
#    dimensões ativas. Referência: mu das 5 000 primeiras imagens de validação. Para comparar, faço a mesma
#    conta com códigos z ~ q(z|x) de outras 5 000 imagens de validação (que não estão na referência).
ref = mu_v[:5000, active]
gen = torch.Generator().manual_seed(1)
with torch.no_grad():
    z_post = model.reparameterize(mu_v[5000:], logvar_v[5000:], gen)   # códigos de imagens reais
    torch.manual_seed(0)
    z_prior = torch.randn(64, D)                                       # os mesmos z da Figure 3
d_post = torch.cdist(z_post[:, active], ref).min(1).values
d_prior = torch.cdist(z_prior[:, active], ref).min(1).values
thr = torch.quantile(d_post, 0.95)
far = d_prior > thr                                                    # amostras "longe" de qualquer imagem real
print(f"E2) distância mediana até o mu real mais próximo: códigos de imagens reais {d_post.median():.2f}, "
      f"amostras da prior {d_prior.median():.2f}")
print(f"    amostras da Figure 3 acima do quantil 95% dos códigos reais ({thr:.2f}): {far.sum().item()} de 64")

# 3) O decoder devolve probabilidades: fração de pixels "cinza" (entre 0.1 e 0.9)
gray = lambda imgs: ((imgs > 0.1) & (imgs < 0.9)).float().mean().item() * 100
print(f"E2) pixels entre 0.1 e 0.9: originais da Figure 2 {gray(x10):.1f}%, reconstruções {gray(rec10):.1f}%, "
      f"amostras da Figure 3 {gray(samples):.1f}%")
print(f"    nas amostras da Figure 3: {gray(samples[far]):.1f}% nas {far.sum().item()} longe, "
      f"{gray(samples[~far]):.1f}% nas outras {(~far).sum().item()}")
far_cells = [(i // 8 + 1, i % 8 + 1) for i in far.nonzero().flatten().tolist()]
print(f"    posições (linha, coluna) das amostras longe: {far_cells}")

# 3b) Parece um dígito? Confiança do classificador auxiliar (digit_judge.py)
judge = train_judge(X_train, y_train)
_, c_rec = judge(rec10)
_, c_smp = judge(samples)
f = far.numpy()
print(f"    confiança média do classificador: reconstruções da Figure 2 {c_rec.mean():.2f}, "
      f"amostras da Figure 3 {c_smp.mean():.2f} (longe {c_smp[f].mean():.2f}, outras {c_smp[~f].mean():.2f})")
print(f"    amostras da Figure 3 com confiança < 0.6: {(c_smp < 0.6).sum()} de 64 ({(c_smp[f] < 0.6).sum()} delas longe)")

# 4) Decodificar mu ou um z sorteado
rec_mu = evaluate(model, X_val, use_mean=True)[0]
rec_z = evaluate(model, X_val)[0]
print(f"E3) reconstrução de validação decodificando mu {rec_mu:.2f} vs z sorteado {rec_z:.2f} nats por imagem")
print(f"    sigma médio no encoder: {torch.exp(0.5 * logvar_v).mean():.3f}; |mu| médio: {mu_v.abs().mean():.3f}")

# Mesmo modelo, dados, seed, beta e orçamento do Exercise 1. Só muda d.
model2, hist2 = train_vae(2, X_train, X_val)
torch.save(model2.state_dict(), DATA / "vae_d2.pt")
model16 = load_or_train(16, X_train, X_val)
model2.eval(); model16.eval()

rec2, kl2 = evaluate(model2, X_test)         # teste só para os números finais
rec16, kl16 = evaluate(model16, X_test)
print("A) teste, nats por imagem:")
print(f"   d = 16: reconstrução {rec16:.2f}, KL {kl16:.2f}, -ELBO {rec16 + kl16:.2f}")
print(f"   d =  2: reconstrução {rec2:.2f}, KL {kl2:.2f}, -ELBO {rec2 + kl2:.2f}")
print(f"   diferença (d=2 - d=16): reconstrução {rec2 - rec16:+.2f}, KL {kl2 - kl16:+.2f}, -ELBO {rec2 + kl2 - rec16 - kl16:+.2f}")
print(f"   d = 2, validação: KL na época 1 = {hist2['val_kl'][0]:.2f}, na época {EPOCHS} = {hist2['val_kl'][-1]:.2f}; "
      f"reconstrução na época {EPOCHS} = {hist2['val_recon'][-1]:.2f}")

with torch.no_grad():
    mu_test = model2.encode(X_test)[0].numpy()
    mu_train = model2.encode(X_train)[0].numpy()

colors = plt.cm.tab10(np.arange(10))
order = np.random.default_rng(0).permutation(len(mu_test))   # ordem aleatória: nenhum dígito fica escondido por baixo
fig, ax = plt.subplots(figsize=(8, 7))
ax.scatter(mu_test[order, 0], mu_test[order, 1], s=3, alpha=0.6, color=colors[y_test.numpy()[order]])
for k in range(10):                                           # pontos vazios só para montar a legenda
    ax.scatter([], [], s=12, color=colors[k], label=f"{k}")
lim = norm.ppf(0.95)
ax.add_patch(plt.Rectangle((-lim, -lim), 2 * lim, 2 * lim, fill=False, ls="--", color="black",
                           label="área da Figure 5"))
ax.set_xlabel("$\\mu_1$"); ax.set_ylabel("$\\mu_2$")
ax.set_title("Figure 4 — Exercise 2: μ das 10 000 imagens de teste (d = 2), cor = dígito")
ax.legend(title="dígito", markerscale=1.5, loc="center left", bbox_to_anchor=(1.01, 0.5))
fig.savefig(FIG / "fig4.png", bbox_inches="tight"); plt.show()

# Sobreposição medida: um k-vizinhos no plano latente, ajustado com o treino e aplicado no teste.
# Os pares que ele mais confunde são os dígitos que se misturam na Figure 4.
knn = KNeighborsClassifier(n_neighbors=15).fit(mu_train, y_train.numpy())
pred = knn.predict(mu_test)
print(f"B) acurácia de um 15-vizinhos no plano latente (teste): {(pred == y_test.numpy()).mean():.4f}")
conf = np.zeros((10, 10), int)
np.add.at(conf, (y_test.numpy(), pred), 1)
pairs = sorted(((conf[i, j] + conf[j, i], i, j) for i in range(10) for j in range(i + 1, 10)), reverse=True)
print("   pares mais confundidos (erros nos dois sentidos):", [(f"{i}-{j}", n) for n, i, j in pairs[:6]])
centers = np.array([np.median(mu_test[y_test.numpy() == k], axis=0) for k in range(10)])
print("   mediana de μ por dígito:", {k: centers[k].round(2).tolist() for k in range(10)})

# Classificador de dígitos de digit_judge.py: diz "isto parece o dígito k" nas imagens geradas.
judge = train_judge(X_train, y_train)
print(f"   classificador auxiliar: acurácia na validação {(judge(X_val)[0] == y_val.numpy()).mean():.4f}")

q = np.linspace(0.05, 0.95, 15)
zs = norm.ppf(q)                             # quantis da N(0, 1): cada faixa tem a mesma massa da prior
print(f"C) valores de z na grade: {zs.round(2).tolist()}")
Z1, Z2 = np.meshgrid(zs, zs[::-1])           # linha de cima = z2 maior, como no gráfico da Figure 4
grid_z = torch.tensor(np.stack([Z1.ravel(), Z2.ravel()], 1), dtype=torch.float32)
with torch.no_grad():
    grid_img = model2.decode(grid_z)
tile = grid_img.reshape(15, 15, 28, 28).permute(0, 2, 1, 3).reshape(15 * 28, 15 * 28)

fig, ax = plt.subplots(figsize=(9, 9.3))
ax.imshow(tile, cmap="gray", vmin=0, vmax=1)
ticks = np.arange(15) * 28 + 13.5
ax.set_xticks(ticks); ax.set_xticklabels([f"{v:.2f}" for v in zs], rotation=90)
ax.set_yticks(ticks); ax.set_yticklabels([f"{v:.2f}" for v in zs[::-1]])
ax.set_xlabel("$z_1 = \\Phi^{-1}(q)$"); ax.set_ylabel("$z_2 = \\Phi^{-1}(q)$")
ax.set_title("Figure 5 — Exercise 2: decoder de d = 2 numa grade 15×15\nnos quantis q = 0.05, ..., 0.95 da N(0, 1)")
fig.savefig(FIG / "fig5.png", bbox_inches="tight"); plt.show()

lab_grid, conf_grid = judge(grid_img)
print("   dígito que o classificador vê em cada célula da Figure 5 (linha de cima = z2 maior):")
for r in range(15):
    print("   " + " ".join(f"{l}{'?' if c < 0.6 else ' '}" for l, c in
                         zip(lab_grid[r * 15:(r + 1) * 15], conf_grid[r * 15:(r + 1) * 15])))
print(f"   células com confiança < 0.6 (marcadas com ?): {(conf_grid < 0.6).sum()} de 225")
print(f"   dígitos que aparecem na grade: {sorted(set(lab_grid.tolist()))}; "
      f"células por dígito: {np.bincount(lab_grid, minlength=10).tolist()}")


def cell_of(z):
    """Célula (linha, coluna) da Figure 5 mais próxima de um ponto z, contando de 1."""
    col = int(np.abs(zs - z[0]).argmin()) + 1
    row = int(np.abs(zs[::-1] - z[1]).argmin()) + 1
    return row, col


print("   célula da Figure 5 mais próxima da mediana de cada dígito:",
      {k: cell_of(centers[k]) for k in range(10)})

with torch.no_grad():
    rec_val = {d: m.decode(m.encode(X_val)[0]) for d, m in [(2, model2), (16, model16)]}
yv = y_val.numpy()
for d in (16, 2):
    acc = (judge(rec_val[d])[0] == yv).mean()
    print(f"E2) d = {d:2d}: o classificador acerta o dígito em {acc:.4f} das reconstruções de validação")


def within_class_spread(imgs):
    """Distância quadrática média de cada imagem até a média do seu dígito (variação de estilo)."""
    return np.mean([((imgs[yv == k] - imgs[yv == k].mean(0)) ** 2).sum(1).mean().item() for k in range(10)])


s_orig = within_class_spread(X_val)
for d in (16, 2):
    print(f"    d = {d:2d}: variação dentro do dígito nas reconstruções = {within_class_spread(rec_val[d]) / s_orig:.1%} "
          f"da variação nas imagens originais")

IA, IB = 3, 2                                # índices de teste: um 0 e um 1
print(f"D) índices de teste: a = {IA} (dígito {y_test[IA].item()}), b = {IB} (dígito {y_test[IB].item()})")
t = torch.linspace(0, 1, 10)[:, None]


def interpolate(xa, xb):
    """Linha de cima: mistura dos pixels. Linha de baixo: mistura dos mu do modelo d = 16, decodificada."""
    pix = (1 - t) * xa + t * xb
    with torch.no_grad():
        mua, mub = model16.encode(torch.stack([xa, xb]))[0]
        lat = model16.decode((1 - t) * mua + t * mub)
    return pix, lat


pix, lat = interpolate(X_test[IA], X_test[IB])

fig, axes = plt.subplots(2, 10, figsize=(13, 3.2))
for j in range(10):
    for row, imgs in enumerate([pix, lat]):
        axes[row, j].imshow(imgs[j].reshape(28, 28), cmap="gray", vmin=0, vmax=1)
        axes[row, j].set_xticks([]); axes[row, j].set_yticks([])
    axes[0, j].set_title(f"t = {t[j].item():.2f}", fontsize=9)
axes[0, 0].set_ylabel("pixels"); axes[1, 0].set_ylabel("latente\n(d = 16)")
fig.supxlabel("t (fração do caminho da imagem a até a imagem b)", y=-0.02)
fig.suptitle(f"Figure 6 — Exercise 2: interpolação da imagem de teste {IA} (dígito {y_test[IA].item()}) "
             f"até a {IB} (dígito {y_test[IB].item()})")
fig.savefig(FIG / "fig6.png", bbox_inches="tight"); plt.show()

gray = lambda imgs: ((imgs > 0.1) & (imgs < 0.9)).float().mean(1).numpy() * 100   # % de pixels "cinza"
for name, imgs in [("pixels", pix), ("latente", lat)]:
    lab, c = judge(imgs)
    print(f"   {name:7s}: dígito visto {lab.tolist()}")
    print(f"            confiança   {np.round(c, 2).tolist()}")
    print(f"            % de pixels entre 0.1 e 0.9 {np.round(gray(imgs), 1).tolist()}")


# O mesmo teste nos 45 pares formados pelas primeiras imagens de teste de cada dígito (as da Figure 2).
# Só os 8 quadros do meio (t de 0.11 a 0.89) entram na média.
first_idx = [int((y_test == k).nonzero()[0]) for k in range(10)]
stats = {"pixels": {"cinza": [], "conf": []}, "latente": {"cinza": [], "conf": []}}
for i in range(10):
    for j in range(i + 1, 10):
        for name, imgs in zip(["pixels", "latente"], interpolate(X_test[first_idx[i]], X_test[first_idx[j]])):
            stats[name]["cinza"].append(gray(imgs[1:-1]).mean())
            stats[name]["conf"].append(judge(imgs[1:-1])[1].mean())
for name in stats:
    print(f"   45 pares, {name:7s}: pixels cinza nos quadros do meio {np.mean(stats[name]['cinza']):.1f}%, "
          f"confiança média do classificador {np.mean(stats[name]['conf']):.2f}")
print(f"   pares em que a linha latente tem menos pixels cinza que a de pixels: "
      f"{(np.array(stats['latente']['cinza']) < np.array(stats['pixels']['cinza'])).sum()} de 45")