1 files changed, 137 insertions, 0 deletions
diff --git a/worldalign/energy.py b/worldalign/energy.py
new file mode 100644
index 0000000..796c41a
--- /dev/null
+++ b/worldalign/energy.py
@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def off_diagonal_mask(size: int, device: torch.device | str) -> torch.Tensor:
+    return ~torch.eye(size, dtype=torch.bool, device=device)
+
+
+def standardized_relation(
+    features: torch.Tensor, mask: torch.Tensor | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Cosine relation field and its standardized off-diagonal values."""
+    features = F.normalize(features.float(), dim=-1)
+    relation = features @ features.T
+    if mask is None:
+        mask = off_diagonal_mask(len(features), features.device)
+    values = relation[mask]
+    standardized = (values - values.mean()) / values.std().clamp_min(1e-6)
+    return relation, standardized
+
+
+def relation_field_energy(
+    visual_relation: torch.Tensor,
+    visual_standardized: torch.Tensor,
+    language_particles: torch.Tensor,
+    temperatures: tuple[float, ...] = (0.03, 0.07, 0.15),
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Second-order and multiscale conditional relation energies."""
+    language_relation, language_standardized = standardized_relation(
+        language_particles
+    )
+    mse = F.mse_loss(language_standardized, visual_standardized)
+    diagonal = torch.eye(
+        len(language_particles),
+        dtype=torch.bool,
+        device=language_particles.device,
+    )
+    conditional_kl = language_particles.new_zeros(())
+    for temperature in temperatures:
+        visual_logits = (visual_relation / temperature).masked_fill(
+            diagonal, -1e4
+        )
+        language_logits = (language_relation / temperature).masked_fill(
+            diagonal, -1e4
+        )
+        visual_probability = F.softmax(visual_logits, dim=-1)
+        conditional_kl = conditional_kl + (
+            visual_probability
+            * (
+                F.log_softmax(visual_logits, dim=-1)
+                - F.log_softmax(language_logits, dim=-1)
+            )
+        ).sum(-1).mean()
+    return mse, conditional_kl
+
+
+def projection_quantile_target(
+    text_features: torch.Tensor,
+    particles: int,
+    projections: int,
+    generator: torch.Generator,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fixed sliced-distribution target from an unpaired text population."""
+    directions = F.normalize(
+        torch.randn(
+            text_features.shape[-1],
+            projections,
+            generator=generator,
+            device=text_features.device,
+        ),
+        dim=0,
+    )
+    projected = (text_features @ directions).sort(dim=0).values
+    quantile_indices = (
+        torch.linspace(
+            0,
+            len(projected) - 1,
+            particles,
+            device=text_features.device,
+        )
+        .round()
+        .long()
+    )
+    return directions, projected[quantile_indices]
+
+
+def sliced_distribution_energy(
+    particles: torch.Tensor,
+    directions: torch.Tensor,
+    target_quantiles: torch.Tensor,
+) -> torch.Tensor:
+    projected = (F.normalize(particles, dim=-1) @ directions).sort(
+        dim=0
+    ).values
+    return F.mse_loss(projected, target_quantiles)
+
+
+def prototype_manifold_energy(
+    particles: torch.Tensor, prototypes: torch.Tensor
+) -> torch.Tensor:
+    particles = F.normalize(particles, dim=-1)
+    prototypes = F.normalize(prototypes, dim=-1)
+    return (1 - (particles @ prototypes.T).max(dim=-1).values).mean()
+
+
+def log_sinkhorn(
+    logits: torch.Tensor, temperature: float, iterations: int = 12
+) -> torch.Tensor:
+    """Doubly stochastic coupling with differentiable log-domain updates."""
+    log_coupling = logits / temperature
+    for _ in range(iterations):
+        log_coupling = log_coupling - torch.logsumexp(
+            log_coupling, dim=1, keepdim=True
+        )
+        log_coupling = log_coupling - torch.logsumexp(
+            log_coupling, dim=0, keepdim=True
+        )
+    return log_coupling.exp()
+
+
+def retrieval_metrics(
+    particles: torch.Tensor, paired_text: torch.Tensor
+) -> dict[str, float]:
+    particles = F.normalize(particles.float(), dim=-1)
+    paired_text = F.normalize(paired_text.float(), dim=-1)
+    similarity = particles @ paired_text.T
+    target = similarity.diagonal()
+    ranks = (similarity > target[:, None]).sum(-1) + 1
+    return {
+        "r@1": float((ranks <= 1).float().mean()),
+        "r@5": float((ranks <= 5).float().mean()),
+        "r@10": float((ranks <= 10).float().mean()),
+        "median_rank": float(ranks.float().median()),
+        "paired_cosine": float(target.mean()),
+    }