378 lines
15 KiB
Python
378 lines
15 KiB
Python
"""
|
|
Core library for the TMC paper:
|
|
"Structured Shared-Private Embedding Multiplexing for Semantic Multiple
|
|
Access in Dynamic Mobile Networks"
|
|
|
|
Builds on the published shared-embedding multiple-access framework
|
|
(Lee, Choi, Lee, IEEE JSAC 2026, doi 10.1109/JSAC.2025.3643816).
|
|
|
|
Pipeline modeled here
|
|
---------------------
|
|
1. Content model (shared-scene decomposition), latent frame:
|
|
z_u = a_u [c; 0] + sqrt(1-a_u^2) [0; p_u] (structured latent)
|
|
c in R^{d_c}: shared scene content, p_u in R^{d-d_c}: private content.
|
|
A frozen foundation encoder outputs RAW embeddings x_u = R z_u with an
|
|
unknown orthogonal mixing R (the encoder's arbitrary basis) -- the
|
|
structure is present but hidden in the coordinates.
|
|
|
|
2. Matched-filter MAC front end (identical to the JSAC/AA-EDMA line):
|
|
tilde_x_u = x_u + sum_{v!=u} beta_uv (h_v/h_u) x_v + n_u,
|
|
Cov(n_u,n_w) = sigma^2 B_uw/(h_u h_w) I_d, sigma^2 = 1/rho.
|
|
|
|
3. Receivers: SR (cancel), SC (combine), B-aware LMMSE (optimal linear,
|
|
isotropic prior), and the decomposition receiver DR that knows the
|
|
shared-subspace basis V_c: BLUE-combining on the shared block +
|
|
SR cancellation and Wiener shrinkage on the private complement.
|
|
|
|
4. Embedding-structure optimization: closed-form spectral recovery of V_c
|
|
from the cross-covariance of paired clean embeddings (GCCA-style), and a
|
|
channel-in-the-loop learned linear adapter refining it.
|
|
|
|
5. Mobility: time-varying a_u(t) from user trajectories around a scene, and
|
|
a decision-directed EWMA affinity tracker.
|
|
"""
|
|
from __future__ import annotations
|
|
import math
|
|
import numpy as np
|
|
|
|
TAU = 0.45 # SER threshold on cosine (same operating definition as prior work)
|
|
|
|
|
|
def set_seed(seed: int):
|
|
np.random.seed(seed)
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Affinity utilities
|
|
# ----------------------------------------------------------------------
|
|
def affinity_matrix(a: np.ndarray) -> np.ndarray:
|
|
"""B_uv = a_u a_v (u != v), B_uu = 1."""
|
|
a = np.asarray(a, dtype=np.float64)
|
|
B = np.outer(a, a)
|
|
np.fill_diagonal(B, 1.0)
|
|
return B
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Content generators
|
|
# ----------------------------------------------------------------------
|
|
def sample_latents_isotropic(batch, U, d, d_c, a, rng):
|
|
"""Structured latents with isotropic random contents (unit norm).
|
|
Returns z of shape (batch, U, d): shared block = first d_c coords."""
|
|
a = np.broadcast_to(np.asarray(a, float), (batch, U)) \
|
|
if np.ndim(a) > 1 or np.ndim(a) == 1 else np.full((batch, U), float(a))
|
|
if a.shape != (batch, U):
|
|
a = np.broadcast_to(np.asarray(a, float), (batch, U))
|
|
z = np.zeros((batch, U, d))
|
|
c = rng.standard_normal((batch, d_c))
|
|
c /= np.linalg.norm(c, axis=1, keepdims=True)
|
|
p = rng.standard_normal((batch, U, d - d_c))
|
|
p /= np.linalg.norm(p, axis=2, keepdims=True)
|
|
z[:, :, :d_c] = a[:, :, None] * c[:, None, :]
|
|
z[:, :, d_c:] = np.sqrt(1.0 - a[:, :, None] ** 2) * p
|
|
return z
|
|
|
|
|
|
class EmbeddingPool:
|
|
"""Real PLM embedding pool (e.g., BERT AG-News, 8000 x 768).
|
|
Centered + unit-normalized; provides PCA coordinates so that structured
|
|
latents can be built from real semantic content."""
|
|
|
|
def __init__(self, X: np.ndarray):
|
|
X = np.asarray(X, dtype=np.float64)
|
|
self.mu = X.mean(axis=0, keepdims=True)
|
|
Xc = X - self.mu
|
|
Xc /= np.linalg.norm(Xc, axis=1, keepdims=True)
|
|
self.X = Xc
|
|
# PCA basis of the (centered, normalized) pool
|
|
_, S, Vt = np.linalg.svd(Xc, full_matrices=False)
|
|
self.pca = Vt # (d, d) rows = principal directions
|
|
self.spectrum = S ** 2 / len(Xc)
|
|
self.N, self.d = Xc.shape
|
|
|
|
def pca_coords(self, idx, k):
|
|
"""Top-k PCA coordinates of pool items idx, renormalized to unit."""
|
|
Y = self.X[idx] @ self.pca[:k].T
|
|
return Y / (np.linalg.norm(Y, axis=1, keepdims=True) + 1e-12)
|
|
|
|
|
|
def sample_latents_pool(pool: EmbeddingPool, batch, U, d, d_c, a, rng,
|
|
idx_pool=None):
|
|
"""Structured latents whose shared/private contents are REAL embeddings:
|
|
shared c = top-d_c PCA coords of one pool sentence, private p_u =
|
|
top-(d-d_c) PCA coords of distinct other sentences.
|
|
|
|
idx_pool: optional index array restricting which pool sentences may be
|
|
drawn (train/holdout partition); None draws from the whole pool."""
|
|
a = np.broadcast_to(np.asarray(a, float), (U,))
|
|
choices = np.arange(pool.N) if idx_pool is None else np.asarray(idx_pool)
|
|
idx = np.array([rng.choice(choices, size=U + 1, replace=False)
|
|
for _ in range(batch)])
|
|
c = pool.pca_coords(idx[:, 0], d_c) # (batch, d_c)
|
|
z = np.zeros((batch, U, d))
|
|
for u in range(U):
|
|
p = pool.pca_coords(idx[:, u + 1], d - d_c)
|
|
z[:, u, :d_c] = a[u] * c
|
|
z[:, u, d_c:] = math.sqrt(1.0 - a[u] ** 2) * p
|
|
return z
|
|
|
|
|
|
def random_orthogonal(d, rng):
|
|
G = rng.standard_normal((d, d))
|
|
Q, Rr = np.linalg.qr(G)
|
|
Q *= np.sign(np.diag(Rr))
|
|
return Q
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Matched-filter MAC front end (JSAC / AA-EDMA convention)
|
|
# ----------------------------------------------------------------------
|
|
def matched_filter(e, B, rho, rng, fading=True):
|
|
"""tilde_u = sum_v (h_v/h_u) B_uv e_v + n_u,
|
|
Cov(n_u,n_w) = (sigma^2 B_uw / (h_u h_w)) I_d. Returns (tilde, h)."""
|
|
batch, U, d = e.shape
|
|
sigma2 = 1.0 / rho
|
|
B = np.asarray(B, dtype=np.float64)
|
|
if B.ndim == 2:
|
|
B = np.broadcast_to(B, (batch, U, U))
|
|
if fading:
|
|
hc = (rng.standard_normal((batch, U)) +
|
|
1j * rng.standard_normal((batch, U))) / math.sqrt(2)
|
|
h = np.clip(np.abs(hc), 0.2, None)
|
|
else:
|
|
h = np.ones((batch, U))
|
|
tilde = np.einsum('buv,bvd,bv,bu->bud', B, e, h, 1.0 / h)
|
|
xi = rng.standard_normal((batch, U, d)) * math.sqrt(sigma2)
|
|
for b in range(batch):
|
|
A = np.linalg.cholesky(B[b])
|
|
tilde[b] += (A @ xi[b]) / h[b][:, None]
|
|
return tilde, h
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Receivers
|
|
# ----------------------------------------------------------------------
|
|
def demux_sr(tilde, B, h):
|
|
"""Similarity-rejecting closed-form demux: (Gamma^{-1} (x) I) tilde."""
|
|
batch, U, d = tilde.shape
|
|
out = np.empty_like(tilde)
|
|
for b in range(batch):
|
|
Gamma = np.diag(1.0 / h[b]) @ B @ np.diag(h[b])
|
|
out[b] = np.linalg.solve(Gamma, tilde[b])
|
|
return out
|
|
|
|
|
|
def demux_sc(tilde, h):
|
|
"""Similarity-combining endpoint: h^2-weighted MRC of all MF outputs."""
|
|
w = h ** 2
|
|
w = w / w.sum(axis=1, keepdims=True)
|
|
comb = np.einsum('bu,bud->bd', w, tilde)
|
|
return np.repeat(comb[:, None, :], tilde.shape[1], axis=1)
|
|
|
|
|
|
def lmmse_matrices(B, h, sigma2, d):
|
|
H = np.diag(h)
|
|
Hi = np.diag(1.0 / h)
|
|
Gamma = Hi @ B @ H
|
|
Cx = B / d
|
|
Cn = sigma2 * (Hi @ B @ Hi)
|
|
S = Gamma @ Cx @ Gamma.T + Cn
|
|
W = np.linalg.solve(S.T, (Cx @ Gamma.T).T).T
|
|
Eerr = Cx - W @ Gamma @ Cx
|
|
return W, Eerr
|
|
|
|
|
|
def demux_lmmse(tilde, B, h, rho):
|
|
"""B-aware LMMSE (optimal linear receiver under isotropic prior).
|
|
Returns (estimates, closed-form per-user total MSE averaged over batch)."""
|
|
batch, U, d = tilde.shape
|
|
sigma2 = 1.0 / rho
|
|
out = np.empty_like(tilde)
|
|
mse_cf = np.zeros(U)
|
|
for b in range(batch):
|
|
W, Eerr = lmmse_matrices(B, h[b], sigma2, d)
|
|
out[b] = W @ tilde[b]
|
|
mse_cf += d * np.diag(Eerr)
|
|
return out, mse_cf / batch
|
|
|
|
|
|
def demux_dr(tilde, B, h, rho, a, Vc):
|
|
"""Decomposition receiver (proposed).
|
|
|
|
Vc: (d, d_c) orthonormal basis of the shared subspace (from the
|
|
embedding-structure optimizer; oracle = true mixing columns).
|
|
Shared block: BLUE-combining of the U looks at the common content,
|
|
followed by Wiener shrinkage. Private complement: SR cancellation +
|
|
per-user Wiener shrinkage. Recombine."""
|
|
batch, U, d = tilde.shape
|
|
d_c = Vc.shape[1]
|
|
sigma2 = 1.0 / rho
|
|
a = np.broadcast_to(np.asarray(a, float), (U,))
|
|
Binv = np.linalg.inv(B)
|
|
out = np.empty_like(tilde)
|
|
Zs = tilde @ Vc # (batch, U, d_c) shared-block obs
|
|
Zp = tilde - (Zs @ Vc.T) # complement part (in ambient frame)
|
|
for b in range(batch):
|
|
hb = h[b]
|
|
Hi = np.diag(1.0 / hb)
|
|
Gamma = Hi @ B @ np.diag(hb)
|
|
Cn = sigma2 * (Hi @ B @ Hi)
|
|
gamma = np.array([
|
|
a[u] + sum(B[u, v] * a[v] * hb[v] / hb[u]
|
|
for v in range(U) if v != u) for u in range(U)])
|
|
Cn_inv = np.linalg.inv(Cn)
|
|
denom = float(gamma @ Cn_inv @ gamma)
|
|
if denom > 1e-12:
|
|
c_hat = (gamma @ Cn_inv @ Zs[b]) / denom
|
|
c_hat *= (1.0 / d_c) / (1.0 / d_c + 1.0 / denom)
|
|
else:
|
|
c_hat = np.zeros(d_c)
|
|
Gp = np.linalg.solve(Gamma, Zp[b]) # SR on the complement
|
|
for u in range(U):
|
|
sig_p = 1.0 - a[u] ** 2
|
|
err_p = (d - d_c) * sigma2 * Binv[u, u] / hb[u] ** 2
|
|
shrink = sig_p / (sig_p + err_p) if sig_p > 0 else 0.0
|
|
out[b, u] = a[u] * (Vc @ c_hat) + shrink * Gp[u]
|
|
return out
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Conventional baselines (orthogonal and power-domain multiple access)
|
|
# ----------------------------------------------------------------------
|
|
def oma_observe(e, h, rho, rng):
|
|
"""Conventional orthogonal MA (OFDMA-style): user u is confined to a
|
|
disjoint d/U-dimensional block and decoded only from that block, so it
|
|
never sees cross-user interference but its recovery is capped by the
|
|
1/U-energy subspace (cosine ceiling ~ sqrt(1/U))."""
|
|
batch, U, d = e.shape
|
|
blk = d // U
|
|
sigma = math.sqrt(1.0 / rho)
|
|
out = np.zeros_like(e)
|
|
for u in range(U):
|
|
sl = slice(u * blk, (u + 1) * blk)
|
|
out[:, u, sl] = e[:, u, sl] + \
|
|
rng.standard_normal((batch, blk)) * sigma / h[:, u][:, None]
|
|
return out
|
|
|
|
|
|
def demux_noma_genie(e, h, rho, rng):
|
|
"""Genie-aided NOMA-SIC upper bound: every user decoded from an
|
|
interference-free observation e_u + n/h_u at per-user SNR rho
|
|
(perfect cancellation, no error propagation)."""
|
|
batch, U, d = e.shape
|
|
sigma = math.sqrt(1.0 / rho)
|
|
return e + rng.standard_normal(e.shape) * sigma / h[:, :, None]
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Embedding-structure optimization
|
|
# ----------------------------------------------------------------------
|
|
def learn_structure_spectral(x_clean, d_c):
|
|
"""Closed-form shared-subspace recovery from N paired CLEAN embeddings.
|
|
|
|
x_clean: (N, U, d) raw (mixed) embeddings of co-located users.
|
|
The averaged symmetrized cross-covariance has column space equal to the
|
|
shared subspace (private parts are independent and average out).
|
|
Returns Vc_hat (d, d_c), eigenvalues (d,)."""
|
|
N, U, d = x_clean.shape
|
|
M = np.zeros((d, d))
|
|
cnt = 0
|
|
for u in range(U):
|
|
for v in range(u + 1, U):
|
|
C = x_clean[:, u, :].T @ x_clean[:, v, :] / N
|
|
M += C + C.T
|
|
cnt += 2
|
|
M /= cnt
|
|
w, V = np.linalg.eigh(M)
|
|
order = np.argsort(w)[::-1]
|
|
return V[:, order[:d_c]], w[order]
|
|
|
|
|
|
def subspace_error(Vhat, Vtrue):
|
|
"""Normalized projection-Frobenius distance in [0,1]."""
|
|
P1 = Vhat @ Vhat.T
|
|
P2 = Vtrue @ Vtrue.T
|
|
k = Vtrue.shape[1]
|
|
return float(np.linalg.norm(P1 - P2) / math.sqrt(2 * k))
|
|
|
|
|
|
def estimate_a_from_clean(x_clean, Vc):
|
|
"""a_u estimate from clean paired data: sqrt(mean shared-block energy)."""
|
|
E = np.linalg.norm(x_clean @ Vc, axis=2) ** 2 # (N, U)
|
|
return np.sqrt(np.clip(E.mean(axis=0), 0.0, 1.0))
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Mobility model and online affinity tracking
|
|
# ----------------------------------------------------------------------
|
|
def mobility_trajectories(U, T, speed, rng, box=60.0, r_scene=28.0,
|
|
a_max=0.95, dt=1.0):
|
|
"""Random-waypoint trajectories around a scene at the origin.
|
|
Returns a_t of shape (T, U): a_u(t) = a_max * exp(-d_u(t)^2 / (2 r^2))."""
|
|
pos = rng.uniform(-box, box, size=(U, 2))
|
|
wp = rng.uniform(-box, box, size=(U, 2))
|
|
a_t = np.zeros((T, U))
|
|
for t in range(T):
|
|
for u in range(U):
|
|
vec = wp[u] - pos[u]
|
|
dist = np.linalg.norm(vec)
|
|
if dist < speed * dt:
|
|
wp[u] = rng.uniform(-box, box, size=2)
|
|
else:
|
|
pos[u] += (speed * dt) * vec / dist
|
|
d2 = (pos ** 2).sum(axis=1)
|
|
a_t[t] = a_max * np.exp(-d2 / (2 * r_scene ** 2))
|
|
return a_t
|
|
|
|
|
|
def pilot_affinity_obs(e_clean, h, rho, rng, a_cap=0.95):
|
|
"""Affinity observation from one orthogonal pilot slot.
|
|
|
|
Each user transmits n_p clean embeddings on an interference-free pilot
|
|
resource; the receiver observes e_u + n/h_u. Off-diagonal Gram entries
|
|
of the pilot observations are unbiased for beta_uv = a_u a_v (independent
|
|
noises), so no bias correction is needed. The share coefficients are
|
|
then the rank-one alternating-least-squares fit of the off-diagonal
|
|
Gram, which uses all pairs jointly."""
|
|
n_p, U, d = e_clean.shape
|
|
sigma = math.sqrt(1.0 / rho)
|
|
y = e_clean + rng.standard_normal(e_clean.shape) * sigma / h[None, :, None]
|
|
G = np.einsum('bud,bvd->buv', y, y).mean(axis=0)
|
|
# rank-1 least-squares fit of the off-diagonal Gram: G_uv ~ a_u a_v.
|
|
# Alternating least squares; uses all pairs jointly (no small-denominator
|
|
# amplification, robust in low-affinity regimes).
|
|
mask = ~np.eye(U, dtype=bool)
|
|
a = np.sqrt(np.clip(np.abs(G[mask]).reshape(U, U - 1).mean(axis=1),
|
|
1e-4, a_cap ** 2))
|
|
for _ in range(20):
|
|
for u in range(U):
|
|
others = [v for v in range(U) if v != u]
|
|
num = sum(G[u, v] * a[v] for v in others)
|
|
den = sum(a[v] ** 2 for v in others) + 1e-9
|
|
a[u] = np.clip(num / den, 0.0, a_cap)
|
|
return a
|
|
|
|
|
|
class AffinityTracker:
|
|
"""EWMA tracker of the per-user share coefficients, driven by sparse
|
|
orthogonal affinity-pilot observations (every K-th slot)."""
|
|
|
|
def __init__(self, U, lam=0.5, a_init=0.3):
|
|
self.a = np.full(U, float(a_init))
|
|
self.lam = lam
|
|
|
|
def update(self, obs):
|
|
self.a = (1 - self.lam) * self.a + self.lam * obs
|
|
return self.a.copy()
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Metrics
|
|
# ----------------------------------------------------------------------
|
|
def metrics(e_hat, e_true, tau=TAU):
|
|
"""(mean cosine, NMSE of raw estimate, SER). e_true unit-norm."""
|
|
nmse = ((e_hat - e_true) ** 2).sum(-1).mean()
|
|
e_n = e_hat / (np.linalg.norm(e_hat, axis=2, keepdims=True) + 1e-12)
|
|
cos = (e_n * e_true).sum(-1)
|
|
return float(cos.mean()), float(nmse), float((cos < tau).mean())
|