├── Example_2_elasticity ├── MLPNet.py ├── KAN_RBF6.py └── Orthotropic_ux.py ├── Example_1_sine_wave ├── MLPNet.py ├── RBF_KAN_Layer.py ├── Sinewave_MLP.py └── Sinewave_RBF.py └── Example_3_transient ├── MLPNet.py ├── KAN_RBF6.py ├── final_mlp_transient_poisson.py └── final_rbfkan_transient_poisson.py /Example_2_elasticity/MLPNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from typing import List, Union 5 | 6 | class MLPNetwork(nn.Module): 7 | def __init__(self, input_dim: int, hidden_layers: List[int], output_dim: int, activation: nn.Module = nn.SiLU(), dtype: torch.dtype = torch.float64): 8 | super(MLPNetwork, self).__init__() 9 | layers = [] 10 | in_dim = input_dim 11 | for hidden_dim in hidden_layers: 12 | layers.append(nn.Linear(in_dim, hidden_dim, dtype=dtype)) 13 | layers.append(activation) 14 | in_dim = hidden_dim 15 | layers.append(nn.Linear(in_dim, output_dim, dtype=dtype)) 16 | self.network = nn.Sequential(*layers) 17 | 18 | def forward(self, x: torch.Tensor) -> torch.Tensor: 19 | return self.network(x) 20 | -------------------------------------------------------------------------------- /Example_1_sine_wave/MLPNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from typing import List, Union 5 | 6 | class MLPNetwork(nn.Module): 7 | def __init__(self, input_dim: int, hidden_layers: List[int], output_dim: int, activation: nn.Module = nn.SiLU(), dtype: torch.dtype = torch.float32): 8 | super(MLPNetwork, self).__init__() 9 | layers = [] 10 | in_dim = input_dim 11 | for hidden_dim in hidden_layers: 12 | layers.append(nn.Linear(in_dim, hidden_dim, dtype=dtype)) 13 | layers.append(activation) 14 | in_dim = hidden_dim 15 | layers.append(nn.Linear(in_dim, output_dim, dtype=dtype)) 16 | self.network = nn.Sequential(*layers) 17 | 18 | self.apply(self._init_weights) 19 | 20 | def forward(self, x: torch.Tensor) -> torch.Tensor: 21 | return self.network(x) 22 | 23 | def _init_weights(self, module): 24 | if isinstance(module, nn.Linear): 25 | # For Xavier Normal Initialization 26 | nn.init.xavier_normal_(module.weight) 27 | if module.bias is not None: 28 | nn.init.constant_(module.bias, 0) 29 | # For Kaiming Normal Initialization 30 | # nn.init.kaiming_normal_(module.weight, mode='fan_in', nonlinearity='relu') 31 | # if module.bias is not None: 32 | # nn.init.constant_(module.bias, 0) 33 | -------------------------------------------------------------------------------- /Example_3_transient/MLPNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from typing import List, Union 5 | 6 | class MLPNetwork(nn.Module): 7 | def __init__(self, input_dim: int, hidden_layers: List[int], output_dim: int, activation: nn.Module = nn.SiLU(), dtype: torch.dtype = torch.float32): 8 | super(MLPNetwork, self).__init__() 9 | layers = [] 10 | in_dim = input_dim 11 | for hidden_dim in hidden_layers: 12 | layers.append(nn.Linear(in_dim, hidden_dim, dtype=dtype)) 13 | layers.append(activation) 14 | in_dim = hidden_dim 15 | layers.append(nn.Linear(in_dim, output_dim, dtype=dtype)) 16 | self.network = nn.Sequential(*layers) 17 | 18 | # Apply custom initialization 19 | self.apply(self._init_weights) 20 | 21 | def forward(self, x: torch.Tensor) -> torch.Tensor: 22 | return self.network(x) 23 | 24 | def _init_weights(self, module): 25 | if isinstance(module, nn.Linear): 26 | # You can use different initializations based on your preference 27 | # For Xavier Normal Initialization 28 | nn.init.xavier_normal_(module.weight) 29 | if module.bias is not None: 30 | nn.init.constant_(module.bias, 0) 31 | # For Kaiming Normal Initialization (use this if you're using ReLU or similar activations) 32 | # nn.init.kaiming_normal_(module.weight, mode='fan_in', nonlinearity='relu') 33 | # if module.bias is not None: 34 | # nn.init.constant_(module.bias, 0) 35 | -------------------------------------------------------------------------------- /Example_2_elasticity/KAN_RBF6.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from typing import List 5 | 6 | class RadialBasisFunctionNetwork(nn.Module): 7 | def __init__(self, hidden_layers: List[int], min_grid: float = -1., max_grid: float = 1., grid_count: int = 5, apply_base_update: bool = False, activation: nn.Module = nn.SiLU(), grid_opt: bool = False, init_scale: float = 0.1, dtype: torch.dtype = torch.float32): 8 | super().__init__() 9 | self.layers = nn.ModuleList() 10 | 11 | self.layers.append(RadialBasisFunctionLayer(hidden_layers[0], hidden_layers[1], min_grid, max_grid, grid_count, apply_base_update, activation, grid_opt, init_scale, dtype)) 12 | 13 | for in_dim, out_dim in zip(hidden_layers[1:-1], hidden_layers[2:]): 14 | self.layers.append(RadialBasisFunctionLayer(in_dim, out_dim, -1., 1., grid_count, apply_base_update, activation, grid_opt, init_scale, dtype)) 15 | 16 | self.dtype = dtype 17 | 18 | def forward(self, x: torch.Tensor) -> torch.Tensor: 19 | x = x.to(self.dtype) 20 | for i, layer in enumerate(self.layers): 21 | x = layer(x) 22 | if i < len(self.layers) - 1: 23 | x = torch.tanh(x) 24 | return x 25 | 26 | class RadialBasisFunctionLayer(nn.Module): 27 | def __init__(self, in_features: int, out_features: int, min_grid: float = -1., max_grid: float = 1., grid_count: int = 5, apply_base_update: bool = False, activation: nn.Module = nn.SiLU(), grid_opt: bool = False, init_scale: float = 0.1, dtype: torch.dtype = torch.float32): 28 | super().__init__() 29 | self.apply_base_update = apply_base_update 30 | self.activation = activation 31 | self.min_grid = min_grid 32 | self.max_grid = max_grid 33 | self.grid_count = grid_count 34 | self.grid = nn.Parameter(torch.linspace(min_grid, max_grid, grid_count, dtype=dtype), requires_grad=grid_opt) 35 | self.rbf_weight = nn.Parameter(torch.randn(in_features * grid_count, out_features, dtype=dtype) * init_scale) 36 | self.base_linear = nn.Linear(in_features, out_features, dtype=dtype) if apply_base_update else None 37 | 38 | def forward(self, x: torch.Tensor) -> torch.Tensor: 39 | x = x.to(self.rbf_weight.dtype) 40 | x_unsqueezed = x.unsqueeze(-1) # Shape: (batch_size, in_features, 1) 41 | rbf_basis = torch.exp(-((x_unsqueezed - self.grid) / ((self.max_grid - self.min_grid) / (self.grid_count - 1))) ** 2) # Shape: (batch_size, in_features, grid_count) 42 | rbf_basis = rbf_basis.view(rbf_basis.size(0), -1) # Shape: (batch_size, in_features * grid_count) 43 | rbf_output = torch.einsum('bi,ij->bj', rbf_basis, self.rbf_weight) # Shape: (batch_size, out_features) 44 | 45 | if self.apply_base_update: 46 | base_output = self.base_linear(self.activation(x)) 47 | rbf_output += base_output 48 | 49 | return rbf_output -------------------------------------------------------------------------------- /Example_3_transient/KAN_RBF6.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from typing import List 5 | 6 | class RadialBasisFunctionNetwork(nn.Module): 7 | def __init__(self, hidden_layers: List[int], min_grid: float = -1., max_grid: float = 1., grid_count: int = 5, apply_base_update: bool = False, activation: nn.Module = nn.SiLU(), grid_opt: bool = False, init_scale: float = 0.1, dtype: torch.dtype = torch.float32): 8 | super().__init__() 9 | self.layers = nn.ModuleList() 10 | 11 | # First hidden layer 12 | self.layers.append(RadialBasisFunctionLayer(hidden_layers[0], hidden_layers[1], min_grid, max_grid, grid_count, apply_base_update, activation, grid_opt, init_scale, dtype)) 13 | 14 | # Remaining hidden layers 15 | for in_dim, out_dim in zip(hidden_layers[1:-1], hidden_layers[2:]): 16 | self.layers.append(RadialBasisFunctionLayer(in_dim, out_dim, -1., 1., grid_count, apply_base_update, activation, grid_opt, init_scale, dtype)) 17 | 18 | self.dtype = dtype 19 | 20 | def forward(self, x: torch.Tensor) -> torch.Tensor: 21 | x = x.to(self.dtype) 22 | for i, layer in enumerate(self.layers): 23 | x = layer(x) 24 | # Apply tanh activation after all hidden layers except the last one 25 | if i < len(self.layers) - 1: 26 | x = torch.tanh(x) 27 | return x 28 | 29 | class RadialBasisFunctionLayer(nn.Module): 30 | def __init__(self, in_features: int, out_features: int, min_grid: float = -1., max_grid: float = 1., grid_count: int = 5, apply_base_update: bool = False, activation: nn.Module = nn.SiLU(), grid_opt: bool = False, init_scale: float = 0.1, dtype: torch.dtype = torch.float32): 31 | super().__init__() 32 | self.apply_base_update = apply_base_update 33 | self.activation = activation 34 | self.min_grid = min_grid 35 | self.max_grid = max_grid 36 | self.grid_count = grid_count 37 | self.grid = nn.Parameter(torch.linspace(min_grid, max_grid, grid_count, dtype=dtype), requires_grad=grid_opt) 38 | self.rbf_weight = nn.Parameter(torch.randn(in_features * grid_count, out_features, dtype=dtype) * init_scale) 39 | self.base_linear = nn.Linear(in_features, out_features, dtype=dtype) if apply_base_update else None 40 | 41 | def forward(self, x: torch.Tensor) -> torch.Tensor: 42 | x = x.to(self.rbf_weight.dtype) 43 | x_unsqueezed = x.unsqueeze(-1) # Shape: (batch_size, in_features, 1) 44 | rbf_basis = torch.exp(-((x_unsqueezed - self.grid) / ((self.max_grid - self.min_grid) / (self.grid_count - 1))) ** 2) # Shape: (batch_size, in_features, grid_count) 45 | rbf_basis = rbf_basis.view(rbf_basis.size(0), -1) # Shape: (batch_size, in_features * grid_count) 46 | rbf_output = torch.einsum('bi,ij->bj', rbf_basis, self.rbf_weight) # Shape: (batch_size, out_features) 47 | 48 | if self.apply_base_update: 49 | base_output = self.base_linear(self.activation(x)) 50 | rbf_output += base_output 51 | 52 | return rbf_output -------------------------------------------------------------------------------- /Example_1_sine_wave/RBF_KAN_Layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from typing import List 5 | 6 | class RadialBasisFunctionNetwork(nn.Module): 7 | def __init__(self, hidden_layers: List[int], min_grid: float = -1., max_grid: float = 1., grid_count: int = 5, apply_base_update: bool = False, activation: nn.Module = nn.SiLU(), grid_opt: bool = False, dtype: torch.dtype = torch.float32, noise_scale: float = 0.1): 8 | super().__init__() 9 | self.layers = nn.ModuleList() 10 | # First layer with specified min_grid and max_grid 11 | self.layers.append(RadialBasisFunctionLayer(hidden_layers[0], hidden_layers[1], min_grid, max_grid, grid_count, apply_base_update, activation, grid_opt, dtype, noise_scale)) 12 | # Subsequent layers with min_grid and max_grid set to -1 and 1 13 | for in_dim, out_dim in zip(hidden_layers[1:-1], hidden_layers[2:]): 14 | self.layers.append(RadialBasisFunctionLayer(in_dim, out_dim, -1., 1., grid_count, apply_base_update, activation, grid_opt, dtype, noise_scale)) 15 | self.dtype = dtype 16 | 17 | def forward(self, x: torch.Tensor) -> torch.Tensor: 18 | x = x.to(self.dtype) 19 | for layer in self.layers[:-1]: 20 | x = torch.tanh(layer(x)) 21 | x = self.layers[-1](x) 22 | return x 23 | 24 | class RadialBasisFunctionLayer(nn.Module): 25 | def __init__(self, in_features: int, out_features: int, min_grid: float = -1., max_grid: float = 1., grid_count: int = 5, apply_base_update: bool = False, activation: nn.Module = nn.SiLU(), grid_opt: bool = False, dtype: torch.dtype = torch.float32, noise_scale: float = 0.1): 26 | super().__init__() 27 | self.apply_base_update = apply_base_update 28 | self.activation = activation 29 | self.min_grid = min_grid 30 | self.max_grid = max_grid 31 | self.grid_count = grid_count 32 | self.grid = nn.Parameter(torch.linspace(min_grid, max_grid, grid_count, dtype=dtype), requires_grad=grid_opt) 33 | self.rbf_weight = nn.Parameter(torch.empty(in_features * grid_count, out_features, dtype=dtype)) 34 | self.scale_base = nn.Parameter(torch.ones(out_features, dtype=dtype)) 35 | self.scale_rbf = nn.Parameter(torch.ones(out_features, dtype=dtype)) 36 | 37 | nn.init.xavier_normal_(self.rbf_weight) 38 | self.rbf_weight.data += torch.randn_like(self.rbf_weight) * noise_scale 39 | 40 | self.base_activation = nn.SiLU() if apply_base_update else None 41 | 42 | def forward(self, x: torch.Tensor) -> torch.Tensor: 43 | x = x.to(self.rbf_weight.dtype) 44 | x_unsqueezed = x.unsqueeze(-1) 45 | rbf_basis = torch.exp(-((x_unsqueezed - self.grid) / ((self.max_grid - self.min_grid) / (self.grid_count - 1))) ** 2) # Shape: (batch_size, in_features, grid_count) 46 | rbf_basis = rbf_basis.view(rbf_basis.size(0), -1) # Shape: (batch_size, in_features * grid_count) 47 | 48 | rbf_output = torch.mm(rbf_basis, self.rbf_weight) # Shape: (batch_size, out_features) 49 | 50 | if self.apply_base_update: 51 | # Compute base activation if required 52 | base_output = self.base_activation(x) 53 | base_output = base_output.mean(dim=-1, keepdim=True) 54 | output = self.scale_base * base_output + self.scale_rbf * rbf_output 55 | else: 56 | output = self.scale_rbf * rbf_output 57 | 58 | return output -------------------------------------------------------------------------------- /Example_1_sine_wave/Sinewave_MLP.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import os 4 | from MLPNet import MLPNetwork as MLPNET 5 | import torch 6 | from torch import nn 7 | import matplotlib.pyplot as plt 8 | from torch import autograd 9 | from tqdm import tqdm 10 | import jax.numpy as jnp 11 | import torch.optim as optim 12 | from torch.optim.lr_scheduler import StepLR 13 | from torch.optim.lr_scheduler import ReduceLROnPlateau 14 | import time 15 | import pandas as pd 16 | import random 17 | from statistics import mean, median 18 | import csv 19 | 20 | ######################################################################################## 21 | # Loading the data 22 | x = np.load('x.npy') 23 | c_train = np.load('c_train.npy') 24 | y_train = np.load('y_train.npy') 25 | c_test = np.load('c_test.npy') 26 | y_test = np.load('y_test.npy') 27 | 28 | dtype = torch.float64 # or torch.float32 or torch.float64 29 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 30 | c_train = torch.tensor(c_train, dtype=dtype, device=device) 31 | y_train = torch.tensor(y_train, dtype=dtype, device=device) 32 | c_test = torch.tensor(c_test, dtype=dtype, device=device) 33 | y_test = torch.tensor(y_test, dtype=dtype, device=device) 34 | x = torch.tensor(x, dtype=dtype, device=device).reshape(len(x),1) 35 | 36 | print(f'c_train shape = {c_train.shape}') 37 | print(f'y_train shape = {y_train.shape}') 38 | print(f'c_test shape = {c_test.shape}') 39 | print(f'y_test shape = {y_test.shape}') 40 | print(f'x shape = {x.shape}') 41 | 42 | print(f'c_train max = {c_train.max()}, c_train min = {c_train.min()}') 43 | print(f'x max = {x.max()}, x min = {x.min()}') 44 | 45 | input_dim_trunk = x.shape[1] 46 | input_dim_branch = c_train.shape[1] 47 | HD = 40 48 | hidden_layers_mlp = [350, 350] 49 | print(f'hidden_layers_mlp = {hidden_layers_mlp}') 50 | activation = nn.Tanh() 51 | 52 | # Hyperparameters 53 | learning_rate = 1e-3 54 | batch_size = 1024 55 | epochs = 20_000 56 | gamma = 0.9 57 | step_size = 500 58 | seed = 84 59 | 60 | c_train = c_train.clone().detach().requires_grad_(True).to(dtype) 61 | y_train = y_train.clone().detach().requires_grad_(True).to(dtype) 62 | c_test = c_test.clone().detach().requires_grad_(True).to(dtype) 63 | y_test = y_test.clone().detach().requires_grad_(True).to(dtype) 64 | 65 | class DeepONet(nn.Module): 66 | def __init__(self): 67 | super(DeepONet, self).__init__() 68 | self.trunk_net = MLPNET(input_dim=input_dim_trunk, hidden_layers=hidden_layers_mlp, activation=activation, output_dim=HD, dtype=dtype).to(device) 69 | self.branch_net = MLPNET(input_dim=input_dim_branch, hidden_layers=hidden_layers_mlp, activation=activation, output_dim=HD, dtype=dtype).to(device) 70 | 71 | def forward(self, x_branch, x_trunk): 72 | y_branch = self.branch_net(x_branch) 73 | y_trunk = self.trunk_net(x_trunk) 74 | 75 | if y_branch.shape[-1] != y_trunk.shape[-1]: 76 | raise AssertionError("Output sizes of branch net and trunk net do not match.") 77 | 78 | Y = torch.einsum('bk,nk->bn', y_branch, y_trunk) 79 | return Y 80 | 81 | model = DeepONet() 82 | 83 | class RMSDLoss(nn.Module): 84 | def __init__(self): 85 | super(RMSDLoss, self).__init__() 86 | 87 | def forward(self, output, target): 88 | if output.shape != target.shape: 89 | raise ValueError("Output and target must have the same shape for RMSD calculation.") 90 | squared_diffs = (output - target).pow(2) 91 | mean_squared_diffs = torch.mean(squared_diffs) 92 | rmsd_value = torch.sqrt(mean_squared_diffs) 93 | return rmsd_value 94 | 95 | criterion = RMSDLoss() 96 | optimizer = optim.Adam(model.parameters(), lr=learning_rate) 97 | scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma) 98 | 99 | # L-BFGS optimizer with line search 100 | # optimizer = optim.LBFGS(model.parameters(), lr=learning_rate, line_search_fn='strong_wolfe') 101 | 102 | def set_seed(seed): 103 | random.seed(seed) 104 | np.random.seed(seed) 105 | torch.manual_seed(seed) 106 | if torch.cuda.is_available(): 107 | torch.cuda.manual_seed(seed) 108 | torch.cuda.manual_seed_all(seed) 109 | set_seed(seed) 110 | 111 | def test_model(model, criterion, c_test, y_test, x, batch_size): 112 | """Evaluates the model on the test dataset in batches.""" 113 | model.eval() 114 | total_loss = 0 115 | num_batches = len(c_test) // batch_size + 1 116 | with torch.no_grad(): 117 | for i in tqdm(range(0, len(c_test), batch_size), desc='Testing'): 118 | batch_loads = c_test[i:i + batch_size] 119 | batch_solutions = y_test[i:i + batch_size] 120 | outputs = model(batch_loads, x) 121 | loss = criterion(outputs, batch_solutions) 122 | total_loss += loss.item() 123 | avg_loss = total_loss / num_batches 124 | return avg_loss 125 | 126 | def count_learnable_parameters(model): 127 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 128 | num_learnable_parameters = count_learnable_parameters(model) 129 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 130 | 131 | train_losses = [] 132 | test_losses = [] 133 | start_time = time.time() 134 | mean_test_loss = 0 135 | 136 | results = [] 137 | 138 | for epoch in range(epochs): 139 | model.train() 140 | total_loss = 0 141 | progress_bar = tqdm(range(0, len(c_train), batch_size), desc=f'Epoch {epoch + 1}/{epochs}') 142 | for i in progress_bar: 143 | batch_loads = c_train[i:i + batch_size] 144 | batch_solutions = y_train[i:i + batch_size] 145 | 146 | def closure(): 147 | optimizer.zero_grad() 148 | outputs = model(batch_loads, x) 149 | loss = criterion(outputs, batch_solutions) 150 | loss.backward() 151 | return loss 152 | 153 | loss = optimizer.step(closure) 154 | total_loss += loss.item() 155 | 156 | progress_bar.set_postfix({'Batch Loss': loss.item()}) 157 | 158 | avg_loss = total_loss / (len(c_train) // batch_size + 1) 159 | train_losses.append(avg_loss) 160 | tqdm.write(f'End of Epoch {epoch + 1}, Average Loss: {avg_loss}') 161 | 162 | test_loss = test_model(model, criterion, c_test, y_test, x, batch_size) 163 | test_losses.append(test_loss) 164 | tqdm.write(f'Test Loss: {test_loss}') 165 | 166 | results.append([epoch + 1, avg_loss, test_loss]) 167 | 168 | scheduler.step() 169 | 170 | end_time = time.time() 171 | training_time = end_time - start_time 172 | print('###################################################################################################') 173 | print('###################################################################################################') 174 | print('###################################################################################################') 175 | print(f'Training time: {training_time:.2f} seconds') 176 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 177 | print('###################################################################################################') 178 | print('###################################################################################################') 179 | print('###################################################################################################') 180 | 181 | model.eval() 182 | y_pred_list = [] 183 | batch_size = 100 184 | with torch.no_grad(): 185 | for i in range(0, len(c_test), batch_size): 186 | batch_loads = c_test[i:i+batch_size].clone().detach().to(device) 187 | batch_pred = model(batch_loads, x) 188 | y_pred_list.append(batch_pred.cpu().numpy()) 189 | 190 | y_pred_deeponet = np.concatenate(y_pred_list, axis=0) 191 | 192 | np.save('y_pred_deeponet.npy', y_pred_deeponet) 193 | 194 | print("Predictions generated and saved successfully.") 195 | print(f"Shape of y_pred_deeponet: {y_pred_deeponet.shape}") 196 | 197 | with open('training_results_deeponet.csv', 'w', newline='') as csvfile: 198 | csv_writer = csv.writer(csvfile) 199 | csv_writer.writerow(['Epoch', 'Train Loss', 'Test Loss']) 200 | csv_writer.writerows(results) 201 | 202 | print("Training results saved in 'training_results_deeponet.csv'") -------------------------------------------------------------------------------- /Example_1_sine_wave/Sinewave_RBF.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import os 4 | from RBF_KAN_Layer import RadialBasisFunctionNetwork as RBFKAN 5 | import torch 6 | from torch import nn 7 | import matplotlib.pyplot as plt 8 | from torch import autograd 9 | from tqdm import tqdm 10 | import jax.numpy as jnp 11 | import torch.optim as optim 12 | from torch.optim.lr_scheduler import StepLR 13 | from torch.optim.lr_scheduler import ReduceLROnPlateau 14 | import time 15 | import pandas as pd 16 | import random 17 | from statistics import mean, median 18 | import csv 19 | 20 | ######################################################################################## 21 | # Loading the data 22 | x = np.load('x.npy') 23 | c_train = np.load('c_train.npy') 24 | y_train = np.load('y_train.npy') 25 | c_test = np.load('c_test.npy') 26 | y_test = np.load('y_test.npy') 27 | 28 | dtype = torch.float64 # or torch.float32 or torch.float64 29 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 30 | c_train = torch.tensor(c_train, device=device) 31 | y_train = torch.tensor(y_train, device=device) 32 | c_test = torch.tensor(c_test, device=device) 33 | y_test = torch.tensor(y_test, device=device) 34 | x = torch.tensor(x, device=device).reshape(len(x),1) 35 | 36 | print(f'c_train shape = {c_train.shape}') 37 | print(f'y_train shape = {y_train.shape}') 38 | print(f'c_test shape = {c_test.shape}') 39 | print(f'y_test shape = {y_test.shape}') 40 | print(f'x shape = {x.shape}') 41 | 42 | print(f'c_train max = {c_train.max()}, c_train min = {c_train.min()}') 43 | print(f'x max = {x.max()}, x min = {x.min()}') 44 | 45 | def define_kan_width(input_dim, W, repeat_hid, output_dim): 46 | width = [input_dim] + [W] * repeat_hid + [output_dim] 47 | return width 48 | 49 | input_dim_trunk = x.shape[1] 50 | input_dim_branch = c_train.shape[1] 51 | HD = 40 52 | hid_trunk = 50 53 | num_layer_trunk = 2 54 | hid_branch = 50 55 | num_layer_branch = 2 56 | trunk_min_grid = x.min() 57 | trunk_max_grid = x.max() 58 | branch_min_grid = c_train.min() 59 | branch_max_grid = c_train.max() 60 | 61 | width_trunk = define_kan_width(input_dim_trunk, hid_trunk, num_layer_trunk, HD) 62 | width_branch = define_kan_width(input_dim_branch, hid_branch, num_layer_branch, HD) 63 | 64 | print(f'hid dim = {hid_trunk}, num_layer = {num_layer_trunk}') 65 | print(f'width_trunk = {width_trunk}, width_branch = {width_branch}') 66 | grid_opt = True 67 | grid_count = 30 68 | init_scale = 0.01 69 | noise_scale = 0.01 70 | print(f'init_scale = {init_scale}') 71 | 72 | # Hyperparameters 73 | learning_rate = 1e-2 74 | batch_size = 1024 75 | epochs = 20_000 76 | gamma = 0.9 77 | step_size = 500 78 | seed = 84 79 | 80 | c_train = c_train.clone().detach().requires_grad_(True).to(dtype) 81 | y_train = y_train.clone().detach().requires_grad_(True).to(dtype) 82 | c_test = c_test.clone().detach().requires_grad_(True).to(dtype) 83 | y_test = y_test.clone().detach().requires_grad_(True).to(dtype) 84 | 85 | class DeepOKAN(nn.Module): 86 | def __init__(self): 87 | super(DeepOKAN, self).__init__() 88 | self.trunk_net = RBFKAN(hidden_layers=width_trunk, dtype=dtype, apply_base_update=False, noise_scale=noise_scale, 89 | min_grid=trunk_min_grid, max_grid=trunk_max_grid, grid_count=grid_count, grid_opt=grid_opt).to(device) 90 | self.branch_net = RBFKAN(hidden_layers=width_branch, dtype=dtype, apply_base_update=False, noise_scale=noise_scale, 91 | min_grid=branch_min_grid, max_grid=branch_max_grid, grid_count=grid_count, grid_opt=grid_opt).to(device) 92 | 93 | def forward(self, x_branch, x_trunk): 94 | y_branch = self.branch_net(x_branch) 95 | y_trunk = self.trunk_net(x_trunk) 96 | 97 | if y_branch.shape[-1] != y_trunk.shape[-1]: 98 | raise AssertionError("Output sizes of branch net and trunk net do not match.") 99 | 100 | Y = torch.einsum('bk,nk->bn', y_branch, y_trunk) 101 | 102 | return Y 103 | 104 | model = DeepOKAN() 105 | 106 | class RMSDLoss(nn.Module): 107 | def __init__(self): 108 | super(RMSDLoss, self).__init__() 109 | 110 | def forward(self, output, target): 111 | if output.shape != target.shape: 112 | raise ValueError("Output and target must have the same shape for RMSD calculation.") 113 | squared_diffs = (output - target).pow(2) 114 | mean_squared_diffs = torch.mean(squared_diffs) 115 | rmsd_value = torch.sqrt(mean_squared_diffs) 116 | return rmsd_value 117 | 118 | criterion = RMSDLoss() 119 | optimizer = optim.Adam(model.parameters(), lr=learning_rate) 120 | scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma) 121 | 122 | # L-BFGS optimizer with line search 123 | # optimizer = optim.LBFGS(model.parameters(), lr=learning_rate, line_search_fn='strong_wolfe') 124 | 125 | def set_seed(seed): 126 | random.seed(seed) 127 | np.random.seed(seed) 128 | torch.manual_seed(seed) 129 | if torch.cuda.is_available(): 130 | torch.cuda.manual_seed(seed) 131 | torch.cuda.manual_seed_all(seed) 132 | set_seed(seed) 133 | 134 | def test_model(model, criterion, c_test, y_test, x, batch_size): 135 | """Evaluates the model on the test dataset in batches.""" 136 | model.eval() 137 | total_loss = 0 138 | num_batches = len(c_test) // batch_size + 1 139 | with torch.no_grad(): 140 | for i in tqdm(range(0, len(c_test), batch_size), desc='Testing'): 141 | batch_loads = c_test[i:i + batch_size] 142 | batch_solutions = y_test[i:i + batch_size] 143 | outputs = model(batch_loads, x) 144 | loss = criterion(outputs, batch_solutions) 145 | total_loss += loss.item() 146 | avg_loss = total_loss / num_batches 147 | return avg_loss 148 | 149 | def count_learnable_parameters(model): 150 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 151 | num_learnable_parameters = count_learnable_parameters(model) 152 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 153 | 154 | import csv 155 | import time 156 | from tqdm import tqdm 157 | import numpy as np 158 | import torch 159 | 160 | train_losses = [] 161 | test_losses = [] 162 | start_time = time.time() 163 | mean_test_loss = 0 164 | 165 | results = [] 166 | 167 | for epoch in range(epochs): 168 | model.train() 169 | total_loss = 0 170 | progress_bar = tqdm(range(0, len(c_train), batch_size), desc=f'Epoch {epoch + 1}/{epochs}') 171 | for i in progress_bar: 172 | batch_loads = c_train[i:i + batch_size] 173 | batch_solutions = y_train[i:i + batch_size] 174 | 175 | def closure(): 176 | optimizer.zero_grad() 177 | outputs = model(batch_loads, x) 178 | loss = criterion(outputs, batch_solutions) 179 | loss.backward() 180 | return loss 181 | 182 | loss = optimizer.step(closure) 183 | total_loss += loss.item() 184 | 185 | progress_bar.set_postfix({'Batch Loss': loss.item()}) 186 | 187 | avg_loss = total_loss / (len(c_train) // batch_size + 1) 188 | train_losses.append(avg_loss) 189 | tqdm.write(f'End of Epoch {epoch + 1}, Average Loss: {avg_loss}') 190 | 191 | test_loss = test_model(model, criterion, c_test, y_test, x, batch_size) 192 | test_losses.append(test_loss) 193 | tqdm.write(f'Test Loss: {test_loss}') 194 | 195 | results.append([epoch + 1, avg_loss, test_loss]) 196 | 197 | scheduler.step() 198 | 199 | end_time = time.time() 200 | training_time = end_time - start_time 201 | print('###################################################################################################') 202 | print('###################################################################################################') 203 | print('###################################################################################################') 204 | print(f'Training time: {training_time:.2f} seconds') 205 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 206 | print('###################################################################################################') 207 | print('###################################################################################################') 208 | print('###################################################################################################') 209 | 210 | model.eval() 211 | y_pred_list = [] 212 | batch_size = 100 213 | with torch.no_grad(): 214 | for i in range(0, len(c_test), batch_size): 215 | batch_loads = c_test[i:i+batch_size].clone().detach().to(device) 216 | batch_pred = model(batch_loads, x) 217 | y_pred_list.append(batch_pred.cpu().numpy()) 218 | 219 | y_pred_deepokan = np.concatenate(y_pred_list, axis=0) 220 | 221 | np.save('y_pred_deepokan.npy', y_pred_deepokan) 222 | 223 | print("Predictions generated and saved successfully.") 224 | print(f"Shape of y_pred_deepokan: {y_pred_deepokan.shape}") 225 | 226 | with open('training_results_deepokan.csv', 'w', newline='') as csvfile: 227 | csv_writer = csv.writer(csvfile) 228 | csv_writer.writerow(['Epoch', 'Train Loss', 'Test Loss']) 229 | csv_writer.writerows(results) 230 | 231 | print("Training results saved in 'training_results_deepokan.csv'") -------------------------------------------------------------------------------- /Example_2_elasticity/Orthotropic_ux.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import os 4 | from KAN_RBF6 import RadialBasisFunctionNetwork as RBFKAN 5 | from MLPNet import MLPNetwork as MLPNET 6 | import torch 7 | from torch import nn 8 | import matplotlib.pyplot as plt 9 | from torch import autograd 10 | from tqdm import tqdm 11 | import torch.optim as optim 12 | from torch.optim.lr_scheduler import StepLR 13 | import scipy.io 14 | import numpy.random as npr 15 | import random 16 | import time 17 | import os 18 | 19 | ######################################################################################## 20 | HD = 5 21 | hid_trunk = 14 22 | repeat_hid_trunk = 1 23 | hid_branch = 14 24 | repeat_hid_branch = 1 25 | 26 | init_scale = 0.01 27 | grid_count = 5 28 | 29 | hidden_layers_mlp = [62] 30 | activation = nn.Tanh() 31 | 32 | learning_rate = 1e-3 33 | batch_size = 64 34 | epochs = 1 35 | gamma = 0.5 36 | step_size_scheduler = 1000 37 | opt = "adam" # "adam" "lbfgs" 38 | 39 | loss_vec_rbf = [] 40 | loss_vec_mlp = [] 41 | test_loss_vec_rbf = [] 42 | test_loss_vec_mlp = [] 43 | dtype = torch.float64 # or torch.float32 or torch.float64 44 | seeding_number = 1 45 | 46 | npr.seed(seeding_number); torch.manual_seed(seeding_number); random.seed(seeding_number) 47 | 48 | ######################################################################################## 49 | def param_and_concatenate_data(filepath): 50 | with h5py.File(filepath, 'r') as hdf: 51 | param_list = [] 52 | disp_list = [] 53 | 54 | param_group = hdf['params'] 55 | for key in sorted(param_group.keys()): 56 | param_list.append(np.array(param_group[key])) 57 | 58 | disp_group = hdf['disp'] 59 | for key in sorted(disp_group.keys()): 60 | disp_list.append(np.array(disp_group[key])) 61 | 62 | params = np.stack(param_list, axis=0) 63 | disps = np.stack(disp_list, axis=0) 64 | 65 | coordinates = np.array(hdf['coordinates']) 66 | 67 | return params, disps, coordinates 68 | 69 | # --------------------------------------------------------------------------------------- 70 | class L2NormLoss(nn.Module): 71 | def __init__(self): 72 | super(L2NormLoss, self).__init__() 73 | 74 | def forward(self, output, target): 75 | return torch.norm(output - target, p=2) 76 | 77 | # --------------------------------------------------------------------------------------- 78 | class RMSDLoss(nn.Module): 79 | def __init__(self): 80 | super(RMSDLoss, self).__init__() 81 | 82 | def forward(self, output, target): 83 | if output.shape != target.shape: 84 | raise ValueError("Output and target must have the same shape for RMSD calculation.") 85 | squared_diffs = (output - target).pow(2) 86 | mean_squared_diffs = torch.mean(squared_diffs) 87 | rmsd_value = torch.sqrt(mean_squared_diffs) 88 | return rmsd_value 89 | 90 | ######################################################################################## 91 | filepath = 'C:/Users/pp2624/Desktop/Research/8. DeepOKANs/Analyses for Arxiv 2 - Orthotropic Elasticity/data_100percent.hdf5' 92 | params, disp, coordinates = param_and_concatenate_data(filepath) 93 | print('max min params') 94 | print(params.max(), params.min()) 95 | print('max min coordinates') 96 | print(coordinates.max(), coordinates.min()) 97 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 98 | print(f'device used is {device}') 99 | 100 | params = torch.tensor(params).to(dtype).to(device) 101 | disp = torch.tensor(disp).to(dtype).to(device) 102 | coordinates = torch.tensor(coordinates).to(dtype).to(device) 103 | 104 | print(f'params shape = {params.shape}') 105 | print(f'disp shape = {disp.shape}') 106 | print(f'coordinates shape = {coordinates.shape}') 107 | 108 | def split_data(params, disp, test_ratio=0.2): 109 | """ Splits the data into training and testing sets. """ 110 | num_data = params.shape[0] 111 | indices = torch.randperm(num_data) 112 | split = int(num_data * (1 - test_ratio)) 113 | 114 | train_indices = indices[:split] 115 | test_indices = indices[split:] 116 | 117 | train_params = params[train_indices] 118 | train_true_ux = disp[train_indices] 119 | test_params = params[test_indices] 120 | test_true_ux = disp[test_indices] 121 | 122 | return train_params, train_true_ux[:,:,0], test_params, test_true_ux[:,:,0] 123 | 124 | train_params, train_true_ux, test_params, test_true_ux = split_data(params, disp) 125 | 126 | ######################################################################################## 127 | def define_kan_width(input_dim, hid, output_dim, repeat_hid): 128 | width = [input_dim] 129 | width += [hid] * repeat_hid 130 | width.append(output_dim) 131 | return width 132 | 133 | width_trunk = define_kan_width(coordinates.shape[1], hid_trunk, HD, repeat_hid_trunk) 134 | width_branch = define_kan_width(train_params.shape[1], hid_branch, HD, repeat_hid_branch) 135 | 136 | print(f'width_trunk = {width_trunk}') 137 | print(f'width_branch = {width_branch}') 138 | 139 | train_params = train_params.clone().detach().requires_grad_(True).to(dtype) 140 | train_true_ux = train_true_ux.clone().detach().requires_grad_(True).to(dtype) 141 | test_params = test_params.clone().detach().requires_grad_(True).to(dtype) 142 | test_true_ux = test_true_ux.clone().detach().requires_grad_(True).to(dtype) 143 | 144 | 145 | ######################################################################################## 146 | class DeepONet_rbf(nn.Module): 147 | def __init__(self): 148 | super(DeepONet_rbf, self).__init__() 149 | self.trunk_net = RBFKAN(hidden_layers=width_trunk, dtype=dtype, apply_base_update=False, min_grid=0, max_grid=1, grid_count=grid_count, grid_opt=False, init_scale=init_scale).to(device) 150 | self.branch_net = RBFKAN(hidden_layers=width_branch, dtype=dtype, apply_base_update=False, min_grid=-1, max_grid=20, grid_count=grid_count, grid_opt=False, init_scale=init_scale).to(device) 151 | 152 | def forward(self, x_branch, x_trunk): 153 | y_branch = self.branch_net(x_branch) 154 | y_trunk = self.trunk_net(x_trunk) 155 | if y_branch.shape[-1] != y_trunk.shape[-1]: 156 | raise AssertionError("Output sizes of branch net and trunk net do not match.") 157 | Y = torch.einsum('bi,ni->bn', y_branch, y_trunk) 158 | return Y 159 | 160 | model_rbf = DeepONet_rbf() 161 | 162 | ######################################################################################## 163 | # Hyperparameters 164 | criterion1 = nn.MSELoss() 165 | criterion2 = L2NormLoss() 166 | criterion3 = RMSDLoss() 167 | 168 | if opt == "adam": 169 | optimizer_rbf = optim.Adam(model_rbf.parameters(), lr=learning_rate) 170 | scheduler_rbf = StepLR(optimizer_rbf, step_size=step_size_scheduler, gamma=gamma) 171 | else: 172 | optimizer_rbf = torch.optim.LBFGS(model_rbf.parameters(), lr=1, max_iter=20, history_size=10, line_search_fn="strong_wolfe") 173 | 174 | def test_model(model, criterion, test_params, test_disp, coordinates): 175 | """Evaluates the model on the test dataset.""" 176 | model.eval() 177 | with torch.no_grad(): 178 | outputs = model(test_params, coordinates) 179 | loss = criterion(outputs, test_disp) 180 | return loss.item() 181 | 182 | # -------------------------------------------------------------------------------------- 183 | start_train_time_rbf = time.time() 184 | for epoch in range(epochs): 185 | total_loss = 0 186 | progress_bar = tqdm(range(0, len(train_params), batch_size), desc=f'Epoch {epoch + 1}/{epochs}') 187 | for i in progress_bar: 188 | batch_params = train_params[i:i + batch_size] 189 | batch_disp = train_true_ux[i:i + batch_size] 190 | def closure_rbf(): 191 | optimizer_rbf.zero_grad() 192 | outputs = model_rbf(batch_params, coordinates) 193 | loss = criterion3(outputs, batch_disp) 194 | loss.backward() 195 | return loss 196 | loss = optimizer_rbf.step(closure_rbf) 197 | total_loss += loss.item() 198 | progress_bar.set_postfix({'Batch Loss': loss.item()}) 199 | 200 | avg_loss_rbf = total_loss / (len(train_params) // batch_size) 201 | loss_vec_rbf.append(avg_loss_rbf) 202 | tqdm.write(f'End of Epoch {epoch + 1}, Average Loss: {avg_loss_rbf}') 203 | 204 | test_loss_rbf = test_model(model_rbf, criterion3, test_params, test_true_ux, coordinates) 205 | test_loss_vec_rbf.append(test_loss_rbf) 206 | tqdm.write(f'Test Loss: {test_loss_rbf}') 207 | 208 | # if (epoch + 1) % 1005 == 0: 209 | # print("Heyyyy") 210 | # with torch.no_grad(): 211 | # test_pred_ux_rbf = model_rbf(test_params, coordinates).cpu().numpy() 212 | # scipy.io.savemat('Results_Orthotropic_RBF_' + str(epoch + 1) + '.mat', {'test_pred_ux_rbf': test_pred_ux_rbf, 'loss_vec_rbf':loss_vec_rbf}) 213 | 214 | if opt == "adam": 215 | scheduler_rbf.step() 216 | 217 | end_training_time_rbf = time.time() - start_train_time_rbf 218 | 219 | ######################################################################################## 220 | ######################################################################################## 221 | ######################################################################################## 222 | class DeepONet_mlp(nn.Module): 223 | def __init__(self): 224 | super(DeepONet_mlp, self).__init__() 225 | self.trunk_net = MLPNET(input_dim=2, hidden_layers=hidden_layers_mlp, activation=activation, output_dim=HD, dtype=dtype).to(device) 226 | self.branch_net = MLPNET(input_dim=6, hidden_layers=hidden_layers_mlp, activation=activation, output_dim=HD, dtype=dtype).to(device) 227 | 228 | def forward(self, x_branch, x_trunk): 229 | y_branch = self.branch_net(x_branch) 230 | y_trunk = self.trunk_net(x_trunk) 231 | if y_branch.shape[-1] != y_trunk.shape[-1]: 232 | raise AssertionError("Output sizes of branch net and trunk net do not match.") 233 | Y = torch.einsum('bi,ni->bn', y_branch, y_trunk) 234 | return Y 235 | 236 | model_mlp = DeepONet_mlp() 237 | 238 | if opt == "adam": 239 | optimizer_mlp = optim.Adam(model_mlp.parameters(), lr=learning_rate) 240 | scheduler_mlp = StepLR(optimizer_mlp, step_size=step_size_scheduler, gamma=gamma) 241 | else: 242 | optimizer_mlp = torch.optim.LBFGS(model_mlp.parameters(), lr=1, max_iter=20, history_size=10, line_search_fn="strong_wolfe") 243 | 244 | # -------------------------------------------------------------------------------------- 245 | start_train_time_mlp = time.time() 246 | for epoch in range(epochs): 247 | total_loss = 0 248 | progress_bar = tqdm(range(0, len(train_params), batch_size), desc=f'Epoch {epoch + 1}/{epochs}') 249 | 250 | for i in progress_bar: 251 | batch_params = train_params[i:i + batch_size] 252 | batch_disp = train_true_ux[i:i + batch_size] 253 | def closure_mlp(): 254 | optimizer_mlp.zero_grad() 255 | outputs = model_mlp(batch_params, coordinates) 256 | loss = criterion3(outputs, batch_disp) 257 | loss.backward() 258 | return loss 259 | loss = optimizer_mlp.step(closure_mlp) 260 | total_loss += loss.item() 261 | progress_bar.set_postfix({'Batch Loss': loss.item()}) 262 | 263 | avg_loss_mlp = total_loss / (len(train_params) // batch_size) 264 | loss_vec_mlp.append(avg_loss_mlp) 265 | tqdm.write(f'End of Epoch {epoch + 1}, Average Loss: {avg_loss_mlp}') 266 | 267 | test_loss_mlp = test_model(model_mlp, criterion3, test_params, test_true_ux, coordinates) 268 | test_loss_vec_mlp.append(test_loss_mlp) 269 | tqdm.write(f'Test Loss: {test_loss_mlp}') 270 | 271 | if opt == "adam": 272 | scheduler_mlp.step() 273 | 274 | # if (epoch + 1) % 1005 == 0: 275 | # print("Heyyyy") 276 | # with torch.no_grad(): 277 | # test_pred_ux_mlp = model_mlp(test_params, coordinates).cpu().numpy() 278 | # scipy.io.savemat('Results_Orthotropic_MLP_' + str(epoch + 1) + '.mat', {'test_pred_ux_mlp': test_pred_ux_mlp, 'loss_vec_mlp':loss_vec_mlp}) 279 | 280 | end_training_time_mlp = time.time() - start_train_time_mlp 281 | 282 | ######################################################################################## 283 | with torch.no_grad(): 284 | coordinates_cpu = coordinates.cpu().numpy() 285 | test_params_cpu = test_params.cpu().numpy() 286 | test_pred_ux_rbf = model_rbf(test_params, coordinates).cpu().numpy() 287 | test_pred_ux_mlp = model_mlp(test_params, coordinates).cpu().numpy() 288 | test_true_ux = test_true_ux.cpu().numpy() 289 | 290 | def count_parameters(model): 291 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 292 | 293 | print(f'Number of learnable parameters in RBF Network: {count_parameters(model_rbf)}') 294 | print(f'Number of learnable parameters in MLP Network: {count_parameters(model_mlp)}') 295 | 296 | ######################################################################################## 297 | scipy.io.savemat('Results_Orthotropic_ModelX.mat', 298 | {'coordinates':coordinates_cpu, 299 | 'test_params':test_params_cpu, 300 | 'pred_ux_rbf':test_pred_ux_rbf, 301 | 'pred_ux_mlp':test_pred_ux_mlp, 302 | 'test_true_ux':test_true_ux, 303 | 'loss_vec_rbf':loss_vec_rbf, 304 | 'loss_vec_mlp':loss_vec_mlp, 305 | 'test_loss_vec_rbf':test_loss_vec_rbf, 306 | 'test_loss_vec_mlp':test_loss_vec_mlp, 307 | 'end_training_time_rbf':end_training_time_rbf, 308 | 'end_training_time_mlp':end_training_time_mlp, 309 | 'parameters_rbf':count_parameters(model_rbf), 310 | 'parameters_mlp':count_parameters(model_mlp)}) 311 | 312 | 313 | 314 | 315 | 316 | 317 | 318 | 319 | 320 | 321 | 322 | -------------------------------------------------------------------------------- /Example_3_transient/final_mlp_transient_poisson.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import os 4 | from MLPNet import MLPNetwork as MLPNET 5 | import torch 6 | from torch import nn 7 | import matplotlib.pyplot as plt 8 | from torch import autograd 9 | from tqdm import tqdm 10 | import jax.numpy as jnp 11 | import torch.optim as optim 12 | from torch.optim.lr_scheduler import StepLR 13 | from torch.optim.lr_scheduler import ReduceLROnPlateau 14 | import time 15 | import pandas as pd 16 | import random 17 | from statistics import mean, median 18 | 19 | time.sleep(5) 20 | dtype = torch.float64 # or torch.float32 or torch.float64 21 | def load_and_concatenate_data(filepath): 22 | with h5py.File(filepath, 'r') as hdf: 23 | # Gather coordinates directly into a numpy array 24 | coordinates = np.array(hdf['coordinates']) 25 | 26 | # Retrieve information about the number of datasets and their shapes 27 | total_datasets = len(hdf['loads'].keys()) 28 | sample_load = hdf['loads']['load_0'] 29 | sample_solution = hdf['solutions']['solution_0'] 30 | 31 | # Determine the shape and datatype for the memmap arrays 32 | loads_shape = (total_datasets, *sample_load.shape) 33 | solutions_shape = (total_datasets, *sample_solution.shape) 34 | 35 | # If memmap files exist, delete them first to avoid shape and dtype conflicts 36 | if os.path.exists('loadsm.dat'): 37 | os.remove('loadsm.dat') 38 | if os.path.exists('solutionsm.dat'): 39 | os.remove('solutionsm.dat') 40 | 41 | # Create memmap arrays for loads and solutions 42 | loads = np.memmap('loadsm.dat', dtype=sample_load.dtype, mode='w+', shape=loads_shape) 43 | solutions = np.memmap('solutionsm.dat', dtype=sample_solution.dtype, mode='w+', shape=solutions_shape) 44 | 45 | # Populate memmap arrays directly from HDF5 without loading entire arrays into memory 46 | for idx in range(total_datasets): 47 | loads[idx] = hdf['loads'][f'load_{idx}'] 48 | solutions[idx] = hdf['solutions'][f'solution_{idx}'] 49 | 50 | return loads, solutions, coordinates 51 | 52 | # Example usage: 53 | filepath = '/scratch/da3205/TransientPoisson_KAN/data_90percent.hdf5' 54 | loads, solutions, coordinates = load_and_concatenate_data(filepath) 55 | solutions = solutions.transpose(0,2,1) 56 | 57 | # Convert to PyTorch tensors 58 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 59 | loads = torch.tensor(loads, dtype=dtype, device=device)[:, :, 1] 60 | solutions = torch.tensor(solutions, dtype=dtype, device=device) 61 | coordinates = torch.tensor(coordinates, dtype=dtype, device=device) 62 | 63 | print(f'loads shape = {loads.shape}') 64 | print(f'solutions shape = {solutions.shape}') 65 | print(f'coordinates shape = {coordinates.shape}') 66 | 67 | def split_data(loads, solutions, test_ratio=0.2): 68 | """ Splits the data into training and testing sets. """ 69 | num_data = loads.shape[0] 70 | indices = torch.randperm(num_data) 71 | split = int(num_data * (1 - test_ratio)) 72 | 73 | train_indices = indices[:split] 74 | test_indices = indices[split:] 75 | 76 | train_loads = loads[train_indices] 77 | train_solutions = solutions[train_indices] 78 | test_loads = loads[test_indices] 79 | test_solutions = solutions[test_indices] 80 | 81 | return train_loads, train_solutions, test_loads, test_solutions 82 | 83 | train_loads, train_solutions, test_loads, test_solutions = split_data(loads, solutions, test_ratio=0.2) 84 | print(f'train_loads shape = {train_loads.shape}, train_solutions shape = {train_solutions.shape}') 85 | print(f'test_loads shape = {test_loads.shape}, test_solutions shape = {test_solutions.shape}') 86 | 87 | input_dim_trunk = coordinates.shape[1] 88 | input_dim_branch = loads.shape[1] 89 | HD = 4 90 | hidden_layers_mlp = [22, 22, 22] 91 | print(f'hidden_layers_mlp = {hidden_layers_mlp}') 92 | activation = nn.Tanh() 93 | 94 | # Hyperparameters 95 | learning_rate = 1e-3 96 | batch_size = 64 97 | epochs = 10_000 98 | gamma = 0.5 99 | step_size = 1_000 100 | seed = 42 101 | 102 | train_loads = train_loads.clone().detach().requires_grad_(True).to(dtype) 103 | train_solutions = train_solutions.clone().detach().requires_grad_(True).to(dtype) 104 | test_loads = test_loads.clone().detach().requires_grad_(True).to(dtype) 105 | test_solutions = test_solutions.clone().detach().requires_grad_(True).to(dtype) 106 | 107 | class DeepONet(nn.Module): 108 | def __init__(self): 109 | super(DeepONet, self).__init__() 110 | self.trunk_net = MLPNET(input_dim=input_dim_trunk, hidden_layers=hidden_layers_mlp, activation=activation, output_dim=HD, dtype=dtype).to(device) 111 | self.branch_net = MLPNET(input_dim=input_dim_branch, hidden_layers=hidden_layers_mlp, activation=activation, output_dim=input_dim_branch*HD, dtype=dtype).to(device) 112 | 113 | def forward(self, x_branch, x_trunk): 114 | y_branch = self.branch_net(x_branch) 115 | y_branch = y_branch.view(-1, train_loads.shape[1], HD) 116 | y_trunk = self.trunk_net(x_trunk) 117 | 118 | if y_branch.shape[-1] != y_trunk.shape[-1]: 119 | raise AssertionError("Output sizes of branch net and trunk net do not match.") 120 | 121 | # Einsum operation to combine outputs into the desired shape 122 | Y = torch.einsum('bik,nk->bni', y_branch, y_trunk) 123 | 124 | return Y 125 | 126 | model = DeepONet() 127 | 128 | class L2NormLoss(nn.Module): 129 | def __init__(self): 130 | super(L2NormLoss, self).__init__() 131 | 132 | def forward(self, output, target): 133 | norm = torch.norm(output - target, p=2) 134 | num_elements = target.shape[0] # This gives the total number of elements in the output tensor 135 | return norm / num_elements 136 | 137 | class RMSDLoss(nn.Module): 138 | def __init__(self): 139 | super(RMSDLoss, self).__init__() 140 | 141 | def forward(self, output, target): 142 | if output.shape != target.shape: 143 | raise ValueError("Output and target must have the same shape for RMSD calculation.") 144 | squared_diffs = (output - target).pow(2) 145 | mean_squared_diffs = torch.mean(squared_diffs) 146 | rmsd_value = torch.sqrt(mean_squared_diffs) 147 | return rmsd_value 148 | 149 | criterion1 = nn.MSELoss() 150 | criterion2 = L2NormLoss() 151 | criterion3 = RMSDLoss() 152 | optimizer = optim.Adam(model.parameters(), lr=learning_rate) 153 | scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma) 154 | 155 | def test_model(model, criterion, test_loads, test_solutions, coordinates, batch_size): 156 | """Evaluates the model on the test dataset in batches.""" 157 | model.eval() # Set the model to evaluation mode 158 | total_loss = 0 159 | num_batches = len(test_loads) // batch_size + 1 160 | with torch.no_grad(): # No need to track gradients for testing 161 | for i in tqdm(range(0, len(test_loads), batch_size), desc='Testing'): 162 | batch_loads = test_loads[i:i + batch_size] 163 | batch_solutions = test_solutions[i:i + batch_size] 164 | outputs = model(batch_loads, coordinates) 165 | loss = criterion(outputs, batch_solutions) 166 | total_loss += loss.item() 167 | avg_loss = total_loss / num_batches 168 | return avg_loss 169 | 170 | def set_seed(seed): 171 | random.seed(seed) 172 | np.random.seed(seed) 173 | torch.manual_seed(seed) 174 | if torch.cuda.is_available(): 175 | torch.cuda.manual_seed(seed) 176 | torch.cuda.manual_seed_all(seed) 177 | set_seed(seed) 178 | 179 | def count_learnable_parameters(model): 180 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 181 | num_learnable_parameters = count_learnable_parameters(model) 182 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 183 | 184 | # Assuming your model, optimizer, loss function, scheduler, and data are defined and set up 185 | train_losses = [] 186 | test_losses = [] 187 | start_time = time.time() 188 | mean_test_loss = 0 189 | for epoch in range(epochs): 190 | model.train() # Set the model to training mode 191 | total_loss = 0 192 | progress_bar = tqdm(range(0, len(train_loads), batch_size), desc=f'Epoch {epoch + 1}/{epochs}') 193 | for i in progress_bar: 194 | batch_loads = train_loads[i:i + batch_size] 195 | batch_solutions = train_solutions[i:i + batch_size] 196 | 197 | # Define the closure for the L-BFGS optimizer 198 | def closure(): 199 | optimizer.zero_grad() 200 | outputs = model(batch_loads, coordinates) 201 | loss = criterion3(outputs, batch_solutions) 202 | loss.backward() 203 | return loss 204 | 205 | # Perform an optimization step 206 | loss = optimizer.step(closure) 207 | total_loss += loss.item() 208 | 209 | # Update the progress bar with the current batch loss 210 | progress_bar.set_postfix({'Batch Loss': loss.item()}) 211 | 212 | # Calculate the average loss for the epoch 213 | avg_loss = total_loss / (len(train_loads) // batch_size + 1) 214 | train_losses.append(avg_loss) 215 | tqdm.write(f'End of Epoch {epoch + 1}, Average Loss: {avg_loss}') 216 | 217 | # Evaluate the model on the testing dataset 218 | test_loss = test_model(model, criterion3, test_loads, test_solutions, coordinates, batch_size) 219 | test_losses.append(test_loss) 220 | tqdm.write(f'Test Loss: {test_loss}') 221 | scheduler.step() 222 | 223 | end_time = time.time() 224 | training_time = end_time - start_time 225 | print('###################################################################################################') 226 | print('###################################################################################################') 227 | print('###################################################################################################') 228 | print(f'Training time: {training_time:.2f} seconds') 229 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 230 | print('###################################################################################################') 231 | print('###################################################################################################') 232 | print('###################################################################################################') 233 | # Create a DataFrame from the losses 234 | loss_df = pd.DataFrame({ 235 | 'Epoch': range(1, epochs + 1), 236 | 'Train Loss': train_losses, 237 | 'Test Loss': test_losses 238 | }) 239 | # Save the DataFrame to a CSV file 240 | loss_df.to_csv('DeepONet_training_testing_losses_seed42.csv', index=False) 241 | print('Training and testing losses have been saved to DeepONet_training_testing_losses.csv') 242 | 243 | # Define a function to get predictions in batches and save directly to file 244 | def get_predictions_and_save(model, test_loads, coordinates, batch_size, filename_pattern): 245 | model.eval() # Set the model to evaluation mode 246 | 247 | # Number of splits 248 | num_splits = 4 249 | split_size = len(test_loads) // num_splits 250 | 251 | # Prepare HDF5 files for writing 252 | for i in range(num_splits): 253 | filename = filename_pattern.format(i + 1) 254 | with h5py.File(filename, 'w') as h5f: 255 | h5f.create_dataset('coordinates', data=coordinates.cpu().numpy(), compression='gzip') 256 | h5f.create_dataset('test_loads', shape=(split_size, *test_loads.shape[1:]), dtype=np.float32, compression='gzip') 257 | h5f.create_dataset('test_pred', shape=(split_size, *model(test_loads[:batch_size], coordinates).shape[1:]), dtype=np.float32, compression='gzip') 258 | h5f.create_dataset('test_solutions', shape=(split_size, *test_solutions.shape[1:]), dtype=np.float32, compression='gzip') 259 | 260 | # Process and save predictions in batches 261 | batch_counter = 0 262 | with torch.no_grad(): # No need to track gradients for prediction 263 | for i in tqdm(range(0, len(test_loads), batch_size), desc='Predicting'): 264 | batch_loads = test_loads[i:i + batch_size] 265 | batch_solutions = test_solutions[i:i + batch_size] 266 | batch_predictions = model(batch_loads, coordinates).cpu().numpy() 267 | 268 | # Determine the corresponding HDF5 file and dataset indices 269 | file_idx = batch_counter // split_size 270 | dataset_idx = batch_counter % split_size 271 | batch_counter += batch_loads.shape[0] 272 | 273 | filename = filename_pattern.format(file_idx + 1) 274 | with h5py.File(filename, 'a') as h5f: 275 | # Determine the actual size to write to avoid broadcasting issues 276 | actual_size = min(split_size - dataset_idx, batch_loads.shape[0]) 277 | h5f['test_loads'][dataset_idx:dataset_idx + actual_size] = batch_loads.cpu().numpy()[:actual_size] 278 | h5f['test_pred'][dataset_idx:dataset_idx + actual_size] = batch_predictions[:actual_size] 279 | h5f['test_solutions'][dataset_idx:dataset_idx + actual_size] = batch_solutions.cpu().numpy()[:actual_size] 280 | 281 | print("Predictions saved successfully.") 282 | 283 | # Assume training is complete 284 | time.sleep(5) 285 | print('Training is complete') 286 | 287 | # Define the filename pattern for saving splits 288 | filename_pattern = 'transient_Poisson_results_mlp_file_seed42_part_{}.h5' 289 | 290 | # Transfer test data to CPU and obtain predictions in batches, saving directly to files 291 | get_predictions_and_save(model, test_loads, coordinates, batch_size, filename_pattern) 292 | 293 | # Print confirmation 294 | time.sleep(5) 295 | print("Numpy arrays are ready and saved to files.") 296 | 297 | # Load and print shapes for verification 298 | for i in range(1, 5): 299 | filename = filename_pattern.format(i) 300 | with h5py.File(filename, 'r') as h5f: 301 | coordinates_cpu = np.array(h5f['coordinates']) 302 | test_loads_cpu = np.array(h5f['test_loads']) 303 | test_pred_cpu = np.array(h5f['test_pred']) 304 | test_solutions_cpu = np.array(h5f['test_solutions']) 305 | print(f'File {filename}:') 306 | print(f'coordinates_cpu shape = {coordinates_cpu.shape}') 307 | print(f'test_loads_cpu shape = {test_loads_cpu.shape}') 308 | print(f'test_pred_cpu shape = {test_pred_cpu.shape}') 309 | print(f'test_solutions_cpu shape = {test_solutions_cpu.shape}') -------------------------------------------------------------------------------- /Example_3_transient/final_rbfkan_transient_poisson.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import os 4 | from KAN_RBF6 import RadialBasisFunctionNetwork as RBFKAN 5 | import torch 6 | from torch import nn 7 | import matplotlib.pyplot as plt 8 | from torch import autograd 9 | from tqdm import tqdm 10 | import jax.numpy as jnp 11 | import torch.optim as optim 12 | from torch.optim.lr_scheduler import StepLR 13 | from torch.optim.lr_scheduler import ReduceLROnPlateau 14 | import time 15 | import pandas as pd 16 | import random 17 | from statistics import mean, median 18 | 19 | time.sleep(5) 20 | dtype = torch.float64 # or torch.float32 or torch.float64 21 | def load_and_concatenate_data(filepath): 22 | with h5py.File(filepath, 'r') as hdf: 23 | # Gather coordinates directly into a numpy array 24 | coordinates = np.array(hdf['coordinates']) 25 | 26 | # Retrieve information about the number of datasets and their shapes 27 | total_datasets = len(hdf['loads'].keys()) 28 | sample_load = hdf['loads']['load_0'] 29 | sample_solution = hdf['solutions']['solution_0'] 30 | 31 | # Determine the shape and datatype for the memmap arrays 32 | loads_shape = (total_datasets, *sample_load.shape) 33 | solutions_shape = (total_datasets, *sample_solution.shape) 34 | 35 | # If memmap files exist, delete them first to avoid shape and dtype conflicts 36 | if os.path.exists('loads20.dat'): 37 | os.remove('loads20.dat') 38 | if os.path.exists('solutions20.dat'): 39 | os.remove('solutions20.dat') 40 | 41 | # Create memmap arrays for loads and solutions 42 | loads = np.memmap('loads20.dat', dtype=sample_load.dtype, mode='w+', shape=loads_shape) 43 | solutions = np.memmap('solutions20.dat', dtype=sample_solution.dtype, mode='w+', shape=solutions_shape) 44 | 45 | # Populate memmap arrays directly from HDF5 without loading entire arrays into memory 46 | for idx in range(total_datasets): 47 | loads[idx] = hdf['loads'][f'load_{idx}'] 48 | solutions[idx] = hdf['solutions'][f'solution_{idx}'] 49 | 50 | return loads, solutions, coordinates 51 | 52 | # Example usage: 53 | filepath = '/scratch/da3205/TransientPoisson_KAN/data_90percent.hdf5' 54 | loads, solutions, coordinates = load_and_concatenate_data(filepath) 55 | solutions = solutions.transpose(0,2,1) 56 | 57 | # Convert to PyTorch tensors 58 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 59 | loads = torch.tensor(loads, device=device)[:,:,1] 60 | solutions = torch.tensor(solutions, device=device) 61 | coordinates = torch.tensor(coordinates, device=device) 62 | 63 | print(f'loads shape = {loads.shape}') 64 | print(f'solutions shape = {solutions.shape}') 65 | print(f'coordinates shape = {coordinates.shape}') 66 | 67 | def split_data(loads, solutions, test_ratio=0.2): 68 | """ Splits the data into training and testing sets. """ 69 | num_data = loads.shape[0] 70 | indices = torch.randperm(num_data) 71 | split = int(num_data * (1 - test_ratio)) 72 | 73 | train_indices = indices[:split] 74 | test_indices = indices[split:] 75 | 76 | train_loads = loads[train_indices] 77 | train_solutions = solutions[train_indices] 78 | test_loads = loads[test_indices] 79 | test_solutions = solutions[test_indices] 80 | 81 | return train_loads, train_solutions, test_loads, test_solutions 82 | 83 | print(f'loads max = {loads.max()}, loads min = {loads.min()}') 84 | print(f'coordinates max = {coordinates.max()}, coordinates min = {coordinates.min()}') 85 | 86 | train_loads, train_solutions, test_loads, test_solutions = split_data(loads, solutions, test_ratio=0.2) 87 | print(f'train_loads shape = {train_loads.shape}, train_solutions shape = {train_solutions.shape}') 88 | print(f'test_loads shape = {test_loads.shape}, test_solutions shape = {test_solutions.shape}') 89 | 90 | def define_kan_width(input_dim, W, repeat_hid, output_dim): 91 | width = [input_dim] + [W] * repeat_hid + [output_dim] 92 | return width 93 | 94 | input_dim_trunk = coordinates.shape[1] 95 | input_dim_branch = loads.shape[1] 96 | HD = 4 97 | hid_trunk = 5 98 | num_layer_trunk = 1 99 | hid_branch = 5 100 | num_layer_branch = 1 101 | trunk_min_grid = coordinates.min() 102 | trunk_max_grid = coordinates.max() 103 | branch_min_grid = loads.min() 104 | branch_max_grid = loads.max() 105 | 106 | width_trunk = define_kan_width(input_dim_trunk, hid_trunk, num_layer_trunk, HD) 107 | width_branch = define_kan_width(input_dim_branch, hid_branch, num_layer_branch, input_dim_branch*HD) 108 | 109 | print(f'hid dim = {hid_trunk}, num_layer = {num_layer_trunk}') 110 | print(f'width_trunk = {width_trunk}, width_branch = {width_branch}') 111 | grid_opt = False 112 | grid_count = 5 113 | init_scale = 0.01 114 | print(f'init_scale = {init_scale}') 115 | 116 | # Hyperparameters 117 | learning_rate = 1e-3 118 | batch_size = 64 119 | epochs = 10_000 120 | gamma = 0.5 121 | step_size = 1_000 122 | seed = 42 123 | 124 | train_loads = train_loads.clone().detach().requires_grad_(True).to(dtype) 125 | train_solutions = train_solutions.clone().detach().requires_grad_(True).to(dtype) 126 | test_loads = test_loads.clone().detach().requires_grad_(True).to(dtype) 127 | test_solutions = test_solutions.clone().detach().requires_grad_(True).to(dtype) 128 | 129 | class DeepONet(nn.Module): 130 | def __init__(self): 131 | super(DeepONet, self).__init__() 132 | self.trunk_net = RBFKAN(hidden_layers=width_trunk, dtype=dtype, apply_base_update=False, init_scale=init_scale, 133 | min_grid=trunk_min_grid, max_grid=trunk_max_grid, grid_count=grid_count, grid_opt=grid_opt).to(device) 134 | self.branch_net = RBFKAN(hidden_layers=width_branch, dtype=dtype, apply_base_update=False, init_scale=init_scale, 135 | min_grid=branch_min_grid, max_grid=branch_max_grid, grid_count=grid_count, grid_opt=grid_opt).to(device) 136 | 137 | def forward(self, x_branch, x_trunk): 138 | y_branch = self.branch_net(x_branch) 139 | y_branch = y_branch.view(-1, train_loads.shape[1], HD) 140 | y_trunk = self.trunk_net(x_trunk) 141 | 142 | if y_branch.shape[-1] != y_trunk.shape[-1]: 143 | raise AssertionError("Output sizes of branch net and trunk net do not match.") 144 | 145 | # Einsum operation to combine outputs into the desired shape 146 | Y = torch.einsum('bik,nk->bni', y_branch, y_trunk) 147 | 148 | return Y 149 | 150 | model = DeepONet() 151 | 152 | class L2NormLoss(nn.Module): 153 | def __init__(self): 154 | super(L2NormLoss, self).__init__() 155 | 156 | def forward(self, output, target): 157 | norm = torch.norm(output - target, p=2) 158 | num_elements = target.shape[0] # This gives the total number of elements in the output tensor 159 | return norm / num_elements 160 | 161 | class RMSDLoss(nn.Module): 162 | def __init__(self): 163 | super(RMSDLoss, self).__init__() 164 | 165 | def forward(self, output, target): 166 | if output.shape != target.shape: 167 | raise ValueError("Output and target must have the same shape for RMSD calculation.") 168 | squared_diffs = (output - target).pow(2) 169 | mean_squared_diffs = torch.mean(squared_diffs) 170 | rmsd_value = torch.sqrt(mean_squared_diffs) 171 | return rmsd_value 172 | 173 | criterion1 = nn.MSELoss() 174 | criterion2 = L2NormLoss() 175 | criterion3 = RMSDLoss() 176 | optimizer = optim.Adam(model.parameters(), lr=learning_rate) 177 | scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma) 178 | 179 | def test_model(model, criterion, test_loads, test_solutions, coordinates, batch_size): 180 | """Evaluates the model on the test dataset in batches.""" 181 | model.eval() # Set the model to evaluation mode 182 | total_loss = 0 183 | num_batches = len(test_loads) // batch_size + 1 184 | with torch.no_grad(): # No need to track gradients for testing 185 | for i in tqdm(range(0, len(test_loads), batch_size), desc='Testing'): 186 | batch_loads = test_loads[i:i + batch_size] 187 | batch_solutions = test_solutions[i:i + batch_size] 188 | outputs = model(batch_loads, coordinates) 189 | loss = criterion(outputs, batch_solutions) 190 | total_loss += loss.item() 191 | avg_loss = total_loss / num_batches 192 | return avg_loss 193 | 194 | def set_seed(seed): 195 | random.seed(seed) 196 | np.random.seed(seed) 197 | torch.manual_seed(seed) 198 | if torch.cuda.is_available(): 199 | torch.cuda.manual_seed(seed) 200 | torch.cuda.manual_seed_all(seed) 201 | set_seed(seed) 202 | 203 | def count_learnable_parameters(model): 204 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 205 | num_learnable_parameters = count_learnable_parameters(model) 206 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 207 | 208 | # Assuming your model, optimizer, loss function, scheduler, and data are defined and set up 209 | train_losses = [] 210 | test_losses = [] 211 | mean_test_loss = 0 212 | start_time = time.time() 213 | for epoch in range(epochs): 214 | model.train() # Set the model to training mode 215 | total_loss = 0 216 | progress_bar = tqdm(range(0, len(train_loads), batch_size), desc=f'Epoch {epoch + 1}/{epochs}') 217 | for i in progress_bar: 218 | batch_loads = train_loads[i:i + batch_size] 219 | batch_solutions = train_solutions[i:i + batch_size] 220 | 221 | # Define the closure for the L-BFGS optimizer 222 | def closure(): 223 | optimizer.zero_grad() 224 | outputs = model(batch_loads, coordinates) 225 | loss = criterion3(outputs, batch_solutions) 226 | loss.backward() 227 | return loss 228 | 229 | # Perform an optimization step 230 | loss = optimizer.step(closure) 231 | total_loss += loss.item() 232 | 233 | # Update the progress bar with the current batch loss 234 | progress_bar.set_postfix({'Batch Loss': loss.item()}) 235 | 236 | # Calculate the average loss for the epoch 237 | avg_loss = total_loss / (len(train_loads) // batch_size + 1) 238 | train_losses.append(avg_loss) 239 | tqdm.write(f'End of Epoch {epoch + 1}, Average Loss: {avg_loss}') 240 | 241 | # Evaluate the model on the testing dataset 242 | test_loss = test_model(model, criterion3, test_loads, test_solutions, coordinates, batch_size) 243 | test_losses.append(test_loss) 244 | tqdm.write(f'Test Loss: {test_loss}') 245 | scheduler.step() 246 | 247 | end_time = time.time() 248 | training_time = end_time - start_time 249 | print('###################################################################################################') 250 | print('###################################################################################################') 251 | print('###################################################################################################') 252 | print(f'Training time: {training_time:.2f} seconds') 253 | print(f'The number of learnable parameters in the model: {num_learnable_parameters}') 254 | print('###################################################################################################') 255 | print('###################################################################################################') 256 | print('###################################################################################################') 257 | # Create a DataFrame from the losses 258 | loss_df = pd.DataFrame({ 259 | 'Epoch': range(1, epochs + 1), 260 | 'Train Loss': train_losses, 261 | 'Test Loss': test_losses 262 | }) 263 | # Save the DataFrame to a CSV file 264 | loss_df.to_csv('DeepOKAN_training_testing_losses_seed42.csv', index=False) 265 | print('Training and testing losses have been saved to DeepOKAN_training_testing_losses.csv') 266 | 267 | # Define a function to get predictions in batches and save directly to file 268 | def get_predictions_and_save(model, test_loads, coordinates, batch_size, filename_pattern): 269 | model.eval() # Set the model to evaluation mode 270 | 271 | # Number of splits 272 | num_splits = 4 273 | split_size = len(test_loads) // num_splits 274 | 275 | # Prepare HDF5 files for writing 276 | for i in range(num_splits): 277 | filename = filename_pattern.format(i + 1) 278 | with h5py.File(filename, 'w') as h5f: 279 | h5f.create_dataset('coordinates', data=coordinates.cpu().numpy(), compression='gzip') 280 | h5f.create_dataset('test_loads', shape=(split_size, *test_loads.shape[1:]), dtype=np.float32, compression='gzip') 281 | h5f.create_dataset('test_pred', shape=(split_size, *model(test_loads[:batch_size], coordinates).shape[1:]), dtype=np.float32, compression='gzip') 282 | h5f.create_dataset('test_solutions', shape=(split_size, *test_solutions.shape[1:]), dtype=np.float32, compression='gzip') 283 | 284 | # Process and save predictions in batches 285 | batch_counter = 0 286 | with torch.no_grad(): # No need to track gradients for prediction 287 | for i in tqdm(range(0, len(test_loads), batch_size), desc='Predicting'): 288 | batch_loads = test_loads[i:i + batch_size] 289 | batch_solutions = test_solutions[i:i + batch_size] 290 | batch_predictions = model(batch_loads, coordinates).cpu().numpy() 291 | 292 | # Determine the corresponding HDF5 file and dataset indices 293 | file_idx = batch_counter // split_size 294 | dataset_idx = batch_counter % split_size 295 | batch_counter += batch_loads.shape[0] 296 | 297 | filename = filename_pattern.format(file_idx + 1) 298 | with h5py.File(filename, 'a') as h5f: 299 | # Determine the actual size to write to avoid broadcasting issues 300 | actual_size = min(split_size - dataset_idx, batch_loads.shape[0]) 301 | h5f['test_loads'][dataset_idx:dataset_idx + actual_size] = batch_loads.cpu().numpy()[:actual_size] 302 | h5f['test_pred'][dataset_idx:dataset_idx + actual_size] = batch_predictions[:actual_size] 303 | h5f['test_solutions'][dataset_idx:dataset_idx + actual_size] = batch_solutions.cpu().numpy()[:actual_size] 304 | 305 | print("Predictions saved successfully.") 306 | 307 | # Assume training is complete 308 | time.sleep(5) 309 | print('Training is complete') 310 | 311 | # Define the filename pattern for saving splits 312 | filename_pattern = 'transient_Poisson_results_rbfkan_file_seed42_part_{}.h5' 313 | 314 | # Transfer test data to CPU and obtain predictions in batches, saving directly to files 315 | get_predictions_and_save(model, test_loads, coordinates, batch_size, filename_pattern) 316 | 317 | # Print confirmation 318 | time.sleep(5) 319 | print("Numpy arrays are ready and saved to files.") 320 | 321 | # Load and print shapes for verification 322 | for i in range(1, 5): 323 | filename = filename_pattern.format(i) 324 | with h5py.File(filename, 'r') as h5f: 325 | coordinates_cpu = np.array(h5f['coordinates']) 326 | test_loads_cpu = np.array(h5f['test_loads']) 327 | test_pred_cpu = np.array(h5f['test_pred']) 328 | test_solutions_cpu = np.array(h5f['test_solutions']) 329 | print(f'File {filename}:') 330 | print(f'coordinates_cpu shape = {coordinates_cpu.shape}') 331 | print(f'test_loads_cpu shape = {test_loads_cpu.shape}') 332 | print(f'test_pred_cpu shape = {test_pred_cpu.shape}') 333 | print(f'test_solutions_cpu shape = {test_solutions_cpu.shape}') --------------------------------------------------------------------------------