1"""Doubled-angle orientation order pooling."""
 2import torch
 3
 4def doubled_angle_pool(energies, eps=1e-8):
 5    """energies: (..., B), nonnegative; bins theta=b*pi/B.
 6    Returns normalized Cartesian order channels (zr, zi, q), plus raw z and mass.
 7    """
 8    B = energies.shape[-1]
 9    theta = torch.arange(B, device=energies.device, dtype=energies.dtype) * torch.pi / B
10    c, s = torch.cos(2 * theta), torch.sin(2 * theta)
11    mass = energies.sum(-1)
12    zr = (energies * c).sum(-1)
13    zi = (energies * s).sum(-1)
14    denom = mass + eps
15    q = torch.sqrt(zr.square() + zi.square()) / denom
16    return torch.stack((zr / denom, zi / denom, q), -1), (zr, zi, q, mass)
17
18def orientation_from_pool(pooled, q_min=1e-4):
19    """Confidence-gated principal orientation in [0, pi)."""
20    zr, zi, q = pooled[..., 0], pooled[..., 1], pooled[..., 2]
21    phi = 0.5 * torch.atan2(zi, zr) % torch.pi
22    return torch.where(q > q_min, phi, torch.zeros_like(phi))