# Copyright 2026, IBM Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import numpy as np
from numpy.linalg import svd
import scipy.sparse as sp
import scipy.sparse.linalg as spla
from typing import List, Optional
logger = logging.getLogger(__name__)
[docs]
def fuse_embeddings_attention_fixed(Zs, k, temperature=1.0, mode="effective_rank"):
"""
FIXED: Attention-weighted SVD fusion with information-content-aware similarity.
The original implementation always produced uniform weights because
similarity = ||Z_norm||_F^2 / n = 1 for any row-normalized Z.
"""
weights_raw = []
for Z in Zs:
Z = np.asarray(Z, dtype=np.float64)
if mode == "effective_rank":
try:
_, s, _ = svd(Z, full_matrices=False)
s2 = s ** 2
if s2.sum() <= 0:
weights_raw.append(0.0)
continue
p = s2 / s2.sum()
entropy = -np.sum(p * np.log(p + 1e-20))
weights_raw.append(float(np.exp(entropy)))
except np.linalg.LinAlgError as exc:
# LAPACK did not converge, or Z holds non-finite values. Weight 1.0
# is a neutral pre-softmax score, so this view neither dominates nor
# is dropped -- but it is not a measured effective rank, so it is
# logged rather than silently mixed in with the measured ones.
logger.warning(
"view weighting (%s): SVD failed (%s); using a neutral "
"weight of 1.0 for this view", "effective_rank", exc,
)
weights_raw.append(1.0)
elif mode == "spectral_norm":
try:
_, s, _ = svd(Z, full_matrices=False)
if s.mean() < 1e-12:
weights_raw.append(0.0)
else:
weights_raw.append(float(s.max() / s.mean()))
except np.linalg.LinAlgError as exc:
# LAPACK did not converge, or Z holds non-finite values. Weight 1.0
# is a neutral pre-softmax score, so this view neither dominates nor
# is dropped -- but it is not a measured spectral-norm ratio, so it is
# logged rather than silently mixed in with the measured ones.
logger.warning(
"view weighting (%s): SVD failed (%s); using a neutral "
"weight of 1.0 for this view", "spectral_norm", exc,
)
weights_raw.append(1.0)
elif mode == "variance":
weights_raw.append(float(np.log1p(Z.var())))
else:
raise ValueError(f"Unknown mode={mode!r}")
weights_raw = np.asarray(weights_raw, dtype=np.float64)
centered = (weights_raw - weights_raw.mean()) / temperature
aw = np.exp(centered) / np.sum(np.exp(centered))
Z_weighted = np.concatenate([w * Z for w, Z in zip(aw, Zs)], axis=1)
U, S, _ = svd(Z_weighted, full_matrices=False)
k_eff = min(k, U.shape[1])
Z_fused = U[:, :k_eff] * S[:k_eff]
if k_eff < k:
Z_fused = np.pad(Z_fused, ((0, 0), (0, k - k_eff)), mode="constant")
return Z_fused, aw
[docs]
def per_node_consensus_gate(Zs):
"""Compute per-node reliability score based on agreement with consensus."""
if not Zs:
raise ValueError("Zs is empty")
Zs_arr = np.stack([np.asarray(Z) for Z in Zs], axis=0)
Zbar = Zs_arr.mean(axis=0)
residuals = np.linalg.norm(Zs_arr - Zbar[None, :, :], axis=2)
logits = -residuals.T
logits = logits - logits.max(axis=1, keepdims=True)
weights = np.exp(logits)
weights = weights / weights.sum(axis=1, keepdims=True)
return weights
[docs]
def fuse_embeddings_consensus_gated(Zs, k):
"""SVD fusion with per-node consensus gating."""
n, d = Zs[0].shape
W = per_node_consensus_gate(Zs)
Z_weighted = np.zeros((n, d), dtype=np.float64)
for v, Z in enumerate(Zs):
Z_weighted += W[:, v:v+1] * Z
U, S, _ = svd(Z_weighted, full_matrices=False)
k_eff = min(k, U.shape[1])
Z_fused = U[:, :k_eff] * S[:k_eff]
if k_eff < k:
Z_fused = np.pad(Z_fused, ((0, 0), (0, k - k_eff)), mode="constant")
return Z_fused, W
[docs]
def calibrate_heat_kernel_fast(
L,
targets,
node_to_idx,
t_grid: Optional[np.ndarray] = None,
loss: str = "l2",
):
"""
Vectorized heat kernel calibration. 3-4x faster than naive loop.
Stacks delta vectors and calls expm_multiply once per t.
"""
if t_grid is None:
t_grid = np.logspace(-2, 2, 40)
if len(targets) == 0:
raise ValueError("targets is empty")
n = L.shape[0]
n_targets = len(targets)
# Stack delta vectors
X = np.zeros((n, n_targets), dtype=np.float64)
target_node_idx: List[List[int]] = []
target_p: List[np.ndarray] = []
for j, item in enumerate(targets):
center = item["center"]
nodes = list(item["nodes"])
idx = [node_to_idx[u] for u in nodes]
p = np.asarray(item.get("pQ", item.get("p")), dtype=np.float64)
p = np.maximum(p, 0.0)
s = p.sum()
if s <= 0:
raise ValueError(f"target {j}: pQ has zero mass")
p = p / s
target_node_idx.append(idx)
target_p.append(p)
X[node_to_idx[center], j] = 1.0
best_loss = np.inf
best_t: Optional[float] = None
for t in t_grid:
Y = spla.expm_multiply(-float(t) * L, X) if sp.issparse(L) \
else _dense_expm_multiply(-float(t) * L, X)
total = 0.0
for j in range(n_targets):
q = np.maximum(Y[target_node_idx[j], j], 0.0)
qsum = q.sum()
if qsum <= 0:
continue
q = q / qsum
p = target_p[j]
if loss == "l2":
total += float(np.sum((q - p) ** 2))
elif loss == "kl":
eps = 1e-12
total += float(np.sum(p * (np.log(p + eps) - np.log(q + eps))))
else:
raise ValueError(f"Unknown loss {loss!r}")
if total < best_loss:
best_loss = total
best_t = float(t)
if best_t is None:
raise RuntimeError("Heat kernel calibration produced no valid t")
return best_loss, best_t
def _dense_expm_multiply(M, X):
"""Fallback for dense Laplacian."""
from scipy.linalg import expm
return expm(np.asarray(M)) @ X
[docs]
def calibrate_polynomial_filter_fast(
L,
targets,
node_to_idx,
K: int = 4,
ridge: float = 1e-5,
) -> np.ndarray:
"""
Vectorized polynomial filter calibration. 2-3x faster.
Builds polynomial basis as block matrix multiplications.
"""
if K < 0:
raise ValueError("K must be non-negative")
if len(targets) == 0:
raise ValueError("targets is empty")
n = L.shape[0]
n_targets = len(targets)
# Stack delta vectors
X = np.zeros((n, n_targets), dtype=np.float64)
target_node_idx: List[List[int]] = []
target_p: List[np.ndarray] = []
for j, item in enumerate(targets):
center = item["center"]
nodes = list(item["nodes"])
idx = [node_to_idx[u] for u in nodes]
p = np.asarray(item.get("pQ", item.get("p")), dtype=np.float64)
p = np.maximum(p, 0.0)
s = p.sum()
if s <= 0:
raise ValueError(f"target {j}: pQ has zero mass")
p = p / s
target_node_idx.append(idx)
target_p.append(p)
X[node_to_idx[center], j] = 1.0
# Build basis tensor
AtA = np.zeros((K + 1, K + 1), dtype=np.float64)
Atb = np.zeros(K + 1, dtype=np.float64)
V = X.copy()
powers = [V.copy()]
for _ in range(K):
V = L @ V
powers.append(np.asarray(V).copy())
# Accumulate normal equations
for j in range(n_targets):
idx = target_node_idx[j]
Phi = np.stack([powers[k][idx, j] for k in range(K + 1)], axis=1)
AtA += Phi.T @ Phi
Atb += Phi.T @ target_p[j]
AtA += float(ridge) * np.eye(K + 1)
try:
coeffs = np.linalg.solve(AtA, Atb)
except np.linalg.LinAlgError:
coeffs = np.linalg.lstsq(AtA, Atb, rcond=None)[0]
if not np.all(np.isfinite(coeffs)):
raise RuntimeError("Polynomial calibration produced non-finite coefficients")
return coeffs
