| """Independent NOTS experiment model; the general FNO is neuraloperator 1.0.2. |
| |
| No source from the authors' NOTS repository is included here. Paper Algorithm 2 |
| requires a newly initialized, fully trained FNO at every posterior draw. |
| """ |
| import math |
|
|
| import torch |
| from torch import nn |
| from neuralop.models import FNO |
|
|
|
|
| class NOTS(nn.Module): |
| """Scalar permeability [B,1,H,W] -> pressure [B,1,H,W].""" |
|
|
| def __init__(self, **model_config): |
| super().__init__() |
| self.model_config = dict(model_config) |
| self.fno = FNO(**model_config) |
| self.reset_posterior_parameters() |
|
|
| @torch.no_grad() |
| def reset_posterior_parameters(self): |
| """He Gaussian weights, zero biases, including complex spectral weights. |
| |
| Complex normal_ uses the requested total variance across real/imaginary |
| parts, and the regularizer uses the modulus squared of those weights. |
| """ |
| for name, parameter in self.named_parameters(): |
| if 'bias' in name: |
| parameter.zero_() |
| else: |
| fan_in = parameter.shape[1] * math.prod(parameter.shape[2:]) if parameter.ndim > 1 else parameter.numel() |
| parameter.normal_(0.0, math.sqrt(2.0 / max(1, fan_in))) |
|
|
| def forward(self, permeability): |
| if permeability.ndim != 4 or permeability.shape[1] != 1: |
| raise ValueError('NOTS expects permeability [batch,1,height,width]') |
| return self.fno(permeability) |
|
|
| def l2_penalty(self): |
| return sum(p.abs().square().sum() for p in self.parameters() if p.requires_grad) |
|
|
