1"""Detailed-balance graph transport primitives and a small learnable layer."""
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6
 7def logarithmic_mean(a, b):
 8    a, b = np.asarray(a, float), np.asarray(b, float)
 9    aa, bb = np.broadcast_arrays(a, b)
10    out = np.empty_like(aa)
11    close = np.isclose(aa, bb, rtol=1e-10, atol=1e-14)
12    out[close] = (aa[close] + bb[close]) / 2
13    out[~close] = (aa[~close] - bb[~close]) / (np.log(aa[~close]) - np.log(bb[~close]))
14    return out
15
16
17def energy(rho, pi):
18    rho, pi = np.asarray(rho), np.asarray(pi)
19    return float(np.sum(rho * np.log(rho / pi)))
20
21
22def rhs(rho, pi, conductance):
23    """Master-equation RHS; rho/pi may be [N] or [N,K]."""
24    q = np.asarray(rho) / np.asarray(pi)
25    if q.ndim == 1:
26        return np.sum(conductance * (q[None, :] - q[:, None]), axis=1)
27    return np.sum(conductance[:, :, None] * (q[None, :, :] - q[:, None, :]), axis=1)
28
29
30def explicit_step(rho, pi, conductance, dt):
31    return np.asarray(rho) + dt * rhs(rho, pi, conductance)
32
33
34def positivity_dt_bound(rho, pi, conductance):
35    out = np.sum(conductance, axis=1)
36    pi = np.asarray(pi); rho = np.asarray(rho)
37    if pi.ndim == 1:
38        out = out / pi
39    else:
40        out = out[:, None] / pi
41    valid = out > 0
42    return float(np.min(rho[valid] / out[valid]))
43
44
45def dissipation(rho, pi, conductance):
46    q = np.asarray(rho) / np.asarray(pi)
47    mu = np.log(q)
48    if q.ndim == 1:
49        lm = logarithmic_mean(q[:, None], q[None, :])
50        return float(.5 * np.sum(conductance * lm * (mu[:, None]-mu[None, :])**2))
51    lm = logarithmic_mean(q[:, None, :], q[None, :, :])
52    return float(.5 * np.sum(conductance[:, :, None] * lm * (mu[None, :, :] - mu[:, None, :])**2))
53
54
55class DetailedBalanceTransport(nn.Module):
56    """K-channel positive graph transport with conservative explicit updates."""
57    def __init__(self, in_dim, channels=2, dt=0.1, eps=1e-6):
58        super().__init__()
59        self.rho_head = nn.Linear(in_dim, channels)
60        self.pi_head = nn.Linear(in_dim, channels)
61        self.edge_head = nn.Linear(3 * in_dim, channels)
62        self.dt, self.eps = dt, eps
63
64    def forward(self, h, edge_index):
65        n = h.shape[0]; src, dst = edge_index
66        rho = torch.nn.functional.softplus(self.rho_head(h)) + self.eps
67        pi = torch.nn.functional.softplus(self.pi_head(h)) + self.eps
68        pi = pi / pi.sum(0, keepdim=True) * rho.sum(0, keepdim=True).detach()
69        z = torch.cat([h[src], h[dst], (h[src]-h[dst]).abs()], dim=1)
70        c0 = torch.nn.functional.softplus(self.edge_head(z)) + self.eps
71        key = torch.minimum(src, dst) * n + torch.maximum(src, dst)
72        cs = torch.zeros((n*n, c0.shape[1]), device=h.device)
73        counts = torch.zeros((n*n, 1), device=h.device)
74        cs.index_add_(0, key, c0); counts.index_add_(0, key, torch.ones_like(c0[:, :1]))
75        c = cs[key] / counts[key].clamp_min(1)
76        delta = rho[src] / pi[src] - rho[dst] / pi[dst]
77        dr = torch.zeros_like(rho)
78        dr.index_add_(0, src, -c * delta); dr.index_add_(0, dst, c * delta)
79        return rho + self.dt * dr, pi, c, rho