File size: 1,587 Bytes
5978445 | 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 | """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)
|