File size: 7,927 Bytes
40c9a4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
class AnalogRouter(nn.Module):
def __init__(self, dim, fp_dim=2048, hidden=64):
super().__init__()
self.dim = int(dim)
if self.dim < 2 or int(fp_dim) < 1 or int(hidden) < 1:
raise ValueError('Invalid spectrum/fingerprint/hidden dimension')
self.encoder = nn.Sequential(nn.Linear(int(fp_dim), int(hidden)), nn.SiLU())
self.context = nn.Sequential(nn.Linear(3 * int(hidden) + 4, int(hidden)), nn.SiLU())
self.router = nn.Sequential(nn.Linear(int(hidden) + 3, 32), nn.SiLU(), nn.Linear(32, 3))
self.register_buffer('mz', torch.arange(self.dim, dtype=torch.float32), persistent=False)
def forward(self, fpq, fpr, ref_l1, massq, massr, sim, mode='transport', return_valid=False):
if mode not in ('transport', 'reweight'):
raise ValueError('Unknown router mode: ' + str(mode))
if ref_l1.ndim != 2 or ref_l1.shape[1] != self.dim:
raise ValueError('ref_l1 must have shape [batch, spectrum dimension]')
batch = ref_l1.shape[0]
massq, massr, sim = (massq.reshape(batch), massr.reshape(batch), sim.reshape(batch))
eq, er = (self.encoder(fpq), self.encoder(fpr))
globals_ = torch.stack((sim, massq / self.dim, massr / self.dim, (massq - massr) / self.dim), dim=1)
context = self.context(torch.cat((eq, er, eq - er, globals_), dim=1))
mz = self.mz.to(dtype=ref_l1.dtype).expand(batch, -1)
ref = ref_l1.clamp_min(0)
ref = ref / ref.sum(1, keepdim=True).clamp_min(1e-12)
peak_features = torch.stack((mz / self.dim, (massr[:, None] - mz) / self.dim, torch.sqrt(ref.clamp_min(0))), dim=2)
logits = self.router(torch.cat((context[:, None, :].expand(-1, self.dim, -1), peak_features), dim=2))
log_weights = F.log_softmax(logits, dim=2)
original = self.mz.to(torch.long).expand(batch, -1)
delta = torch.round(massq).to(torch.long) - torch.round(massr).to(torch.long)
shifted = original + delta[:, None] if mode == 'transport' else original
upper = torch.minimum(torch.round(massq).to(torch.long) + 3, torch.full_like(delta, self.dim - 1))
source_valid = original >= 1
copy_valid = source_valid & (original <= upper[:, None])
shift_valid = source_valid & (shifted >= 1) & (shifted <= upper[:, None]) & (shifted < self.dim)
positive = ref > 0
safe_ref = torch.where(positive, ref, torch.ones_like(ref))
log_ref = torch.log(safe_ref)
retained = torch.cat((copy_valid & positive, shift_valid & positive), dim=1)
log_mass = torch.cat((log_ref + log_weights[:, :, 0], log_ref + log_weights[:, :, 1]), dim=1)
log_mass = log_mass.masked_fill(~retained, -float('inf'))
valid = retained.any(dim=1)
safe_log_mass = torch.where(valid[:, None], log_mass, torch.zeros_like(log_mass))
centered_log_mass = safe_log_mass - safe_log_mass.amax(dim=1, keepdim=True)
log_total = torch.logsumexp(centered_log_mass, dim=1, keepdim=True)
normalized_mass = torch.where(retained, torch.exp(centered_log_mass - log_total), torch.zeros_like(safe_log_mass))
copy_mass, shift_mass = normalized_mass.split(self.dim, dim=1)
output = torch.zeros_like(ref)
output.scatter_add_(1, original, copy_mass)
output.scatter_add_(1, shifted.clamp(0, self.dim - 1), shift_mass)
return (output, valid) if return_valid else output
def _router_data(data):
required = ('features', 'masses', 'spectra', 'neighbors', 'similarities')
missing = [key for key in required if key not in data]
if missing:
raise ValueError('Missing router data fields: ' + ', '.join(missing))
features = np.asarray(data['features'], dtype=np.float32)
masses = np.asarray(data['masses'], dtype=np.float32)
spectra = np.asarray(data['spectra'], dtype=np.float32)
neighbors = np.asarray(data['neighbors'], dtype=np.int64)
similarities = np.asarray(data['similarities'], dtype=np.float32)
n = len(features)
if features.ndim != 2 or features.shape[1] < 1 or masses.shape != (n,):
raise ValueError('Invalid fingerprint/mass shapes')
if spectra.ndim != 2 or spectra.shape[0] != n or spectra.shape[1] < 2:
raise ValueError('Invalid spectrum shape')
if neighbors.ndim != 2 or neighbors.shape[0] != n or similarities.shape != neighbors.shape:
raise ValueError('Invalid neighbour/similarity shapes')
if not all((np.isfinite(x).all() for x in (features, masses, spectra, similarities))):
raise ValueError('Router data contain non-finite values')
if (spectra < 0).any() or (masses <= 0).any():
raise ValueError('Spectra must be nonnegative and masses positive')
if (neighbors < -1).any() or (neighbors >= n).any():
raise ValueError('Neighbour indices must be -1 or valid molecular indices')
if (similarities < 0).any() or (similarities > 1 + 1e-06).any():
raise ValueError('Similarities must lie in [0, 1]')
return dict(features=features, masses=masses, spectra=spectra, neighbors=neighbors, similarities=similarities)
def _router_indices(indices, n, label):
x = np.asarray(indices, dtype=np.int64).reshape(-1)
if (x < 0).any() or (x >= n).any() or len(np.unique(x)) != len(x):
raise ValueError(label + ' must contain distinct valid molecular indices')
return x
def _router_eligible(data, indices, min_similarity):
refs = data['neighbors'][indices]
sims = data['similarities'][indices]
eligible = (refs >= 0) & (sims >= float(min_similarity))
eligible &= refs != indices[:, None]
safe = np.maximum(refs, 0)
eligible &= data['spectra'][safe].sum(2) > 1e-12
return eligible
def _router_forward_pairs(model, data, queries, refs, sims, device, mode):
tensor = lambda x: torch.as_tensor(x, dtype=torch.float32, device=device)
return model(tensor(data['features'][queries]), tensor(data['features'][refs]), tensor(data['spectra'][refs]), tensor(data['masses'][queries]), tensor(data['masses'][refs]), tensor(sims), mode=mode, return_valid=True)
def predict_router(model, data, query_indices, device='cpu', mode='transport', batch_size=64, min_similarity=0.35):
data = _router_data(data)
queries = _router_indices(query_indices, len(data['features']), 'query_indices')
if int(batch_size) < 1:
raise ValueError('batch_size must be positive')
eligible = _router_eligible(data, queries, min_similarity)
where_q, where_k = np.nonzero(eligible)
result = np.zeros((len(queries), data['spectra'].shape[1]), dtype=np.float32)
weight_sum = np.zeros(len(queries), dtype=np.float64)
previous_training = model.training
model.to(device)
model.eval()
try:
with torch.no_grad():
for start in range(0, len(where_q), int(batch_size)):
positions = where_q[start:start + int(batch_size)]
columns = where_k[start:start + int(batch_size)]
q = queries[positions]
r = data['neighbors'][q, columns]
sims = data['similarities'][q, columns]
p, valid = _router_forward_pairs(model, data, q, r, sims, device, mode)
p = p.cpu().numpy()
weights = np.power(sims.astype(np.float64), 4) * valid.cpu().numpy()
np.add.at(result, positions, p * weights[:, None].astype(np.float32))
np.add.at(weight_sum, positions, weights)
finally:
model.train(previous_training)
covered = weight_sum > 1e-12
result[covered] /= weight_sum[covered, None].astype(np.float32)
totals = result.sum(1, dtype=np.float64)
covered &= totals > 1e-12
result[covered] /= totals[covered, None].astype(np.float32)
result[~covered] = 0
return (result, covered)
|