| import torch |
| import math |
| import torch.nn as nn |
| import numpy as np |
| import torch.nn.functional as F |
| from timm.layers import trunc_normal_ |
| from onescience.modules.fourier.MultiWaveletTransform import ( |
| MultiWaveletTransform1D, |
| MultiWaveletTransform2D, |
| MultiWaveletTransform3D, |
| ) |
| from onescience.modules.mlp.MLP import StandardMLP |
| from onescience.modules.embedding import timestep_embedding, unified_pos_embedding |
| from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI |
|
|
| MultiWaveletTransformList = [ |
| None, |
| MultiWaveletTransform1D, |
| MultiWaveletTransform2D, |
| MultiWaveletTransform3D, |
| ] |
|
|
| class Model(nn.Module): |
| |
| def __init__( |
| self, args, device, alpha=2, L=0, c=1, base="legendre", s1=128, s2=128 |
| ): |
| super(Model, self).__init__() |
| self.__name__ = "MWT" |
| self.args = args |
| self.k = args.mwt_k |
| self.WMT_dim = c * self.k**2 |
| if args.geotype == "structured_1D": |
| self.WMT_dim = c * self.k |
| self.c = c |
| self.s1 = s1 |
| self.s2 = s2 |
| |
| |
| if ( |
| args.unified_pos and args.geotype != "unstructured" |
| ): |
| self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device) |
| self.preprocess = StandardMLP( |
| input_dim=args.fun_dim + args.ref ** len(args.shapelist), |
| output_dim=self.WMT_dim, |
| hidden_dims=[args.n_hidden * 2], |
| activation=args.act, |
| use_bias=True |
| ) |
| else: |
| self.preprocess = StandardMLP( |
| input_dim=args.fun_dim + args.space_dim, |
| output_dim=self.WMT_dim, |
| hidden_dims=[args.n_hidden * 2], |
| activation=args.act, |
| use_bias=True |
| ) |
| |
| if args.time_input: |
| self.time_fc = nn.Sequential( |
| nn.Linear(self.WMT_dim, args.n_hidden), |
| nn.SiLU(), |
| nn.Linear(args.n_hidden, self.WMT_dim), |
| ) |
| |
| |
| if self.args.geotype == "unstructured": |
| self.fftproject_in = GeoSpectralConv2d( |
| in_channels=self.WMT_dim, |
| out_channels=self.WMT_dim, |
| modes1=args.modes, |
| modes2=args.modes, |
| s1=s1, |
| s2=s2 |
| ) |
| self.fftproject_out = GeoSpectralConv2d( |
| in_channels=self.WMT_dim, |
| out_channels=self.WMT_dim, |
| modes1=args.modes, |
| modes2=args.modes, |
| s1=s1, |
| s2=s2 |
| ) |
| self.iphi = IPHI() |
| self.augmented_resolution = [s1, s2] |
| self.padding = [(16 - size % 16) % 16 for size in [s1, s2]] |
| else: |
| target = 2 ** (math.ceil(np.log2(max(args.shapelist)))) |
| self.padding = [(target - size) for size in args.shapelist] |
| self.augmented_resolution = [target for _ in range(len(self.padding))] |
|
|
| dim = len(self.padding) |
| transform_class = MultiWaveletTransformList[dim] |
|
|
| self.spectral_layers = nn.ModuleList( |
| [ |
| transform_class( |
| k=self.k, |
| alpha=alpha, |
| L=L, |
| c=c, |
| base=base |
| ) |
| for _ in range(args.n_layers) |
| ] |
| ) |
| |
| |
| self.fc1 = nn.Linear(self.WMT_dim, args.n_hidden) |
| self.fc2 = nn.Linear(args.n_hidden, args.out_dim) |
|
|
| def structured_geo(self, x, fx, T=None): |
| B, N, _ = x.shape |
| if self.args.unified_pos: |
| x = self.pos.repeat(x.shape[0], 1, 1) |
| if fx is not None: |
| fx = torch.cat((x, fx), -1) |
| fx = self.preprocess(fx) |
| else: |
| fx = self.preprocess(x) |
|
|
| if T is not None: |
| Time_emb = timestep_embedding(T, self.WMT_dim) |
| Time_emb = self.time_fc(Time_emb) |
| if Time_emb.ndim == 2: |
| Time_emb = Time_emb.unsqueeze(1) |
| fx = fx + Time_emb |
| |
| x = fx.permute(0, 2, 1).reshape(B, self.WMT_dim, *self.args.shapelist) |
| if not all(item == 0 for item in self.padding): |
| if len(self.args.shapelist) == 2: |
| x = F.pad(x, [0, self.padding[1], 0, self.padding[0]]) |
| elif len(self.args.shapelist) == 3: |
| x = F.pad( |
| x, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]] |
| ) |
| x = ( |
| x.reshape(B, self.WMT_dim, -1) |
| .permute(0, 2, 1) |
| .contiguous() |
| .reshape( |
| B, |
| *self.augmented_resolution, |
| self.c, |
| self.k**2 if self.args.geotype != "structured_1D" else self.k |
| ) |
| ) |
| for i in range(self.args.n_layers): |
| x = self.spectral_layers[i](x) |
| if i < self.args.n_layers - 1: |
| x = F.gelu(x) |
| x = ( |
| x.reshape(B, -1, self.WMT_dim) |
| .permute(0, 2, 1) |
| .contiguous() |
| .reshape(B, self.WMT_dim, *self.augmented_resolution) |
| ) |
| if not all(item == 0 for item in self.padding): |
| if len(self.args.shapelist) == 2: |
| x = x[..., : -self.padding[0], : -self.padding[1]] |
| elif len(self.args.shapelist) == 3: |
| x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]] |
| x = x.reshape(B, self.WMT_dim, -1).permute(0, 2, 1) |
| x = self.fc1(x) |
| x = F.gelu(x) |
| x = self.fc2(x) |
| return x |
|
|
| def unstructured_geo(self, x, fx, T=None): |
| B, N, _ = x.shape |
| original_pos = x |
| if fx is not None: |
| fx = torch.cat((x, fx), -1) |
| fx = self.preprocess(fx) |
| else: |
| fx = self.preprocess(x) |
|
|
| if T is not None: |
| Time_emb = timestep_embedding(T, self.WMT_dim) |
| Time_emb = self.time_fc(Time_emb) |
| if Time_emb.ndim == 2: |
| Time_emb = Time_emb.unsqueeze(1) |
| fx = fx + Time_emb |
|
|
| x = self.fftproject_in( |
| fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None |
| ) |
| x = ( |
| x.reshape(B, self.WMT_dim, -1) |
| .permute(0, 2, 1) |
| .contiguous() |
| .reshape(B, *self.augmented_resolution, self.c, self.k**2) |
| ) |
| for i in range(self.args.n_layers): |
| x = self.spectral_layers[i](x) |
| if i < self.args.n_layers - 1: |
| x = F.gelu(x) |
| x = ( |
| x.reshape(B, -1, self.WMT_dim) |
| .permute(0, 2, 1) |
| .contiguous() |
| .reshape(B, self.WMT_dim, *self.augmented_resolution) |
| ) |
| x = self.fftproject_out( |
| x, x_out=original_pos, iphi=self.iphi, code=None |
| ).permute(0, 2, 1) |
| x = self.fc1(x) |
| x = F.gelu(x) |
| x = self.fc2(x) |
| return x |
|
|
| def forward(self, x, fx, T=None, geo=None): |
| if self.args.geotype == "unstructured": |
| return self.unstructured_geo(x, fx, T) |
| else: |
| return self.structured_geo(x, fx, T) |
|
|
