1"""Small differentiable Hodge-dual electrostatic loss on a periodic 2-D grid."""
 2import torch
 3
 4
 5def periodic_grad(u, h):
 6    return torch.stack(((torch.roll(u, -1, -2)-torch.roll(u, 1, -2))/(2*h),
 7                        (torch.roll(u, -1, -1)-torch.roll(u, 1, -1))/(2*h)), dim=-1)
 8
 9
10def curl_2d(A, h):
11    """2-D scalar-vector-potential curl: (d_y A_z, -d_x A_z)."""
12    az = A[..., 2]
13    dy = (torch.roll(az, -1, -2)-torch.roll(az, 1, -2))/(2*h)
14    dx = (torch.roll(az, -1, -1)-torch.roll(az, 1, -1))/(2*h)
15    return torch.stack((dy, -dx), dim=-1)
16
17
18def divergence_2d(p, h):
19    px, py = p[..., 0], p[..., 1]
20    dx = (torch.roll(px, -1, -1)-torch.roll(px, 1, -1))/(2*h)
21    dy = (torch.roll(py, -1, -2)-torch.roll(py, 1, -2))/(2*h)
22    return dx + dy
23
24
25def dielectric(n, eps_perp=2.0, eps_a=0.5):
26    n = n / (torch.linalg.vector_norm(n, dim=-1, keepdim=True) + 1e-8)
27    I = torch.eye(3, device=n.device, dtype=n.dtype)
28    return eps_perp*I + eps_a*n[..., :, None]*n[..., None, :]
29
30
31def dual_electrostatic_loss(p0, A, n, h, eps_perp=2.0, eps_a=0.5):
32    """Positive dual density plus optional elastic term is left to the caller."""
33    p = p0 + curl_2d(A, h)
34    E = dielectric(n, eps_perp, eps_a)
35    density = 0.5 * torch.einsum('...i,...ij,...j->...', p, torch.linalg.inv(E), p)
36    return density.mean(), p