├── README.md ├── allen-cahn ├── E1 │ ├── data │ │ ├── generate_data.py │ │ └── initial_field.py │ └── mcnp100 │ │ ├── initial_field.py │ │ ├── main.py │ │ ├── mc_loss.py │ │ ├── model.py │ │ ├── tools.py │ │ └── train.py ├── E2 │ ├── data │ │ ├── generate_data.py │ │ └── initial_field.py │ └── mcnp100 │ │ ├── initial_field.py │ │ ├── main.py │ │ ├── mc_loss.py │ │ ├── model.py │ │ ├── tools.py │ │ └── train.py ├── E3 │ ├── data │ │ ├── generate_data.py │ │ └── initial_field.py │ └── mcnp100 │ │ ├── initial_field.py │ │ ├── main.py │ │ ├── mc_loss.py │ │ ├── model.py │ │ ├── tools.py │ │ └── train.py └── E4 │ ├── data │ ├── generate_data.py │ └── initial_field.py │ └── mcnp100 │ ├── initial_field.py │ ├── main.py │ ├── mc_loss.py │ ├── model.py │ ├── tools.py │ └── train.py ├── convection_diffusion ├── N10 │ ├── data │ │ ├── config_pde.py │ │ └── generate_data.py │ └── mcnp10 │ │ ├── config_pde.py │ │ ├── fk_loss.py │ │ ├── main.py │ │ ├── model.py │ │ ├── tools.py │ │ └── train.py └── N5 │ ├── data │ ├── config_pde.py │ └── generate_data.py │ └── mcnp10 │ ├── config_pde.py │ ├── main.py │ ├── mc_loss.py │ ├── model.py │ ├── tools.py │ └── train.py └── navier-stokes ├── E1 ├── data │ ├── generate_data.py │ └── nse.py └── mcnp10 │ ├── main.py │ ├── mc_loss.py │ ├── model.py │ ├── tools.py │ └── train.py ├── E2 ├── data │ ├── generate_data.py │ └── nse.py └── mcnp10 │ ├── main.py │ ├── mc_loss.py │ ├── model.py │ ├── tools.py │ └── train.py ├── E3 ├── data │ ├── generate_data.py │ └── nse.py └── mcnp10 │ ├── main.py │ ├── mc_loss.py │ ├── model.py │ ├── tools.py │ └── train.py └── E4 ├── data ├── generate_data.py └── nse.py └── mcnp10 ├── main.py ├── mc_loss.py ├── model.py ├── tools.py └── train.py /README.md: -------------------------------------------------------------------------------- 1 | # Monte Carlo Neural PDE Solver (MCNP Solver) 2 | 3 | This repository contains the code for the paper 4 | - Monte Carlo Neural PDE Solver for Learning PDEs via Probabilistic Representation (IEEE TPAMI 2025) 5 | 6 | ## Quick Start 7 | 8 | To generate the data-set, use, 9 | ```bash 10 | python generate_data.py 11 | ``` 12 | To train MCNP, use 13 | ```bash 14 | python main.py 15 | ``` 16 | 17 | ## Citation 18 | 19 | If you find our work useful in your research, please consider citing: 20 | ``` 21 | @ARTICLE{10916840, 22 | author={Zhang, Rui and Meng, Qi and Zhu, Rongchan and Wang, Yue and Shi, Wenlei and Zhang, Shihua and Ma, Zhi-Ming and Liu, Tie-Yan}, 23 | journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 24 | title={Monte Carlo Neural PDE Solver for Learning PDEs Via Probabilistic Representation}, 25 | year={2025}, 26 | volume={}, 27 | number={}, 28 | pages={1-18}, 29 | doi={10.1109/TPAMI.2025.3548673}} 30 | ``` 31 | 32 | If you have any questions, please feel free to contact me via: rayzhang@amss.ac.cn 33 | -------------------------------------------------------------------------------- /allen-cahn/E1/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import os 4 | import random 5 | import numpy as np 6 | from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph 7 | from initial_field import GaussianRF 8 | 9 | 10 | def setup_seed(seed): 11 | torch.manual_seed(seed) 12 | torch.cuda.manual_seed(seed) 13 | torch.cuda.manual_seed_all(seed) 14 | np.random.seed(seed) 15 | random.seed(seed) 16 | os.environ['PYTHONHASHSEED'] = str(seed) 17 | torch.backends.cudnn.deterministic = True 18 | torch.backends.cudnn.benchmark = False 19 | torch.backends.cudnn.enabled = True 20 | 21 | 22 | setup_seed(0) 23 | size_x = 1024 24 | size_t = 200 25 | T = 1.0 26 | delta_t = 1e-6 27 | num_test = 200 28 | num_val = 200 29 | 30 | def generate_u0(b_size, N, num_grid): 31 | ''' 32 | Input: 33 | b_size: int, num of pre-generated data 34 | N: int, highest frequency fourier basis; 35 | num_grid: int, number of grids; 36 | Output: 37 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 38 | ''' 39 | B_n = torch.rand(b_size, N).reshape(b_size, N, 1, 1) 40 | grid = torch.tensor(np.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid)).reshape(1, 1, -1, 1) 41 | n = 2 * torch.Tensor(range(1, N + 1)).reshape(1, -1, 1, 1) 42 | sin_grid = torch.sin(torch.pi * n * grid) 43 | return (B_n * sin_grid).sum(axis=1) 44 | 45 | setup_seed(1) 46 | N = 5 47 | size_x = 65 48 | size_t = 100 49 | T = 1.0 50 | data_test = np.zeros([num_test, size_t+1, size_x]) 51 | initial = GaussianRF(1024) 52 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 53 | field = ScalarField(grid, 2) 54 | for i in range(num_test): 55 | maxvalue = np.nan 56 | while np.isnan(maxvalue) or maxvalue>10: 57 | bc_x_left = {"value": 0} 58 | bc_x_right = {"value": 0} 59 | term_1 = f"laplace(c) * 0.01 + c - c**3" 60 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 61 | b = generate_u0(1, N, 1024) 62 | field.data = np.array(b).reshape(-1,) 63 | storage = MemoryStorage() # store intermediate information of the simulation 64 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 65 | a = torch.tensor(storage.data) 66 | a = (a[:, 15:-1:16] + a[:, 16::16])/2 67 | data_test[i, :, 1:-1] = a 68 | maxvalue = np.abs(data_test[i]).max() 69 | print(i, np.abs(data_test[i]).max()) 70 | 71 | data_test = torch.tensor(data_test) 72 | torch.save(data_test, 'dataset/data_test') 73 | 74 | 75 | setup_seed(2) 76 | N = 5 77 | size_x = 65 78 | size_t = 100 79 | T = 1.0 80 | data_val = np.zeros([num_val, size_t+1, size_x]) 81 | initial = GaussianRF(1024) 82 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 83 | field = ScalarField(grid, 2) 84 | for i in range(num_val): 85 | maxvalue = np.nan 86 | while np.isnan(maxvalue) or maxvalue>10: 87 | bc_x_left = {"value": 0} 88 | bc_x_right = {"value": 0} 89 | term_1 = f"laplace(c) * 0.01 + c - c**3" 90 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 91 | b = generate_u0(1, N, 1024) 92 | field.data = np.array(b).reshape(-1,) 93 | storage = MemoryStorage() # store intermediate information of the simulation 94 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 95 | a = torch.tensor(storage.data) 96 | a = (a[:, 15:-1:16] + a[:, 16::16])/2 97 | data_val[i, :, 1:-1] = a 98 | maxvalue = np.abs(data_val[i]).max() 99 | print(i, np.abs(data_val[i]).max()) 100 | 101 | data_val = torch.tensor(data_val) 102 | torch.save(data_val, 'dataset/data_val') -------------------------------------------------------------------------------- /allen-cahn/E1/data/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E1/mcnp100/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E1/mcnp100/main.py: -------------------------------------------------------------------------------- 1 | 2 | from model import FNO 3 | from train import train 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def setup_seed(seed): 21 | torch.manual_seed(seed) 22 | torch.cuda.manual_seed(seed) 23 | torch.cuda.manual_seed_all(seed) 24 | np.random.seed(seed) 25 | random.seed(seed) 26 | os.environ['PYTHONHASHSEED'] = str(seed) 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = True 29 | torch.backends.cudnn.enabled = True 30 | 31 | 32 | 33 | def main(cfg): 34 | if not os.path.exists(f'log'): 35 | os.mkdir(f'log') 36 | if not os.path.exists(f'model'): 37 | os.mkdir(f'model') 38 | setup_seed(cfg.seed) 39 | dateTimeObj = datetime.now() 40 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}_{dateTimeObj.time().second}' 41 | cfg.timestring = timestring 42 | logfile = f'log/z0seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{cfg.weight_decay}_{timestring}.csv' 43 | 44 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 45 | json.dump(cfg.__dict__, f, indent=2) 46 | 47 | sys.stdout = open(logfile, 'w') 48 | 49 | print('--------args----------') 50 | for k in list(vars(cfg).keys()): 51 | print('%s: %s' % (k, vars(cfg)[k])) 52 | print('--------args----------\n') 53 | 54 | sys.stdout.flush() 55 | 56 | net = FNO(cfg.modes, 32) 57 | 58 | cfg.delta_t = cfg.T/cfg.time_steps 59 | train(cfg, net) 60 | torch.save(net.state_dict(), f'model/net_seed_{cfg.seed}_{timestring}.pt') 61 | 62 | 63 | if __name__ == "__main__": 64 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 65 | 66 | parser.add_argument('--data_path', type=str, 67 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 68 | help='path of data') 69 | 70 | parser.add_argument('--device', type=str, default='cuda:3', 71 | help='Used device') 72 | 73 | parser.add_argument('--seed', type=int, default=0, 74 | help='seed') 75 | 76 | parser.add_argument('--batch_size', type=int, default=200, 77 | help='batchsize of the operator learning') 78 | 79 | parser.add_argument('--step_size', type=int, default=2000, 80 | help='step_size of optim') 81 | 82 | parser.add_argument('--gamma', type=float, default=0.5, 83 | help='gamma of optim') 84 | 85 | parser.add_argument('--lr', type=float, default=0.01, 86 | help='lr of optim') 87 | 88 | parser.add_argument('--weight_decay', type=float, default=0.0, 89 | help='lr of optim') 90 | 91 | parser.add_argument('--num_iterations', type=int, default=20000, 92 | help='num_iterations of optim') 93 | 94 | parser.add_argument('--size', type=int, default=64, 95 | help='data spatial size') 96 | 97 | parser.add_argument('--T', type=float, default=1.0, 98 | help='final time') 99 | 100 | parser.add_argument('--time_steps', type=int, default=100, 101 | help='number of time_steps in data') 102 | 103 | parser.add_argument('--EPS', type=float, default=1e-8, 104 | help='epsilon') 105 | 106 | parser.add_argument('--lamb', type=float, default=1, 107 | help='lamb') 108 | 109 | parser.add_argument('--modes', type=int, default=16, 110 | help='modes in fno') 111 | 112 | cfg = parser.parse_args() 113 | for seed in [0, 1, 2]: 114 | cfg.seed = seed 115 | main(cfg) 116 | -------------------------------------------------------------------------------- /allen-cahn/E1/mcnp100/mc_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import math 3 | import copy 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from scipy.stats import levy_stable 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def mc_loss(u_0, model, p, config): 17 | device = config.device 18 | size = config.size 19 | batch_size = config.batch_size 20 | time_steps = config.time_steps 21 | delta_t = config.delta_t 22 | 23 | u_input = u_0.reshape(u_0.shape[0], 1, size+1).repeat(1, time_steps, 1).reshape(-1, size+1, 1) 24 | t = delta_t * torch.tensor(range(1, time_steps+1)).to(device).reshape(1, -1).repeat(u_0.shape[0], 1).reshape(-1,) 25 | u = model(u_input, t).reshape(batch_size, -1, size+1) 26 | u = torch.concat([u_0.reshape(batch_size, 1, size+1), u], dim=1) 27 | 28 | u_0, u_1 = u[:, :-1, :], u[:, 1:, :] 29 | f_temp = 0.5*(u_0+u_1) 30 | f_temp = f_temp - f_temp**3 31 | u_hat = torch.einsum('bts, gs->btg', u_0 + f_temp* delta_t, p) 32 | return torch.sqrt(torch.mean(torch.square(u_hat - u_1))) -------------------------------------------------------------------------------- /allen-cahn/E1/mcnp100/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.parameter import Parameter 6 | import matplotlib.pyplot as plt 7 | 8 | import operator 9 | from functools import reduce 10 | from functools import partial 11 | from timeit import default_timer 12 | 13 | 14 | ################################################################ 15 | # 1d fourier layer 16 | ################################################################ 17 | class SpectralConv1d(nn.Module): 18 | def __init__(self, in_channels, out_channels, modes1): 19 | super(SpectralConv1d, self).__init__() 20 | 21 | """ 22 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 23 | """ 24 | 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 28 | 29 | self.scale = (1 / (in_channels*out_channels)) 30 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 31 | 32 | # Complex multiplication 33 | def compl_mul1d(self, input, weights): 34 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 35 | return torch.einsum("bix,iox->box", input, weights) 36 | 37 | def forward(self, x): 38 | batchsize = x.shape[0] 39 | #Compute Fourier coeffcients up to factor of e^(- something constant) 40 | x_ft = torch.fft.rfft(x) 41 | 42 | # Multiply relevant Fourier modes 43 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 44 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 45 | 46 | #Return to physical space 47 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 48 | return x 49 | 50 | class FNO(nn.Module): 51 | def __init__(self, modes, width): 52 | super(FNO, self).__init__() 53 | 54 | """ 55 | The overall network. It contains 4 layers of the Fourier layer. 56 | 1. Lift the input to the desire channel dimension by self.fc0 . 57 | 2. 4 layers of the integral operators u' = (W + K)(u). 58 | W defined by self.w; K defined by self.conv . 59 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 60 | 61 | input: the solution of the initial condition and location (a(x), x) 62 | input shape: (batchsize, x=s, c=2) 63 | output: the solution of a later timestep 64 | output shape: (batchsize, x=s, c=1) 65 | """ 66 | 67 | self.modes1 = modes 68 | self.width = width 69 | self.fc0 = nn.Linear(3, self.width) # input channel is 2: (a(x), x) 70 | 71 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 74 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 75 | self.w0 = nn.Conv1d(self.width, self.width, 1) 76 | self.w1 = nn.Conv1d(self.width, self.width, 1) 77 | self.w2 = nn.Conv1d(self.width, self.width, 1) 78 | self.w3 = nn.Conv1d(self.width, self.width, 1) 79 | 80 | self.fc1 = nn.Linear(self.width, 128) 81 | self.fc2 = nn.Linear(128, 1) 82 | 83 | def forward(self, x, t): 84 | grid = self.get_grid(x.shape, x.device) 85 | time = t.reshape(-1, 1, 1).repeat(1, x.shape[1], 1) 86 | x = torch.cat((x, grid), dim=-1) 87 | x = torch.concat([x, time], dim=-1) 88 | x = self.fc0(x) 89 | x = x.permute(0, 2, 1) 90 | 91 | x1 = self.conv0(x) 92 | x2 = self.w0(x) 93 | x = x1 + x2 94 | x = F.gelu(x) 95 | 96 | x1 = self.conv1(x) 97 | x2 = self.w1(x) 98 | x = x1 + x2 99 | x = F.gelu(x) 100 | 101 | x1 = self.conv2(x) 102 | x2 = self.w2(x) 103 | x = x1 + x2 104 | x = F.gelu(x) 105 | 106 | x1 = self.conv3(x) 107 | x2 = self.w3(x) 108 | x = x1 + x2 109 | x = x.permute(0, 2, 1) 110 | x = self.fc1(x) 111 | x = F.gelu(x) 112 | x = self.fc2(x) 113 | return x 114 | 115 | def get_grid(self, shape, device): 116 | batchsize, size_x = shape[0], shape[1] 117 | gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) 118 | gridx = gridx.reshape(1, size_x, 1).repeat([batchsize, 1, 1]) 119 | return gridx.to(device) -------------------------------------------------------------------------------- /allen-cahn/E1/mcnp100/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_sol(sol): 11 | ''' 12 | sol: torch.Tensor, size = (num_grid, ) 13 | return: figs 14 | ''' 15 | temp = sol.detach().cpu().numpy().reshape(-1, ) 16 | num_grid = temp.shape[0] - 1 17 | x_axis = (1 / num_grid) * torch.Tensor(range(num_grid + 1)) 18 | plt.figure() 19 | plt.plot(x_axis.numpy(), temp.numpy()) 20 | plt.show() 21 | 22 | 23 | def setup_seed(seed): 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | np.random.seed(seed) 28 | random.seed(seed) 29 | os.environ['PYTHONHASHSEED'] = str(seed) 30 | torch.backends.cudnn.deterministic = True 31 | torch.backends.cudnn.benchmark = False 32 | torch.backends.cudnn.enabled = True -------------------------------------------------------------------------------- /allen-cahn/E1/mcnp100/train.py: -------------------------------------------------------------------------------- 1 | from mc_loss import mc_loss 2 | from initial_field import GaussianRF 3 | import sys 4 | import math 5 | import copy 6 | import random 7 | import numpy as np 8 | import torch 9 | import torch.optim as optim 10 | import matplotlib.pyplot as plt 11 | from scipy.stats import norm 12 | import os 13 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 14 | plt.rcParams["animation.html"] = "jshtml" 15 | 16 | 17 | def test(config, net, test_data): 18 | device = config.device 19 | num = test_data.shape[0] 20 | delta_steps = 10 21 | delta_t = config.T / delta_steps 22 | u_0 = test_data[:, 0, :].to(device)[:, :, None] 23 | rela_err = torch.zeros(delta_steps) 24 | rela_err_1 = torch.zeros(delta_steps) 25 | rela_err_max = torch.zeros(delta_steps) 26 | for time_step in range(1, 11): 27 | net.eval() 28 | t = (time_step*delta_t) * torch.ones(num).to(device) 29 | u = net(u_0, t).detach() 30 | u_t = test_data[:, delta_steps * time_step, :].to(device)[:, :, None] 31 | rela_err[time_step-1] = (torch.norm((u - u_t).reshape(u.shape[0], -1), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), dim=1)).mean() 32 | rela_err_1[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=1, dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=1, dim=1)).mean() 33 | rela_err_max[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=float('inf'), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=float('inf'), dim=1)).mean() 34 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 35 | print('mean relative l_2 error', rela_err.mean().item()) 36 | print('mean relative l_1 error', rela_err_1.mean().item()) 37 | print('mean relative l_inf error', rela_err_max.mean().item()) 38 | return rela_err.mean().item() 39 | 40 | 41 | def p_matrix(config): 42 | size_x = config.size 43 | dt = 1/100 44 | p = np.zeros([size_x+1, size_x+1]) 45 | x = np.linspace(0, 1, size_x+1) 46 | dx = 1/size_x 47 | kappa = 0.01 48 | sigma = math.sqrt(2 * kappa * dt) 49 | for i in range(1, size_x): 50 | if x[i] <= 2/3: 51 | d = x[i] 52 | p[i, :] = (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-x+d)**2)/(2 * sigma**2))) * dx - (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-x+d-2*d)**2)/(2 * sigma**2))) * dx 53 | if x[i] > 2/3: 54 | p[i, :] = np.flip(p[size_x-i, :]) 55 | return torch.tensor(p).float() 56 | 57 | 58 | def generate_u0(b_size, N, num_grid, device): 59 | ''' 60 | Input: 61 | b_size: int, num of pre-generated data 62 | N: int, highest frequency fourier basis; 63 | num_grid: int, number of grids; 64 | Output: 65 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 66 | ''' 67 | B_n = torch.rand(b_size, N, device=device).reshape(b_size, N, 1, 1) 68 | grid = torch.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid, device=device).reshape(1, 1, -1, 1) 69 | n = 2 * torch.range(1, N, device=device).reshape(1, -1, 1, 1) 70 | sin_grid = torch.sin(torch.pi * n * grid) 71 | return (B_n * sin_grid).sum(axis=1) 72 | 73 | 74 | def train(config, net): 75 | N = 5 76 | device = config.device 77 | size = config.size 78 | batch_size = config.batch_size 79 | net = net.to(device) 80 | optimizer = optim.Adam(net.parameters(), config.lr) 81 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 82 | val_data = torch.load(config.data_path + 'data_test').to(device).float() 83 | test_data = torch.load(config.data_path + 'data_test').to(device).float() 84 | val_loss = 1e50 85 | p = p_matrix(config).to(device) 86 | for step in range(config.num_iterations+1): 87 | u_0 = generate_u0(batch_size, N, 1024, device) 88 | u_0 = (u_0[:, 15:-1:16] + u_0[:, 16::16])/2 89 | u_0 = torch.concat([torch.zeros(batch_size, 1, 1, device=device), u_0, torch.zeros(batch_size, 1, 1, device=device)], dim=1) 90 | net.train() 91 | loss = mc_loss(u_0, net, p, config) 92 | loss.backward() 93 | optimizer.step() 94 | optimizer.zero_grad(set_to_none=True) 95 | scheduler.step() 96 | if step % 50 == 0: 97 | print('########################') 98 | print('training loss', step, loss.detach().item()) 99 | print('VAL_LOSS') 100 | print('########################') 101 | temp = test(config, net, val_data) 102 | if temp < val_loss: 103 | val_loss = temp 104 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 105 | print('TEST_LOSS') 106 | print('########################') 107 | test(config, net, test_data) 108 | sys.stdout.flush() 109 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 110 | print('FINAL_LOSS') 111 | print('########################') 112 | test(config, net, test_data) 113 | sys.stdout.flush() -------------------------------------------------------------------------------- /allen-cahn/E2/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import os 4 | import random 5 | import numpy as np 6 | from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph 7 | from initial_field import GaussianRF 8 | 9 | 10 | def setup_seed(seed): 11 | torch.manual_seed(seed) 12 | torch.cuda.manual_seed(seed) 13 | torch.cuda.manual_seed_all(seed) 14 | np.random.seed(seed) 15 | random.seed(seed) 16 | os.environ['PYTHONHASHSEED'] = str(seed) 17 | torch.backends.cudnn.deterministic = True 18 | torch.backends.cudnn.benchmark = False 19 | torch.backends.cudnn.enabled = True 20 | 21 | 22 | setup_seed(0) 23 | size_x = 1024 24 | size_t = 200 25 | T = 1.0 26 | delta_t = 1e-6 27 | num_test = 200 28 | num_val = 200 29 | 30 | def generate_u0(b_size, N, num_grid): 31 | ''' 32 | Input: 33 | b_size: int, num of pre-generated data 34 | N: int, highest frequency fourier basis; 35 | num_grid: int, number of grids; 36 | Output: 37 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 38 | ''' 39 | B_n = torch.rand(b_size, N).reshape(b_size, N, 1, 1) 40 | grid = torch.tensor(np.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid)).reshape(1, 1, -1, 1) 41 | n = 2 * torch.Tensor(range(1, N + 1)).reshape(1, -1, 1, 1) 42 | sin_grid = torch.sin(torch.pi * n * grid) 43 | return (B_n * sin_grid).sum(axis=1) 44 | 45 | setup_seed(1) 46 | N = 10 47 | size_x = 65 48 | size_t = 100 49 | T = 1.0 50 | data_test = np.zeros([num_test, size_t+1, size_x]) 51 | initial = GaussianRF(1024) 52 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 53 | field = ScalarField(grid, 2) 54 | for i in range(num_test): 55 | maxvalue = np.nan 56 | while np.isnan(maxvalue) or maxvalue>10: 57 | bc_x_left = {"value": 0} 58 | bc_x_right = {"value": 0} 59 | term_1 = f"laplace(c) * 0.01 + c - c**3" 60 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 61 | b = generate_u0(1, N, 1024) 62 | field.data = np.array(b).reshape(-1,) 63 | storage = MemoryStorage() # store intermediate information of the simulation 64 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 65 | a = torch.tensor(storage.data) 66 | a = (a[:, 15:-1:16] + a[:, 16::16])/2 67 | data_test[i, :, 1:-1] = a 68 | maxvalue = np.abs(data_test[i]).max() 69 | print(i, np.abs(data_test[i]).max()) 70 | 71 | data_test = torch.tensor(data_test) 72 | torch.save(data_test, 'dataset/data_test') 73 | 74 | 75 | setup_seed(2) 76 | N = 10 77 | size_x = 65 78 | size_t = 100 79 | T = 1.0 80 | data_val = np.zeros([num_val, size_t+1, size_x]) 81 | initial = GaussianRF(1024) 82 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 83 | field = ScalarField(grid, 2) 84 | for i in range(num_val): 85 | maxvalue = np.nan 86 | while np.isnan(maxvalue) or maxvalue>10: 87 | bc_x_left = {"value": 0} 88 | bc_x_right = {"value": 0} 89 | term_1 = f"laplace(c) * 0.01 + c - c**3" 90 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 91 | b = generate_u0(1, N, 1024) 92 | field.data = np.array(b).reshape(-1,) 93 | storage = MemoryStorage() # store intermediate information of the simulation 94 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 95 | a = torch.tensor(storage.data) 96 | a = (a[:, 15:-1:16] + a[:, 16::16])/2 97 | data_val[i, :, 1:-1] = a 98 | maxvalue = np.abs(data_val[i]).max() 99 | print(i, np.abs(data_val[i]).max()) 100 | 101 | data_val = torch.tensor(data_val) 102 | torch.save(data_val, 'dataset/data_val') -------------------------------------------------------------------------------- /allen-cahn/E2/data/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E2/mcnp100/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E2/mcnp100/main.py: -------------------------------------------------------------------------------- 1 | 2 | from model import FNO 3 | from train import train 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def setup_seed(seed): 21 | torch.manual_seed(seed) 22 | torch.cuda.manual_seed(seed) 23 | torch.cuda.manual_seed_all(seed) 24 | np.random.seed(seed) 25 | random.seed(seed) 26 | os.environ['PYTHONHASHSEED'] = str(seed) 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = True 29 | torch.backends.cudnn.enabled = True 30 | 31 | 32 | 33 | def main(cfg): 34 | if not os.path.exists(f'log'): 35 | os.mkdir(f'log') 36 | if not os.path.exists(f'model'): 37 | os.mkdir(f'model') 38 | setup_seed(cfg.seed) 39 | dateTimeObj = datetime.now() 40 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}_{dateTimeObj.time().second}' 41 | cfg.timestring = timestring 42 | logfile = f'log/z0seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{cfg.weight_decay}_{timestring}.csv' 43 | 44 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 45 | json.dump(cfg.__dict__, f, indent=2) 46 | 47 | sys.stdout = open(logfile, 'w') 48 | 49 | print('--------args----------') 50 | for k in list(vars(cfg).keys()): 51 | print('%s: %s' % (k, vars(cfg)[k])) 52 | print('--------args----------\n') 53 | 54 | sys.stdout.flush() 55 | 56 | net = FNO(cfg.modes, 32) 57 | 58 | cfg.delta_t = cfg.T/cfg.time_steps 59 | train(cfg, net) 60 | torch.save(net.state_dict(), f'model/net_seed_{cfg.seed}_{timestring}.pt') 61 | 62 | 63 | if __name__ == "__main__": 64 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 65 | 66 | parser.add_argument('--data_path', type=str, 67 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 68 | help='path of data') 69 | 70 | parser.add_argument('--device', type=str, default='cuda:3', 71 | help='Used device') 72 | 73 | parser.add_argument('--seed', type=int, default=0, 74 | help='seed') 75 | 76 | parser.add_argument('--batch_size', type=int, default=200, 77 | help='batchsize of the operator learning') 78 | 79 | parser.add_argument('--step_size', type=int, default=2000, 80 | help='step_size of optim') 81 | 82 | parser.add_argument('--gamma', type=float, default=0.5, 83 | help='gamma of optim') 84 | 85 | parser.add_argument('--lr', type=float, default=0.01, 86 | help='lr of optim') 87 | 88 | parser.add_argument('--weight_decay', type=float, default=0.0, 89 | help='lr of optim') 90 | 91 | parser.add_argument('--num_iterations', type=int, default=20000, 92 | help='num_iterations of optim') 93 | 94 | parser.add_argument('--size', type=int, default=64, 95 | help='data spatial size') 96 | 97 | parser.add_argument('--T', type=float, default=1.0, 98 | help='final time') 99 | 100 | parser.add_argument('--time_steps', type=int, default=100, 101 | help='number of time_steps in data') 102 | 103 | parser.add_argument('--EPS', type=float, default=1e-8, 104 | help='epsilon') 105 | 106 | parser.add_argument('--lamb', type=float, default=1, 107 | help='lamb') 108 | 109 | parser.add_argument('--modes', type=int, default=16, 110 | help='modes in fno') 111 | 112 | cfg = parser.parse_args() 113 | for seed in [0, 1, 2]: 114 | cfg.seed = seed 115 | main(cfg) 116 | -------------------------------------------------------------------------------- /allen-cahn/E2/mcnp100/mc_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import math 3 | import copy 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from scipy.stats import levy_stable 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def mc_loss(u_0, model, p, config): 17 | device = config.device 18 | size = config.size 19 | batch_size = config.batch_size 20 | time_steps = config.time_steps 21 | delta_t = config.delta_t 22 | 23 | u_input = u_0.reshape(u_0.shape[0], 1, size+1).repeat(1, time_steps, 1).reshape(-1, size+1, 1) 24 | t = delta_t * torch.tensor(range(1, time_steps+1)).to(device).reshape(1, -1).repeat(u_0.shape[0], 1).reshape(-1,) 25 | u = model(u_input, t).reshape(batch_size, -1, size+1) 26 | u = torch.concat([u_0.reshape(batch_size, 1, size+1), u], dim=1) 27 | 28 | u_0, u_1 = u[:, :-1, :], u[:, 1:, :] 29 | f_temp = 0.5*(u_0+u_1) 30 | f_temp = f_temp - f_temp**3 31 | u_hat = torch.einsum('bts, gs->btg', u_0 + f_temp* delta_t, p) 32 | return torch.sqrt(torch.mean(torch.square(u_hat - u_1))) -------------------------------------------------------------------------------- /allen-cahn/E2/mcnp100/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.parameter import Parameter 6 | import matplotlib.pyplot as plt 7 | 8 | import operator 9 | from functools import reduce 10 | from functools import partial 11 | from timeit import default_timer 12 | 13 | 14 | ################################################################ 15 | # 1d fourier layer 16 | ################################################################ 17 | class SpectralConv1d(nn.Module): 18 | def __init__(self, in_channels, out_channels, modes1): 19 | super(SpectralConv1d, self).__init__() 20 | 21 | """ 22 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 23 | """ 24 | 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 28 | 29 | self.scale = (1 / (in_channels*out_channels)) 30 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 31 | 32 | # Complex multiplication 33 | def compl_mul1d(self, input, weights): 34 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 35 | return torch.einsum("bix,iox->box", input, weights) 36 | 37 | def forward(self, x): 38 | batchsize = x.shape[0] 39 | #Compute Fourier coeffcients up to factor of e^(- something constant) 40 | x_ft = torch.fft.rfft(x) 41 | 42 | # Multiply relevant Fourier modes 43 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 44 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 45 | 46 | #Return to physical space 47 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 48 | return x 49 | 50 | class FNO(nn.Module): 51 | def __init__(self, modes, width): 52 | super(FNO, self).__init__() 53 | 54 | """ 55 | The overall network. It contains 4 layers of the Fourier layer. 56 | 1. Lift the input to the desire channel dimension by self.fc0 . 57 | 2. 4 layers of the integral operators u' = (W + K)(u). 58 | W defined by self.w; K defined by self.conv . 59 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 60 | 61 | input: the solution of the initial condition and location (a(x), x) 62 | input shape: (batchsize, x=s, c=2) 63 | output: the solution of a later timestep 64 | output shape: (batchsize, x=s, c=1) 65 | """ 66 | 67 | self.modes1 = modes 68 | self.width = width 69 | self.fc0 = nn.Linear(3, self.width) # input channel is 2: (a(x), x) 70 | 71 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 74 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 75 | self.w0 = nn.Conv1d(self.width, self.width, 1) 76 | self.w1 = nn.Conv1d(self.width, self.width, 1) 77 | self.w2 = nn.Conv1d(self.width, self.width, 1) 78 | self.w3 = nn.Conv1d(self.width, self.width, 1) 79 | 80 | self.fc1 = nn.Linear(self.width, 128) 81 | self.fc2 = nn.Linear(128, 1) 82 | 83 | def forward(self, x, t): 84 | grid = self.get_grid(x.shape, x.device) 85 | time = t.reshape(-1, 1, 1).repeat(1, x.shape[1], 1) 86 | x = torch.cat((x, grid), dim=-1) 87 | x = torch.concat([x, time], dim=-1) 88 | x = self.fc0(x) 89 | x = x.permute(0, 2, 1) 90 | 91 | x1 = self.conv0(x) 92 | x2 = self.w0(x) 93 | x = x1 + x2 94 | x = F.gelu(x) 95 | 96 | x1 = self.conv1(x) 97 | x2 = self.w1(x) 98 | x = x1 + x2 99 | x = F.gelu(x) 100 | 101 | x1 = self.conv2(x) 102 | x2 = self.w2(x) 103 | x = x1 + x2 104 | x = F.gelu(x) 105 | 106 | x1 = self.conv3(x) 107 | x2 = self.w3(x) 108 | x = x1 + x2 109 | x = x.permute(0, 2, 1) 110 | x = self.fc1(x) 111 | x = F.gelu(x) 112 | x = self.fc2(x) 113 | return x 114 | 115 | def get_grid(self, shape, device): 116 | batchsize, size_x = shape[0], shape[1] 117 | gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) 118 | gridx = gridx.reshape(1, size_x, 1).repeat([batchsize, 1, 1]) 119 | return gridx.to(device) -------------------------------------------------------------------------------- /allen-cahn/E2/mcnp100/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_sol(sol): 11 | ''' 12 | sol: torch.Tensor, size = (num_grid, ) 13 | return: figs 14 | ''' 15 | temp = sol.detach().cpu().numpy().reshape(-1, ) 16 | num_grid = temp.shape[0] - 1 17 | x_axis = (1 / num_grid) * torch.Tensor(range(num_grid + 1)) 18 | plt.figure() 19 | plt.plot(x_axis.numpy(), temp.numpy()) 20 | plt.show() 21 | 22 | 23 | def setup_seed(seed): 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | np.random.seed(seed) 28 | random.seed(seed) 29 | os.environ['PYTHONHASHSEED'] = str(seed) 30 | torch.backends.cudnn.deterministic = True 31 | torch.backends.cudnn.benchmark = False 32 | torch.backends.cudnn.enabled = True -------------------------------------------------------------------------------- /allen-cahn/E2/mcnp100/train.py: -------------------------------------------------------------------------------- 1 | from mc_loss import mc_loss 2 | from initial_field import GaussianRF 3 | import sys 4 | import math 5 | import copy 6 | import random 7 | import numpy as np 8 | import torch 9 | import torch.optim as optim 10 | import matplotlib.pyplot as plt 11 | from scipy.stats import norm 12 | import os 13 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 14 | plt.rcParams["animation.html"] = "jshtml" 15 | 16 | 17 | def test(config, net, test_data): 18 | device = config.device 19 | num = test_data.shape[0] 20 | delta_steps = 10 21 | delta_t = config.T / delta_steps 22 | u_0 = test_data[:, 0, :].to(device)[:, :, None] 23 | rela_err = torch.zeros(delta_steps) 24 | rela_err_1 = torch.zeros(delta_steps) 25 | rela_err_max = torch.zeros(delta_steps) 26 | for time_step in range(1, 11): 27 | net.eval() 28 | t = (time_step*delta_t) * torch.ones(num).to(device) 29 | u = net(u_0, t).detach() 30 | u_t = test_data[:, delta_steps * time_step, :].to(device)[:, :, None] 31 | rela_err[time_step-1] = (torch.norm((u - u_t).reshape(u.shape[0], -1), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), dim=1)).mean() 32 | rela_err_1[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=1, dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=1, dim=1)).mean() 33 | rela_err_max[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=float('inf'), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=float('inf'), dim=1)).mean() 34 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 35 | print('mean relative l_2 error', rela_err.mean().item()) 36 | print('mean relative l_1 error', rela_err_1.mean().item()) 37 | print('mean relative l_inf error', rela_err_max.mean().item()) 38 | return rela_err.mean().item() 39 | 40 | 41 | def p_matrix(config): 42 | size_x = config.size 43 | dt = 1/100 44 | p = np.zeros([size_x+1, size_x+1]) 45 | x = np.linspace(0, 1, size_x+1) 46 | dx = 1/size_x 47 | kappa = 0.01 48 | sigma = math.sqrt(2 * kappa * dt) 49 | for i in range(1, size_x): 50 | if x[i] <= 2/3: 51 | d = x[i] 52 | p[i, :] = (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-x+d)**2)/(2 * sigma**2))) * dx - (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-x+d-2*d)**2)/(2 * sigma**2))) * dx 53 | if x[i] > 2/3: 54 | p[i, :] = np.flip(p[size_x-i, :]) 55 | return torch.tensor(p).float() 56 | 57 | 58 | def generate_u0(b_size, N, num_grid, device): 59 | ''' 60 | Input: 61 | b_size: int, num of pre-generated data 62 | N: int, highest frequency fourier basis; 63 | num_grid: int, number of grids; 64 | Output: 65 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 66 | ''' 67 | B_n = torch.rand(b_size, N, device=device).reshape(b_size, N, 1, 1) 68 | grid = torch.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid, device=device).reshape(1, 1, -1, 1) 69 | n = 2 * torch.range(1, N, device=device).reshape(1, -1, 1, 1) 70 | sin_grid = torch.sin(torch.pi * n * grid) 71 | return (B_n * sin_grid).sum(axis=1) 72 | 73 | 74 | def train(config, net): 75 | N = 10 76 | device = config.device 77 | size = config.size 78 | batch_size = config.batch_size 79 | net = net.to(device) 80 | optimizer = optim.Adam(net.parameters(), config.lr) 81 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 82 | val_data = torch.load(config.data_path + 'data_test').to(device).float() 83 | test_data = torch.load(config.data_path + 'data_test').to(device).float() 84 | val_loss = 1e50 85 | p = p_matrix(config).to(device) 86 | for step in range(config.num_iterations+1): 87 | u_0 = generate_u0(batch_size, N, 1024, device) 88 | u_0 = (u_0[:, 15:-1:16] + u_0[:, 16::16])/2 89 | u_0 = torch.concat([torch.zeros(batch_size, 1, 1, device=device), u_0, torch.zeros(batch_size, 1, 1, device=device)], dim=1) 90 | net.train() 91 | loss = mc_loss(u_0, net, p, config) 92 | loss.backward() 93 | optimizer.step() 94 | optimizer.zero_grad(set_to_none=True) 95 | scheduler.step() 96 | if step % 50 == 0: 97 | print('########################') 98 | print('training loss', step, loss.detach().item()) 99 | print('VAL_LOSS') 100 | print('########################') 101 | temp = test(config, net, val_data) 102 | if temp < val_loss: 103 | val_loss = temp 104 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 105 | print('TEST_LOSS') 106 | print('########################') 107 | test(config, net, test_data) 108 | sys.stdout.flush() 109 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 110 | print('FINAL_LOSS') 111 | print('########################') 112 | test(config, net, test_data) 113 | sys.stdout.flush() -------------------------------------------------------------------------------- /allen-cahn/E3/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import os 4 | import random 5 | import numpy as np 6 | from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph 7 | from initial_field import GaussianRF 8 | 9 | 10 | def setup_seed(seed): 11 | torch.manual_seed(seed) 12 | torch.cuda.manual_seed(seed) 13 | torch.cuda.manual_seed_all(seed) 14 | np.random.seed(seed) 15 | random.seed(seed) 16 | os.environ['PYTHONHASHSEED'] = str(seed) 17 | torch.backends.cudnn.deterministic = True 18 | torch.backends.cudnn.benchmark = False 19 | torch.backends.cudnn.enabled = True 20 | 21 | 22 | setup_seed(0) 23 | size_x = 1024 24 | size_t = 200 25 | T = 1.0 26 | delta_t = 1e-6 27 | num_test = 200 28 | num_val = 200 29 | 30 | def generate_u0(b_size, N, num_grid): 31 | ''' 32 | Input: 33 | b_size: int, num of pre-generated data 34 | N: int, highest frequency fourier basis; 35 | num_grid: int, number of grids; 36 | Output: 37 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 38 | ''' 39 | B_n = torch.rand(b_size, N).reshape(b_size, N, 1, 1) 40 | grid = torch.tensor(np.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid)).reshape(1, 1, -1, 1) 41 | n = 2 * torch.Tensor(range(1, N + 1)).reshape(1, -1, 1, 1) 42 | sin_grid = torch.sin(torch.pi * n * grid) 43 | return (B_n * sin_grid).sum(axis=1) 44 | 45 | 46 | setup_seed(1) 47 | N = 5 48 | size_x = 65 49 | size_t = 100 50 | T = 1.0 51 | data_test = np.zeros([num_test, size_t+1, size_x]) 52 | initial = GaussianRF(1024) 53 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 54 | field = ScalarField(grid, 2) 55 | for i in range(num_test): 56 | maxvalue = np.nan 57 | while np.isnan(maxvalue) or maxvalue>10: 58 | bc_x_left = {"derivative": 0} 59 | bc_x_right = {"derivative": 0} 60 | term_1 = f"laplace(c) * 0.01 + c - c**3" 61 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 62 | b = generate_u0(1, N, 1024) 63 | field.data = np.array(b).reshape(-1,) 64 | storage = MemoryStorage() # store intermediate information of the simulation 65 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 66 | a = torch.tensor(storage.data) 67 | a_in = (a[:, 15:-1:16] + a[:, 16::16])/2 68 | data_test[i, :, 1:-1] = a_in 69 | data_test[i, :, 0], data_test[i, :, -1] = a[:, 0], a[:, -1] 70 | maxvalue = np.abs(data_test[i]).max() 71 | print(i, np.abs(data_test[i]).max()) 72 | 73 | data_test = torch.tensor(data_test) 74 | torch.save(data_test, 'dataset/data_test') 75 | 76 | 77 | setup_seed(2) 78 | N = 5 79 | size_x = 65 80 | size_t = 100 81 | T = 1.0 82 | data_val = np.zeros([num_val, size_t+1, size_x]) 83 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 84 | field = ScalarField(grid, 2) 85 | for i in range(num_val): 86 | maxvalue = np.nan 87 | while np.isnan(maxvalue) or maxvalue>10: 88 | bc_x_left = {"derivative": 0} 89 | bc_x_right = {"derivative": 0} 90 | term_1 = f"laplace(c) * 0.01 + c - c**3" 91 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 92 | b = generate_u0(1, N, 1024) 93 | field.data = np.array(b).reshape(-1,) 94 | storage = MemoryStorage() # store intermediate information of the simulation 95 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 96 | a = torch.tensor(storage.data) 97 | a_in = (a[:, 15:-1:16] + a[:, 16::16])/2 98 | data_val[i, :, 1:-1] = a_in 99 | data_val[i, :, 0], data_val[i, :, -1] = a[:, 0], a[:, -1] 100 | maxvalue = np.abs(data_val[i]).max() 101 | print(i, np.abs(data_val[i]).max()) 102 | 103 | data_val = torch.tensor(data_val) 104 | torch.save(data_val, 'dataset/data_val') -------------------------------------------------------------------------------- /allen-cahn/E3/data/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E3/mcnp100/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E3/mcnp100/main.py: -------------------------------------------------------------------------------- 1 | 2 | from model import FNO 3 | from train import train 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def setup_seed(seed): 21 | torch.manual_seed(seed) 22 | torch.cuda.manual_seed(seed) 23 | torch.cuda.manual_seed_all(seed) 24 | np.random.seed(seed) 25 | random.seed(seed) 26 | os.environ['PYTHONHASHSEED'] = str(seed) 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = True 29 | torch.backends.cudnn.enabled = True 30 | 31 | 32 | 33 | def main(cfg): 34 | if not os.path.exists(f'log'): 35 | os.mkdir(f'log') 36 | if not os.path.exists(f'model'): 37 | os.mkdir(f'model') 38 | setup_seed(cfg.seed) 39 | dateTimeObj = datetime.now() 40 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}_{dateTimeObj.time().second}' 41 | cfg.timestring = timestring 42 | logfile = f'log/z0seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{cfg.weight_decay}_{timestring}.csv' 43 | 44 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 45 | json.dump(cfg.__dict__, f, indent=2) 46 | 47 | sys.stdout = open(logfile, 'w') 48 | 49 | print('--------args----------') 50 | for k in list(vars(cfg).keys()): 51 | print('%s: %s' % (k, vars(cfg)[k])) 52 | print('--------args----------\n') 53 | 54 | sys.stdout.flush() 55 | 56 | net = FNO(cfg.modes, 32) 57 | 58 | cfg.delta_t = cfg.T/cfg.time_steps 59 | train(cfg, net) 60 | torch.save(net.state_dict(), f'model/net_seed_{cfg.seed}_{timestring}.pt') 61 | 62 | 63 | if __name__ == "__main__": 64 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 65 | 66 | parser.add_argument('--data_path', type=str, 67 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 68 | help='path of data') 69 | 70 | parser.add_argument('--device', type=str, default='cuda:1', 71 | help='Used device') 72 | 73 | parser.add_argument('--seed', type=int, default=0, 74 | help='seed') 75 | 76 | parser.add_argument('--batch_size', type=int, default=200, 77 | help='batchsize of the operator learning') 78 | 79 | parser.add_argument('--step_size', type=int, default=2000, 80 | help='step_size of optim') 81 | 82 | parser.add_argument('--gamma', type=float, default=0.5, 83 | help='gamma of optim') 84 | 85 | parser.add_argument('--lr', type=float, default=0.01, 86 | help='lr of optim') 87 | 88 | parser.add_argument('--weight_decay', type=float, default=0.0, 89 | help='lr of optim') 90 | 91 | parser.add_argument('--num_iterations', type=int, default=20000, 92 | help='num_iterations of optim') 93 | 94 | parser.add_argument('--size', type=int, default=64, 95 | help='data spatial size') 96 | 97 | parser.add_argument('--T', type=float, default=1.0, 98 | help='final time') 99 | 100 | parser.add_argument('--time_steps', type=int, default=100, 101 | help='number of time_steps in data') 102 | 103 | parser.add_argument('--EPS', type=float, default=1e-8, 104 | help='epsilon') 105 | 106 | parser.add_argument('--lamb', type=float, default=1, 107 | help='lamb') 108 | 109 | parser.add_argument('--modes', type=int, default=16, 110 | help='modes in fno') 111 | 112 | # cfg = parser.parse_args() 113 | # for lamb in [1, 10, 100]: 114 | # cfg.lamb = lamb 115 | # main(cfg) 116 | 117 | # cfg = parser.parse_args() 118 | # for modes in [8, 12, 16]: 119 | # cfg.modes = modes 120 | # main(cfg) 121 | 122 | # cfg = parser.parse_args() 123 | # for modes in [20]: 124 | # cfg.modes = modes 125 | # main(cfg) 126 | 127 | # cfg = parser.parse_args() 128 | # for step_size in [2000]: 129 | # cfg.step_size = step_size 130 | # main(cfg) 131 | 132 | cfg = parser.parse_args() 133 | for seed in [0, 1, 2]: 134 | cfg.seed = seed 135 | main(cfg) 136 | 137 | # cfg = parser.parse_args() 138 | # for seed in [1, 2]: 139 | # cfg.seed = seed 140 | # main(cfg) -------------------------------------------------------------------------------- /allen-cahn/E3/mcnp100/mc_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import math 3 | import copy 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from scipy.stats import levy_stable 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def mc_loss(u_0, model, p, config): 17 | device = config.device 18 | size = config.size 19 | batch_size = config.batch_size 20 | time_steps = config.time_steps 21 | delta_t = config.delta_t 22 | 23 | u_input = u_0.reshape(u_0.shape[0], 1, size+1).repeat(1, time_steps, 1).reshape(-1, size+1, 1) 24 | t = delta_t * torch.tensor(range(1, time_steps+1)).to(device).reshape(1, -1).repeat(u_0.shape[0], 1).reshape(-1,) 25 | u = model(u_input, t).reshape(batch_size, -1, size+1) 26 | u = torch.concat([u_0.reshape(batch_size, 1, size+1), u], dim=1) 27 | 28 | u_0, u_1 = u[:, :-1, :], u[:, 1:, :] 29 | f_temp = 0.5*(u_0+u_1) 30 | f_temp = f_temp - f_temp**3 31 | u_hat = torch.einsum('bts, gs->btg', u_0 + f_temp* delta_t, p) 32 | return torch.sqrt(torch.mean(torch.square(u_hat - u_1))) -------------------------------------------------------------------------------- /allen-cahn/E3/mcnp100/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.parameter import Parameter 6 | import matplotlib.pyplot as plt 7 | 8 | import operator 9 | from functools import reduce 10 | from functools import partial 11 | from timeit import default_timer 12 | 13 | 14 | ################################################################ 15 | # 1d fourier layer 16 | ################################################################ 17 | class SpectralConv1d(nn.Module): 18 | def __init__(self, in_channels, out_channels, modes1): 19 | super(SpectralConv1d, self).__init__() 20 | 21 | """ 22 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 23 | """ 24 | 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 28 | 29 | self.scale = (1 / (in_channels*out_channels)) 30 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 31 | 32 | # Complex multiplication 33 | def compl_mul1d(self, input, weights): 34 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 35 | return torch.einsum("bix,iox->box", input, weights) 36 | 37 | def forward(self, x): 38 | batchsize = x.shape[0] 39 | #Compute Fourier coeffcients up to factor of e^(- something constant) 40 | x_ft = torch.fft.rfft(x) 41 | 42 | # Multiply relevant Fourier modes 43 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 44 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 45 | 46 | #Return to physical space 47 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 48 | return x 49 | 50 | class FNO(nn.Module): 51 | def __init__(self, modes, width): 52 | super(FNO, self).__init__() 53 | 54 | """ 55 | The overall network. It contains 4 layers of the Fourier layer. 56 | 1. Lift the input to the desire channel dimension by self.fc0 . 57 | 2. 4 layers of the integral operators u' = (W + K)(u). 58 | W defined by self.w; K defined by self.conv . 59 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 60 | 61 | input: the solution of the initial condition and location (a(x), x) 62 | input shape: (batchsize, x=s, c=2) 63 | output: the solution of a later timestep 64 | output shape: (batchsize, x=s, c=1) 65 | """ 66 | 67 | self.modes1 = modes 68 | self.width = width 69 | self.fc0 = nn.Linear(3, self.width) # input channel is 2: (a(x), x) 70 | 71 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 74 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 75 | self.w0 = nn.Conv1d(self.width, self.width, 1) 76 | self.w1 = nn.Conv1d(self.width, self.width, 1) 77 | self.w2 = nn.Conv1d(self.width, self.width, 1) 78 | self.w3 = nn.Conv1d(self.width, self.width, 1) 79 | 80 | self.fc1 = nn.Linear(self.width, 128) 81 | self.fc2 = nn.Linear(128, 1) 82 | 83 | def forward(self, x, t): 84 | grid = self.get_grid(x.shape, x.device) 85 | time = t.reshape(-1, 1, 1).repeat(1, x.shape[1], 1) 86 | x = torch.cat((x, grid), dim=-1) 87 | x = torch.concat([x, time], dim=-1) 88 | x = self.fc0(x) 89 | x = x.permute(0, 2, 1) 90 | 91 | x1 = self.conv0(x) 92 | x2 = self.w0(x) 93 | x = x1 + x2 94 | x = F.gelu(x) 95 | 96 | x1 = self.conv1(x) 97 | x2 = self.w1(x) 98 | x = x1 + x2 99 | x = F.gelu(x) 100 | 101 | x1 = self.conv2(x) 102 | x2 = self.w2(x) 103 | x = x1 + x2 104 | x = F.gelu(x) 105 | 106 | x1 = self.conv3(x) 107 | x2 = self.w3(x) 108 | x = x1 + x2 109 | x = x.permute(0, 2, 1) 110 | x = self.fc1(x) 111 | x = F.gelu(x) 112 | x = self.fc2(x) 113 | return x 114 | 115 | def get_grid(self, shape, device): 116 | batchsize, size_x = shape[0], shape[1] 117 | gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) 118 | gridx = gridx.reshape(1, size_x, 1).repeat([batchsize, 1, 1]) 119 | return gridx.to(device) -------------------------------------------------------------------------------- /allen-cahn/E3/mcnp100/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_sol(sol): 11 | ''' 12 | sol: torch.Tensor, size = (num_grid, ) 13 | return: figs 14 | ''' 15 | temp = sol.detach().cpu().numpy().reshape(-1, ) 16 | num_grid = temp.shape[0] - 1 17 | x_axis = (1 / num_grid) * torch.Tensor(range(num_grid + 1)) 18 | plt.figure() 19 | plt.plot(x_axis.numpy(), temp.numpy()) 20 | plt.show() 21 | 22 | 23 | def setup_seed(seed): 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | np.random.seed(seed) 28 | random.seed(seed) 29 | os.environ['PYTHONHASHSEED'] = str(seed) 30 | torch.backends.cudnn.deterministic = True 31 | torch.backends.cudnn.benchmark = False 32 | torch.backends.cudnn.enabled = True -------------------------------------------------------------------------------- /allen-cahn/E3/mcnp100/train.py: -------------------------------------------------------------------------------- 1 | from mc_loss import mc_loss 2 | from initial_field import GaussianRF 3 | import sys 4 | import math 5 | import copy 6 | import random 7 | import numpy as np 8 | import torch 9 | import torch.optim as optim 10 | import matplotlib.pyplot as plt 11 | from scipy.stats import norm 12 | import os 13 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 14 | plt.rcParams["animation.html"] = "jshtml" 15 | 16 | 17 | def test(config, net, test_data): 18 | device = config.device 19 | num = test_data.shape[0] 20 | delta_steps = 10 21 | delta_t = config.T / delta_steps 22 | u_0 = test_data[:, 0, :].to(device)[:, :, None] 23 | rela_err = torch.zeros(delta_steps) 24 | rela_err_1 = torch.zeros(delta_steps) 25 | rela_err_max = torch.zeros(delta_steps) 26 | for time_step in range(1, 11): 27 | net.eval() 28 | t = (time_step*delta_t) * torch.ones(num).to(device) 29 | u = net(u_0, t).detach() 30 | u_t = test_data[:, delta_steps * time_step, :].to(device)[:, :, None] 31 | rela_err[time_step-1] = (torch.norm((u - u_t).reshape(u.shape[0], -1), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), dim=1)).mean() 32 | rela_err_1[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=1, dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=1, dim=1)).mean() 33 | rela_err_max[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=float('inf'), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=float('inf'), dim=1)).mean() 34 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 35 | print('mean relative l_2 error', rela_err.mean().item()) 36 | print('mean relative l_1 error', rela_err_1.mean().item()) 37 | print('mean relative l_inf error', rela_err_max.mean().item()) 38 | return rela_err.mean().item() 39 | 40 | 41 | def p_matrix(config): 42 | size_x = config.size 43 | dt = config.delta_t 44 | p = np.zeros([size_x+1, size_x+1]) 45 | x = np.linspace(0, 1, size_x+1) 46 | dx = 1/size_x 47 | kappa = 0.01 48 | sigma = math.sqrt(2 * kappa * dt) 49 | for i in range(size_x+1): 50 | if x[i] <= 2/3: 51 | d = x[i] 52 | p[i, :] = (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((x-d)**2)/(2 * sigma**2))) * dx + (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-x-d)**2)/(2 * sigma**2))) * dx 53 | p[i, 0] = (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-d)**2)/(2 * sigma**2))) * dx 54 | if x[i] > 2/3: 55 | p[i, :] = np.flip(p[size_x-i, :]) 56 | return torch.tensor(p).float() 57 | 58 | 59 | def generate_u0(b_size, N, num_grid, device): 60 | ''' 61 | Input: 62 | b_size: int, num of pre-generated data 63 | N: int, highest frequency fourier basis; 64 | num_grid: int, number of grids; 65 | Output: 66 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 67 | ''' 68 | B_n = torch.rand(b_size, N, device=device).reshape(b_size, N, 1, 1) 69 | grid = torch.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid, device=device).reshape(1, 1, -1, 1) 70 | n = 2 * torch.range(1, N, device=device).reshape(1, -1, 1, 1) 71 | sin_grid = torch.sin(torch.pi * n * grid) 72 | return (B_n * sin_grid).sum(axis=1) 73 | 74 | 75 | def train(config, net): 76 | N = 5 77 | device = config.device 78 | size = config.size 79 | batch_size = config.batch_size 80 | net = net.to(device) 81 | optimizer = optim.Adam(net.parameters(), config.lr) 82 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 83 | val_data = torch.load(config.data_path + 'data_val').to(device).float() 84 | test_data = torch.load(config.data_path + 'data_test').to(device).float() 85 | val_loss = 1e50 86 | p = p_matrix(config).to(device) 87 | for step in range(config.num_iterations+1): 88 | u = generate_u0(batch_size, N, 1024, device) 89 | u_0 = (u[:, 15:-1:16] + u[:, 16::16])/2 90 | u_0 = torch.concat([u[:, 0, None], u_0, u[:, -1, None]], dim=1) 91 | net.train() 92 | loss = mc_loss(u_0, net, p, config) 93 | loss.backward() 94 | optimizer.step() 95 | optimizer.zero_grad(set_to_none=True) 96 | scheduler.step() 97 | if step % 50 == 0: 98 | print('########################') 99 | print('training loss', step, loss.detach().item()) 100 | print('VAL_LOSS') 101 | print('########################') 102 | temp = test(config, net, val_data) 103 | if temp < val_loss: 104 | val_loss = temp 105 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 106 | print('TEST_LOSS') 107 | print('########################') 108 | test(config, net, test_data) 109 | sys.stdout.flush() 110 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 111 | print('FINAL_LOSS') 112 | print('########################') 113 | test(config, net, test_data) 114 | sys.stdout.flush() -------------------------------------------------------------------------------- /allen-cahn/E4/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import os 4 | import random 5 | import numpy as np 6 | from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph 7 | from initial_field import GaussianRF 8 | 9 | 10 | def setup_seed(seed): 11 | torch.manual_seed(seed) 12 | torch.cuda.manual_seed(seed) 13 | torch.cuda.manual_seed_all(seed) 14 | np.random.seed(seed) 15 | random.seed(seed) 16 | os.environ['PYTHONHASHSEED'] = str(seed) 17 | torch.backends.cudnn.deterministic = True 18 | torch.backends.cudnn.benchmark = False 19 | torch.backends.cudnn.enabled = True 20 | 21 | 22 | setup_seed(0) 23 | size_x = 1024 24 | size_t = 200 25 | T = 1.0 26 | delta_t = 1e-6 27 | num_test = 200 28 | num_val = 200 29 | 30 | def generate_u0(b_size, N, num_grid): 31 | ''' 32 | Input: 33 | b_size: int, num of pre-generated data 34 | N: int, highest frequency fourier basis; 35 | num_grid: int, number of grids; 36 | Output: 37 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 38 | ''' 39 | B_n = torch.rand(b_size, N).reshape(b_size, N, 1, 1) 40 | grid = torch.tensor(np.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid)).reshape(1, 1, -1, 1) 41 | n = 2 * torch.Tensor(range(1, N + 1)).reshape(1, -1, 1, 1) 42 | sin_grid = torch.sin(torch.pi * n * grid) 43 | return (B_n * sin_grid).sum(axis=1) 44 | 45 | 46 | setup_seed(1) 47 | N = 10 48 | size_x = 65 49 | size_t = 100 50 | T = 1.0 51 | data_test = np.zeros([num_test, size_t+1, size_x]) 52 | initial = GaussianRF(1024) 53 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 54 | field = ScalarField(grid, 2) 55 | for i in range(num_test): 56 | maxvalue = np.nan 57 | while np.isnan(maxvalue) or maxvalue>10: 58 | bc_x_left = {"derivative": 0} 59 | bc_x_right = {"derivative": 0} 60 | term_1 = f"laplace(c) * 0.01 + c - c**3" 61 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 62 | b = generate_u0(1, N, 1024) 63 | field.data = np.array(b).reshape(-1,) 64 | storage = MemoryStorage() # store intermediate information of the simulation 65 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 66 | a = torch.tensor(storage.data) 67 | a_in = (a[:, 15:-1:16] + a[:, 16::16])/2 68 | data_test[i, :, 1:-1] = a_in 69 | data_test[i, :, 0], data_test[i, :, -1] = a[:, 0], a[:, -1] 70 | maxvalue = np.abs(data_test[i]).max() 71 | print(i, np.abs(data_test[i]).max()) 72 | 73 | data_test = torch.tensor(data_test) 74 | torch.save(data_test, 'dataset/data_test') 75 | 76 | 77 | setup_seed(2) 78 | N = 10 79 | size_x = 65 80 | size_t = 100 81 | T = 1.0 82 | data_val = np.zeros([num_val, size_t+1, size_x]) 83 | grid = CartesianGrid([[0, 1]], 1024) # generate grid 84 | field = ScalarField(grid, 2) 85 | for i in range(num_val): 86 | maxvalue = np.nan 87 | while np.isnan(maxvalue) or maxvalue>10: 88 | bc_x_left = {"derivative": 0} 89 | bc_x_right = {"derivative": 0} 90 | term_1 = f"laplace(c) * 0.01 + c - c**3" 91 | eq = PDE({"c": f"{term_1}"}, bc=[bc_x_left, bc_x_right]) 92 | b = generate_u0(1, N, 1024) 93 | field.data = np.array(b).reshape(-1,) 94 | storage = MemoryStorage() # store intermediate information of the simulation 95 | res = eq.solve(field, T, dt=delta_t, tracker=storage.tracker(T/size_t)) # solve the PDE 96 | a = torch.tensor(storage.data) 97 | a_in = (a[:, 15:-1:16] + a[:, 16::16])/2 98 | data_val[i, :, 1:-1] = a_in 99 | data_val[i, :, 0], data_val[i, :, -1] = a[:, 0], a[:, -1] 100 | maxvalue = np.abs(data_val[i]).max() 101 | print(i, np.abs(data_val[i]).max()) 102 | 103 | data_val = torch.tensor(data_val) 104 | torch.save(data_val, 'dataset/data_val') -------------------------------------------------------------------------------- /allen-cahn/E4/data/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E4/mcnp100/initial_field.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | class GaussianRF(object): 6 | 7 | def __init__(self, size, alpha=2, tau=3, sigma=None, boundary="periodic", device=None): 8 | self.device = device 9 | 10 | sigma = tau**(0.5*(2*alpha - 1)) 11 | 12 | k_max = size//2 13 | 14 | k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 15 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0) 16 | 17 | self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0)) 18 | self.sqrt_eig[0] = 0.0 19 | 20 | self.size = [] 21 | self.size.append(size) 22 | 23 | self.size = tuple(self.size) 24 | 25 | def sample(self, N): 26 | 27 | coeff = torch.randn(N, *self.size, 2, device=self.device) 28 | 29 | coeff[...,0] = self.sqrt_eig*coeff[...,0] 30 | coeff[...,1] = self.sqrt_eig*coeff[...,1] 31 | 32 | u = torch.fft.ifft(torch.view_as_complex(coeff)).real 33 | 34 | return u -------------------------------------------------------------------------------- /allen-cahn/E4/mcnp100/main.py: -------------------------------------------------------------------------------- 1 | 2 | from model import FNO 3 | from train import train 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def setup_seed(seed): 21 | torch.manual_seed(seed) 22 | torch.cuda.manual_seed(seed) 23 | torch.cuda.manual_seed_all(seed) 24 | np.random.seed(seed) 25 | random.seed(seed) 26 | os.environ['PYTHONHASHSEED'] = str(seed) 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = True 29 | torch.backends.cudnn.enabled = True 30 | 31 | 32 | 33 | def main(cfg): 34 | if not os.path.exists(f'log'): 35 | os.mkdir(f'log') 36 | if not os.path.exists(f'model'): 37 | os.mkdir(f'model') 38 | setup_seed(cfg.seed) 39 | dateTimeObj = datetime.now() 40 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}_{dateTimeObj.time().second}' 41 | cfg.timestring = timestring 42 | logfile = f'log/z1seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{cfg.weight_decay}_{timestring}.csv' 43 | 44 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 45 | json.dump(cfg.__dict__, f, indent=2) 46 | 47 | sys.stdout = open(logfile, 'w') 48 | 49 | print('--------args----------') 50 | for k in list(vars(cfg).keys()): 51 | print('%s: %s' % (k, vars(cfg)[k])) 52 | print('--------args----------\n') 53 | 54 | sys.stdout.flush() 55 | 56 | net = FNO(cfg.modes, 32) 57 | 58 | cfg.delta_t = cfg.T/cfg.time_steps 59 | train(cfg, net) 60 | torch.save(net.state_dict(), f'model/net_seed_{cfg.seed}_{timestring}.pt') 61 | 62 | 63 | if __name__ == "__main__": 64 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 65 | 66 | parser.add_argument('--data_path', type=str, 67 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 68 | help='path of data') 69 | 70 | parser.add_argument('--device', type=str, default='cuda:1', 71 | help='Used device') 72 | 73 | parser.add_argument('--seed', type=int, default=0, 74 | help='seed') 75 | 76 | parser.add_argument('--batch_size', type=int, default=200, 77 | help='batchsize of the operator learning') 78 | 79 | parser.add_argument('--step_size', type=int, default=2000, 80 | help='step_size of optim') 81 | 82 | parser.add_argument('--gamma', type=float, default=0.5, 83 | help='gamma of optim') 84 | 85 | parser.add_argument('--lr', type=float, default=0.01, 86 | help='lr of optim') 87 | 88 | parser.add_argument('--weight_decay', type=float, default=0.0, 89 | help='lr of optim') 90 | 91 | parser.add_argument('--num_iterations', type=int, default=20000, 92 | help='num_iterations of optim') 93 | 94 | parser.add_argument('--size', type=int, default=64, 95 | help='data spatial size') 96 | 97 | parser.add_argument('--T', type=float, default=1.0, 98 | help='final time') 99 | 100 | parser.add_argument('--time_steps', type=int, default=100, 101 | help='number of time_steps in data') 102 | 103 | parser.add_argument('--EPS', type=float, default=1e-8, 104 | help='epsilon') 105 | 106 | parser.add_argument('--lamb', type=float, default=1, 107 | help='lamb') 108 | 109 | parser.add_argument('--modes', type=int, default=16, 110 | help='modes in fno') 111 | 112 | # cfg = parser.parse_args() 113 | # for lamb in [1, 10, 100]: 114 | # cfg.lamb = lamb 115 | # main(cfg) 116 | 117 | # cfg = parser.parse_args() 118 | # for modes in [8, 12, 16]: 119 | # cfg.modes = modes 120 | # main(cfg) 121 | 122 | # cfg = parser.parse_args() 123 | # for modes in [20]: 124 | # cfg.modes = modes 125 | # main(cfg) 126 | 127 | # cfg = parser.parse_args() 128 | # for step_size in [2000]: 129 | # cfg.step_size = step_size 130 | # main(cfg) 131 | 132 | cfg = parser.parse_args() 133 | for seed in [0, 1, 2]: 134 | cfg.seed = seed 135 | main(cfg) 136 | 137 | # cfg = parser.parse_args() 138 | # for seed in [1, 2]: 139 | # cfg.seed = seed 140 | # main(cfg) -------------------------------------------------------------------------------- /allen-cahn/E4/mcnp100/mc_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import math 3 | import copy 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from scipy.stats import levy_stable 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def mc_loss(u_0, model, p, config): 17 | device = config.device 18 | size = config.size 19 | batch_size = config.batch_size 20 | time_steps = config.time_steps 21 | delta_t = config.delta_t 22 | 23 | u_input = u_0.reshape(u_0.shape[0], 1, size+1).repeat(1, time_steps, 1).reshape(-1, size+1, 1) 24 | t = delta_t * torch.tensor(range(1, time_steps+1)).to(device).reshape(1, -1).repeat(u_0.shape[0], 1).reshape(-1,) 25 | u = model(u_input, t).reshape(batch_size, -1, size+1) 26 | u = torch.concat([u_0.reshape(batch_size, 1, size+1), u], dim=1) 27 | 28 | u_0, u_1 = u[:, :-1, :], u[:, 1:, :] 29 | f_temp = 0.5*(u_0+u_1) 30 | f_temp = f_temp - f_temp**3 31 | u_hat = torch.einsum('bts, gs->btg', u_0 + f_temp* delta_t, p) 32 | return torch.sqrt(torch.mean(torch.square(u_hat - u_1))) -------------------------------------------------------------------------------- /allen-cahn/E4/mcnp100/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.parameter import Parameter 6 | import matplotlib.pyplot as plt 7 | 8 | import operator 9 | from functools import reduce 10 | from functools import partial 11 | from timeit import default_timer 12 | 13 | 14 | ################################################################ 15 | # 1d fourier layer 16 | ################################################################ 17 | class SpectralConv1d(nn.Module): 18 | def __init__(self, in_channels, out_channels, modes1): 19 | super(SpectralConv1d, self).__init__() 20 | 21 | """ 22 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 23 | """ 24 | 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 28 | 29 | self.scale = (1 / (in_channels*out_channels)) 30 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 31 | 32 | # Complex multiplication 33 | def compl_mul1d(self, input, weights): 34 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 35 | return torch.einsum("bix,iox->box", input, weights) 36 | 37 | def forward(self, x): 38 | batchsize = x.shape[0] 39 | #Compute Fourier coeffcients up to factor of e^(- something constant) 40 | x_ft = torch.fft.rfft(x) 41 | 42 | # Multiply relevant Fourier modes 43 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 44 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 45 | 46 | #Return to physical space 47 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 48 | return x 49 | 50 | class FNO(nn.Module): 51 | def __init__(self, modes, width): 52 | super(FNO, self).__init__() 53 | 54 | """ 55 | The overall network. It contains 4 layers of the Fourier layer. 56 | 1. Lift the input to the desire channel dimension by self.fc0 . 57 | 2. 4 layers of the integral operators u' = (W + K)(u). 58 | W defined by self.w; K defined by self.conv . 59 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 60 | 61 | input: the solution of the initial condition and location (a(x), x) 62 | input shape: (batchsize, x=s, c=2) 63 | output: the solution of a later timestep 64 | output shape: (batchsize, x=s, c=1) 65 | """ 66 | 67 | self.modes1 = modes 68 | self.width = width 69 | self.fc0 = nn.Linear(3, self.width) # input channel is 2: (a(x), x) 70 | 71 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 74 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 75 | self.w0 = nn.Conv1d(self.width, self.width, 1) 76 | self.w1 = nn.Conv1d(self.width, self.width, 1) 77 | self.w2 = nn.Conv1d(self.width, self.width, 1) 78 | self.w3 = nn.Conv1d(self.width, self.width, 1) 79 | 80 | self.fc1 = nn.Linear(self.width, 128) 81 | self.fc2 = nn.Linear(128, 1) 82 | 83 | def forward(self, x, t): 84 | grid = self.get_grid(x.shape, x.device) 85 | time = t.reshape(-1, 1, 1).repeat(1, x.shape[1], 1) 86 | x = torch.cat((x, grid), dim=-1) 87 | x = torch.concat([x, time], dim=-1) 88 | x = self.fc0(x) 89 | x = x.permute(0, 2, 1) 90 | 91 | x1 = self.conv0(x) 92 | x2 = self.w0(x) 93 | x = x1 + x2 94 | x = F.gelu(x) 95 | 96 | x1 = self.conv1(x) 97 | x2 = self.w1(x) 98 | x = x1 + x2 99 | x = F.gelu(x) 100 | 101 | x1 = self.conv2(x) 102 | x2 = self.w2(x) 103 | x = x1 + x2 104 | x = F.gelu(x) 105 | 106 | x1 = self.conv3(x) 107 | x2 = self.w3(x) 108 | x = x1 + x2 109 | x = x.permute(0, 2, 1) 110 | x = self.fc1(x) 111 | x = F.gelu(x) 112 | x = self.fc2(x) 113 | return x 114 | 115 | def get_grid(self, shape, device): 116 | batchsize, size_x = shape[0], shape[1] 117 | gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) 118 | gridx = gridx.reshape(1, size_x, 1).repeat([batchsize, 1, 1]) 119 | return gridx.to(device) -------------------------------------------------------------------------------- /allen-cahn/E4/mcnp100/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_sol(sol): 11 | ''' 12 | sol: torch.Tensor, size = (num_grid, ) 13 | return: figs 14 | ''' 15 | temp = sol.detach().cpu().numpy().reshape(-1, ) 16 | num_grid = temp.shape[0] - 1 17 | x_axis = (1 / num_grid) * torch.Tensor(range(num_grid + 1)) 18 | plt.figure() 19 | plt.plot(x_axis.numpy(), temp.numpy()) 20 | plt.show() 21 | 22 | 23 | def setup_seed(seed): 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | np.random.seed(seed) 28 | random.seed(seed) 29 | os.environ['PYTHONHASHSEED'] = str(seed) 30 | torch.backends.cudnn.deterministic = True 31 | torch.backends.cudnn.benchmark = False 32 | torch.backends.cudnn.enabled = True -------------------------------------------------------------------------------- /allen-cahn/E4/mcnp100/train.py: -------------------------------------------------------------------------------- 1 | from mc_loss import mc_loss 2 | from initial_field import GaussianRF 3 | import sys 4 | import math 5 | import copy 6 | import random 7 | import numpy as np 8 | import torch 9 | import torch.optim as optim 10 | import matplotlib.pyplot as plt 11 | from scipy.stats import norm 12 | import os 13 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 14 | plt.rcParams["animation.html"] = "jshtml" 15 | 16 | 17 | def test(config, net, test_data): 18 | device = config.device 19 | num = test_data.shape[0] 20 | delta_steps = 10 21 | delta_t = config.T / delta_steps 22 | u_0 = test_data[:, 0, :].to(device)[:, :, None] 23 | rela_err = torch.zeros(delta_steps) 24 | rela_err_1 = torch.zeros(delta_steps) 25 | rela_err_max = torch.zeros(delta_steps) 26 | for time_step in range(1, 11): 27 | net.eval() 28 | t = (time_step*delta_t) * torch.ones(num).to(device) 29 | u = net(u_0, t).detach() 30 | u_t = test_data[:, delta_steps * time_step, :].to(device)[:, :, None] 31 | rela_err[time_step-1] = (torch.norm((u - u_t).reshape(u.shape[0], -1), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), dim=1)).mean() 32 | rela_err_1[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=1, dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=1, dim=1)).mean() 33 | rela_err_max[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=float('inf'), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=float('inf'), dim=1)).mean() 34 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 35 | print('mean relative l_2 error', rela_err.mean().item()) 36 | print('mean relative l_1 error', rela_err_1.mean().item()) 37 | print('mean relative l_inf error', rela_err_max.mean().item()) 38 | return rela_err.mean().item() 39 | 40 | 41 | def p_matrix(config): 42 | size_x = config.size 43 | dt = config.delta_t 44 | p = np.zeros([size_x+1, size_x+1]) 45 | x = np.linspace(0, 1, size_x+1) 46 | dx = 1/size_x 47 | kappa = 0.01 48 | sigma = math.sqrt(2 * kappa * dt) 49 | for i in range(size_x+1): 50 | if x[i] <= 2/3: 51 | d = x[i] 52 | p[i, :] = (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((x-d)**2)/(2 * sigma**2))) * dx + (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-x-d)**2)/(2 * sigma**2))) * dx 53 | p[i, 0] = (1/(math.sqrt(2*torch.pi) * sigma) * np.exp(-((-d)**2)/(2 * sigma**2))) * dx 54 | if x[i] > 2/3: 55 | p[i, :] = np.flip(p[size_x-i, :]) 56 | return torch.tensor(p).float() 57 | 58 | 59 | def generate_u0(b_size, N, num_grid, device): 60 | ''' 61 | Input: 62 | b_size: int, num of pre-generated data 63 | N: int, highest frequency fourier basis; 64 | num_grid: int, number of grids; 65 | Output: 66 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 67 | ''' 68 | B_n = torch.rand(b_size, N, device=device).reshape(b_size, N, 1, 1) 69 | grid = torch.linspace(0.5/num_grid, 1 - 0.5/num_grid, num_grid, device=device).reshape(1, 1, -1, 1) 70 | n = 2 * torch.range(1, N, device=device).reshape(1, -1, 1, 1) 71 | sin_grid = torch.sin(torch.pi * n * grid) 72 | return (B_n * sin_grid).sum(axis=1) 73 | 74 | 75 | def train(config, net): 76 | N = 10 77 | device = config.device 78 | size = config.size 79 | batch_size = config.batch_size 80 | net = net.to(device) 81 | optimizer = optim.Adam(net.parameters(), config.lr) 82 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 83 | val_data = torch.load(config.data_path + 'data_val').to(device).float() 84 | test_data = torch.load(config.data_path + 'data_test').to(device).float() 85 | val_loss = 1e50 86 | p = p_matrix(config).to(device) 87 | for step in range(config.num_iterations+1): 88 | u = generate_u0(batch_size, N, 1024, device) 89 | u_0 = (u[:, 15:-1:16] + u[:, 16::16])/2 90 | u_0 = torch.concat([u[:, 0, None], u_0, u[:, -1, None]], dim=1) 91 | net.train() 92 | loss = mc_loss(u_0, net, p, config) 93 | loss.backward() 94 | optimizer.step() 95 | optimizer.zero_grad(set_to_none=True) 96 | scheduler.step() 97 | if step % 50 == 0: 98 | print('########################') 99 | print('training loss', step, loss.detach().item()) 100 | print('VAL_LOSS') 101 | print('########################') 102 | temp = test(config, net, val_data) 103 | if temp < val_loss: 104 | val_loss = temp 105 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 106 | print('TEST_LOSS') 107 | print('########################') 108 | test(config, net, test_data) 109 | sys.stdout.flush() 110 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 111 | print('FINAL_LOSS') 112 | print('########################') 113 | test(config, net, test_data) 114 | sys.stdout.flush() -------------------------------------------------------------------------------- /convection_diffusion/N10/data/config_pde.py: -------------------------------------------------------------------------------- 1 | class Config_E1(object): 2 | device = 'cuda:2' 3 | num_test = 200 4 | num_val = 200 5 | num_train = 1000 6 | total_time = 2 7 | record_steps = 200 8 | b = 0.1 9 | N = 10 10 | kappa = 0.005 11 | size = 64 12 | sup_size = 1024 13 | 14 | class Config_E2(object): 15 | device = 'cuda:2' 16 | num_test = 200 17 | num_val = 200 18 | num_train = 1000 19 | total_time = 2 20 | record_steps = 200 21 | b = 0.1 22 | N = 10 23 | kappa = 0.01 24 | size = 64 25 | sup_size = 1024 -------------------------------------------------------------------------------- /convection_diffusion/N10/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | import matplotlib.pyplot as plt 5 | import math 6 | import os 7 | from config_pde import * 8 | 9 | 10 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 11 | 12 | 13 | 14 | def setup_seed(seed): 15 | torch.manual_seed(seed) 16 | torch.cuda.manual_seed(seed) 17 | torch.cuda.manual_seed_all(seed) 18 | np.random.seed(seed) 19 | random.seed(seed) 20 | os.environ['PYTHONHASHSEED'] = str(seed) 21 | torch.backends.cudnn.deterministic = True 22 | torch.backends.cudnn.benchmark = False 23 | torch.backends.cudnn.enabled = True 24 | 25 | 26 | def generate_u0(b_size, N, num_grid, device='cuda'): 27 | ''' 28 | Input: 29 | b_size: int, num of pre-generated data 30 | N: int, highest frequency fourier basis; 31 | num_grid: int, number of grids; 32 | Output: 33 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization 34 | ''' 35 | B_n = torch.rand(b_size, N).to(device).reshape(b_size, N, 1, 1) 36 | grid = torch.tensor(np.linspace(0, 1-1/num_grid, num_grid)).to(device).reshape(1, 1, -1, 1) 37 | n = 2 * torch.Tensor(range(1, N + 1)).to(device).reshape(1, -1, 1, 1) 38 | sin_grid = torch.sin(torch.pi * n * grid) 39 | return (B_n * sin_grid).sum(axis=1) 40 | 41 | 42 | def psm_cde(u_0, kappa, b, T, record_steps, dt = 1e-5): 43 | ''' 44 | Input: 45 | u_0: torch.Tensor, size = (b_size, num_grid, 1), initialization; 46 | kappa: float, positive constant; 47 | b: float, convection term; 48 | T: final time 49 | Output: 50 | u: torch.Tensor, size = (b_size, num_grid, num_steps+1) 51 | ''' 52 | total_steps = int(T/dt) 53 | b_size, num_grid = u_0.shape[0], u_0.shape[1] 54 | u_psm = u_0 55 | u = u_0 56 | sub_t = int(total_steps // record_steps) 57 | for time_step in range(0, total_steps): 58 | u_h = torch.fft.fft(u, dim=1) 59 | k_max = num_grid//2 60 | k_x = torch.cat((torch.arange(start=0, end=k_max, step=1, device=u.device), 61 | torch.arange(start=-k_max, end=0, step=1, device=u.device)), 0).reshape(1,num_grid,1) 62 | ux_h = 2j *np.pi*k_x*u_h 63 | ux = torch.fft.irfft(ux_h[:, :, :k_max+1], dim=1, n=num_grid) 64 | uxx_h = 2j *np.pi*k_x*ux_h 65 | uxx = torch.fft.irfft(uxx_h[:, :, :k_max+1], dim=1, n=num_grid) 66 | u_1 = u + b * ux * dt + kappa * dt * uxx 67 | 68 | u_h = torch.fft.fft(u_1, dim=1) 69 | k_max = num_grid//2 70 | k_x = torch.cat((torch.arange(start=0, end=k_max, step=1, device=u.device), 71 | torch.arange(start=-k_max, end=0, step=1, device=u.device)), 0).reshape(1,num_grid,1) 72 | ux_h = 2j *np.pi*k_x*u_h 73 | ux = torch.fft.irfft(ux_h[:, :, :k_max+1], dim=1, n=num_grid) 74 | uxx_h = 2j *np.pi*k_x*ux_h 75 | uxx = torch.fft.irfft(uxx_h[:, :, :k_max+1], dim=1, n=num_grid) 76 | u_2 = u + b * ux * dt + kappa * dt * uxx 77 | u = (u_1+u_2)/2 78 | if (time_step+1) % sub_t == 0: 79 | u_psm = torch.concat([u_psm, u], dim=-1) 80 | return u_psm 81 | 82 | 83 | 84 | def check_directory() -> None: 85 | """ 86 | Check if log directory exists within experiments 87 | """ 88 | if not os.path.exists(f'dataset'): 89 | os.mkdir(f'dataset') 90 | 91 | 92 | 93 | check_directory() 94 | 95 | 96 | 97 | for experiment in ['E1', 'E2']: 98 | setup_seed(0) 99 | cfg = eval('Config_' + experiment + '()') 100 | device = cfg.device 101 | kappa = cfg.kappa 102 | b = cfg.b 103 | b_size = cfg.num_train 104 | N = cfg.N 105 | num_grid = cfg.sup_size 106 | sub_x = cfg.sup_size // cfg.size 107 | total_time = cfg.total_time 108 | record_steps = cfg.record_steps 109 | 110 | u_0 = generate_u0(b_size, N, num_grid, device) 111 | u = psm_cde(u_0, kappa, b, total_time, record_steps) 112 | u = u[:, ::sub_x, :] 113 | print(u.shape, u.max()) 114 | torch.save(u.float().permute(0, 2, 1), './dataset/' + experiment + '_train_data') 115 | 116 | 117 | for experiment in ['E1', 'E2']: 118 | setup_seed(1) 119 | cfg = eval('Config_' + experiment + '()') 120 | device = cfg.device 121 | kappa = cfg.kappa 122 | b = cfg.b 123 | b_size = cfg.num_test 124 | N = cfg.N 125 | num_grid = cfg.sup_size 126 | sub_x = cfg.sup_size // cfg.size 127 | total_time = cfg.total_time 128 | record_steps = cfg.record_steps 129 | 130 | u_0 = generate_u0(b_size, N, num_grid, device) 131 | u = psm_cde(u_0, kappa, b, total_time, record_steps) 132 | u = u[:, ::sub_x, :] 133 | print(u.shape, u.max()) 134 | torch.save(u.float().permute(0, 2, 1), './dataset/' + experiment + '_test_data') 135 | 136 | 137 | for experiment in ['E1', 'E2']: 138 | setup_seed(2) 139 | cfg = eval('Config_' + experiment + '()') 140 | device = cfg.device 141 | kappa = cfg.kappa 142 | b = cfg.b 143 | b_size = cfg.num_val 144 | N = cfg.N 145 | num_grid = cfg.sup_size 146 | sub_x = cfg.sup_size // cfg.size 147 | total_time = cfg.total_time 148 | 149 | u_0 = generate_u0(b_size, N, num_grid, device) 150 | u = psm_cde(u_0, kappa, b, total_time, record_steps) 151 | u = u[:, ::sub_x, :] 152 | print(u.shape, u.max()) 153 | torch.save(u.float().permute(0, 2, 1), './dataset/' + experiment + '_val_data') 154 | 155 | -------------------------------------------------------------------------------- /convection_diffusion/N10/mcnp10/config_pde.py: -------------------------------------------------------------------------------- 1 | class Config_E1(object): 2 | device = 'cuda:2' 3 | num_test = 200 4 | num_val = 200 5 | num_train = 1000 6 | total_time = 2 7 | record_steps = 200 8 | b = 0.1 9 | N = 10 10 | kappa = 0.005 11 | size = 64 12 | sup_size = 1024 13 | 14 | class Config_E2(object): 15 | device = 'cuda:2' 16 | num_test = 200 17 | num_val = 200 18 | num_train = 1000 19 | total_time = 2 20 | record_steps = 200 21 | b = 0.1 22 | N = 10 23 | kappa = 0.01 24 | size = 64 25 | sup_size = 1024 -------------------------------------------------------------------------------- /convection_diffusion/N10/mcnp10/fk_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import math 3 | import copy 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from scipy.stats import levy_stable 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def mc_loss(u_0, model, config): 17 | device = config.device 18 | size = config.size 19 | batch_size = config.batch_size 20 | 21 | u_input = u_0.reshape(u_0.shape[0], 1, size).repeat(1, config.time_steps, 1).reshape(-1, size, 1) 22 | t = config.delta_t * torch.arange(1, config.time_steps+1, device=device).reshape(1, -1).repeat(u_0.shape[0], 1).reshape(-1,) 23 | u = model(u_input, t).reshape(batch_size, -1, size) 24 | u = torch.concat([u_0.reshape(batch_size, 1, size), u], dim=1) 25 | 26 | kappa = config.kappa 27 | b = config.b 28 | pad = config.pad 29 | u_0, u_1 = u[:, :-1, :], u[:, 1:, :] 30 | delta_t = config.T / config.time_steps 31 | sigma = math.sqrt(2 * kappa * delta_t) 32 | x = torch.linspace(-pad, 1+pad-1/size, (2*pad+1) * size, device=device).reshape((2*pad+1), -1, 1) 33 | mu = torch.linspace(0, 1-1/size, size, device=device) + b * delta_t 34 | mu = (mu - mu.floor()).reshape(1, 1, -1) 35 | w = (1/(math.sqrt(2*torch.pi) * sigma) * torch.exp(-((x-mu)**2)/(2 * sigma**2)).sum(axis=0)) /size 36 | u_hat = torch.einsum('bts, sw->btw', u_0, w) 37 | return torch.sqrt(torch.mean(torch.square(u_hat - u_1))) -------------------------------------------------------------------------------- /convection_diffusion/N10/mcnp10/main.py: -------------------------------------------------------------------------------- 1 | 2 | from model import FNO 3 | from train import train 4 | from config_pde import * 5 | 6 | import json 7 | import sys 8 | import copy 9 | from datetime import datetime 10 | import random 11 | import argparse 12 | import numpy as np 13 | import torch 14 | import torch.optim as optim 15 | import matplotlib.pyplot as plt 16 | import os 17 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 18 | plt.rcParams["animation.html"] = "jshtml" 19 | 20 | 21 | def setup_seed(seed): 22 | torch.manual_seed(seed) 23 | torch.cuda.manual_seed(seed) 24 | torch.cuda.manual_seed_all(seed) 25 | np.random.seed(seed) 26 | random.seed(seed) 27 | os.environ['PYTHONHASHSEED'] = str(seed) 28 | torch.backends.cudnn.deterministic = True 29 | torch.backends.cudnn.benchmark = True 30 | torch.backends.cudnn.enabled = True 31 | 32 | 33 | 34 | def main(cfg): 35 | if not os.path.exists(f'log'): 36 | os.mkdir(f'log') 37 | if not os.path.exists(f'model'): 38 | os.mkdir(f'model') 39 | setup_seed(cfg.seed) 40 | dateTimeObj = datetime.now() 41 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}_{dateTimeObj.time().second}' 42 | cfg.timestring = timestring 43 | logfile = f'log/20_{cfg.experiment}_seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{timestring}.csv' 44 | 45 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 46 | json.dump(cfg.__dict__, f, indent=2) 47 | 48 | sys.stdout = open(logfile, 'w') 49 | 50 | print('--------args----------') 51 | for k in list(vars(cfg).keys()): 52 | print('%s: %s' % (k, vars(cfg)[k])) 53 | print('--------args----------\n') 54 | 55 | sys.stdout.flush() 56 | 57 | net = FNO(cfg.modes, 32) 58 | 59 | cfg.delta_t = cfg.T/cfg.time_steps 60 | train(cfg, net) 61 | torch.save(net.state_dict(), f'model/net_seed_{cfg.seed}_{timestring}.pt') 62 | 63 | 64 | if __name__ == "__main__": 65 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 66 | 67 | parser.add_argument('--data_path', type=str, 68 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 69 | help='path of data') 70 | 71 | parser.add_argument('--device', type=str, default='cuda:2', 72 | help='Used device') 73 | 74 | parser.add_argument('--seed', type=int, default=0, 75 | help='seed') 76 | 77 | parser.add_argument('--batch_size', type=int, default=200, 78 | help='batchsize of the operator learning') 79 | 80 | parser.add_argument('--step_size', type=int, default=1000, 81 | help='step_size of optim') 82 | 83 | parser.add_argument('--gamma', type=float, default=0.5, 84 | help='gamma of optim') 85 | 86 | parser.add_argument('--lr', type=float, default=0.01, 87 | help='lr of optim') 88 | 89 | parser.add_argument('--weight_decay', type=float, default=0.0, 90 | help='lr of optim') 91 | 92 | parser.add_argument('--num_iterations', type=int, default=10000, 93 | help='num_iterations of optim') 94 | 95 | parser.add_argument('--size', type=int, default=64, 96 | help='data spatial size') 97 | 98 | parser.add_argument('--T', type=float, default=2.0, 99 | help='final time') 100 | 101 | parser.add_argument('--time_steps', type=int, default=10, 102 | help='number of time_steps in data') 103 | 104 | parser.add_argument('--EPS', type=float, default=1e-8, 105 | help='epsilon') 106 | 107 | parser.add_argument('--experiment', type=str, default='E2', 108 | help='index of the experiments') 109 | 110 | parser.add_argument('--pad', type=int, default=1, 111 | help='num_copy in fk loss') 112 | 113 | parser.add_argument('--modes', type=int, default=20, 114 | help='modes in fno') 115 | 116 | for seed in [0, 1, 2]: 117 | for experiment in ['E2', 'E1']: 118 | cfg = parser.parse_args() 119 | cfg.seed = seed 120 | cfg.experiment = experiment 121 | cfg_pde = eval('Config_' + cfg.experiment + '()') 122 | cfg.N, cfg.kappa, cfg.b = cfg_pde.N, cfg_pde.kappa, cfg_pde.b 123 | main(cfg) 124 | -------------------------------------------------------------------------------- /convection_diffusion/N10/mcnp10/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.parameter import Parameter 6 | import matplotlib.pyplot as plt 7 | 8 | import operator 9 | from functools import reduce 10 | from functools import partial 11 | from timeit import default_timer 12 | 13 | 14 | ################################################################ 15 | # 1d fourier layer 16 | ################################################################ 17 | class SpectralConv1d(nn.Module): 18 | def __init__(self, in_channels, out_channels, modes1): 19 | super(SpectralConv1d, self).__init__() 20 | 21 | """ 22 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 23 | """ 24 | 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 28 | 29 | self.scale = (1 / (in_channels*out_channels)) 30 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 31 | 32 | # Complex multiplication 33 | def compl_mul1d(self, input, weights): 34 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 35 | return torch.einsum("bix,iox->box", input, weights) 36 | 37 | def forward(self, x): 38 | batchsize = x.shape[0] 39 | #Compute Fourier coeffcients up to factor of e^(- something constant) 40 | x_ft = torch.fft.rfft(x) 41 | 42 | # Multiply relevant Fourier modes 43 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 44 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 45 | 46 | #Return to physical space 47 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 48 | return x 49 | 50 | class FNO(nn.Module): 51 | def __init__(self, modes, width): 52 | super(FNO, self).__init__() 53 | 54 | """ 55 | The overall network. It contains 4 layers of the Fourier layer. 56 | 1. Lift the input to the desire channel dimension by self.fc0 . 57 | 2. 4 layers of the integral operators u' = (W + K)(u). 58 | W defined by self.w; K defined by self.conv . 59 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 60 | 61 | input: the solution of the initial condition and location (a(x), x) 62 | input shape: (batchsize, x=s, c=2) 63 | output: the solution of a later timestep 64 | output shape: (batchsize, x=s, c=1) 65 | """ 66 | 67 | self.modes1 = modes 68 | self.width = width 69 | self.fc0 = nn.Linear(3, self.width) # input channel is 2: (a(x), x) 70 | 71 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 74 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 75 | self.w0 = nn.Conv1d(self.width, self.width, 1) 76 | self.w1 = nn.Conv1d(self.width, self.width, 1) 77 | self.w2 = nn.Conv1d(self.width, self.width, 1) 78 | self.w3 = nn.Conv1d(self.width, self.width, 1) 79 | 80 | self.fc1 = nn.Linear(self.width, 128) 81 | self.fc2 = nn.Linear(128, 1) 82 | 83 | def forward(self, x, t): 84 | grid = self.get_grid(x.shape, x.device) 85 | time = t.reshape(-1, 1, 1).repeat(1, x.shape[1], 1) 86 | x = torch.cat((x, grid), dim=-1) 87 | x = torch.concat([x, time], dim=-1) 88 | x = self.fc0(x) 89 | x = x.permute(0, 2, 1) 90 | 91 | x1 = self.conv0(x) 92 | x2 = self.w0(x) 93 | x = x1 + x2 94 | x = F.gelu(x) 95 | 96 | x1 = self.conv1(x) 97 | x2 = self.w1(x) 98 | x = x1 + x2 99 | x = F.gelu(x) 100 | 101 | x1 = self.conv2(x) 102 | x2 = self.w2(x) 103 | x = x1 + x2 104 | x = F.gelu(x) 105 | 106 | x1 = self.conv3(x) 107 | x2 = self.w3(x) 108 | x = x1 + x2 109 | x = x.permute(0, 2, 1) 110 | x = self.fc1(x) 111 | x = F.gelu(x) 112 | x = self.fc2(x) 113 | return x 114 | 115 | def get_grid(self, shape, device): 116 | batchsize, size_x = shape[0], shape[1] 117 | gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) 118 | gridx = gridx.reshape(1, size_x, 1).repeat([batchsize, 1, 1]) 119 | return gridx.to(device) -------------------------------------------------------------------------------- /convection_diffusion/N10/mcnp10/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_sol(sol): 11 | ''' 12 | sol: torch.Tensor, size = (num_grid, ) 13 | return: figs 14 | ''' 15 | temp = sol.detach().cpu().numpy().reshape(-1, ) 16 | num_grid = temp.shape[0] - 1 17 | x_axis = (1 / num_grid) * torch.Tensor(range(num_grid + 1)) 18 | plt.figure() 19 | plt.plot(x_axis.numpy(), temp.numpy()) 20 | plt.show() 21 | 22 | 23 | def setup_seed(seed): 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | np.random.seed(seed) 28 | random.seed(seed) 29 | os.environ['PYTHONHASHSEED'] = str(seed) 30 | torch.backends.cudnn.deterministic = True 31 | torch.backends.cudnn.benchmark = False 32 | torch.backends.cudnn.enabled = True -------------------------------------------------------------------------------- /convection_diffusion/N10/mcnp10/train.py: -------------------------------------------------------------------------------- 1 | from mc_loss import mc_loss 2 | 3 | import sys 4 | import math 5 | import copy 6 | import random 7 | import numpy as np 8 | import torch 9 | import torch.optim as optim 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def test(config, net, test_data): 17 | device = config.device 18 | num = test_data.shape[0] 19 | delta_steps = 10 20 | delta_t = config.T / delta_steps 21 | u_0 = test_data[:, 0, :].to(device)[:, :, None] 22 | rela_err = torch.zeros(delta_steps) 23 | rela_err_1 = torch.zeros(delta_steps) 24 | rela_err_max = torch.zeros(delta_steps) 25 | for time_step in range(1, 11): 26 | net.eval() 27 | t = (time_step*delta_t) * torch.ones(num).to(device) 28 | u = net(u_0, t).detach() 29 | u_t = test_data[:, 2 * delta_steps * time_step, :].to(device)[:, :, None] 30 | rela_err[time_step-1] = (torch.norm((u - u_t).reshape(u.shape[0], -1), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), dim=1)).mean() 31 | rela_err_1[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=1, dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=1, dim=1)).mean() 32 | rela_err_max[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=float('inf'), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=float('inf'), dim=1)).mean() 33 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 34 | print('mean relative l_2 error', rela_err.mean().item()) 35 | print('mean relative l_1 error', rela_err_1.mean().item()) 36 | print('mean relative l_inf error', rela_err_max.mean().item()) 37 | return rela_err.mean().item() 38 | 39 | 40 | 41 | def train(config, net): 42 | device = config.device 43 | size = config.size 44 | batch_size = config.batch_size 45 | N = config.N 46 | grid = torch.tensor(np.linspace(0, 1-1/size, size), dtype=torch.float).to(device).reshape(1, 1, -1, 1) 47 | n = 2 * torch.Tensor(range(1, N + 1)).to(device).reshape(1, -1, 1, 1) 48 | sin_grid = torch.sin(torch.pi * n * grid) 49 | net = net.to(device) 50 | optimizer = optim.Adam(net.parameters(), config.lr) 51 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 52 | val_data = torch.load(config.data_path + config.experiment + '_val_data').to(device) 53 | test_data = torch.load(config.data_path + config.experiment + '_test_data').to(device) 54 | val_loss = 1e50 55 | for step in range(config.num_iterations+1): 56 | B_n = torch.rand(batch_size, N, 1, 1).to(device) 57 | u_0 = (B_n * sin_grid).sum(axis=1) 58 | net.train() 59 | loss = mc_loss(u_0, net, config) 60 | loss.backward() 61 | optimizer.step() 62 | optimizer.zero_grad(set_to_none=True) 63 | scheduler.step() 64 | if step % 50 == 0: 65 | print('########################') 66 | print('training loss', step, loss.detach().item()) 67 | print('VAL_LOSS') 68 | print('########################') 69 | temp = test(config, net, val_data) 70 | if temp < val_loss: 71 | val_loss = temp 72 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 73 | print('TEST_LOSS') 74 | print('########################') 75 | test(config, net, test_data) 76 | sys.stdout.flush() 77 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 78 | print('FINAL_LOSS') 79 | print('########################') 80 | test(config, net, test_data) 81 | sys.stdout.flush() -------------------------------------------------------------------------------- /convection_diffusion/N5/data/config_pde.py: -------------------------------------------------------------------------------- 1 | class Config_E1(object): 2 | device = 'cuda:2' 3 | num_test = 200 4 | num_val = 200 5 | num_train = 1000 6 | total_time = 2 7 | record_steps = 200 8 | b = 0.1 9 | N = 5 10 | kappa = 0.005 11 | size = 64 12 | sup_size = 1024 13 | 14 | class Config_E2(object): 15 | device = 'cuda:2' 16 | num_test = 200 17 | num_val = 200 18 | num_train = 1000 19 | total_time = 2 20 | record_steps = 200 21 | b = 0.1 22 | N = 5 23 | kappa = 0.01 24 | size = 64 25 | sup_size = 1024 -------------------------------------------------------------------------------- /convection_diffusion/N5/mcnp10/config_pde.py: -------------------------------------------------------------------------------- 1 | class Config_E1(object): 2 | device = 'cuda:2' 3 | num_test = 200 4 | num_val = 200 5 | num_train = 1000 6 | total_time = 2 7 | record_steps = 200 8 | b = 0.1 9 | N = 5 10 | kappa = 0.005 11 | size = 64 12 | sup_size = 1024 13 | 14 | class Config_E2(object): 15 | device = 'cuda:2' 16 | num_test = 200 17 | num_val = 200 18 | num_train = 1000 19 | total_time = 2 20 | record_steps = 200 21 | b = 0.1 22 | N = 5 23 | kappa = 0.01 24 | size = 64 25 | sup_size = 1024 -------------------------------------------------------------------------------- /convection_diffusion/N5/mcnp10/main.py: -------------------------------------------------------------------------------- 1 | 2 | from model import FNO 3 | from train import train 4 | from config_pde import * 5 | 6 | import json 7 | import sys 8 | import copy 9 | from datetime import datetime 10 | import random 11 | import argparse 12 | import numpy as np 13 | import torch 14 | import torch.optim as optim 15 | import matplotlib.pyplot as plt 16 | import os 17 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 18 | plt.rcParams["animation.html"] = "jshtml" 19 | 20 | 21 | def setup_seed(seed): 22 | torch.manual_seed(seed) 23 | torch.cuda.manual_seed(seed) 24 | torch.cuda.manual_seed_all(seed) 25 | np.random.seed(seed) 26 | random.seed(seed) 27 | os.environ['PYTHONHASHSEED'] = str(seed) 28 | torch.backends.cudnn.deterministic = True 29 | torch.backends.cudnn.benchmark = True 30 | torch.backends.cudnn.enabled = True 31 | 32 | 33 | 34 | def main(cfg): 35 | if not os.path.exists(f'log'): 36 | os.mkdir(f'log') 37 | if not os.path.exists(f'model'): 38 | os.mkdir(f'model') 39 | setup_seed(cfg.seed) 40 | dateTimeObj = datetime.now() 41 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}_{dateTimeObj.time().second}' 42 | cfg.timestring = timestring 43 | logfile = f'log/{cfg.experiment}_seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{timestring}.csv' 44 | 45 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 46 | json.dump(cfg.__dict__, f, indent=2) 47 | 48 | sys.stdout = open(logfile, 'w') 49 | 50 | print('--------args----------') 51 | for k in list(vars(cfg).keys()): 52 | print('%s: %s' % (k, vars(cfg)[k])) 53 | print('--------args----------\n') 54 | 55 | sys.stdout.flush() 56 | 57 | net = FNO(cfg.modes, 32) 58 | 59 | cfg.delta_t = cfg.T/cfg.time_steps 60 | train(cfg, net) 61 | torch.save(net.state_dict(), f'model/net_seed_{cfg.seed}_{timestring}.pt') 62 | 63 | 64 | if __name__ == "__main__": 65 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 66 | 67 | parser.add_argument('--data_path', type=str, 68 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 69 | help='path of data') 70 | 71 | parser.add_argument('--device', type=str, default='cuda:2', 72 | help='Used device') 73 | 74 | parser.add_argument('--seed', type=int, default=0, 75 | help='seed') 76 | 77 | parser.add_argument('--batch_size', type=int, default=200, 78 | help='batchsize of the operator learning') 79 | 80 | parser.add_argument('--step_size', type=int, default=1000, 81 | help='step_size of optim') 82 | 83 | parser.add_argument('--gamma', type=float, default=0.5, 84 | help='gamma of optim') 85 | 86 | parser.add_argument('--lr', type=float, default=0.01, 87 | help='lr of optim') 88 | 89 | parser.add_argument('--weight_decay', type=float, default=0.0, 90 | help='lr of optim') 91 | 92 | parser.add_argument('--num_iterations', type=int, default=10000, 93 | help='num_iterations of optim') 94 | 95 | parser.add_argument('--size', type=int, default=64, 96 | help='data spatial size') 97 | 98 | parser.add_argument('--T', type=float, default=2.0, 99 | help='final time') 100 | 101 | parser.add_argument('--time_steps', type=int, default=10, 102 | help='number of time_steps in data') 103 | 104 | parser.add_argument('--EPS', type=float, default=1e-8, 105 | help='epsilon') 106 | 107 | parser.add_argument('--experiment', type=str, default='E1', 108 | help='index of the experiments') 109 | 110 | parser.add_argument('--pad', type=int, default=1, 111 | help='num_copy in fk loss') 112 | 113 | parser.add_argument('--modes', type=int, default=16, 114 | help='modes in fno') 115 | 116 | 117 | for seed in [0, 1, 2]: 118 | for experiment in ['E2', 'E1']: 119 | cfg = parser.parse_args() 120 | cfg.seed = seed 121 | cfg.experiment = experiment 122 | cfg_pde = eval('Config_' + cfg.experiment + '()') 123 | cfg.N, cfg.kappa, cfg.b = cfg_pde.N, cfg_pde.kappa, cfg_pde.b 124 | main(cfg) -------------------------------------------------------------------------------- /convection_diffusion/N5/mcnp10/mc_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import math 3 | import copy 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from scipy.stats import levy_stable 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def mc_loss(u_0, model, config): 17 | device = config.device 18 | size = config.size 19 | batch_size = config.batch_size 20 | 21 | u_input = u_0.reshape(u_0.shape[0], 1, size).repeat(1, config.time_steps, 1).reshape(-1, size, 1) 22 | t = config.delta_t * torch.arange(1, config.time_steps+1, device=device).reshape(1, -1).repeat(u_0.shape[0], 1).reshape(-1,) 23 | u = model(u_input, t).reshape(batch_size, -1, size) 24 | u = torch.concat([u_0.reshape(batch_size, 1, size), u], dim=1) 25 | 26 | kappa = config.kappa 27 | b = config.b 28 | pad = config.pad 29 | u_0, u_1 = u[:, :-1, :], u[:, 1:, :] 30 | delta_t = config.T / config.time_steps 31 | sigma = math.sqrt(2 * kappa * delta_t) 32 | x = torch.linspace(-pad, 1+pad-1/size, (2*pad+1) * size, device=device).reshape((2*pad+1), -1, 1) 33 | mu = torch.linspace(0, 1-1/size, size, device=device) + b * delta_t 34 | mu = (mu - mu.floor()).reshape(1, 1, -1) 35 | w = (1/(math.sqrt(2*torch.pi) * sigma) * torch.exp(-((x-mu)**2)/(2 * sigma**2)).sum(axis=0)) /size 36 | u_hat = torch.einsum('bts, sw->btw', u_0, w) 37 | return torch.sqrt(torch.mean(torch.square(u_hat - u_1))) -------------------------------------------------------------------------------- /convection_diffusion/N5/mcnp10/model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.parameter import Parameter 6 | import matplotlib.pyplot as plt 7 | 8 | import operator 9 | from functools import reduce 10 | from functools import partial 11 | from timeit import default_timer 12 | 13 | 14 | ################################################################ 15 | # 1d fourier layer 16 | ################################################################ 17 | class SpectralConv1d(nn.Module): 18 | def __init__(self, in_channels, out_channels, modes1): 19 | super(SpectralConv1d, self).__init__() 20 | 21 | """ 22 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 23 | """ 24 | 25 | self.in_channels = in_channels 26 | self.out_channels = out_channels 27 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 28 | 29 | self.scale = (1 / (in_channels*out_channels)) 30 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 31 | 32 | # Complex multiplication 33 | def compl_mul1d(self, input, weights): 34 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 35 | return torch.einsum("bix,iox->box", input, weights) 36 | 37 | def forward(self, x): 38 | batchsize = x.shape[0] 39 | #Compute Fourier coeffcients up to factor of e^(- something constant) 40 | x_ft = torch.fft.rfft(x) 41 | 42 | # Multiply relevant Fourier modes 43 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 44 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 45 | 46 | #Return to physical space 47 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 48 | return x 49 | 50 | class FNO(nn.Module): 51 | def __init__(self, modes, width): 52 | super(FNO, self).__init__() 53 | 54 | """ 55 | The overall network. It contains 4 layers of the Fourier layer. 56 | 1. Lift the input to the desire channel dimension by self.fc0 . 57 | 2. 4 layers of the integral operators u' = (W + K)(u). 58 | W defined by self.w; K defined by self.conv . 59 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 60 | 61 | input: the solution of the initial condition and location (a(x), x) 62 | input shape: (batchsize, x=s, c=2) 63 | output: the solution of a later timestep 64 | output shape: (batchsize, x=s, c=1) 65 | """ 66 | 67 | self.modes1 = modes 68 | self.width = width 69 | self.fc0 = nn.Linear(3, self.width) # input channel is 2: (a(x), x) 70 | 71 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 74 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 75 | self.w0 = nn.Conv1d(self.width, self.width, 1) 76 | self.w1 = nn.Conv1d(self.width, self.width, 1) 77 | self.w2 = nn.Conv1d(self.width, self.width, 1) 78 | self.w3 = nn.Conv1d(self.width, self.width, 1) 79 | 80 | self.fc1 = nn.Linear(self.width, 128) 81 | self.fc2 = nn.Linear(128, 1) 82 | 83 | def forward(self, x, t): 84 | grid = self.get_grid(x.shape, x.device) 85 | time = t.reshape(-1, 1, 1).repeat(1, x.shape[1], 1) 86 | x = torch.cat((x, grid), dim=-1) 87 | x = torch.concat([x, time], dim=-1) 88 | x = self.fc0(x) 89 | x = x.permute(0, 2, 1) 90 | 91 | x1 = self.conv0(x) 92 | x2 = self.w0(x) 93 | x = x1 + x2 94 | x = F.gelu(x) 95 | 96 | x1 = self.conv1(x) 97 | x2 = self.w1(x) 98 | x = x1 + x2 99 | x = F.gelu(x) 100 | 101 | x1 = self.conv2(x) 102 | x2 = self.w2(x) 103 | x = x1 + x2 104 | x = F.gelu(x) 105 | 106 | x1 = self.conv3(x) 107 | x2 = self.w3(x) 108 | x = x1 + x2 109 | x = x.permute(0, 2, 1) 110 | x = self.fc1(x) 111 | x = F.gelu(x) 112 | x = self.fc2(x) 113 | return x 114 | 115 | def get_grid(self, shape, device): 116 | batchsize, size_x = shape[0], shape[1] 117 | gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) 118 | gridx = gridx.reshape(1, size_x, 1).repeat([batchsize, 1, 1]) 119 | return gridx.to(device) -------------------------------------------------------------------------------- /convection_diffusion/N5/mcnp10/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_sol(sol): 11 | ''' 12 | sol: torch.Tensor, size = (num_grid, ) 13 | return: figs 14 | ''' 15 | temp = sol.detach().cpu().numpy().reshape(-1, ) 16 | num_grid = temp.shape[0] - 1 17 | x_axis = (1 / num_grid) * torch.Tensor(range(num_grid + 1)) 18 | plt.figure() 19 | plt.plot(x_axis.numpy(), temp.numpy()) 20 | plt.show() 21 | 22 | 23 | def setup_seed(seed): 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | np.random.seed(seed) 28 | random.seed(seed) 29 | os.environ['PYTHONHASHSEED'] = str(seed) 30 | torch.backends.cudnn.deterministic = True 31 | torch.backends.cudnn.benchmark = False 32 | torch.backends.cudnn.enabled = True -------------------------------------------------------------------------------- /convection_diffusion/N5/mcnp10/train.py: -------------------------------------------------------------------------------- 1 | from mc_loss import mc_loss 2 | 3 | import sys 4 | import math 5 | import copy 6 | import random 7 | import numpy as np 8 | import torch 9 | import torch.optim as optim 10 | import matplotlib.pyplot as plt 11 | import os 12 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 13 | plt.rcParams["animation.html"] = "jshtml" 14 | 15 | 16 | def test(config, net, test_data): 17 | device = config.device 18 | num = test_data.shape[0] 19 | delta_steps = 10 20 | delta_t = config.T / delta_steps 21 | u_0 = test_data[:, 0, :].to(device)[:, :, None] 22 | rela_err = torch.zeros(delta_steps) 23 | rela_err_1 = torch.zeros(delta_steps) 24 | rela_err_max = torch.zeros(delta_steps) 25 | for time_step in range(1, 11): 26 | net.eval() 27 | t = (time_step*delta_t) * torch.ones(num).to(device) 28 | u = net(u_0, t).detach() 29 | u_t = test_data[:, 2 * delta_steps * time_step, :].to(device)[:, :, None] 30 | rela_err[time_step-1] = (torch.norm((u - u_t).reshape(u.shape[0], -1), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), dim=1)).mean() 31 | rela_err_1[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=1, dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=1, dim=1)).mean() 32 | rela_err_max[time_step-1] = (torch.norm((u- u_t).reshape(u.shape[0], -1), p=float('inf'), dim=1) / torch.norm(u_t.reshape(u.shape[0], -1), p=float('inf'), dim=1)).mean() 33 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 34 | print('mean relative l_2 error', rela_err.mean().item()) 35 | print('mean relative l_1 error', rela_err_1.mean().item()) 36 | print('mean relative l_inf error', rela_err_max.mean().item()) 37 | return rela_err.mean().item() 38 | 39 | 40 | 41 | def train(config, net): 42 | device = config.device 43 | size = config.size 44 | batch_size = config.batch_size 45 | N = config.N 46 | grid = torch.tensor(np.linspace(0, 1-1/size, size), dtype=torch.float).to(device).reshape(1, 1, -1, 1) 47 | n = 2 * torch.Tensor(range(1, N + 1)).to(device).reshape(1, -1, 1, 1) 48 | sin_grid = torch.sin(torch.pi * n * grid) 49 | net = net.to(device) 50 | optimizer = optim.Adam(net.parameters(), config.lr) 51 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 52 | val_data = torch.load(config.data_path + config.experiment + '_val_data').to(device) 53 | test_data = torch.load(config.data_path + config.experiment + '_test_data').to(device) 54 | val_loss = 1e50 55 | for step in range(config.num_iterations+1): 56 | B_n = torch.rand(batch_size, N, 1, 1).to(device) 57 | u_0 = (B_n * sin_grid).sum(axis=1) 58 | net.train() 59 | loss = mc_loss(u_0, net, config) 60 | loss.backward() 61 | optimizer.step() 62 | optimizer.zero_grad(set_to_none=True) 63 | scheduler.step() 64 | if step % 50 == 0: 65 | print('########################') 66 | print('training loss', step, loss.detach().item()) 67 | print('VAL_LOSS') 68 | print('########################') 69 | temp = test(config, net, val_data) 70 | if temp < val_loss: 71 | val_loss = temp 72 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 73 | print('TEST_LOSS') 74 | print('########################') 75 | test(config, net, test_data) 76 | sys.stdout.flush() 77 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 78 | print('FINAL_LOSS') 79 | print('########################') 80 | test(config, net, test_data) 81 | sys.stdout.flush() -------------------------------------------------------------------------------- /navier-stokes/E1/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | import math 5 | import os 6 | import time 7 | from nse import navier_stokes_2d, GaussianRF 8 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 9 | 10 | 11 | def check_directory() -> None: 12 | """ 13 | Check if log directory exists within experiments 14 | """ 15 | if not os.path.exists(f'dataset'): 16 | os.mkdir(f'dataset') 17 | 18 | 19 | def setup_seed(seed): 20 | torch.manual_seed(seed) 21 | torch.cuda.manual_seed(seed) 22 | torch.cuda.manual_seed_all(seed) 23 | np.random.seed(seed) 24 | random.seed(seed) 25 | os.environ['PYTHONHASHSEED'] = str(seed) 26 | torch.backends.cudnn.deterministic = True 27 | torch.backends.cudnn.benchmark = False 28 | torch.backends.cudnn.enabled = True 29 | 30 | 31 | 32 | check_directory() 33 | device = torch.device('cuda:1') 34 | 35 | s = 256 36 | sub = 4 37 | 38 | # Set up 2d GRF with covariance parameters 39 | GRF = GaussianRF(s, device=device) 40 | 41 | # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y))) 42 | t = torch.linspace(0, 1, s+1, device=device) 43 | t = t[0: -1] 44 | 45 | X, Y = torch.meshgrid(t, t, indexing='ij') 46 | 47 | # Number of snapshots from solution 48 | 49 | 50 | setup_seed(0) 51 | record_steps = 200 52 | T = 20 53 | N = 200 54 | bsize = 200 55 | c = 0 56 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 57 | f = 0.1*(torch.sin(2*math.pi*(X + Y)) + torch.cos(2*math.pi*(X + Y))) 58 | 59 | 60 | for j in range(N//bsize): 61 | # Sample random feilds 62 | w0 = GRF(bsize) 63 | visc = 10 ** (- 4 * torch.ones(bsize).to(device)) 64 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 65 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 66 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 67 | u[c:(c+bsize),...] = sol 68 | c += bsize 69 | print(j, c) 70 | torch.save(u, './dataset/data_test') 71 | 72 | setup_seed(1) 73 | record_steps = 200 74 | T = 20 75 | N = 200 76 | bsize = 200 77 | c = 0 78 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 79 | f = 0.1*(torch.sin(2*math.pi*(X + Y)) + torch.cos(2*math.pi*(X + Y))) 80 | 81 | 82 | for j in range(N//bsize): 83 | # Sample random feilds 84 | w0 = GRF(bsize) 85 | visc = 10 ** (- 4 * torch.ones(bsize).to(device)) 86 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 87 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 88 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 89 | u[c:(c+bsize),...] = sol 90 | c += bsize 91 | print(j, c) 92 | torch.save(u, './dataset/data_val') 93 | -------------------------------------------------------------------------------- /navier-stokes/E1/data/nse.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import os 4 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 5 | 6 | class GaussianRF(object): 7 | 8 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 9 | self.dim = 2 10 | self.device = device 11 | 12 | if sigma is None: 13 | sigma = tau**(0.5*(2*alpha - self.dim)) 14 | k_max = size//2 15 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 16 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 17 | 18 | k_x = wavenumers.transpose(0,1) 19 | k_y = wavenumers 20 | 21 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 22 | self.sqrt_eig[0,0] = 0.0 23 | 24 | self.size = [] 25 | for j in range(self.dim): 26 | self.size.append(size) 27 | 28 | self.size = tuple(self.size) 29 | 30 | def __call__(self, N): 31 | 32 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 33 | coeff = self.sqrt_eig * coeff 34 | 35 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 36 | 37 | 38 | def force(x): 39 | z = x.sum(axis=1).reshape(x.shape[0], 1) 40 | return 0.1 * torch.sin(2 * math.pi * z) + 0.1 * torch.cos(2 * math.pi * z) 41 | 42 | 43 | def navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1): 44 | 45 | # Grid size - must be power of 2 46 | N = w0.size()[-1] 47 | 48 | # Maximum frequency 49 | k_max = math.floor(N/2.0) 50 | 51 | # Number of steps to final time 52 | steps = math.ceil(T/delta_t) 53 | 54 | # Initial vorticity to Fourier space 55 | w_h = torch.fft.rfft2(w0) 56 | 57 | # Forcing to Fourier space 58 | f_h = torch.fft.rfft2(f) 59 | 60 | #If same forcing for the whole batch 61 | if len(f_h.size()) < len(w_h.size()): 62 | f_h = torch.unsqueeze(f_h, 0) 63 | 64 | # Record solution every this number of steps 65 | record_time = math.floor(steps/record_steps) 66 | 67 | # Wavenumbers in y-direction 68 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device), torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N,1) 69 | # Wavenumbers in x-direction 70 | k_x = k_y.transpose(0,1) 71 | 72 | # Truncate redundant modes 73 | k_x = k_x[..., :k_max + 1] 74 | k_y = k_y[..., :k_max + 1] 75 | 76 | # Negative Laplacian in Fourier space 77 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 78 | lap[0,0] = 1.0 79 | # Dealiasing mask 80 | dealias = torch.unsqueeze(torch.logical_and(torch.abs(k_y) <= (2.0/3.0)*k_max, torch.abs(k_x) <= (2.0/3.0)*k_max).float(), 0) 81 | 82 | # Saving solution and time 83 | sol = torch.zeros(*w0.size(), record_steps, device=w0.device) 84 | sol_t = torch.zeros(record_steps, device=w0.device) 85 | 86 | # Record counter 87 | c = 0 88 | # Physical time 89 | t = 0.0 90 | for j in range(steps): 91 | # Stream function in Fourier space: solve Poisson equation 92 | psi_h = w_h / lap 93 | 94 | # Velocity field in x-direction = psi_y 95 | q = 2. * math.pi * k_y * 1j * psi_h 96 | q = torch.fft.irfft2(q, s=(N, N)) 97 | 98 | # Velocity field in y-direction = -psi_x 99 | v = -2. * math.pi * k_x * 1j * psi_h 100 | v = torch.fft.irfft2(v, s=(N, N)) 101 | 102 | # Partial x of vorticity 103 | w_x = 2. * math.pi * k_x * 1j * w_h 104 | w_x = torch.fft.irfft2(w_x, s=(N, N)) 105 | 106 | # Partial y of vorticity 107 | w_y = 2. * math.pi * k_y * 1j * w_h 108 | w_y = torch.fft.irfft2(w_y, s=(N, N)) 109 | 110 | # Non-linear term (u.grad(w)): compute in physical space then back to Fourier space 111 | F_h = torch.fft.rfft2(q*w_x + v*w_y) 112 | 113 | # Dealias 114 | F_h = dealias* F_h 115 | 116 | # Crank-Nicolson update 117 | w_h = (-delta_t*F_h + delta_t*f_h + (1.0 - 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...])*w_h)/(1.0 + 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...]) 118 | 119 | # Update real time (used only for recording) 120 | t += delta_t 121 | 122 | if (j+1) % record_time == 0: 123 | # Solution in physical space 124 | w = torch.fft.irfft2(w_h, s=(N, N)) 125 | 126 | # Record solution and time 127 | sol[...,c] = w 128 | sol_t[c] = t 129 | 130 | c += 1 131 | return sol, sol_t 132 | 133 | -------------------------------------------------------------------------------- /navier-stokes/E1/mcnp10/main.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from train import train 3 | from tools import setup_seed 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def main(cfg): 21 | if not os.path.exists(f'log'): 22 | os.mkdir(f'log') 23 | if not os.path.exists(f'model'): 24 | os.mkdir(f'model') 25 | setup_seed(cfg.seed) 26 | dateTimeObj = datetime.now() 27 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}' 28 | logfile = f'log/seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{timestring}.csv' 29 | 30 | cfg.timestring = timestring 31 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 32 | json.dump(cfg.__dict__, f, indent=2) 33 | 34 | sys.stdout = open(logfile, 'w') 35 | 36 | print('--------args----------') 37 | for k in list(vars(cfg).keys()): 38 | print('%s: %s' % (k, vars(cfg)[k])) 39 | print('--------args----------\n') 40 | 41 | sys.stdout.flush() 42 | cfg.delta_t = cfg.T/cfg.time_steps 43 | net = FNO(cfg.NUM_TRUNCATION, cfg.NUM_TRUNCATION, cfg.num_channel) 44 | train(cfg, net) 45 | 46 | 47 | if __name__ == "__main__": 48 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 49 | 50 | parser.add_argument('--data_path', type=str, 51 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 52 | help='path of data') 53 | 54 | parser.add_argument('--device', type=str, default='cuda:0', 55 | help='Used device') 56 | 57 | parser.add_argument('--seed', type=int, default=0, 58 | help='seed') 59 | 60 | parser.add_argument('--batch_size', type=int, default=10, 61 | help='batchsize of the operator learning') 62 | 63 | parser.add_argument('--step_size', type=int, default=2000, 64 | help='step_size of optim') 65 | 66 | parser.add_argument('--gamma', type=float, default=0.8, 67 | help='gamma of optim') 68 | 69 | parser.add_argument('--lr', type=float, default=0.01, 70 | help='lr of optim') 71 | 72 | parser.add_argument('--weight_decay', type=float, default=0.0, 73 | help='lr of optim') 74 | 75 | parser.add_argument('--num_iterations', type=int, default=20000, 76 | help='num_iterations of optim') 77 | 78 | parser.add_argument('--nu', type=float, default=1/10000, 79 | help='nu in NSE') 80 | 81 | parser.add_argument('--size', type=int, default=64, 82 | help='data spatial size') 83 | 84 | parser.add_argument('--T', type=float, default=10.0, 85 | help='final time') 86 | 87 | parser.add_argument('--time_steps', type=int, default=10, 88 | help='number of time_steps in data') 89 | 90 | parser.add_argument('--EPS', type=float, default=1e-8, 91 | help='epsilon') 92 | 93 | parser.add_argument('--NUM_TRUNCATION', type=int, default=16, 94 | help='NUM_TRUNCATION in FNO') 95 | 96 | parser.add_argument('--num_channel', type=int, default=36, 97 | help='num_channel in FNO') 98 | 99 | parser.add_argument('--sup_u', type=int, default=2, 100 | help='sup_u') 101 | 102 | parser.add_argument('--sup_w', type=int, default=1, 103 | help='sup_w') 104 | 105 | parser.add_argument('--k', type=int, default=3, 106 | help='k') 107 | 108 | cfg = parser.parse_args() 109 | for seed in [0, 1, 2]: 110 | cfg.seed = seed 111 | main(cfg) -------------------------------------------------------------------------------- /navier-stokes/E1/mcnp10/mc_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import numpy as np 4 | from tools import w2v_g 5 | 6 | 7 | 8 | def mc_loss(w_0, x, f_1, model, config): 9 | device = config.device 10 | size = config.size 11 | sup_w, sup_u = config.sup_w, config.sup_u 12 | sup_size_w, sup_size_u = sup_w * size, sup_u * size 13 | batch_size = config.batch_size 14 | time_steps = config.time_steps 15 | delta_t = config.delta_t 16 | k = config.k 17 | nu = config.nu 18 | sigma = math.sqrt(2 * nu * delta_t) 19 | w_0 = w_0.reshape([1, batch_size, size, size, 1]) 20 | 21 | w_input = w_0.repeat(time_steps, 1, 1, 1, 1).reshape(-1, size, size, 1) 22 | t = delta_t * torch.arange(1, time_steps+1, device=device).reshape(-1, 1).repeat(1,batch_size).reshape(-1,) 23 | w = model(w_input, t).reshape(time_steps, batch_size, size, size, 1) 24 | w = torch.concat([w_0.reshape(1, -1, size, size, 1), w], dim=0) 25 | 26 | w_0, w_1 = w[:-1, ...], w[1:, ...] 27 | v = w2v_g(w.reshape(-1, size, size, 1), size).reshape(time_steps+1, batch_size, size, size, 2) 28 | v_0, v_1 = v[:-1, ...], v[1:, ...] 29 | 30 | if sup_u > 1: 31 | v_0 = torch.fft.irfft(torch.fft.rfft(v_0, dim=2), dim=2, n=sup_size_u) 32 | v_0 = sup_u ** 2 * torch.fft.irfft(torch.fft.rfft(v_0, dim=3), dim=3, n=sup_size_u) 33 | mu_1 = x - v_1 * delta_t 34 | index = mu_1 - mu_1.floor() 35 | index = (index * (sup_size_u)).round().reshape(-1, 2) 36 | index[index==sup_size_u] = 0 37 | index = index.long() 38 | index_batch = torch.arange(batch_size * time_steps, device=device).reshape(-1, 1).repeat(1, size**2).reshape(-1,).long() 39 | index = index_batch * sup_size_u ** 2 + index[:, 0] * sup_size_u + index[:, 1] 40 | 41 | v_2 = torch.concat([torch.take(v_0[..., 0], index.long()).reshape(time_steps, batch_size, size, size, 1), 42 | torch.take(v_0[..., 1], index.long()).reshape(time_steps, batch_size, size, size, 1)], dim=-1) 43 | 44 | mu = x - 0.5*(v_1+v_2) * delta_t 45 | f_2 = 0.1* (torch.sin(2*math.pi*(mu).sum(axis=-1)) + torch.cos(2*math.pi*(mu).sum(axis=-1))) 46 | 47 | if sup_w > 1: 48 | w_0 = torch.fft.irfft(torch.fft.rfft(w_0, dim=2), dim=2, n=sup_size_w) 49 | w_0 = sup_w ** 2 * torch.fft.irfft(torch.fft.rfft(w_0, dim=3), dim=3, n=sup_size_w) 50 | 51 | delta_gridx = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 52 | delta_gridx = delta_gridx.reshape(-1, 1, 1).repeat([1, 2*k+1, 1]) 53 | delta_gridy = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 54 | delta_gridy = delta_gridy.reshape(1, -1, 1).repeat([2*k+1, 1, 1]) 55 | delta_grid = torch.cat((delta_gridx, delta_gridy), dim=-1) 56 | delta_grid = delta_grid[delta_grid.norm(dim=-1) <= (k)/sup_size_w] 57 | delta_size = delta_grid.shape[0] 58 | 59 | loc_p = mu.reshape(time_steps, batch_size, size, size, 1, 2) + delta_grid.reshape(1, 1, 1, 1, -1, 2) 60 | loc_p = (loc_p * sup_size_w).round()/sup_size_w 61 | loc_density = 1/(2*torch.pi*sigma**2) * torch.exp(-((loc_p-mu.reshape(time_steps, batch_size, size, size, 1, 2))**2).sum(axis=-1)/(2 * sigma**2)) * (1/sup_size_w**2) 62 | loc_density = loc_density/loc_density.detach().sum(dim=-1, keepdim=True) 63 | w_0 = w_0.reshape(-1, sup_size_w, sup_size_w) 64 | loc_p = loc_p - loc_p.floor() 65 | loc_p = (loc_p * (sup_size_w)).round().reshape(-1, 2) 66 | loc_p[loc_p==sup_size_w] = 0 67 | loc_p = loc_p.long() 68 | index_batch = torch.arange(batch_size * time_steps, device = device).reshape(-1, 1).repeat(1, delta_size * size**2).reshape(-1,).long() 69 | loc_p = index_batch * sup_size_w ** 2 + loc_p[:, 0] * sup_size_w + loc_p[:, 1] 70 | w_hat = torch.take(w_0, loc_p.long()).reshape(time_steps, batch_size, size, size, delta_size, 1) 71 | f = 0.5 * (f_1.reshape(1, 1, size, size, 1) + f_2[..., None]) 72 | w_hat = (w_hat * loc_density[..., None]).sum(axis=-2) + f * delta_t 73 | 74 | return torch.sqrt(torch.mean(torch.square(w_hat - w_1))) -------------------------------------------------------------------------------- /navier-stokes/E1/mcnp10/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_contour(w): 11 | ''' 12 | :param w: shape: s * s 13 | :return: figs 14 | ''' 15 | temp = w.detach().cpu().numpy() 16 | s = w.shape[0] 17 | plt.figure() 18 | plt.contour(temp.reshape(s, s), 50) 19 | plt.show() 20 | 21 | 22 | def get_grid(shape, device): 23 | batchsize, size_x, size_y = shape[0], shape[1], shape[2] 24 | gridx = torch.tensor(np.linspace(0, 1 - 1 / size_x, size_x), dtype=torch.float) 25 | gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1]) 26 | gridy = torch.tensor(np.linspace(0, 1 - 1 / size_y, size_y), dtype=torch.float) 27 | gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1]) 28 | return torch.cat((gridx, gridy), dim=-1).to(device) 29 | 30 | 31 | def kernel_x(x, config): 32 | device = x.device 33 | xkt = 2 * torch.pi * torch.mm(x, config.K.t().to(device)) 34 | return torch.sin(xkt), torch.cos(xkt) 35 | 36 | 37 | def w2v_g(w, size): 38 | ''' 39 | input: vortex field w, size = [bw, Batch_grid, 1] 40 | output: velocity field u, size = [bw, Batch_size, 2] 41 | ''' 42 | device = w.device 43 | N = size 44 | bw = w.shape[0] 45 | #Maximum frequency 46 | k_max = math.floor(N/2.0) 47 | w0 = w.reshape(bw, N, N) 48 | 49 | #Initial vorticity to Fourier space 50 | w_h = torch.fft.rfft2(w0) 51 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(N,1) 52 | #Wavenumbers in x-direction 53 | k_x = k_y.transpose(0,1) 54 | #Truncate redundant modes 55 | k_x = k_x[..., :k_max + 1] 56 | k_y = k_y[..., :k_max + 1] 57 | #Negative Laplacian in Fourier space 58 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 59 | lap[0,0] = 1.0 60 | psi_h = w_h / lap 61 | 62 | #Velocity field in x-direction = psi_y 63 | q = 2. * math.pi * k_y * 1j * psi_h 64 | q = torch.fft.irfft2(q, s=(N, N)) 65 | 66 | #Velocity field in y-direction = -psi_x 67 | v = -2. * math.pi * k_x * 1j * psi_h 68 | v = torch.fft.irfft2(v, s=(N, N)) 69 | return torch.concat([q[:,:,:,None], v[:,:,:,None]], dim=-1).reshape(bw, -1, 2) 70 | 71 | 72 | def setup_seed(seed): 73 | torch.manual_seed(seed) 74 | torch.cuda.manual_seed(seed) 75 | torch.cuda.manual_seed_all(seed) 76 | np.random.seed(seed) 77 | random.seed(seed) 78 | os.environ['PYTHONHASHSEED'] = str(seed) 79 | torch.backends.cudnn.deterministic = True 80 | torch.backends.cudnn.benchmark = True 81 | torch.backends.cudnn.enabled = True 82 | 83 | 84 | class GaussianRF(object): 85 | 86 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 87 | self.dim = 2 88 | self.device = device 89 | 90 | if sigma is None: 91 | sigma = tau**(0.5*(2*alpha - self.dim)) 92 | k_max = size//2 93 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 94 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 95 | 96 | k_x = wavenumers.transpose(0,1) 97 | k_y = wavenumers 98 | 99 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 100 | self.sqrt_eig[0,0] = 0.0 101 | 102 | self.size = [] 103 | for j in range(self.dim): 104 | self.size.append(size) 105 | 106 | self.size = tuple(self.size) 107 | 108 | def __call__(self, N): 109 | 110 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 111 | coeff = self.sqrt_eig * coeff 112 | 113 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 114 | -------------------------------------------------------------------------------- /navier-stokes/E1/mcnp10/train.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from tools import GaussianRF 3 | from mc_loss import mc_loss 4 | import math 5 | 6 | 7 | 8 | import sys 9 | import time 10 | import copy 11 | import random 12 | import numpy as np 13 | import torch 14 | import torch.optim as optim 15 | import matplotlib.pyplot as plt 16 | import os 17 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 18 | plt.rcParams["animation.html"] = "jshtml" 19 | 20 | 21 | def test(config, net, test_data): 22 | num = test_data.shape[0] 23 | device = config.device 24 | delta_steps = 10 25 | 26 | delta_t = config.T / delta_steps 27 | w_0 = test_data[:, :, :, 0].to(device)[:, :, :, None] 28 | rela_err = torch.zeros(delta_steps) 29 | rela_err_1 = torch.zeros(delta_steps) 30 | rela_err_max = torch.zeros(delta_steps) 31 | for time_step in range(1, delta_steps+1): 32 | net.eval() 33 | t = (delta_t * time_step) * torch.ones(num).to(device) 34 | w = net(w_0, t).detach() 35 | w_t = test_data[..., delta_steps*time_step].to(device)[:, :, :, None] 36 | rela_err[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), dim=1)).mean() 37 | rela_err_1[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=1, dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=1, dim=1)).mean() 38 | rela_err_max[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=float('inf'), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=float('inf'), dim=1)).mean() 39 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 40 | print('mean relative l_2 error', rela_err.mean().item()) 41 | print('mean relative l_1 error', rela_err_1.mean().item()) 42 | print('mean relative l_inf error', rela_err_max.mean().item()) 43 | return rela_err.mean().item() 44 | 45 | 46 | def train_epoch(net, GRF, config, x, f, optimizer, scheduler): 47 | batch_size, size = config.batch_size, config.size 48 | net.train() 49 | w_0 = GRF(batch_size).reshape(batch_size, size, size, 1) 50 | loss = mc_loss(w_0, x, f, net, config) 51 | loss.backward() 52 | optimizer.step() 53 | optimizer.zero_grad(set_to_none=True) 54 | scheduler.step() 55 | 56 | 57 | 58 | def train(config, net): 59 | device = config.device 60 | size = config.size 61 | batch_size = config.batch_size 62 | net = net.to(device) 63 | optimizer = optim.AdamW(net.parameters(), config.lr, weight_decay=config.weight_decay) 64 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 65 | val_data = torch.load(config.data_path + 'data_val').to(device).float()[..., :int(10*config.T)+1] 66 | test_data = torch.load(config.data_path+'data_test').to(device).float()[..., :int(10*config.T)+1] 67 | val_loss = 1e50 68 | 69 | gridx = torch.linspace(0, 1 - 1 / size, size, device=device) 70 | gridx = gridx.reshape(size, 1, 1).repeat([1, size, 1]) 71 | gridy = torch.linspace(0, 1 - 1 / size, size, device=device) 72 | gridy = gridy.reshape(1, size, 1).repeat([size, 1, 1]) 73 | x = torch.cat((gridx, gridy), dim=-1) 74 | f = 0.1* (torch.sin(2*math.pi*(x).sum(axis=-1)) + torch.cos(2*math.pi*(x).sum(axis=-1))) 75 | GRF = GaussianRF(size, device=device) 76 | for step in range(config.num_iterations+1): 77 | train_epoch(net, GRF, config, x, f, optimizer, scheduler) 78 | if step % 50 == 0: 79 | print('########################') 80 | print('training loss', step) 81 | print('VAL_LOSS') 82 | print('########################') 83 | temp = test(config, net, val_data) 84 | if temp < val_loss: 85 | val_loss = temp 86 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 87 | print('TEST_LOSS') 88 | print('########################') 89 | test(config, net, test_data) 90 | sys.stdout.flush() 91 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 92 | print('FINAL_LOSS') 93 | print('########################') 94 | test(config, net, test_data) 95 | sys.stdout.flush() -------------------------------------------------------------------------------- /navier-stokes/E2/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | import math 5 | import time 6 | import os 7 | 8 | from nse import navier_stokes_2d, GaussianRF 9 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 10 | 11 | 12 | def check_directory() -> None: 13 | """ 14 | Check if log directory exists within experiments 15 | """ 16 | if not os.path.exists(f'dataset'): 17 | os.mkdir(f'dataset') 18 | 19 | 20 | def setup_seed(seed): 21 | torch.manual_seed(seed) 22 | torch.cuda.manual_seed(seed) 23 | torch.cuda.manual_seed_all(seed) 24 | np.random.seed(seed) 25 | random.seed(seed) 26 | os.environ['PYTHONHASHSEED'] = str(seed) 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = False 29 | torch.backends.cudnn.enabled = True 30 | 31 | 32 | setup_seed(0) 33 | check_directory() 34 | device = torch.device('cuda') 35 | 36 | s = 256 37 | sub = 4 38 | 39 | # Set up 2d GRF with covariance parameters 40 | GRF = GaussianRF(s, device=device) 41 | 42 | # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y))) 43 | t = torch.linspace(0, 1, s+1, device=device) 44 | t = t[0: -1] 45 | 46 | X, Y = torch.meshgrid(t, t, indexing='ij') 47 | 48 | 49 | record_steps = 200 50 | T = 20 51 | N = 200 52 | bsize = 200 53 | c = 0 54 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 55 | 56 | 57 | for j in range(N//bsize): 58 | # Sample random feilds 59 | w0 = GRF(bsize) 60 | visc = 10 ** (- 5 * torch.ones(bsize).to(device)) 61 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 62 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 63 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 64 | u[c:(c+bsize),...] = sol 65 | c += bsize 66 | print(j, c) 67 | torch.save(u, './dataset/data_test') 68 | 69 | setup_seed(1) 70 | record_steps = 200 71 | T = 20 72 | N = 200 73 | bsize = 200 74 | c = 0 75 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 76 | f = 0.1*(torch.sin(2*math.pi*(X + Y)) + torch.cos(2*math.pi*(X + Y))) 77 | 78 | 79 | for j in range(N//bsize): 80 | # Sample random feilds 81 | w0 = GRF(bsize) 82 | visc = 10 ** (- 5 * torch.ones(bsize).to(device)) 83 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 84 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 85 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 86 | u[c:(c+bsize),...] = sol 87 | c += bsize 88 | print(j, c) 89 | torch.save(u, './dataset/data_val') 90 | 91 | -------------------------------------------------------------------------------- /navier-stokes/E2/data/nse.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import os 4 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 5 | 6 | class GaussianRF(object): 7 | 8 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 9 | self.dim = 2 10 | self.device = device 11 | 12 | if sigma is None: 13 | sigma = tau**(0.5*(2*alpha - self.dim)) 14 | k_max = size//2 15 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 16 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 17 | 18 | k_x = wavenumers.transpose(0,1) 19 | k_y = wavenumers 20 | 21 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 22 | self.sqrt_eig[0,0] = 0.0 23 | 24 | self.size = [] 25 | for j in range(self.dim): 26 | self.size.append(size) 27 | 28 | self.size = tuple(self.size) 29 | 30 | def __call__(self, N): 31 | 32 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 33 | coeff = self.sqrt_eig * coeff 34 | 35 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 36 | 37 | 38 | def force(x): 39 | z = x.sum(axis=1).reshape(x.shape[0], 1) 40 | return 0.1 * torch.sin(2 * math.pi * z) + 0.1 * torch.cos(2 * math.pi * z) 41 | 42 | 43 | def navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1): 44 | 45 | # Grid size - must be power of 2 46 | N = w0.size()[-1] 47 | 48 | # Maximum frequency 49 | k_max = math.floor(N/2.0) 50 | 51 | # Number of steps to final time 52 | steps = math.ceil(T/delta_t) 53 | 54 | # Initial vorticity to Fourier space 55 | w_h = torch.fft.rfft2(w0) 56 | 57 | # Forcing to Fourier space 58 | f_h = torch.fft.rfft2(f) 59 | 60 | #If same forcing for the whole batch 61 | if len(f_h.size()) < len(w_h.size()): 62 | f_h = torch.unsqueeze(f_h, 0) 63 | 64 | # Record solution every this number of steps 65 | record_time = math.floor(steps/record_steps) 66 | 67 | # Wavenumbers in y-direction 68 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device), torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N,1) 69 | # Wavenumbers in x-direction 70 | k_x = k_y.transpose(0,1) 71 | 72 | # Truncate redundant modes 73 | k_x = k_x[..., :k_max + 1] 74 | k_y = k_y[..., :k_max + 1] 75 | 76 | # Negative Laplacian in Fourier space 77 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 78 | lap[0,0] = 1.0 79 | # Dealiasing mask 80 | dealias = torch.unsqueeze(torch.logical_and(torch.abs(k_y) <= (2.0/3.0)*k_max, torch.abs(k_x) <= (2.0/3.0)*k_max).float(), 0) 81 | 82 | # Saving solution and time 83 | sol = torch.zeros(*w0.size(), record_steps, device=w0.device) 84 | sol_t = torch.zeros(record_steps, device=w0.device) 85 | 86 | # Record counter 87 | c = 0 88 | # Physical time 89 | t = 0.0 90 | for j in range(steps): 91 | # Stream function in Fourier space: solve Poisson equation 92 | psi_h = w_h / lap 93 | 94 | # Velocity field in x-direction = psi_y 95 | q = 2. * math.pi * k_y * 1j * psi_h 96 | q = torch.fft.irfft2(q, s=(N, N)) 97 | 98 | # Velocity field in y-direction = -psi_x 99 | v = -2. * math.pi * k_x * 1j * psi_h 100 | v = torch.fft.irfft2(v, s=(N, N)) 101 | 102 | # Partial x of vorticity 103 | w_x = 2. * math.pi * k_x * 1j * w_h 104 | w_x = torch.fft.irfft2(w_x, s=(N, N)) 105 | 106 | # Partial y of vorticity 107 | w_y = 2. * math.pi * k_y * 1j * w_h 108 | w_y = torch.fft.irfft2(w_y, s=(N, N)) 109 | 110 | # Non-linear term (u.grad(w)): compute in physical space then back to Fourier space 111 | F_h = torch.fft.rfft2(q*w_x + v*w_y) 112 | 113 | # Dealias 114 | F_h = dealias* F_h 115 | 116 | # Crank-Nicolson update 117 | w_h = (-delta_t*F_h + delta_t*f_h + (1.0 - 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...])*w_h)/(1.0 + 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...]) 118 | 119 | # Update real time (used only for recording) 120 | t += delta_t 121 | 122 | if (j+1) % record_time == 0: 123 | # Solution in physical space 124 | w = torch.fft.irfft2(w_h, s=(N, N)) 125 | 126 | # Record solution and time 127 | sol[...,c] = w 128 | sol_t[c] = t 129 | 130 | c += 1 131 | return sol, sol_t 132 | 133 | -------------------------------------------------------------------------------- /navier-stokes/E2/mcnp10/main.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from train import train 3 | from tools import setup_seed 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def main(cfg): 21 | if not os.path.exists(f'log'): 22 | os.mkdir(f'log') 23 | if not os.path.exists(f'model'): 24 | os.mkdir(f'model') 25 | setup_seed(cfg.seed) 26 | dateTimeObj = datetime.now() 27 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}' 28 | logfile = f'log/seed_{cfg.seed}_k_{cfg.k}_u_{cfg.sup_u}_w_{cfg.sup_w}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{timestring}.csv' 29 | 30 | cfg.timestring = timestring 31 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 32 | json.dump(cfg.__dict__, f, indent=2) 33 | 34 | sys.stdout = open(logfile, 'w') 35 | 36 | print('--------args----------') 37 | for k in list(vars(cfg).keys()): 38 | print('%s: %s' % (k, vars(cfg)[k])) 39 | print('--------args----------\n') 40 | 41 | sys.stdout.flush() 42 | cfg.delta_t = cfg.T/cfg.time_steps 43 | net = FNO(cfg.NUM_TRUNCATION, cfg.NUM_TRUNCATION, cfg.num_channel) 44 | train(cfg, net) 45 | 46 | 47 | if __name__ == "__main__": 48 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 49 | 50 | parser.add_argument('--data_path', type=str, 51 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 52 | help='path of data') 53 | 54 | parser.add_argument('--seed', type=int, default=0, 55 | help='seed') 56 | 57 | parser.add_argument('--batch_size', type=int, default=10, 58 | help='batchsize of the operator learning') 59 | 60 | parser.add_argument('--step_size', type=int, default=2000, 61 | help='step_size of optim') 62 | 63 | parser.add_argument('--gamma', type=float, default=0.8, 64 | help='gamma of optim') 65 | 66 | parser.add_argument('--lr', type=float, default=0.01, 67 | help='lr of optim') 68 | 69 | parser.add_argument('--weight_decay', type=float, default=0.0, 70 | help='lr of optim') 71 | 72 | parser.add_argument('--num_iterations', type=int, default=20000, 73 | help='num_iterations of optim') 74 | 75 | parser.add_argument('--nu', type=float, default=1/100000, 76 | help='nu in NSE') 77 | 78 | parser.add_argument('--size', type=int, default=64, 79 | help='data spatial size') 80 | 81 | parser.add_argument('--T', type=float, default=10.0, 82 | help='final time') 83 | 84 | parser.add_argument('--time_steps', type=int, default=10, 85 | help='number of time_steps in data') 86 | 87 | parser.add_argument('--EPS', type=float, default=1e-8, 88 | help='epsilon') 89 | 90 | parser.add_argument('--NUM_TRUNCATION', type=int, default=16, 91 | help='NUM_TRUNCATION in FNO') 92 | 93 | parser.add_argument('--num_channel', type=int, default=36, 94 | help='num_channel in FNO') 95 | 96 | parser.add_argument('--sup_u', type=int, default=4, 97 | help='sup_u') 98 | 99 | parser.add_argument('--sup_w', type=int, default=2, 100 | help='sup_w') 101 | 102 | parser.add_argument('--k', type=int, default=2, 103 | help='k') 104 | 105 | parser.add_argument('--device', type=str, default='cuda:3', 106 | help='Used device') 107 | 108 | cfg = parser.parse_args() 109 | for seed in [0, 1, 2]: 110 | cfg.seed = seed 111 | main(cfg) -------------------------------------------------------------------------------- /navier-stokes/E2/mcnp10/mc_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import numpy as np 4 | from tools import w2v_g 5 | 6 | 7 | 8 | def mc_loss(w_0, x, f_1, model, config): 9 | device = config.device 10 | size = config.size 11 | sup_w, sup_u = config.sup_w, config.sup_u 12 | sup_size_w, sup_size_u = sup_w * size, sup_u * size 13 | batch_size = config.batch_size 14 | time_steps = config.time_steps 15 | delta_t = config.delta_t 16 | k = config.k 17 | nu = config.nu 18 | sigma = math.sqrt(2 * nu * delta_t) 19 | w_0 = w_0.reshape([1, batch_size, size, size, 1]) 20 | 21 | w_input = w_0.repeat(time_steps, 1, 1, 1, 1).reshape(-1, size, size, 1) 22 | t = delta_t * torch.arange(1, time_steps+1, device=device).reshape(-1, 1).repeat(1,batch_size).reshape(-1,) 23 | w = model(w_input, t).reshape(time_steps, batch_size, size, size, 1) 24 | w = torch.concat([w_0.reshape(1, -1, size, size, 1), w], dim=0) 25 | 26 | w_0, w_1 = w[:-1, ...], w[1:, ...] 27 | v = w2v_g(w.reshape(-1, size, size, 1), size).reshape(time_steps+1, batch_size, size, size, 2) 28 | v_0, v_1 = v[:-1, ...], v[1:, ...] 29 | 30 | if sup_u > 1: 31 | v_0 = torch.fft.irfft(torch.fft.rfft(v_0, dim=2), dim=2, n=sup_size_u) 32 | v_0 = sup_u ** 2 * torch.fft.irfft(torch.fft.rfft(v_0, dim=3), dim=3, n=sup_size_u) 33 | mu_1 = x - v_1 * delta_t 34 | index = mu_1 - mu_1.floor() 35 | index = (index * (sup_size_u)).round().reshape(-1, 2) 36 | index[index==sup_size_u] = 0 37 | index = index.long() 38 | index_batch = torch.arange(batch_size * time_steps, device=device).reshape(-1, 1).repeat(1, size**2).reshape(-1,).long() 39 | index = index_batch * sup_size_u ** 2 + index[:, 0] * sup_size_u + index[:, 1] 40 | 41 | v_2 = torch.concat([torch.take(v_0[..., 0], index.long()).reshape(time_steps, batch_size, size, size, 1), 42 | torch.take(v_0[..., 1], index.long()).reshape(time_steps, batch_size, size, size, 1)], dim=-1) 43 | 44 | mu = x - 0.5*(v_1+v_2) * delta_t 45 | f_2 = 0.1* (torch.sin(2*math.pi*(mu).sum(axis=-1)) + torch.cos(2*math.pi*(mu).sum(axis=-1))) 46 | 47 | if sup_w > 1: 48 | w_0 = torch.fft.irfft(torch.fft.rfft(w_0, dim=2), dim=2, n=sup_size_w) 49 | w_0 = sup_w ** 2 * torch.fft.irfft(torch.fft.rfft(w_0, dim=3), dim=3, n=sup_size_w) 50 | 51 | delta_gridx = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 52 | delta_gridx = delta_gridx.reshape(-1, 1, 1).repeat([1, 2*k+1, 1]) 53 | delta_gridy = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 54 | delta_gridy = delta_gridy.reshape(1, -1, 1).repeat([2*k+1, 1, 1]) 55 | delta_grid = torch.cat((delta_gridx, delta_gridy), dim=-1) 56 | delta_grid = delta_grid[delta_grid.norm(dim=-1) <= (k)/sup_size_w] 57 | delta_size = delta_grid.shape[0] 58 | 59 | loc_p = mu.reshape(time_steps, batch_size, size, size, 1, 2) + delta_grid.reshape(1, 1, 1, 1, -1, 2) 60 | loc_p = (loc_p * sup_size_w).round()/sup_size_w 61 | loc_density = 1/(2*torch.pi*sigma**2) * torch.exp(-((loc_p-mu.reshape(time_steps, batch_size, size, size, 1, 2))**2).sum(axis=-1)/(2 * sigma**2)) * (1/sup_size_w**2) 62 | loc_density = loc_density/loc_density.detach().sum(dim=-1, keepdim=True) 63 | w_0 = w_0.reshape(-1, sup_size_w, sup_size_w) 64 | loc_p = loc_p - loc_p.floor() 65 | loc_p = (loc_p * (sup_size_w)).round().reshape(-1, 2) 66 | loc_p[loc_p==sup_size_w] = 0 67 | loc_p = loc_p.long() 68 | index_batch = torch.arange(batch_size * time_steps, device = device).reshape(-1, 1).repeat(1, delta_size * size**2).reshape(-1,).long() 69 | loc_p = index_batch * sup_size_w ** 2 + loc_p[:, 0] * sup_size_w + loc_p[:, 1] 70 | w_hat = torch.take(w_0, loc_p.long()).reshape(time_steps, batch_size, size, size, delta_size, 1) 71 | f = 0.5 * (f_1.reshape(1, 1, size, size, 1) + f_2[..., None]) 72 | w_hat = (w_hat * loc_density[..., None]).sum(axis=-2) + f * delta_t 73 | 74 | return torch.sqrt(torch.mean(torch.square(w_hat - w_1))) -------------------------------------------------------------------------------- /navier-stokes/E2/mcnp10/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_contour(w): 11 | ''' 12 | :param w: shape: s * s 13 | :return: figs 14 | ''' 15 | temp = w.detach().cpu().numpy() 16 | s = w.shape[0] 17 | plt.figure() 18 | plt.contour(temp.reshape(s, s), 50) 19 | plt.show() 20 | 21 | 22 | def get_grid(shape, device): 23 | batchsize, size_x, size_y = shape[0], shape[1], shape[2] 24 | gridx = torch.tensor(np.linspace(0, 1 - 1 / size_x, size_x), dtype=torch.float) 25 | gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1]) 26 | gridy = torch.tensor(np.linspace(0, 1 - 1 / size_y, size_y), dtype=torch.float) 27 | gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1]) 28 | return torch.cat((gridx, gridy), dim=-1).to(device) 29 | 30 | 31 | def kernel_x(x, config): 32 | device = x.device 33 | xkt = 2 * torch.pi * torch.mm(x, config.K.t().to(device)) 34 | return torch.sin(xkt), torch.cos(xkt) 35 | 36 | 37 | def w2v_g(w, size): 38 | ''' 39 | input: vortex field w, size = [bw, Batch_grid, 1] 40 | output: velocity field u, size = [bw, Batch_size, 2] 41 | ''' 42 | device = w.device 43 | N = size 44 | bw = w.shape[0] 45 | #Maximum frequency 46 | k_max = math.floor(N/2.0) 47 | w0 = w.reshape(bw, N, N) 48 | 49 | #Initial vorticity to Fourier space 50 | w_h = torch.fft.rfft2(w0) 51 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(N,1) 52 | #Wavenumbers in x-direction 53 | k_x = k_y.transpose(0,1) 54 | #Truncate redundant modes 55 | k_x = k_x[..., :k_max + 1] 56 | k_y = k_y[..., :k_max + 1] 57 | #Negative Laplacian in Fourier space 58 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 59 | lap[0,0] = 1.0 60 | psi_h = w_h / lap 61 | 62 | #Velocity field in x-direction = psi_y 63 | q = 2. * math.pi * k_y * 1j * psi_h 64 | q = torch.fft.irfft2(q, s=(N, N)) 65 | 66 | #Velocity field in y-direction = -psi_x 67 | v = -2. * math.pi * k_x * 1j * psi_h 68 | v = torch.fft.irfft2(v, s=(N, N)) 69 | return torch.concat([q[:,:,:,None], v[:,:,:,None]], dim=-1).reshape(bw, -1, 2) 70 | 71 | 72 | def setup_seed(seed): 73 | torch.manual_seed(seed) 74 | torch.cuda.manual_seed(seed) 75 | torch.cuda.manual_seed_all(seed) 76 | np.random.seed(seed) 77 | random.seed(seed) 78 | os.environ['PYTHONHASHSEED'] = str(seed) 79 | torch.backends.cudnn.deterministic = True 80 | torch.backends.cudnn.benchmark = True 81 | torch.backends.cudnn.enabled = True 82 | 83 | 84 | class GaussianRF(object): 85 | 86 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 87 | self.dim = 2 88 | self.device = device 89 | 90 | if sigma is None: 91 | sigma = tau**(0.5*(2*alpha - self.dim)) 92 | k_max = size//2 93 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 94 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 95 | 96 | k_x = wavenumers.transpose(0,1) 97 | k_y = wavenumers 98 | 99 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 100 | self.sqrt_eig[0,0] = 0.0 101 | 102 | self.size = [] 103 | for j in range(self.dim): 104 | self.size.append(size) 105 | 106 | self.size = tuple(self.size) 107 | 108 | def __call__(self, N): 109 | 110 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 111 | coeff = self.sqrt_eig * coeff 112 | 113 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 114 | -------------------------------------------------------------------------------- /navier-stokes/E2/mcnp10/train.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from tools import GaussianRF 3 | from mc_loss import mc_loss 4 | import math 5 | 6 | 7 | 8 | import sys 9 | import time 10 | import copy 11 | import random 12 | import numpy as np 13 | import torch 14 | import torch.optim as optim 15 | import matplotlib.pyplot as plt 16 | import os 17 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 18 | plt.rcParams["animation.html"] = "jshtml" 19 | 20 | 21 | def test(config, net, test_data): 22 | num = test_data.shape[0] 23 | device = config.device 24 | delta_steps = 10 25 | 26 | delta_t = config.T / delta_steps 27 | w_0 = test_data[:, :, :, 0].to(device)[:, :, :, None] 28 | rela_err = torch.zeros(delta_steps) 29 | rela_err_1 = torch.zeros(delta_steps) 30 | rela_err_max = torch.zeros(delta_steps) 31 | for time_step in range(1, delta_steps+1): 32 | net.eval() 33 | t = (delta_t * time_step) * torch.ones(num).to(device) 34 | w = net(w_0, t).detach() 35 | w_t = test_data[..., delta_steps*time_step].to(device)[:, :, :, None] 36 | rela_err[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), dim=1)).mean() 37 | rela_err_1[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=1, dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=1, dim=1)).mean() 38 | rela_err_max[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=float('inf'), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=float('inf'), dim=1)).mean() 39 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 40 | print('mean relative l_2 error', rela_err.mean().item()) 41 | print('mean relative l_1 error', rela_err_1.mean().item()) 42 | print('mean relative l_inf error', rela_err_max.mean().item()) 43 | return rela_err.mean().item() 44 | 45 | 46 | def train_epoch(net, GRF, config, x, f, optimizer, scheduler): 47 | batch_size, size = config.batch_size, config.size 48 | net.train() 49 | w_0 = GRF(batch_size).reshape(batch_size, size, size, 1) 50 | loss = mc_loss(w_0, x, f, net, config) 51 | loss.backward() 52 | optimizer.step() 53 | optimizer.zero_grad(set_to_none=True) 54 | scheduler.step() 55 | 56 | 57 | 58 | def train(config, net): 59 | device = config.device 60 | size = config.size 61 | batch_size = config.batch_size 62 | net = net.to(device) 63 | optimizer = optim.AdamW(net.parameters(), config.lr, weight_decay=config.weight_decay) 64 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 65 | val_data = torch.load(config.data_path + 'data_val').to(device).float()[..., :int(10*config.T)+1] 66 | test_data = torch.load(config.data_path+'data_test').to(device).float()[..., :int(10*config.T)+1] 67 | val_loss = 1e50 68 | 69 | gridx = torch.linspace(0, 1 - 1 / size, size, device=device) 70 | gridx = gridx.reshape(size, 1, 1).repeat([1, size, 1]) 71 | gridy = torch.linspace(0, 1 - 1 / size, size, device=device) 72 | gridy = gridy.reshape(1, size, 1).repeat([size, 1, 1]) 73 | x = torch.cat((gridx, gridy), dim=-1) 74 | f = 0.1* (torch.sin(2*math.pi*(x).sum(axis=-1)) + torch.cos(2*math.pi*(x).sum(axis=-1))) 75 | GRF = GaussianRF(size, device=device) 76 | for step in range(config.num_iterations+1): 77 | train_epoch(net, GRF, config, x, f, optimizer, scheduler) 78 | if step % 50 == 0: 79 | print('########################') 80 | print('training loss', step) 81 | print('VAL_LOSS') 82 | print('########################') 83 | temp = test(config, net, val_data) 84 | if temp < val_loss: 85 | val_loss = temp 86 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 87 | print('TEST_LOSS') 88 | print('########################') 89 | test(config, net, test_data) 90 | sys.stdout.flush() 91 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 92 | print('FINAL_LOSS') 93 | print('########################') 94 | test(config, net, test_data) 95 | sys.stdout.flush() -------------------------------------------------------------------------------- /navier-stokes/E3/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | import math 5 | import os 6 | 7 | from nse import navier_stokes_2d, GaussianRF 8 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 9 | 10 | 11 | def check_directory() -> None: 12 | """ 13 | Check if log directory exists within experiments 14 | """ 15 | if not os.path.exists(f'dataset'): 16 | os.mkdir(f'dataset') 17 | 18 | 19 | def setup_seed(seed): 20 | torch.manual_seed(seed) 21 | torch.cuda.manual_seed(seed) 22 | torch.cuda.manual_seed_all(seed) 23 | np.random.seed(seed) 24 | random.seed(seed) 25 | os.environ['PYTHONHASHSEED'] = str(seed) 26 | torch.backends.cudnn.deterministic = True 27 | torch.backends.cudnn.benchmark = False 28 | torch.backends.cudnn.enabled = True 29 | 30 | 31 | setup_seed(0) 32 | check_directory() 33 | device = torch.device('cuda') 34 | 35 | s = 256 36 | sub = 4 37 | 38 | # Set up 2d GRF with covariance parameters 39 | GRF = GaussianRF(s, device=device) 40 | 41 | # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y))) 42 | t = torch.linspace(0, 1, s+1, device=device) 43 | t = t[0: -1] 44 | 45 | X, Y = torch.meshgrid(t, t, indexing='ij') 46 | 47 | # Number of snapshots from solution 48 | 49 | 50 | 51 | ###################################################################### 52 | 53 | record_steps = 100 54 | T = 10 55 | N = 200 56 | bsize = 200 57 | c = 0 58 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 59 | f = 0.1*torch.cos(8*math.pi*(X)) 60 | 61 | 62 | for j in range(N//bsize): 63 | # Sample random feilds 64 | w0 = GRF(bsize) 65 | visc = 10 ** (- 4 * torch.ones(bsize).to(device)) 66 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 67 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 68 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 69 | u[c:(c+bsize),...] = sol 70 | c += bsize 71 | print(j, c) 72 | torch.save(u, './dataset/data_test') 73 | 74 | setup_seed(1) 75 | record_steps = 100 76 | T = 10 77 | N = 200 78 | bsize = 200 79 | c = 0 80 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 81 | f = 0.1*torch.cos(8*math.pi*(X)) 82 | 83 | 84 | for j in range(N//bsize): 85 | # Sample random feilds 86 | w0 = GRF(bsize) 87 | visc = 10 ** (- 4 * torch.ones(bsize).to(device)) 88 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 89 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 90 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 91 | u[c:(c+bsize),...] = sol 92 | c += bsize 93 | print(j, c) 94 | torch.save(u, './dataset/data_val') 95 | -------------------------------------------------------------------------------- /navier-stokes/E3/data/nse.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import os 4 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 5 | 6 | class GaussianRF(object): 7 | 8 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 9 | self.dim = 2 10 | self.device = device 11 | 12 | if sigma is None: 13 | sigma = tau**(0.5*(2*alpha - self.dim)) 14 | k_max = size//2 15 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 16 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 17 | 18 | k_x = wavenumers.transpose(0,1) 19 | k_y = wavenumers 20 | 21 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 22 | self.sqrt_eig[0,0] = 0.0 23 | 24 | self.size = [] 25 | for j in range(self.dim): 26 | self.size.append(size) 27 | 28 | self.size = tuple(self.size) 29 | 30 | def __call__(self, N): 31 | 32 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 33 | coeff = self.sqrt_eig * coeff 34 | 35 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 36 | 37 | 38 | def force(x): 39 | z = x.sum(axis=1).reshape(x.shape[0], 1) 40 | return 0.1 * torch.sin(2 * math.pi * z) + 0.1 * torch.cos(2 * math.pi * z) 41 | 42 | 43 | def navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1): 44 | 45 | # Grid size - must be power of 2 46 | N = w0.size()[-1] 47 | 48 | # Maximum frequency 49 | k_max = math.floor(N/2.0) 50 | 51 | # Number of steps to final time 52 | steps = math.ceil(T/delta_t) 53 | 54 | # Initial vorticity to Fourier space 55 | w_h = torch.fft.rfft2(w0) 56 | 57 | # Forcing to Fourier space 58 | f_h = torch.fft.rfft2(f) 59 | 60 | #If same forcing for the whole batch 61 | if len(f_h.size()) < len(w_h.size()): 62 | f_h = torch.unsqueeze(f_h, 0) 63 | 64 | # Record solution every this number of steps 65 | record_time = math.floor(steps/record_steps) 66 | 67 | # Wavenumbers in y-direction 68 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device), torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N,1) 69 | # Wavenumbers in x-direction 70 | k_x = k_y.transpose(0,1) 71 | 72 | # Truncate redundant modes 73 | k_x = k_x[..., :k_max + 1] 74 | k_y = k_y[..., :k_max + 1] 75 | 76 | # Negative Laplacian in Fourier space 77 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 78 | lap[0,0] = 1.0 79 | # Dealiasing mask 80 | dealias = torch.unsqueeze(torch.logical_and(torch.abs(k_y) <= (2.0/3.0)*k_max, torch.abs(k_x) <= (2.0/3.0)*k_max).float(), 0) 81 | 82 | # Saving solution and time 83 | sol = torch.zeros(*w0.size(), record_steps, device=w0.device) 84 | sol_t = torch.zeros(record_steps, device=w0.device) 85 | 86 | # Record counter 87 | c = 0 88 | # Physical time 89 | t = 0.0 90 | for j in range(steps): 91 | # Stream function in Fourier space: solve Poisson equation 92 | psi_h = w_h / lap 93 | 94 | # Velocity field in x-direction = psi_y 95 | q = 2. * math.pi * k_y * 1j * psi_h 96 | q = torch.fft.irfft2(q, s=(N, N)) 97 | 98 | # Velocity field in y-direction = -psi_x 99 | v = -2. * math.pi * k_x * 1j * psi_h 100 | v = torch.fft.irfft2(v, s=(N, N)) 101 | 102 | # Partial x of vorticity 103 | w_x = 2. * math.pi * k_x * 1j * w_h 104 | w_x = torch.fft.irfft2(w_x, s=(N, N)) 105 | 106 | # Partial y of vorticity 107 | w_y = 2. * math.pi * k_y * 1j * w_h 108 | w_y = torch.fft.irfft2(w_y, s=(N, N)) 109 | 110 | # Non-linear term (u.grad(w)): compute in physical space then back to Fourier space 111 | F_h = torch.fft.rfft2(q*w_x + v*w_y) 112 | 113 | # Dealias 114 | F_h = dealias* F_h 115 | 116 | # Crank-Nicolson update 117 | w_h = (-delta_t*F_h + delta_t*f_h + (1.0 - 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...])*w_h)/(1.0 + 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...]) 118 | 119 | # Update real time (used only for recording) 120 | t += delta_t 121 | 122 | if (j+1) % record_time == 0: 123 | # Solution in physical space 124 | w = torch.fft.irfft2(w_h, s=(N, N)) 125 | 126 | # Record solution and time 127 | sol[...,c] = w 128 | sol_t[c] = t 129 | 130 | c += 1 131 | return sol, sol_t 132 | 133 | -------------------------------------------------------------------------------- /navier-stokes/E3/mcnp10/main.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from train import train 3 | from tools import setup_seed 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def main(cfg): 21 | if not os.path.exists(f'log'): 22 | os.mkdir(f'log') 23 | if not os.path.exists(f'model'): 24 | os.mkdir(f'model') 25 | setup_seed(cfg.seed) 26 | dateTimeObj = datetime.now() 27 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}' 28 | logfile = f'log/seed_{cfg.seed}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{timestring}.csv' 29 | 30 | cfg.timestring = timestring 31 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 32 | json.dump(cfg.__dict__, f, indent=2) 33 | 34 | sys.stdout = open(logfile, 'w') 35 | 36 | print('--------args----------') 37 | for k in list(vars(cfg).keys()): 38 | print('%s: %s' % (k, vars(cfg)[k])) 39 | print('--------args----------\n') 40 | 41 | sys.stdout.flush() 42 | cfg.delta_t = cfg.T/cfg.time_steps 43 | net = FNO(cfg.NUM_TRUNCATION, cfg.NUM_TRUNCATION, cfg.num_channel) 44 | train(cfg, net) 45 | 46 | 47 | if __name__ == "__main__": 48 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 49 | 50 | parser.add_argument('--data_path', type=str, 51 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 52 | help='path of data') 53 | 54 | parser.add_argument('--device', type=str, default='cuda:3', 55 | help='Used device') 56 | 57 | parser.add_argument('--seed', type=int, default=0, 58 | help='seed') 59 | 60 | parser.add_argument('--batch_size', type=int, default=10, 61 | help='batchsize of the operator learning') 62 | 63 | parser.add_argument('--step_size', type=int, default=2000, 64 | help='step_size of optim') 65 | 66 | parser.add_argument('--gamma', type=float, default=0.8, 67 | help='gamma of optim') 68 | 69 | parser.add_argument('--lr', type=float, default=0.01, 70 | help='lr of optim') 71 | 72 | parser.add_argument('--weight_decay', type=float, default=0.0, 73 | help='lr of optim') 74 | 75 | parser.add_argument('--num_iterations', type=int, default=20000, 76 | help='num_iterations of optim') 77 | 78 | parser.add_argument('--nu', type=float, default=1/10000, 79 | help='nu in NSE') 80 | 81 | parser.add_argument('--size', type=int, default=64, 82 | help='data spatial size') 83 | 84 | parser.add_argument('--T', type=float, default=10.0, 85 | help='final time') 86 | 87 | parser.add_argument('--time_steps', type=int, default=10, 88 | help='number of time_steps in data') 89 | 90 | parser.add_argument('--EPS', type=float, default=1e-8, 91 | help='epsilon') 92 | 93 | parser.add_argument('--NUM_TRUNCATION', type=int, default=16, 94 | help='NUM_TRUNCATION in FNO') 95 | 96 | parser.add_argument('--num_channel', type=int, default=36, 97 | help='num_channel in FNO') 98 | 99 | parser.add_argument('--sup_u', type=int, default=2, 100 | help='sup_u') 101 | 102 | parser.add_argument('--sup_w', type=int, default=1, 103 | help='sup_w') 104 | 105 | parser.add_argument('--k', type=int, default=3, 106 | help='k') 107 | 108 | cfg = parser.parse_args() 109 | for seed in [0, 1, 2]: 110 | cfg.seed = seed 111 | main(cfg) -------------------------------------------------------------------------------- /navier-stokes/E3/mcnp10/mc_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import numpy as np 4 | from tools import w2v_g 5 | 6 | 7 | 8 | def mc_loss(w_0, x, f_1, model, config): 9 | device = config.device 10 | size = config.size 11 | sup_w, sup_u = config.sup_w, config.sup_u 12 | sup_size_w, sup_size_u = sup_w * size, sup_u * size 13 | batch_size = config.batch_size 14 | time_steps = config.time_steps 15 | delta_t = config.delta_t 16 | k = config.k 17 | nu = config.nu 18 | sigma = math.sqrt(2 * nu * delta_t) 19 | w_0 = w_0.reshape([1, batch_size, size, size, 1]) 20 | 21 | w_input = w_0.repeat(time_steps, 1, 1, 1, 1).reshape(-1, size, size, 1) 22 | t = delta_t * torch.arange(1, time_steps+1, device=device).reshape(-1, 1).repeat(1,batch_size).reshape(-1,) 23 | w = model(w_input, t).reshape(time_steps, batch_size, size, size, 1) 24 | w = torch.concat([w_0.reshape(1, -1, size, size, 1), w], dim=0) 25 | 26 | w_0, w_1 = w[:-1, ...], w[1:, ...] 27 | v = w2v_g(w.reshape(-1, size, size, 1), size).reshape(time_steps+1, batch_size, size, size, 2) 28 | v_0, v_1 = v[:-1, ...], v[1:, ...] 29 | 30 | if sup_u > 1: 31 | v_0 = torch.fft.irfft(torch.fft.rfft(v_0, dim=2), dim=2, n=sup_size_u) 32 | v_0 = sup_u ** 2 * torch.fft.irfft(torch.fft.rfft(v_0, dim=3), dim=3, n=sup_size_u) 33 | mu_1 = x - v_1 * delta_t 34 | index = mu_1 - mu_1.floor() 35 | index = (index * (sup_size_u)).round().reshape(-1, 2) 36 | index[index==sup_size_u] = 0 37 | index = index.long() 38 | index_batch = torch.arange(batch_size * time_steps, device=device).reshape(-1, 1).repeat(1, size**2).reshape(-1,).long() 39 | index = index_batch * sup_size_u ** 2 + index[:, 0] * sup_size_u + index[:, 1] 40 | 41 | v_2 = torch.concat([torch.take(v_0[..., 0], index.long()).reshape(time_steps, batch_size, size, size, 1), 42 | torch.take(v_0[..., 1], index.long()).reshape(time_steps, batch_size, size, size, 1)], dim=-1) 43 | 44 | mu = x - 0.5*(v_1+v_2) * delta_t 45 | f_2 = 0.1* torch.cos(8*math.pi*(mu[..., 0])) 46 | 47 | if sup_w > 1: 48 | w_0 = torch.fft.irfft(torch.fft.rfft(w_0, dim=2), dim=2, n=sup_size_w) 49 | w_0 = sup_w ** 2 * torch.fft.irfft(torch.fft.rfft(w_0, dim=3), dim=3, n=sup_size_w) 50 | 51 | delta_gridx = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 52 | delta_gridx = delta_gridx.reshape(-1, 1, 1).repeat([1, 2*k+1, 1]) 53 | delta_gridy = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 54 | delta_gridy = delta_gridy.reshape(1, -1, 1).repeat([2*k+1, 1, 1]) 55 | delta_grid = torch.cat((delta_gridx, delta_gridy), dim=-1) 56 | delta_grid = delta_grid[delta_grid.norm(dim=-1) <= (k)/sup_size_w] 57 | delta_size = delta_grid.shape[0] 58 | 59 | loc_p = mu.reshape(time_steps, batch_size, size, size, 1, 2) + delta_grid.reshape(1, 1, 1, 1, -1, 2) 60 | loc_p = (loc_p * sup_size_w).round()/sup_size_w 61 | loc_density = 1/(2*torch.pi*sigma**2) * torch.exp(-((loc_p-mu.reshape(time_steps, batch_size, size, size, 1, 2))**2).sum(axis=-1)/(2 * sigma**2)) * (1/sup_size_w**2) 62 | loc_density = loc_density/loc_density.detach().sum(dim=-1, keepdim=True) 63 | w_0 = w_0.reshape(-1, sup_size_w, sup_size_w) 64 | loc_p = loc_p - loc_p.floor() 65 | loc_p = (loc_p * (sup_size_w)).round().reshape(-1, 2) 66 | loc_p[loc_p==sup_size_w] = 0 67 | loc_p = loc_p.long() 68 | index_batch = torch.arange(batch_size * time_steps, device = device).reshape(-1, 1).repeat(1, delta_size * size**2).reshape(-1,).long() 69 | loc_p = index_batch * sup_size_w ** 2 + loc_p[:, 0] * sup_size_w + loc_p[:, 1] 70 | w_hat = torch.take(w_0, loc_p.long()).reshape(time_steps, batch_size, size, size, delta_size, 1) 71 | f = 0.5 * (f_1.reshape(1, 1, size, size, 1) + f_2[..., None]) 72 | w_hat = (w_hat * loc_density[..., None]).sum(axis=-2) + f * delta_t 73 | 74 | return torch.sqrt(torch.mean(torch.square(w_hat - w_1))) -------------------------------------------------------------------------------- /navier-stokes/E3/mcnp10/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_contour(w): 11 | ''' 12 | :param w: shape: s * s 13 | :return: figs 14 | ''' 15 | temp = w.detach().cpu().numpy() 16 | s = w.shape[0] 17 | plt.figure() 18 | plt.contour(temp.reshape(s, s), 50) 19 | plt.show() 20 | 21 | 22 | def get_grid(shape, device): 23 | batchsize, size_x, size_y = shape[0], shape[1], shape[2] 24 | gridx = torch.tensor(np.linspace(0, 1 - 1 / size_x, size_x), dtype=torch.float) 25 | gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1]) 26 | gridy = torch.tensor(np.linspace(0, 1 - 1 / size_y, size_y), dtype=torch.float) 27 | gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1]) 28 | return torch.cat((gridx, gridy), dim=-1).to(device) 29 | 30 | 31 | def kernel_x(x, config): 32 | device = x.device 33 | xkt = 2 * torch.pi * torch.mm(x, config.K.t().to(device)) 34 | return torch.sin(xkt), torch.cos(xkt) 35 | 36 | 37 | def w2v_g(w, size): 38 | ''' 39 | input: vortex field w, size = [bw, Batch_grid, 1] 40 | output: velocity field u, size = [bw, Batch_size, 2] 41 | ''' 42 | device = w.device 43 | N = size 44 | bw = w.shape[0] 45 | #Maximum frequency 46 | k_max = math.floor(N/2.0) 47 | w0 = w.reshape(bw, N, N) 48 | 49 | #Initial vorticity to Fourier space 50 | w_h = torch.fft.rfft2(w0) 51 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(N,1) 52 | #Wavenumbers in x-direction 53 | k_x = k_y.transpose(0,1) 54 | #Truncate redundant modes 55 | k_x = k_x[..., :k_max + 1] 56 | k_y = k_y[..., :k_max + 1] 57 | #Negative Laplacian in Fourier space 58 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 59 | lap[0,0] = 1.0 60 | psi_h = w_h / lap 61 | 62 | #Velocity field in x-direction = psi_y 63 | q = 2. * math.pi * k_y * 1j * psi_h 64 | q = torch.fft.irfft2(q, s=(N, N)) 65 | 66 | #Velocity field in y-direction = -psi_x 67 | v = -2. * math.pi * k_x * 1j * psi_h 68 | v = torch.fft.irfft2(v, s=(N, N)) 69 | return torch.concat([q[:,:,:,None], v[:,:,:,None]], dim=-1).reshape(bw, -1, 2) 70 | 71 | 72 | def setup_seed(seed): 73 | torch.manual_seed(seed) 74 | torch.cuda.manual_seed(seed) 75 | torch.cuda.manual_seed_all(seed) 76 | np.random.seed(seed) 77 | random.seed(seed) 78 | os.environ['PYTHONHASHSEED'] = str(seed) 79 | torch.backends.cudnn.deterministic = True 80 | torch.backends.cudnn.benchmark = True 81 | torch.backends.cudnn.enabled = True 82 | 83 | 84 | class GaussianRF(object): 85 | 86 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 87 | self.dim = 2 88 | self.device = device 89 | 90 | if sigma is None: 91 | sigma = tau**(0.5*(2*alpha - self.dim)) 92 | k_max = size//2 93 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 94 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 95 | 96 | k_x = wavenumers.transpose(0,1) 97 | k_y = wavenumers 98 | 99 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 100 | self.sqrt_eig[0,0] = 0.0 101 | 102 | self.size = [] 103 | for j in range(self.dim): 104 | self.size.append(size) 105 | 106 | self.size = tuple(self.size) 107 | 108 | def __call__(self, N): 109 | 110 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 111 | coeff = self.sqrt_eig * coeff 112 | 113 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 114 | -------------------------------------------------------------------------------- /navier-stokes/E3/mcnp10/train.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from tools import GaussianRF 3 | from mc_loss import mc_loss 4 | import math 5 | 6 | 7 | 8 | import sys 9 | import time 10 | import copy 11 | import random 12 | import numpy as np 13 | import torch 14 | import torch.optim as optim 15 | import matplotlib.pyplot as plt 16 | import os 17 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 18 | plt.rcParams["animation.html"] = "jshtml" 19 | 20 | 21 | def test(config, net, test_data): 22 | num = test_data.shape[0] 23 | device = config.device 24 | delta_steps = 10 25 | 26 | delta_t = config.T / delta_steps 27 | w_0 = test_data[:, :, :, 0].to(device)[:, :, :, None] 28 | rela_err = torch.zeros(delta_steps) 29 | rela_err_1 = torch.zeros(delta_steps) 30 | rela_err_max = torch.zeros(delta_steps) 31 | for time_step in range(1, delta_steps+1): 32 | net.eval() 33 | t = (delta_t * time_step) * torch.ones(num).to(device) 34 | w = net(w_0, t).detach() 35 | w_t = test_data[..., delta_steps*time_step].to(device)[:, :, :, None] 36 | rela_err[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), dim=1)).mean() 37 | rela_err_1[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=1, dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=1, dim=1)).mean() 38 | rela_err_max[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=float('inf'), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=float('inf'), dim=1)).mean() 39 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 40 | print('mean relative l_2 error', rela_err.mean().item()) 41 | print('mean relative l_1 error', rela_err_1.mean().item()) 42 | print('mean relative l_inf error', rela_err_max.mean().item()) 43 | return rela_err.mean().item() 44 | 45 | 46 | def train_epoch(net, GRF, config, x, f, optimizer, scheduler): 47 | batch_size, size = config.batch_size, config.size 48 | net.train() 49 | w_0 = GRF(batch_size).reshape(batch_size, size, size, 1) 50 | loss = mc_loss(w_0, x, f, net, config) 51 | loss.backward() 52 | optimizer.step() 53 | optimizer.zero_grad(set_to_none=True) 54 | scheduler.step() 55 | 56 | 57 | def train(config, net): 58 | device = config.device 59 | size = config.size 60 | batch_size = config.batch_size 61 | net = net.to(device) 62 | optimizer = optim.AdamW(net.parameters(), config.lr, weight_decay=config.weight_decay) 63 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 64 | val_data = torch.load(config.data_path + 'data_val').to(device).float()[..., :int(10*config.T)+1] 65 | test_data = torch.load(config.data_path+'data_test').to(device).float()[..., :int(10*config.T)+1] 66 | val_loss = 1e50 67 | 68 | gridx = torch.linspace(0, 1 - 1 / size, size, device=device) 69 | gridx = gridx.reshape(size, 1, 1).repeat([1, size, 1]) 70 | gridy = torch.linspace(0, 1 - 1 / size, size, device=device) 71 | gridy = gridy.reshape(1, size, 1).repeat([size, 1, 1]) 72 | x = torch.cat((gridx, gridy), dim=-1) 73 | f = 0.1*torch.cos(8*math.pi*(x[..., 0])) 74 | GRF = GaussianRF(size, device=device) 75 | for step in range(config.num_iterations+1): 76 | train_epoch(net, GRF, config, x, f, optimizer, scheduler) 77 | if step % 50 == 0: 78 | print('########################') 79 | print('training loss', step) 80 | print('VAL_LOSS') 81 | print('########################') 82 | temp = test(config, net, val_data) 83 | if temp < val_loss: 84 | val_loss = temp 85 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 86 | print('TEST_LOSS') 87 | print('########################') 88 | test(config, net, test_data) 89 | sys.stdout.flush() 90 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 91 | print('FINAL_LOSS') 92 | print('########################') 93 | test(config, net, test_data) 94 | sys.stdout.flush() -------------------------------------------------------------------------------- /navier-stokes/E4/data/generate_data.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | import math 5 | import os 6 | 7 | from nse import navier_stokes_2d, GaussianRF 8 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 9 | 10 | 11 | def check_directory() -> None: 12 | """ 13 | Check if log directory exists within experiments 14 | """ 15 | if not os.path.exists(f'dataset'): 16 | os.mkdir(f'dataset') 17 | 18 | 19 | def setup_seed(seed): 20 | torch.manual_seed(seed) 21 | torch.cuda.manual_seed(seed) 22 | torch.cuda.manual_seed_all(seed) 23 | np.random.seed(seed) 24 | random.seed(seed) 25 | os.environ['PYTHONHASHSEED'] = str(seed) 26 | torch.backends.cudnn.deterministic = True 27 | torch.backends.cudnn.benchmark = False 28 | torch.backends.cudnn.enabled = True 29 | 30 | 31 | setup_seed(0) 32 | check_directory() 33 | device = torch.device('cuda:1') 34 | 35 | s = 256 36 | sub = 4 37 | 38 | # Set up 2d GRF with covariance parameters 39 | GRF = GaussianRF(s, device=device) 40 | 41 | # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y))) 42 | t = torch.linspace(0, 1, s+1, device=device) 43 | t = t[0: -1] 44 | 45 | X, Y = torch.meshgrid(t, t, indexing='ij') 46 | 47 | # Number of snapshots from solution 48 | 49 | 50 | 51 | ###################################################################### 52 | 53 | record_steps = 100 54 | T = 10 55 | N = 200 56 | bsize = 200 57 | c = 0 58 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 59 | f = 0.1*torch.cos(8*math.pi*(X)) 60 | 61 | 62 | for j in range(N//bsize): 63 | # Sample random feilds 64 | w0 = GRF(bsize) 65 | visc = 10 ** (- 5 * torch.ones(bsize).to(device)) 66 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 67 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 68 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 69 | u[c:(c+bsize),...] = sol 70 | c += bsize 71 | print(j, c) 72 | torch.save(u, './dataset/data_test') 73 | 74 | setup_seed(1) 75 | record_steps = 100 76 | T = 10 77 | N = 200 78 | bsize = 200 79 | c = 0 80 | u = torch.zeros(N, s//sub, s//sub, record_steps+1) 81 | f = 0.1*torch.cos(8*math.pi*(X)) 82 | 83 | 84 | for j in range(N//bsize): 85 | # Sample random feilds 86 | w0 = GRF(bsize) 87 | visc = 10 ** (- 5 * torch.ones(bsize).to(device)) 88 | sol, sol_t = navier_stokes_2d(w0, f, visc, T, 1e-4, record_steps) 89 | w0 = w0[:, ::sub, ::sub].reshape(-1, s//sub, s//sub, 1) 90 | sol = torch.concat([w0, sol[:, ::sub, ::sub, :]], dim=3) 91 | u[c:(c+bsize),...] = sol 92 | c += bsize 93 | print(j, c) 94 | torch.save(u, './dataset/data_val') 95 | -------------------------------------------------------------------------------- /navier-stokes/E4/data/nse.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import os 4 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 5 | 6 | class GaussianRF(object): 7 | 8 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 9 | self.dim = 2 10 | self.device = device 11 | 12 | if sigma is None: 13 | sigma = tau**(0.5*(2*alpha - self.dim)) 14 | k_max = size//2 15 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 16 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 17 | 18 | k_x = wavenumers.transpose(0,1) 19 | k_y = wavenumers 20 | 21 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 22 | self.sqrt_eig[0,0] = 0.0 23 | 24 | self.size = [] 25 | for j in range(self.dim): 26 | self.size.append(size) 27 | 28 | self.size = tuple(self.size) 29 | 30 | def __call__(self, N): 31 | 32 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 33 | coeff = self.sqrt_eig * coeff 34 | 35 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 36 | 37 | 38 | def force(x): 39 | z = x.sum(axis=1).reshape(x.shape[0], 1) 40 | return 0.1 * torch.sin(2 * math.pi * z) + 0.1 * torch.cos(2 * math.pi * z) 41 | 42 | 43 | def navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1): 44 | 45 | # Grid size - must be power of 2 46 | N = w0.size()[-1] 47 | 48 | # Maximum frequency 49 | k_max = math.floor(N/2.0) 50 | 51 | # Number of steps to final time 52 | steps = math.ceil(T/delta_t) 53 | 54 | # Initial vorticity to Fourier space 55 | w_h = torch.fft.rfft2(w0) 56 | 57 | # Forcing to Fourier space 58 | f_h = torch.fft.rfft2(f) 59 | 60 | #If same forcing for the whole batch 61 | if len(f_h.size()) < len(w_h.size()): 62 | f_h = torch.unsqueeze(f_h, 0) 63 | 64 | # Record solution every this number of steps 65 | record_time = math.floor(steps/record_steps) 66 | 67 | # Wavenumbers in y-direction 68 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device), torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N,1) 69 | # Wavenumbers in x-direction 70 | k_x = k_y.transpose(0,1) 71 | 72 | # Truncate redundant modes 73 | k_x = k_x[..., :k_max + 1] 74 | k_y = k_y[..., :k_max + 1] 75 | 76 | # Negative Laplacian in Fourier space 77 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 78 | lap[0,0] = 1.0 79 | # Dealiasing mask 80 | dealias = torch.unsqueeze(torch.logical_and(torch.abs(k_y) <= (2.0/3.0)*k_max, torch.abs(k_x) <= (2.0/3.0)*k_max).float(), 0) 81 | 82 | # Saving solution and time 83 | sol = torch.zeros(*w0.size(), record_steps, device=w0.device) 84 | sol_t = torch.zeros(record_steps, device=w0.device) 85 | 86 | # Record counter 87 | c = 0 88 | # Physical time 89 | t = 0.0 90 | for j in range(steps): 91 | # Stream function in Fourier space: solve Poisson equation 92 | psi_h = w_h / lap 93 | 94 | # Velocity field in x-direction = psi_y 95 | q = 2. * math.pi * k_y * 1j * psi_h 96 | q = torch.fft.irfft2(q, s=(N, N)) 97 | 98 | # Velocity field in y-direction = -psi_x 99 | v = -2. * math.pi * k_x * 1j * psi_h 100 | v = torch.fft.irfft2(v, s=(N, N)) 101 | 102 | # Partial x of vorticity 103 | w_x = 2. * math.pi * k_x * 1j * w_h 104 | w_x = torch.fft.irfft2(w_x, s=(N, N)) 105 | 106 | # Partial y of vorticity 107 | w_y = 2. * math.pi * k_y * 1j * w_h 108 | w_y = torch.fft.irfft2(w_y, s=(N, N)) 109 | 110 | # Non-linear term (u.grad(w)): compute in physical space then back to Fourier space 111 | F_h = torch.fft.rfft2(q*w_x + v*w_y) 112 | 113 | # Dealias 114 | F_h = dealias* F_h 115 | 116 | # Crank-Nicolson update 117 | w_h = (-delta_t*F_h + delta_t*f_h + (1.0 - 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...])*w_h)/(1.0 + 0.5*delta_t*visc.reshape(-1, 1, 1)*lap[None,...]) 118 | 119 | # Update real time (used only for recording) 120 | t += delta_t 121 | 122 | if (j+1) % record_time == 0: 123 | # Solution in physical space 124 | w = torch.fft.irfft2(w_h, s=(N, N)) 125 | 126 | # Record solution and time 127 | sol[...,c] = w 128 | sol_t[c] = t 129 | 130 | c += 1 131 | return sol, sol_t 132 | 133 | -------------------------------------------------------------------------------- /navier-stokes/E4/mcnp10/main.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from train import train 3 | from tools import setup_seed 4 | 5 | import json 6 | import sys 7 | import copy 8 | from datetime import datetime 9 | import random 10 | import argparse 11 | import numpy as np 12 | import torch 13 | import torch.optim as optim 14 | import matplotlib.pyplot as plt 15 | import os 16 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 17 | plt.rcParams["animation.html"] = "jshtml" 18 | 19 | 20 | def main(cfg): 21 | if not os.path.exists(f'log'): 22 | os.mkdir(f'log') 23 | if not os.path.exists(f'model'): 24 | os.mkdir(f'model') 25 | setup_seed(cfg.seed) 26 | dateTimeObj = datetime.now() 27 | timestring = f'{dateTimeObj.date().month}_{dateTimeObj.date().day}_{dateTimeObj.time().hour}_{dateTimeObj.time().minute}' 28 | logfile = f'log/seed_{cfg.seed}_k_{cfg.k}_u_{cfg.sup_u}_w_{cfg.sup_w}_Adam_{cfg.num_iterations}_{cfg.step_size}_{cfg.gamma}_{cfg.lr}_{timestring}.csv' 29 | 30 | cfg.timestring = timestring 31 | with open('log/cfg_'+ timestring +'.txt', 'w') as f: 32 | json.dump(cfg.__dict__, f, indent=2) 33 | 34 | sys.stdout = open(logfile, 'w') 35 | 36 | print('--------args----------') 37 | for k in list(vars(cfg).keys()): 38 | print('%s: %s' % (k, vars(cfg)[k])) 39 | print('--------args----------\n') 40 | 41 | sys.stdout.flush() 42 | cfg.delta_t = cfg.T/cfg.time_steps 43 | net = FNO(cfg.NUM_TRUNCATION, cfg.NUM_TRUNCATION, cfg.num_channel) 44 | train(cfg, net) 45 | 46 | 47 | if __name__ == "__main__": 48 | parser = argparse.ArgumentParser(description='Hyper-parameters of pretraining') 49 | 50 | parser.add_argument('--data_path', type=str, 51 | default=os.path.abspath(os.path.join(os.getcwd(), "../")) + '/data/dataset/', 52 | help='path of data') 53 | 54 | parser.add_argument('--seed', type=int, default=0, 55 | help='seed') 56 | 57 | parser.add_argument('--batch_size', type=int, default=10, 58 | help='batchsize of the operator learning') 59 | 60 | parser.add_argument('--step_size', type=int, default=2000, 61 | help='step_size of optim') 62 | 63 | parser.add_argument('--gamma', type=float, default=0.8, 64 | help='gamma of optim') 65 | 66 | parser.add_argument('--lr', type=float, default=0.01, 67 | help='lr of optim') 68 | 69 | parser.add_argument('--weight_decay', type=float, default=0.0, 70 | help='lr of optim') 71 | 72 | parser.add_argument('--num_iterations', type=int, default=20000, 73 | help='num_iterations of optim') 74 | 75 | parser.add_argument('--nu', type=float, default=1/100000, 76 | help='nu in NSE') 77 | 78 | parser.add_argument('--size', type=int, default=64, 79 | help='data spatial size') 80 | 81 | parser.add_argument('--T', type=float, default=10.0, 82 | help='final time') 83 | 84 | parser.add_argument('--time_steps', type=int, default=10, 85 | help='number of time_steps in data') 86 | 87 | parser.add_argument('--EPS', type=float, default=1e-8, 88 | help='epsilon') 89 | 90 | parser.add_argument('--NUM_TRUNCATION', type=int, default=16, 91 | help='NUM_TRUNCATION in FNO') 92 | 93 | parser.add_argument('--num_channel', type=int, default=36, 94 | help='num_channel in FNO') 95 | 96 | parser.add_argument('--sup_u', type=int, default=4, 97 | help='sup_u') 98 | 99 | parser.add_argument('--sup_w', type=int, default=2, 100 | help='sup_w') 101 | 102 | parser.add_argument('--k', type=int, default=2, 103 | help='k') 104 | 105 | parser.add_argument('--device', type=str, default='cuda:0', 106 | help='Used device') 107 | 108 | cfg = parser.parse_args() 109 | for seed in [0, 1, 2]: 110 | cfg.seed = seed 111 | main(cfg) -------------------------------------------------------------------------------- /navier-stokes/E4/mcnp10/mc_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import numpy as np 4 | from tools import w2v_g 5 | 6 | 7 | 8 | def mc_loss(w_0, x, f_1, model, config): 9 | device = config.device 10 | size = config.size 11 | sup_w, sup_u = config.sup_w, config.sup_u 12 | sup_size_w, sup_size_u = sup_w * size, sup_u * size 13 | batch_size = config.batch_size 14 | time_steps = config.time_steps 15 | delta_t = config.delta_t 16 | k = config.k 17 | nu = config.nu 18 | sigma = math.sqrt(2 * nu * delta_t) 19 | w_0 = w_0.reshape([1, batch_size, size, size, 1]) 20 | 21 | w_input = w_0.repeat(time_steps, 1, 1, 1, 1).reshape(-1, size, size, 1) 22 | t = delta_t * torch.arange(1, time_steps+1, device=device).reshape(-1, 1).repeat(1,batch_size).reshape(-1,) 23 | w = model(w_input, t).reshape(time_steps, batch_size, size, size, 1) 24 | w = torch.concat([w_0.reshape(1, -1, size, size, 1), w], dim=0) 25 | 26 | w_0, w_1 = w[:-1, ...], w[1:, ...] 27 | v = w2v_g(w.reshape(-1, size, size, 1), size).reshape(time_steps+1, batch_size, size, size, 2) 28 | v_0, v_1 = v[:-1, ...], v[1:, ...] 29 | 30 | if sup_u > 1: 31 | v_0 = torch.fft.irfft(torch.fft.rfft(v_0, dim=2), dim=2, n=sup_size_u) 32 | v_0 = sup_u ** 2 * torch.fft.irfft(torch.fft.rfft(v_0, dim=3), dim=3, n=sup_size_u) 33 | mu_1 = x - v_1 * delta_t 34 | index = mu_1 - mu_1.floor() 35 | index = (index * (sup_size_u)).round().reshape(-1, 2) 36 | index[index==sup_size_u] = 0 37 | index = index.long() 38 | index_batch = torch.arange(batch_size * time_steps, device=device).reshape(-1, 1).repeat(1, size**2).reshape(-1,).long() 39 | index = index_batch * sup_size_u ** 2 + index[:, 0] * sup_size_u + index[:, 1] 40 | 41 | v_2 = torch.concat([torch.take(v_0[..., 0], index.long()).reshape(time_steps, batch_size, size, size, 1), 42 | torch.take(v_0[..., 1], index.long()).reshape(time_steps, batch_size, size, size, 1)], dim=-1) 43 | 44 | mu = x - 0.5*(v_1+v_2) * delta_t 45 | f_2 = 0.1* torch.cos(8*math.pi*(mu[..., 0])) 46 | 47 | if sup_w > 1: 48 | w_0 = torch.fft.irfft(torch.fft.rfft(w_0, dim=2), dim=2, n=sup_size_w) 49 | w_0 = sup_w ** 2 * torch.fft.irfft(torch.fft.rfft(w_0, dim=3), dim=3, n=sup_size_w) 50 | 51 | delta_gridx = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 52 | delta_gridx = delta_gridx.reshape(-1, 1, 1).repeat([1, 2*k+1, 1]) 53 | delta_gridy = torch.linspace(-k/sup_size_w, k/sup_size_w, 2*k+1, device=device) 54 | delta_gridy = delta_gridy.reshape(1, -1, 1).repeat([2*k+1, 1, 1]) 55 | delta_grid = torch.cat((delta_gridx, delta_gridy), dim=-1) 56 | delta_grid = delta_grid[delta_grid.norm(dim=-1) <= (k)/sup_size_w] 57 | delta_size = delta_grid.shape[0] 58 | 59 | loc_p = mu.reshape(time_steps, batch_size, size, size, 1, 2) + delta_grid.reshape(1, 1, 1, 1, -1, 2) 60 | loc_p = (loc_p * sup_size_w).round()/sup_size_w 61 | loc_density = 1/(2*torch.pi*sigma**2) * torch.exp(-((loc_p-mu.reshape(time_steps, batch_size, size, size, 1, 2))**2).sum(axis=-1)/(2 * sigma**2)) * (1/sup_size_w**2) 62 | loc_density = loc_density/loc_density.detach().sum(dim=-1, keepdim=True) 63 | w_0 = w_0.reshape(-1, sup_size_w, sup_size_w) 64 | loc_p = loc_p - loc_p.floor() 65 | loc_p = (loc_p * (sup_size_w)).round().reshape(-1, 2) 66 | loc_p[loc_p==sup_size_w] = 0 67 | loc_p = loc_p.long() 68 | index_batch = torch.arange(batch_size * time_steps, device = device).reshape(-1, 1).repeat(1, delta_size * size**2).reshape(-1,).long() 69 | loc_p = index_batch * sup_size_w ** 2 + loc_p[:, 0] * sup_size_w + loc_p[:, 1] 70 | w_hat = torch.take(w_0, loc_p.long()).reshape(time_steps, batch_size, size, size, delta_size, 1) 71 | f = 0.5 * (f_1.reshape(1, 1, size, size, 1) + f_2[..., None]) 72 | w_hat = (w_hat * loc_density[..., None]).sum(axis=-2) + f * delta_t 73 | 74 | return torch.sqrt(torch.mean(torch.square(w_hat - w_1))) -------------------------------------------------------------------------------- /navier-stokes/E4/mcnp10/tools.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import torch 3 | import numpy as np 4 | import math 5 | import random 6 | import os 7 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 8 | 9 | 10 | def plot_contour(w): 11 | ''' 12 | :param w: shape: s * s 13 | :return: figs 14 | ''' 15 | temp = w.detach().cpu().numpy() 16 | s = w.shape[0] 17 | plt.figure() 18 | plt.contour(temp.reshape(s, s), 50) 19 | plt.show() 20 | 21 | 22 | def get_grid(shape, device): 23 | batchsize, size_x, size_y = shape[0], shape[1], shape[2] 24 | gridx = torch.tensor(np.linspace(0, 1 - 1 / size_x, size_x), dtype=torch.float) 25 | gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1]) 26 | gridy = torch.tensor(np.linspace(0, 1 - 1 / size_y, size_y), dtype=torch.float) 27 | gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1]) 28 | return torch.cat((gridx, gridy), dim=-1).to(device) 29 | 30 | 31 | def kernel_x(x, config): 32 | device = x.device 33 | xkt = 2 * torch.pi * torch.mm(x, config.K.t().to(device)) 34 | return torch.sin(xkt), torch.cos(xkt) 35 | 36 | 37 | def w2v_g(w, size): 38 | ''' 39 | input: vortex field w, size = [bw, Batch_grid, 1] 40 | output: velocity field u, size = [bw, Batch_size, 2] 41 | ''' 42 | device = w.device 43 | N = size 44 | bw = w.shape[0] 45 | #Maximum frequency 46 | k_max = math.floor(N/2.0) 47 | w0 = w.reshape(bw, N, N) 48 | 49 | #Initial vorticity to Fourier space 50 | w_h = torch.fft.rfft2(w0) 51 | k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(N,1) 52 | #Wavenumbers in x-direction 53 | k_x = k_y.transpose(0,1) 54 | #Truncate redundant modes 55 | k_x = k_x[..., :k_max + 1] 56 | k_y = k_y[..., :k_max + 1] 57 | #Negative Laplacian in Fourier space 58 | lap = 4*(math.pi**2)*(k_x**2 + k_y**2) 59 | lap[0,0] = 1.0 60 | psi_h = w_h / lap 61 | 62 | #Velocity field in x-direction = psi_y 63 | q = 2. * math.pi * k_y * 1j * psi_h 64 | q = torch.fft.irfft2(q, s=(N, N)) 65 | 66 | #Velocity field in y-direction = -psi_x 67 | v = -2. * math.pi * k_x * 1j * psi_h 68 | v = torch.fft.irfft2(v, s=(N, N)) 69 | return torch.concat([q[:,:,:,None], v[:,:,:,None]], dim=-1).reshape(bw, -1, 2) 70 | 71 | 72 | def setup_seed(seed): 73 | torch.manual_seed(seed) 74 | torch.cuda.manual_seed(seed) 75 | torch.cuda.manual_seed_all(seed) 76 | np.random.seed(seed) 77 | random.seed(seed) 78 | os.environ['PYTHONHASHSEED'] = str(seed) 79 | torch.backends.cudnn.deterministic = True 80 | torch.backends.cudnn.benchmark = True 81 | torch.backends.cudnn.enabled = True 82 | 83 | 84 | class GaussianRF(object): 85 | 86 | def __init__(self, size, alpha=2.5, tau=7, sigma=None, boundary="periodic", device=None): 87 | self.dim = 2 88 | self.device = device 89 | 90 | if sigma is None: 91 | sigma = tau**(0.5*(2*alpha - self.dim)) 92 | k_max = size//2 93 | wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \ 94 | torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1) 95 | 96 | k_x = wavenumers.transpose(0,1) 97 | k_y = wavenumers 98 | 99 | self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) 100 | self.sqrt_eig[0,0] = 0.0 101 | 102 | self.size = [] 103 | for j in range(self.dim): 104 | self.size.append(size) 105 | 106 | self.size = tuple(self.size) 107 | 108 | def __call__(self, N): 109 | 110 | coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device) 111 | coeff = self.sqrt_eig * coeff 112 | 113 | return torch.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real 114 | -------------------------------------------------------------------------------- /navier-stokes/E4/mcnp10/train.py: -------------------------------------------------------------------------------- 1 | from model import FNO 2 | from tools import GaussianRF 3 | from mc_loss import mc_loss 4 | import math 5 | 6 | 7 | 8 | import sys 9 | import time 10 | import copy 11 | import random 12 | import numpy as np 13 | import torch 14 | import torch.optim as optim 15 | import matplotlib.pyplot as plt 16 | import os 17 | os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" 18 | plt.rcParams["animation.html"] = "jshtml" 19 | 20 | 21 | def test(config, net, test_data): 22 | num = test_data.shape[0] 23 | device = config.device 24 | delta_steps = 10 25 | 26 | delta_t = config.T / delta_steps 27 | w_0 = test_data[:, :, :, 0].to(device)[:, :, :, None] 28 | rela_err = torch.zeros(delta_steps) 29 | rela_err_1 = torch.zeros(delta_steps) 30 | rela_err_max = torch.zeros(delta_steps) 31 | for time_step in range(1, delta_steps+1): 32 | net.eval() 33 | t = (delta_t * time_step) * torch.ones(num).to(device) 34 | w = net(w_0, t).detach() 35 | w_t = test_data[..., delta_steps*time_step].to(device)[:, :, :, None] 36 | rela_err[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), dim=1)).mean() 37 | rela_err_1[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=1, dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=1, dim=1)).mean() 38 | rela_err_max[time_step-1] = (torch.norm((w- w_t).reshape(w.shape[0], -1), p=float('inf'), dim=1) / torch.norm(w_t.reshape(w.shape[0], -1), p=float('inf'), dim=1)).mean() 39 | print(time_step, 'relative l_2 error', rela_err[time_step-1].item(), 'relative l_1 error', rela_err_1[time_step-1].item(), 'relative l_inf error', rela_err_max[time_step-1].item()) 40 | print('mean relative l_2 error', rela_err.mean().item()) 41 | print('mean relative l_1 error', rela_err_1.mean().item()) 42 | print('mean relative l_inf error', rela_err_max.mean().item()) 43 | return rela_err.mean().item() 44 | 45 | 46 | def train_epoch(net, GRF, config, x, f, optimizer, scheduler): 47 | batch_size, size = config.batch_size, config.size 48 | net.train() 49 | w_0 = GRF(batch_size).reshape(batch_size, size, size, 1) 50 | loss = mc_loss(w_0, x, f, net, config) 51 | loss.backward() 52 | optimizer.step() 53 | optimizer.zero_grad(set_to_none=True) 54 | scheduler.step() 55 | 56 | 57 | 58 | def train(config, net): 59 | device = config.device 60 | size = config.size 61 | batch_size = config.batch_size 62 | net = net.to(device) 63 | optimizer = optim.AdamW(net.parameters(), config.lr, weight_decay=config.weight_decay) 64 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma) 65 | val_data = torch.load(config.data_path + 'data_val').to(device).float()[..., :int(10*config.T)+1] 66 | test_data = torch.load(config.data_path+'data_test').to(device).float()[..., :int(10*config.T)+1] 67 | val_loss = 1e50 68 | 69 | gridx = torch.linspace(0, 1 - 1 / size, size, device=device) 70 | gridx = gridx.reshape(size, 1, 1).repeat([1, size, 1]) 71 | gridy = torch.linspace(0, 1 - 1 / size, size, device=device) 72 | gridy = gridy.reshape(1, size, 1).repeat([size, 1, 1]) 73 | x = torch.cat((gridx, gridy), dim=-1) 74 | f = 0.1*torch.cos(8*math.pi*(x[..., 0])) 75 | GRF = GaussianRF(size, device=device) 76 | for step in range(config.num_iterations+1): 77 | train_epoch(net, GRF, config, x, f, optimizer, scheduler) 78 | if step % 50 == 0: 79 | print('########################') 80 | print('training loss', step) 81 | print('VAL_LOSS') 82 | print('########################') 83 | temp = test(config, net, val_data) 84 | if temp < val_loss: 85 | val_loss = temp 86 | torch.save(net.state_dict(), f'model/net_seed_{config.seed}_{config.timestring}.pt') 87 | print('TEST_LOSS') 88 | print('########################') 89 | test(config, net, test_data) 90 | sys.stdout.flush() 91 | net.load_state_dict(torch.load(f'model/net_seed_{config.seed}_{config.timestring}.pt')) 92 | print('FINAL_LOSS') 93 | print('########################') 94 | test(config, net, test_data) 95 | sys.stdout.flush() --------------------------------------------------------------------------------