src/embed_hierarchy.py

"""
Symmetric Poincaré vs. Euclidean hierarchy embedding.

Design choice, stated explicitly (deviates from how the rest of this
codebase trains hyperbolic components, on purpose, for this experiment):
the hierarchical field-predictor (model.py) parameterizes points in the
tangent space and only ever calls expmap0 at the boundary, trained with
plain Euclidean Adam. That's a reasonable simplification for a predictor
whose main job is spatiotemporal forecasting. For THIS experiment — whose
entire point is to test whether hyperbolic geometry helps hierarchy
recovery — using genuine Riemannian optimization (geoopt.ManifoldParameter
+ geoopt.optim.RiemannianAdam, verified working in this session before
being used here) is more faithful to the original Nickel & Kiela / Sala
et al. methodology this whole test is trying to reproduce in miniature.
Using a weaker approximation here would bias the comparison against the
Poincaré model and undermine the point of running the experiment at all.

The Euclidean baseline is a plain nn.Embedding trained with ordinary
Adam — the natural, undiluted comparison point.
"""
from __future__ import annotations
from typing import Dict, Optional, Tuple

import torch
import torch.nn as nn
import geoopt

from .data_pbdb_taxonomy import TaxonomyEdgeDataset


class HierarchyEmbedding(nn.Module):
    """
    Shared interface over the two geometries so training/eval code never
    needs an if/else on `geometry` — everything goes through .points(),
    .distance(), and .parameters()/.embed_optimizer().
    """
    def __init__(
        self,
        num_nodes: int,
        dim: int = 8,
        geometry: str = "poincare",
        c: float = 1.0,
        learnable_c: bool = False,
        init_scale: float = 1e-3,
    ):
        super().__init__()
        if geometry not in ("poincare", "euclidean"):
            raise ValueError(f"geometry must be 'poincare' or 'euclidean', got {geometry!r}")
        self.geometry = geometry
        self.dim = dim
        self.num_nodes = num_nodes

        init = torch.randn(num_nodes, dim) * init_scale
        if geometry == "poincare":
            self.manifold = geoopt.PoincareBall(c=c, learnable=learnable_c)
            self.emb = geoopt.ManifoldParameter(init, manifold=self.manifold)
        else:
            self.manifold = None
            self.emb = nn.Parameter(init)

    def points(self, idx: torch.Tensor) -> torch.Tensor:
        return self.emb[idx]

    def distance(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        if self.geometry == "poincare":
            return self.manifold.dist(a, b)
        return (a - b).pow(2).sum(-1).clamp_min(1e-12).sqrt()

    def make_optimizer(self, lr: float, curvature_lr_mult: float = 0.1, optimizer_type: str = "radam"):

        if self.geometry == "poincare":
            groups = [{"params": [self.emb], "lr": lr}]
            if hasattr(self.manifold, "isp_c"):
                groups.append({"params": [self.manifold.isp_c], "lr": lr * curvature_lr_mult})
            if optimizer_type == "radam":
                return geoopt.optim.RiemannianAdam(groups)
            elif optimizer_type == "rsgd":
                return geoopt.optim.RiemannianSGD(groups, lr=lr)
            else:
                raise ValueError(f"optimizer_type must be 'radam' or 'rsgd', got {optimizer_type!r}")
        return torch.optim.Adam(self.parameters(), lr=lr) if optimizer_type == "radam" \
            else torch.optim.SGD(self.parameters(), lr=lr)

    @torch.no_grad()
    def clip_to_ball(self, margin: float = 0.95):
        if self.geometry != "poincare":
            return
        self.emb.data = self.manifold.projx(self.emb.data)
        max_norm = (1.0 / self.manifold.c.clamp_min(1e-8).sqrt()) * margin
        norms = self.emb.data.norm(dim=-1, keepdim=True).clamp_min(1e-12)
        factor = torch.clamp(max_norm / norms, max=1.0)
        self.emb.data = self.emb.data * factor

    @torch.no_grad()
    def clamp_curvature(self, c_min: float = 0.1, c_max: float = 3.0):
        if self.geometry != "poincare" or not hasattr(self.manifold, "isp_c"):
            return
        lo = torch.log(torch.expm1(torch.tensor(c_min, dtype=self.manifold.isp_c.dtype)))
        hi = torch.log(torch.expm1(torch.tensor(c_max, dtype=self.manifold.isp_c.dtype)))
        self.manifold.isp_c.clamp_(lo.item(), hi.item())

    def radii(self) -> torch.Tensor:
        return self.emb.detach().norm(dim=-1)


def negative_sample(
    child_idx: torch.Tensor,
    true_parent_idx: torch.Tensor,
    num_nodes: int,
    k: int,
    generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
    B = child_idx.shape[0]
    neg = torch.randint(0, num_nodes, (B, k), generator=generator)
    collision = neg.eq(true_parent_idx.unsqueeze(1)) | neg.eq(child_idx.unsqueeze(1))
    while collision.any():
        resample = torch.randint(0, num_nodes, (int(collision.sum().item()),), generator=generator)
        neg[collision] = resample
        collision = neg.eq(true_parent_idx.unsqueeze(1)) | neg.eq(child_idx.unsqueeze(1))
    return neg


def ranking_loss(
    model: HierarchyEmbedding,
    child_idx: torch.Tensor,
    parent_idx: torch.Tensor,
    neg_idx: torch.Tensor,
    margin: float = 1.0,
) -> torch.Tensor:
    child_pts = model.points(child_idx)                      # (B, D)
    parent_pts = model.points(parent_idx)                    # (B, D)
    B, K = neg_idx.shape
    neg_pts = model.points(neg_idx.reshape(-1)).reshape(B, K, -1)

    d_pos = model.distance(child_pts, parent_pts)             # (B,)
    d_neg = model.distance(
        child_pts.unsqueeze(1).expand(-1, K, -1).reshape(-1, model.dim),
        neg_pts.reshape(-1, model.dim),
    ).reshape(B, K)                                            # (B, K)

    loss = torch.relu(margin + d_pos.unsqueeze(1) - d_neg)
    return loss.mean()


def softmax_ranking_loss(
    model: HierarchyEmbedding,
    child_idx: torch.Tensor,
    parent_idx: torch.Tensor,
    neg_idx: torch.Tensor,
) -> torch.Tensor:
    child_pts = model.points(child_idx)                       # (B, D)
    parent_pts = model.points(parent_idx)                     # (B, D)
    B, K = neg_idx.shape
    neg_pts = model.points(neg_idx.reshape(-1)).reshape(B, K, -1)

    d_pos = model.distance(child_pts, parent_pts)              # (B,)
    d_neg = model.distance(
        child_pts.unsqueeze(1).expand(-1, K, -1).reshape(-1, model.dim),
        neg_pts.reshape(-1, model.dim),
    ).reshape(B, K)                                             # (B, K)

    logits = torch.cat([(-d_pos).unsqueeze(1), -d_neg], dim=1)  # (B, K+1); true parent at index 0
    target = torch.zeros(B, dtype=torch.long, device=logits.device)
    return torch.nn.functional.cross_entropy(logits, target)


def train_hierarchy_embedding(
    dataset: TaxonomyEdgeDataset,
    geometry: str = "poincare",
    dim: int = 8,
    epochs: int = 50,
    batch_size: int = 256,
    lr: float = 1e-3,
    neg_samples: int = 10,
    margin: float = 1.0,
    c: float = 1.0,
    learnable_c: bool = False,
    loss_type: str = "margin",
    burn_in_epochs: int = 0,
    burn_in_lr_mult: float = 0.1,
    curvature_lr_mult: float = 0.1,
    c_min: float = 0.1,
    c_max: float = 3.0,
    optimizer_type: str = "radam",
    device: str = "cpu",
    seed: int = 0,
) -> Tuple[HierarchyEmbedding, Dict[str, float]]:
    if loss_type not in ("margin", "softmax"):
        raise ValueError(f"loss_type must be 'margin' or 'softmax', got {loss_type!r}")

    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)

    model = HierarchyEmbedding(
        num_nodes=dataset.num_nodes, dim=dim, geometry=geometry,
        c=c, learnable_c=learnable_c,
    ).to(device)
    opt = model.make_optimizer(lr, curvature_lr_mult=curvature_lr_mult, optimizer_type=optimizer_type)
    for group in opt.param_groups:
        group["lr"] = lr * burn_in_lr_mult if burn_in_epochs > 0 else lr

    edge_idx = torch.tensor(dataset.edge_idx, dtype=torch.long)  # (E, 2) = (child, parent)
    n_edges = edge_idx.shape[0]
    loss_history = []
    c_history = []
    max_norm_history = []

    for epoch in range(epochs):
        if burn_in_epochs > 0 and epoch == burn_in_epochs:
            for group in opt.param_groups:
                group["lr"] = lr

        perm = torch.randperm(n_edges, generator=gen)
        epoch_loss, n_batches = 0.0, 0
        for start in range(0, n_edges, batch_size):
            batch_idx = perm[start : start + batch_size]
            batch = edge_idx[batch_idx].to(device)
            child_idx, parent_idx = batch[:, 0], batch[:, 1]
            neg_idx = negative_sample(
                child_idx, parent_idx, dataset.num_nodes, neg_samples, generator=gen
            ).to(device)

            if loss_type == "margin":
                loss = ranking_loss(model, child_idx, parent_idx, neg_idx, margin=margin)
            else:
                loss = softmax_ranking_loss(model, child_idx, parent_idx, neg_idx)

            if not torch.isfinite(loss):
                raise RuntimeError(
                    f"[NON_FINITE_LOSS] loss={loss.item()} at epoch {epoch+1}, "
                    f"geometry={geometry}, loss_type={loss_type} — stopping "
                    f"rather than continuing with a corrupted embedding."
                )
            opt.zero_grad()
            loss.backward()
            opt.step()
            model.clamp_curvature(c_min=c_min, c_max=c_max)
            model.clip_to_ball()
            epoch_loss += loss.item()
            n_batches += 1

        mean_loss = epoch_loss / max(n_batches, 1)
        loss_history.append(mean_loss)
        if geometry == "poincare":
            c_history.append(model.manifold.c.item())
            max_norm_history.append(model.emb.detach().norm(dim=-1).max().item())

    metrics = {"final_loss": loss_history[-1], "loss_history": loss_history}
    if geometry == "poincare":
        metrics["c_history"] = c_history
        metrics["max_norm_history"] = max_norm_history
    return model, metrics