| """Independent DeepONet and exponential vRBA implementation from the paper.""" |
| import math |
| import torch |
| from torch import nn |
|
|
|
|
| class VRBADeepONet(nn.Module): |
| def __init__(self, sensors=200, width=100, hidden_layers=4, latent_dim=100): |
| super().__init__() |
| def network(inputs): |
| layers = [] |
| for _ in range(hidden_layers): |
| layers.extend([nn.Linear(inputs, width), nn.GELU()]) |
| inputs = width |
| layers.append(nn.Linear(width, latent_dim)) |
| return nn.Sequential(*layers) |
| self.branch = network(sensors) |
| self.trunk = network(1) |
|
|
| def forward(self, pressure, time): |
| """Normalized pressure [B,M], normalized time [N,1] -> radius [B,N].""" |
| return self.branch(pressure) @ self.trunk(time).transpose(0, 1) |
|
|
|
|
| class ExponentialVRBA: |
| """Eqs. 19, 40--44: detached EMA weights and function marginal sampling. |
| |
| phi mixes the max-normalized exponential tilt with a uniform floor. |
| The exponential divided by its maximum is evaluated as exp((r-rmax)/eps). |
| """ |
| def __init__(self, functions, points, device, eta=0.1, gamma=0.99, |
| phi=0.9, c=1.0, initial_weight=1.0): |
| self.weights = torch.full((functions, points), initial_weight, device=device) |
| self.eta, self.gamma, self.phi, self.c = eta, gamma, phi, c |
|
|
| @torch.no_grad() |
| def probabilities(self): |
| scores = self.weights.sum(dim=1) |
| return scores / scores.sum() |
|
|
| @torch.no_grad() |
| def update(self, indices, residual, step): |
| r = residual.detach().abs() |
| maximum = r.amax(dim=1, keepdim=True) |
| epsilon = self.c * maximum / math.log(step + 2.0) |
| |
| safe = epsilon.clamp_min(torch.finfo(r.dtype).tiny) |
| tilt = torch.exp((r - maximum) / safe) |
| |
| |
| tilt = self.phi * tilt + (1.0 - self.phi) * tilt.mean(dim=1, keepdim=True) |
| updated = self.gamma * self.weights[indices] + self.eta * tilt |
| self.weights[indices] = updated |
| return updated |
|
|
