| 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) |
|
|
