diff --git a/ThermalSolver/autoregressive_module.py b/ThermalSolver/autoregressive_module.py new file mode 100644 index 0000000..67fbbad --- /dev/null +++ b/ThermalSolver/autoregressive_module.py @@ -0,0 +1,298 @@ +import torch +from lightning import LightningModule +from torch_geometric.data import Batch +import importlib +from matplotlib import pyplot as plt +from matplotlib.tri import Triangulation +from .model.finite_difference import FiniteDifferenceStep +import os + +def import_class(class_path: str): + module_path, class_name = class_path.rsplit(".", 1) # split last dot + module = importlib.import_module(module_path) # import the module + cls = getattr(module, class_name) # get the class + return cls + + +def _plot_mesh(pos, y, y_pred, y_true ,batch, i, batch_idx): + + idx = batch == 0 + y = y[idx].detach().cpu() + y_pred = y_pred[idx].detach().cpu() + pos = pos[idx].detach().cpu() + y_true = y_true[idx].detach().cpu() + # print(torch.max(y_true), torch.min(y_true)) + folder = f"{batch_idx:02d}_images" + if os.path.exists(folder) is False: + os.makedirs(folder) + pos = pos.detach().cpu() + tria = Triangulation(pos[:, 0], pos[:, 1]) + plt.figure(figsize=(18, 5)) + plt.subplot(1, 3, 1) + plt.tricontourf(tria, y.squeeze().numpy(), levels=14) + plt.colorbar() + plt.title("Step t-1") + plt.subplot(1, 3, 2) + plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=14) + plt.colorbar() + plt.title("Step t Predicted") + plt.subplot(1, 3, 3) + plt.tricontourf(tria, y_true.squeeze().numpy(), levels=14) + plt.colorbar() + plt.title("t True") + plt.suptitle("GNO", fontsize=16) + name = f"{folder}/graph_iter_{i:04d}.png" + plt.savefig(name, dpi=72) + plt.close() + +def _plot_losses(losses, batch_idx): + folder = f"{batch_idx:02d}_images" + plt.figure() + plt.plot(losses) + plt.yscale("log") + plt.xlabel("Iteration") + plt.ylabel("Loss") + plt.title("Test Loss over Iterations") + plt.grid(True) + file_name = f"{folder}/test_loss.png" + plt.savefig(file_name, dpi=300) + plt.close() + + +class GraphSolver(LightningModule): + def __init__( + self, + model_class_path: str, + model_init_args: dict = {}, + loss: torch.nn.Module = None, + start_unrolling_steps: int = 1, + increase_every: int = 20, + increase_rate: float = 2, + max_unrolling_steps: int = 100, + max_inference_iters: int = 1000, + inner_steps: int = 16, + ): + super().__init__() + self.model = import_class(model_class_path)(**model_init_args) + # for param in self.model.parameters(): + # print(f"Param: {param.shape}, Grad: {param.grad}") + # print(f"Param: {param[0]}") + self.fd_net = FiniteDifferenceStep() + self.loss = loss if loss is not None else torch.nn.MSELoss() + self.start_unrolling = start_unrolling_steps + self.current_unrolling_steps = self.start_unrolling + self.increase_every = increase_every + self.increase_rate = increase_rate + self.max_unrolling_steps = max_unrolling_steps + self.max_inference_iters = max_inference_iters + self.threshold = 1e-4 + self.inner_steps = inner_steps + + def _compute_deg(self, edge_index, edge_attr, num_nodes): + deg = torch.zeros(num_nodes, device=edge_index.device) + deg = deg.scatter_add(0, edge_index[1], edge_attr) + return deg + 1e-7 + + def _compute_loss(self, x, y): + return self.loss(x, y) + + def _log_loss(self, loss, batch, stage: str): + self.log( + f"{stage}/loss", + loss, + on_step=True, + on_epoch=True, + prog_bar=True, + batch_size=int(batch.num_graphs), + ) + return loss + + @staticmethod + def _compute_c_ij(c, edge_index): + """ + TODO: add docstring. + """ + return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze() + + def _compute_model_steps( + self, x, edge_index, edge_attr, boundary_mask, boundary_values + ): + + out = x + self.model(x, edge_index, edge_attr) + # out[boundary_mask] = boundary_values.unsqueeze(-1) + plt.figure() + return out + + def _check_convergence(self, out, x): + residual_norm = torch.norm(out - x) + if residual_norm < self.threshold * torch.norm(x): + return True + return False + + def _preprocess_batch(self, batch: Batch): + x, y, c, edge_index, edge_attr = ( + batch.x, + batch.y, + batch.c, + batch.edge_index, + batch.edge_attr, + ) + # edge_attr = 1 / edge_attr + c_ij = self._compute_c_ij(c, edge_index) + edge_attr = edge_attr * (c_ij) # / 100) + return x, y, edge_index, edge_attr + + def training_step(self, batch: Batch): + x, y, edge_index, edge_attr = self._preprocess_batch(batch) + # deg = self._compute_deg(edge_index, edge_attr, x.size(0)) + losses = [] + # print(x.shape, y.shape) + # # print(torch.max(edge_index), torch.min(edge_index)) + # plt.figure() + # plt.subplot(2,3,1) + # plt.scatter(batch.pos[:,0].cpu(), batch.pos[:,1].cpu(), c=x.squeeze().cpu()) + # plt.subplot(2,3,2) + # plt.scatter(batch.pos[:,0].cpu(), batch.pos[:,1].cpu(), c=y[:,0,:].squeeze().cpu()) + # plt.subplot(2,3,3) + # plt.scatter(batch.pos[:,0].cpu(), batch.pos[:,1].cpu(), c=y[:,1,:].squeeze().cpu()) + # plt.subplot(2,3,4) + # plt.scatter(batch.pos[:,0].cpu(), batch.pos[:,1].cpu(), c=y[:,2,:].squeeze().cpu()) + # plt.subplot(2,3,5) + # plt.scatter(batch.pos[:,0].cpu(), batch.pos[:,1].cpu(), c=y[:,3,:].squeeze().cpu()) + # plt.subplot(2,3,6) + # plt.scatter(batch.pos[:,0].cpu(), batch.pos[:,1].cpu(), c=y[:,4,:].squeeze().cpu()) + # plt.suptitle("Training Batch Visualization", fontsize=16) + # plt.savefig("training_batch_visualization.png", dpi=300) + # plt.close() + # y = z + pos = batch.pos + boundary_mask = batch.boundary_mask + boundary_values = batch.boundary_values + # plt.scatter(pos[boundary_mask,0].cpu(), pos[boundary_mask,1].cpu(), c=boundary_values.cpu(), s=1) + # plt.savefig("boundary_nodes.png", dpi=300) + # y = z + print(y.shape) + for i in range(self.current_unrolling_steps * self.inner_steps): + out = self._compute_model_steps( + # torch.cat([x,pos], dim=-1), + x, + edge_index, + edge_attr, + # deg, + batch.boundary_mask, + batch.boundary_values, + ) + x = out + # print(out.shape, y[:, i, :].shape) + losses.append(self.loss(out.flatten(), y[:, i, :].flatten())) + + print(losses) + + loss = torch.stack(losses).mean() + # for param in self.model.parameters(): + # print(f"Param: {param.shape}, Grad: {param.grad}") + # print(f"Param: {param[0]}") + self._log_loss(loss, batch, "train") + return loss + + # def on_train_epoch_start(self): + # print(f"Current unrolling steps: {self.current_unrolling_steps}, dataset unrolling steps: {self.trainer.datamodule.train_dataset.unrolling_steps}") + # return super().on_train_epoch_start() + + def on_train_epoch_end(self): + if ( + (self.current_epoch + 1) % self.increase_every == 0 + and self.current_epoch > 0 + ): + dm = self.trainer.datamodule + self.current_unrolling_steps = min( + int(self.current_unrolling_steps * self.increase_rate), + self.max_unrolling_steps + ) + dm.unrolling_steps = self.current_unrolling_steps + return super().on_train_epoch_end() + + def validation_step(self, batch: Batch, _): + # x, y, edge_index, edge_attr = self._preprocess_batch(batch) + + # deg = self._compute_deg(edge_index, edge_attr, x.size(0)) + # for i in range(self.max_inference_iters * self.inner_steps): + # out = self._compute_model_steps( + # x, + # edge_index, + # edge_attr, + # deg, + # batch.boundary_mask, + # batch.boundary_values, + # ) + # converged = self._check_convergence(out, x) + # x = out + # if converged: + # break + # print(y.shape, out.shape) + # loss = self.loss(out, y[:,-1,:]) + # self._log_loss(loss, batch, "val") + # self.log("val/iterations", i + 1, on_step=False, on_epoch=True, prog_bar=True, batch_size=int(batch.num_graphs),) + # return loss + + x, y, edge_index, edge_attr = self._preprocess_batch(batch) + # deg = self._compute_deg(edge_index, edge_attr, x.size(0)) + losses = [] + pos = batch.pos + for i in range(self.current_unrolling_steps * self.inner_steps): + out = self._compute_model_steps( + # torch.cat([x,pos], dim=-1), + x, + edge_index, + edge_attr, + # deg, + batch.boundary_mask, + batch.boundary_values, + ) + _plot_mesh(batch.pos, x, out, y[:, i, :], batch.batch, i, self.current_epoch) + x = out + losses.append(self.loss(out, y[:, i, :])) + + loss = torch.stack(losses).mean() + self._log_loss(loss, batch, "val") + return loss + + def test_step(self, batch: Batch, batch_idx): + x, y, edge_index, edge_attr = self._preprocess_batch(batch) + + deg = self._compute_deg(edge_index, edge_attr, x.size(0)) + losses = [] + for i in range(self.max_iters): + out = self._compute_model_steps( + x, + edge_index, + edge_attr.unsqueeze(-1), + deg, + batch.boundary_mask, + batch.boundary_values, + ) + converged = self._check_convergence(out, x) + # _plot_mesh(batch.pos, y, out, batch.batch, i, batch_idx) + losses.append(self.loss(out, y).item()) + if converged: + break + x = out + loss = self.loss(out, y) + # _plot_losses(losses, batch_idx) + self._log_loss(loss, batch, "test") + self.log( + "test/iterations", + i + 1, + on_step=False, + on_epoch=True, + prog_bar=True, + batch_size=int(batch.num_graphs), + ) + + def configure_optimizers(self): + optimizer = torch.optim.AdamW(self.parameters(), lr=1e-2) + return optimizer + + def _impose_bc(self, x: torch.Tensor, data: Batch): + x[data.boundary_mask] = data.boundary_values + return x diff --git a/ThermalSolver/graph_datamodule_unsteady.py b/ThermalSolver/graph_datamodule_unsteady.py new file mode 100644 index 0000000..b3d3976 --- /dev/null +++ b/ThermalSolver/graph_datamodule_unsteady.py @@ -0,0 +1,219 @@ +import torch +from tqdm import tqdm +from lightning import LightningDataModule +from datasets import load_dataset +from torch_geometric.data import Data +from torch_geometric.loader import DataLoader +from torch_geometric.utils import to_undirected +from .mesh_data import MeshData +# from torch.utils.data import Dataset + +class GraphDataModule(LightningDataModule): + def __init__( + self, + hf_repo: str, + split_name: str, + train_size: float = 0.2, + val_size: float = 0.1, + test_size: float = 0.1, + batch_size: int = 32, + remove_boundary_edges: bool = False, + build_radial_graph: bool = False, + radius: float = None, + start_unrolling_steps: int = 1, + ): + super().__init__() + self.hf_repo = hf_repo + self.split_name = split_name + self.dataset_dict = {} + self.train_dataset, self.val_dataset, self.test_dataset = None, None, None + self.unrolling_steps = start_unrolling_steps + self.geometry_dict = {} + self.train_size = train_size + self.val_size = val_size + self.test_size = test_size + self.batch_size = batch_size + self.remove_boundary_edges = remove_boundary_edges + self.build_radial_graph = build_radial_graph + self.radius = radius + + def prepare_data(self): + dataset = load_dataset(self.hf_repo, name="snapshots")[self.split_name] + geometry = load_dataset(self.hf_repo, name="geometry")[self.split_name] + + total_len = len(dataset) + train_len = int(self.train_size * total_len) + valid_len = int(self.val_size * total_len) + self.dataset_dict = { + "train": dataset.select(range(0, train_len)), + "val": dataset.select(range(train_len, train_len + valid_len)), + "test": dataset.select(range(train_len + valid_len, total_len)), + } + self.geometry_dict = { + "train": geometry.select(range(0, train_len)), + "val": geometry.select(range(train_len, train_len + valid_len)), + "test": geometry.select(range(train_len + valid_len, total_len)), + } + + def _compute_boundary_mask( + self, bottom_ids, right_ids, top_ids, left_ids, temperature + ): + left_ids = left_ids[~torch.isin(left_ids, bottom_ids)] + right_ids = right_ids[~torch.isin(right_ids, bottom_ids)] + left_ids = left_ids[~torch.isin(left_ids, top_ids)] + right_ids = right_ids[~torch.isin(right_ids, top_ids)] + + bottom_bc = temperature[bottom_ids].median() + bottom_bc_mask = torch.ones(len(bottom_ids)) * bottom_bc + left_bc = temperature[left_ids].median() + left_bc_mask = torch.ones(len(left_ids)) * left_bc + right_bc = temperature[right_ids].median() + right_bc_mask = torch.ones(len(right_ids)) * right_bc + + boundary_values = torch.cat( + [bottom_bc_mask, right_bc_mask, left_bc_mask], dim=0 + ) + boundary_mask = torch.cat([bottom_ids, right_ids, left_ids], dim=0) + + return boundary_mask, boundary_values + + def _build_dataset( + self, + snapshot: dict, + geometry: dict, + ) -> Data: + conductivity = torch.tensor( + geometry["conductivity"], dtype=torch.float32 + ) + temperatures = torch.tensor(snapshot["temperatures"], dtype=torch.float32)[:2] + times = torch.tensor(snapshot["times"], dtype=torch.float32) + + pos = torch.tensor(geometry["points"], dtype=torch.float32)[:, :2] + + bottom_ids = torch.tensor( + geometry["bottom_boundary_ids"], dtype=torch.long + ) + top_ids = torch.tensor(geometry["top_boundary_ids"], dtype=torch.long) + left_ids = torch.tensor(geometry["left_boundary_ids"], dtype=torch.long) + right_ids = torch.tensor( + geometry["right_boundary_ids"], dtype=torch.long + ) + + if self.build_radial_graph: + from pina.graph import RadiusGraph + + if self.radius is None: + raise ValueError("Radius must be specified for radial graph.") + edge_index = RadiusGraph.compute_radius_graph( + pos, radius=self.radius + ) + from torch_geometric.utils import remove_self_loops + + edge_index, _ = remove_self_loops(edge_index) + else: + edge_index = torch.tensor( + geometry["edge_index"], dtype=torch.int64 + ).T + edge_index = to_undirected(edge_index, num_nodes=pos.size(0)) + + boundary_mask, boundary_values = self._compute_boundary_mask( + bottom_ids, right_ids, top_ids, left_ids, temperatures[0,:] + ) + + if self.remove_boundary_edges: + boundary_idx = torch.unique(boundary_mask) + edge_index_mask = ~torch.isin(edge_index[1], boundary_idx) + edge_index = edge_index[:, edge_index_mask] + + edge_attr = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1) + + n_data = temperatures.size(0) - self.unrolling_steps + data = [] + for i in range(n_data): + x = temperatures[i, :].unsqueeze(-1) + print(x.shape) + y = temperatures[i + 1 : i + 1 + self.unrolling_steps, :].unsqueeze(-1).permute(1,0,2) + # print(y.shape) + data.append(MeshData( + x=x, + y=y, + c=conductivity.unsqueeze(-1), + edge_index=edge_index, + pos=pos, + edge_attr=edge_attr, + boundary_mask=boundary_mask, + boundary_values=boundary_values, + )) + return data + + def setup(self, stage: str = None): + if stage == "fit" or stage is None: + self.train_data = [ + self._build_dataset(snap, geom) + for snap, geom in tqdm( + zip( + self.dataset_dict["train"], self.geometry_dict["train"] + ), + desc="Building train graphs", + total=len(self.dataset_dict["train"]), + ) + ] + self.val_data = [ + self._build_dataset(snap, geom) + for snap, geom in tqdm( + zip(self.dataset_dict["val"], self.geometry_dict["val"]), + desc="Building val graphs", + total=len(self.dataset_dict["val"]), + ) + ] + if stage == "test" or stage is None: + self.test_data = [ + self._build_dataset(snap, geom) + for snap, geom in tqdm( + zip(self.dataset_dict["test"], self.geometry_dict["test"]), + desc="Building test graphs", + total=len(self.dataset_dict["test"]), + ) + ] + + # def create_autoregressive_datasets(self, dataset: str, no_unrolling: bool = False): + # if dataset == "train": + # return AutoregressiveDataset(self.train_data, self.unrolling_steps, no_unrolling) + # if dataset == "val": + # return AutoregressiveDataset(self.val_data, self.unrolling_steps, no_unrolling) + # if dataset == "test": + # return AutoregressiveDataset(self.test_data, self.unrolling_steps, no_unrolling) + + def train_dataloader(self): + # ds = self.create_autoregressive_datasets(dataset="train") + # self.train_dataset = ds + print(type(self.train_data[0])) + ds = [i for data in self.train_data for i in data] + print(type(ds[0])) + return DataLoader( + ds, + batch_size=self.batch_size, + shuffle=True, + num_workers=8, + pin_memory=True, + ) + + def val_dataloader(self): + ds = [i for data in self.val_data for i in data] + return DataLoader( + ds, + batch_size=self.batch_size, + shuffle=False, + num_workers=8, + pin_memory=True, + ) + + def test_dataloader(self): + ds = self.create_autoregressive_datasets(dataset="test", no_unrolling=True) + return DataLoader( + ds, + batch_size=self.batch_size, + shuffle=False, + num_workers=8, + pin_memory=True, + ) diff --git a/ThermalSolver/graph_module.py b/ThermalSolver/graph_module.py index 88b5fcd..b60ec73 100644 --- a/ThermalSolver/graph_module.py +++ b/ThermalSolver/graph_module.py @@ -5,7 +5,7 @@ import importlib from matplotlib import pyplot as plt from matplotlib.tri import Triangulation from .model.finite_difference import FiniteDifferenceStep - +import os def import_class(class_path: str): module_path, class_name = class_path.rsplit(".", 1) # split last dot @@ -14,13 +14,15 @@ def import_class(class_path: str): return cls -def _plot_mesh(pos, y, y_pred, batch, i): +def _plot_mesh(pos, y, y_pred, batch, i, batch_idx): idx = batch == 0 y = y[idx].detach().cpu() y_pred = y_pred[idx].detach().cpu() pos = pos[idx].detach().cpu() - + folder = f"{batch_idx:02d}_images" + if os.path.exists(folder) is False: + os.makedirs(folder) pos = pos.detach().cpu() tria = Triangulation(pos[:, 0], pos[:, 1]) plt.figure(figsize=(18, 5)) @@ -37,10 +39,23 @@ def _plot_mesh(pos, y, y_pred, batch, i): plt.colorbar() plt.title("Error") plt.suptitle("GNO", fontsize=16) - name = f"images/graph_iter_{i:04d}.png" + name = f"{folder}/graph_iter_{i:04d}.png" plt.savefig(name, dpi=72) plt.close() +def _plot_losses(losses, batch_idx): + folder = f"{batch_idx:02d}_images" + plt.figure() + plt.plot(losses) + plt.yscale("log") + plt.xlabel("Iteration") + plt.ylabel("Loss") + plt.title("Test Loss over Iterations") + plt.grid(True) + file_name = f"{folder}/test_loss.png" + plt.savefig(file_name, dpi=300) + plt.close() + class GraphSolver(LightningModule): def __init__( @@ -231,7 +246,6 @@ class GraphSolver(LightningModule): x, y, edge_index, edge_attr = self._preprocess_batch(batch) deg = self._compute_deg(edge_index, edge_attr, x.size(0)) - for i in range(self.current_iters): out = self._compute_model_steps( x, @@ -257,36 +271,8 @@ class GraphSolver(LightningModule): batch_size=int(batch.num_graphs), ) - def test_step(self, batch: Batch, _): - x, y, edge_index, edge_attr = self._preprocess_batch(batch) - - deg = self._compute_deg(edge_index, edge_attr, x.size(0)) - - for i in range(self.max_iters): - out = self._compute_model_steps( - x, - edge_index, - edge_attr.unsqueeze(-1), - deg, - batch.boundary_mask, - batch.boundary_values, - ) - converged = self._check_convergence(out, x) - # _plot_mesh(batch.pos, y, out, batch.batch, i) - if converged: - break - x = out - loss = self.loss(out, y) - - self._log_loss(loss, batch, "test") - self.log( - "test/iterations", - i + 1, - on_step=False, - on_epoch=True, - prog_bar=True, - batch_size=int(batch.num_graphs), - ) + def test_step(self, batch: Batch, batch_idx): + pass def configure_optimizers(self): optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3) diff --git a/ThermalSolver/model/learnable_finite_difference.py b/ThermalSolver/model/learnable_finite_difference.py index f78fff0..a606a70 100644 --- a/ThermalSolver/model/learnable_finite_difference.py +++ b/ThermalSolver/model/learnable_finite_difference.py @@ -1,119 +1,22 @@ -# import torch -# import torch.nn as nn -# from torch_geometric.nn import MessagePassing -# from torch.nn.utils import spectral_norm - -# class GCNConvLayer(MessagePassing): -# def __init__(self, in_channels, out_channels): -# super().__init__(aggr="add") -# self.lin_l = spectral_norm(nn.Linear(in_channels, out_channels, bias=False)) -# self.lin_r = spectral_norm(nn.Linear(in_channels, out_channels, bias=False)) - -# def forward(self, x, edge_index, edge_attr, deg): -# out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg) -# out = self.lin_l(out) -# return out - -# def message(self, x_j, edge_attr): -# return x_j * edge_attr - -# def aggregate(self, inputs, index, deg): -# """ -# TODO: add docstring. -# """ -# out = super().aggregate(inputs, index) -# deg = deg + 1e-7 -# return out / deg.view(-1, 1) - - -# class CorrectionNet(nn.Module): -# def __init__(self, hidden_dim=8, n_layers=1): -# super().__init__() -# # self.enc = GCNConvLayer(1, hidden_dim) -# self.enc = nn.Sequential( -# spectral_norm(nn.Linear(1, hidden_dim//2)), -# nn.GELU(), -# spectral_norm(nn.Linear(hidden_dim//2, hidden_dim)), -# ) -# self.layers = torch.nn.ModuleList([GCNConvLayer(hidden_dim, hidden_dim) for _ in range(n_layers)]) -# self.relu = nn.GELU() - -# self.dec = nn.Sequential( -# spectral_norm(nn.Linear(hidden_dim, hidden_dim//2)), -# nn.GELU(), -# spectral_norm(nn.Linear(hidden_dim//2, 1)), -# ) - -# def forward(self, x, edge_index, edge_attr, deg,): -# # h = self.enc(x, edge_index, edge_attr, deg) -# # h = self.relu(self.enc(x)) -# h = self.enc(x) -# for layer in self.layers: -# h = layer(h, edge_index, edge_attr, deg) -# # h = self.norm(h) -# h = self.relu(h) -# # out = self.dec(h, edge_index, edge_attr, deg) -# out = self.dec(h) -# return out - - import torch import torch.nn as nn from torch_geometric.nn import MessagePassing from torch.nn.utils import spectral_norm +from torch_geometric.nn.conv import GCNConv - -class CorrectionNet(MessagePassing): - """ - TODO: add docstring. - """ - - def __init__(self, hidden_dim=16): +class GCNConvLayer(MessagePassing): + def __init__(self, in_channels, out_channels): super().__init__(aggr="add") - self.in_net = nn.Sequential( - spectral_norm(nn.Linear(1, hidden_dim // 2)), - nn.GELU(), - spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)), - ) - - self.out_net = nn.Sequential( - spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)), - nn.GELU(), - spectral_norm(nn.Linear(hidden_dim // 2, 1)), - ) - - self.lin_msg = spectral_norm( - nn.Linear(hidden_dim, hidden_dim, bias=False) - ) - self.lin_update = spectral_norm( - nn.Linear(hidden_dim, hidden_dim, bias=False) - ) - self.alpha = nn.Parameter(torch.tensor(0.0)) - self.beta = nn.Parameter(torch.tensor(0.0)) + self.lin_l = nn.Linear(in_channels, out_channels, bias=True) + # self.lin_r = spectral_norm(nn.Linear(in_channels, out_channels, bias=False)) def forward(self, x, edge_index, edge_attr, deg): - """ - TODO: add docstring. - """ - x = self.in_net(x) out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg) - return self.out_net(out) + out = self.lin_l(out) + return out def message(self, x_j, edge_attr): - """ - TODO: add docstring. - """ - alpha = torch.sigmoid(self.alpha) - msg = x_j * edge_attr - msg = (1 - alpha) * msg + alpha * self.lin_msg(msg) - return msg - - def update(self, aggr_out, x): - """ - TODO: add docstring. - """ - beta = torch.sigmoid(self.beta) - return aggr_out * (1 - beta) + self.lin_msg(x) * beta + return x_j * edge_attr.view(-1, 1) def aggregate(self, inputs, index, deg): """ @@ -122,3 +25,45 @@ class CorrectionNet(MessagePassing): out = super().aggregate(inputs, index) deg = deg + 1e-7 return out / deg.view(-1, 1) + + +class CorrectionNet(nn.Module): + def __init__(self, input_dim=1, output_dim=1, hidden_dim=8, n_layers=8): + super().__init__() + self.enc = nn.Linear(input_dim, hidden_dim, bias=False) + # self.layers = n_layers + # self.l = GCNConv(hidden_dim, hidden_dim, aggr="mean") + self.layers = torch.nn.ModuleList( + [GCNConv(hidden_dim, hidden_dim, aggr="mean", bias=False) for _ in range(n_layers)] + ) + self.dec = nn.Linear(hidden_dim, output_dim) + + def forward(self, x, edge_index, edge_attr,): + h = self.enc(x) + # h = self.relu(h) + for l in self.layers: + # print(f"Forward pass layer {_}") + h = l(h, edge_index, edge_attr) + # h = self.relu(h) + out = self.dec(h) + return out + + +class MLPNet(nn.Module): + def __init__(self, input_dim=1, output_dim=1, hidden_dim=8, n_layers=1): + super().__init__() + layers = [] + func = torch.nn.ReLU + + self.network = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + func(), + nn.Linear(hidden_dim, hidden_dim), + func(), + nn.Linear(hidden_dim, hidden_dim), + func(), + nn.Linear(hidden_dim, output_dim), + ) + + def forward(self, x, edge_index=None, edge_attr=None): + return self.network(x) \ No newline at end of file diff --git a/experiments/config.yaml b/experiments/config.yaml index c88c524..8216fc9 100644 --- a/experiments/config.yaml +++ b/experiments/config.yaml @@ -20,7 +20,7 @@ trainer: mode: min patience: 10 verbose: false - max_epochs: 200 + max_epochs: 2000 min_epochs: null max_steps: -1 min_steps: null diff --git a/experiments/config_01.yaml b/experiments/config_01.yaml index 4604b8a..686dc87 100644 --- a/experiments/config_01.yaml +++ b/experiments/config_01.yaml @@ -9,8 +9,8 @@ trainer: logger: - class_path: lightning.pytorch.loggers.TensorBoardLogger init_args: - save_dir: lightning_logs - name: "01" + save_dir: logs + name: "test" version: null callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint @@ -33,26 +33,21 @@ trainer: log_every_n_steps: null inference_mode: true default_root_dir: null - accumulate_grad_batches: 6 - # gradient_clip_val: 1.0 + accumulate_grad_batches: 4 + gradient_clip_val: 1.0 model: - class_path: ThermalSolver.module.GraphSolver + class_path: ThermalSolver.graph_module.GraphSolver init_args: - model_class_path: ThermalSolver.model.local_gno.GatingGNO + model_class_path: ThermalSolver.model.LearnableGraphFiniteDifference model_init_args: - x_ch_node: 1 - f_ch_node: 1 - hidden: 16 - layers: 1 - edge_ch: 3 - out_ch: 1 + max_iters: 250 unrolling_steps: 64 data: - class_path: ThermalSolver.data_module.GraphDataModule + class_path: ThermalSolver.graph_datamodule.GraphDataModule init_args: hf_repo: "SISSAmathLab/thermal-conduction" - split_name: "2000_ref_1" - batch_size: 4 + split_name: "1000_40x30" + batch_size: 8 train_size: 0.8 test_size: 0.1 test_size: 0.1 diff --git a/experiments/config_64.yaml b/experiments/config_64.yaml index 617b8ac..6765a5d 100644 --- a/experiments/config_64.yaml +++ b/experiments/config_64.yaml @@ -33,11 +33,11 @@ trainer: log_every_n_steps: null inference_mode: true default_root_dir: null - accumulate_grad_batches: 6 + accumulate_grad_batches: 2 gradient_clip_val: 1.0 model: - class_path: ThermalSolver.module.GraphSolver + class_path: ThermalSolver.graph_module.GraphSolver init_args: model_class_path: ThermalSolver.model.local_gno.GatingGNO model_init_args: @@ -49,11 +49,11 @@ model: out_ch: 1 unrolling_steps: 1 data: - class_path: ThermalSolver.data_module.GraphDataModule + class_path: ThermalSolver.graph_datamodule.GraphDataModule init_args: hf_repo: "SISSAmathLab/thermal-conduction" - split_name: "2000" - batch_size: 4 + split_name: "2000_ref_1" + batch_size: 10 train_size: 0.8 test_size: 0.1 test_size: 0.1 diff --git a/experiments/config_autoregressive.yaml b/experiments/config_autoregressive.yaml new file mode 100644 index 0000000..ec9e94a --- /dev/null +++ b/experiments/config_autoregressive.yaml @@ -0,0 +1,70 @@ +# lightning.pytorch==2.5.5 +seed_everything: 1999 +trainer: + accelerator: gpu + strategy: auto + devices: 1 + num_nodes: 1 + precision: null + logger: + - class_path: lightning.pytorch.loggers.TensorBoardLogger + init_args: + save_dir: logs.autoregressive + name: "test" + version: null + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val/loss + mode: min + save_top_k: 1 + filename: best-checkpoint + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val/loss + mode: min + patience: 50 + verbose: false + max_epochs: 1000 + min_epochs: null + max_steps: -1 + min_steps: null + overfit_batches: 0.0 + log_every_n_steps: null + accumulate_grad_batches: 1 + # reload_dataloaders_every_n_epochs: 1 + default_root_dir: null + +model: + class_path: ThermalSolver.autoregressive_module.GraphSolver + init_args: + model_class_path: ThermalSolver.model.learnable_finite_difference.CorrectionNet + model_init_args: + input_dim: 1 + hidden_dim: 24 + # output_dim: 1 + n_layers: 1 + start_unrolling_steps: 1 + increase_every: 100000 + increase_rate: 2 + max_inference_iters: 300 + max_unrolling_steps: 40 + inner_steps: 1 + +data: + class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule + init_args: + hf_repo: "SISSAmathLab/thermal-conduction-unsteady" + split_name: "50_samples_easy" + batch_size: 64 + train_size: 0.02 + val_size: 0.02 + test_size: 0.96 + build_radial_graph: true + radius: 0.5 + remove_boundary_edges: true + start_unrolling_steps: 1 + +optimizer: null +lr_scheduler: null +# ckpt_path: logs/test/version_0/checkpoints/best-checkpoint.ckpt diff --git a/experiments/config_fd.yaml b/experiments/config_fd.yaml new file mode 100644 index 0000000..bb7a2b3 --- /dev/null +++ b/experiments/config_fd.yaml @@ -0,0 +1,58 @@ +# lightning.pytorch==2.5.5 +seed_everything: 1999 +trainer: + accelerator: gpu + strategy: auto + devices: 1 + num_nodes: 1 + precision: null + logger: + - class_path: lightning.pytorch.loggers.TensorBoardLogger + init_args: + save_dir: logs + name: "fd" + version: null + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val/loss + mode: min + save_top_k: 1 + filename: best-checkpoint + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val/loss + mode: min + patience: 2 + verbose: false + max_epochs: 1000 + min_epochs: null + max_steps: -1 + min_steps: null + overfit_batches: 0.0 + log_every_n_steps: null + inference_mode: true + default_root_dir: null + accumulate_grad_batches: 4 + gradient_clip_val: 1.0 +model: + class_path: ThermalSolver.graph_module.GraphSolver + init_args: + model_class_path: ThermalSolver.model.GraphFiniteDifference + # model_init_args: + max_iters: 10000 + # unrolling_steps: 64 +data: + class_path: ThermalSolver.graph_datamodule.GraphDataModule + init_args: + hf_repo: "SISSAmathLab/thermal-conduction" + split_name: "1000_1_40x30" + batch_size: 8 + train_size: 0.8 + test_size: 0.1 + test_size: 0.1 + # build_radial_graph: true + # radius: 1.5 +optimizer: null +lr_scheduler: null +# ckpt_path: lightning_logs/01/version_0/checkpoints/best-checkpoint.ckpt \ No newline at end of file diff --git a/experiments/config_gino.yaml b/experiments/config_gino.yaml new file mode 100644 index 0000000..50f63b3 --- /dev/null +++ b/experiments/config_gino.yaml @@ -0,0 +1,74 @@ +# lightning.pytorch==2.5.5 +seed_everything: 1999 +trainer: + accelerator: gpu + strategy: auto + devices: 1 + num_nodes: 1 + precision: null + logger: + - class_path: lightning.pytorch.loggers.TensorBoardLogger + init_args: + save_dir: logs + name: "test" + version: null + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val/loss + mode: min + save_top_k: 1 + filename: best-checkpoint + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val/loss + mode: min + patience: 15 + verbose: false + max_epochs: 1000 + min_epochs: null + max_steps: -1 + min_steps: null + overfit_batches: 0.0 + log_every_n_steps: null + inference_mode: true + default_root_dir: null + # accumulate_grad_batches: 2 + # gradient_clip_val: 1.0 +model: + class_path: ThermalSolver.graph_module.GraphSolver + init_args: + model_class_path: neuralop.models import GINO + model_init_args: + in_channels: 3 # Es: coordinate (x, y, z) + valore della conducibilità k + out_channels: 1 # Es: temperatura T + + # Parametri per l'encoder e il decoder GNO + gno_coord_features=3, # Dimensionalità delle coordinate per GNO (es. 3D) + gno_n_layers=2, # Numero di layer GNO nell'encoder e nel decoder + gno_hidden_channels=64, # Canali nascosti per i layer GNO + + # Parametri per il processore FNO + fno_n_modes=(16, 16, 16), # Numero di modi di Fourier per ogni dimensione + fno_n_layers=4, # Numero di layer FNO + fno_hidden_channels=64, # Canali nascosti per i layer FNO + + # Canali per il lifting e la proiezione + lifting_channels=256, # Dimensione dello spazio latente dopo il lifting iniziale + projection_channels=256, # Dimensione prima della proiezione finale + + # Padding del dominio per il processore FNO + domain_padding=0.05 + # unrolling_steps: 64 +data: + class_path: ThermalSolver.graph_datamodule.GraphDataModule + init_args: + hf_repo: "SISSAmathLab/thermal-conduction" + split_name: "2000_ref_1" + batch_size: 64 + train_size: 0.8 + test_size: 0.1 + test_size: 0.1 +optimizer: null +lr_scheduler: null +# ckpt_path: lightning_logs/01/version_0/checkpoints/best-checkpoint.ckpt \ No newline at end of file diff --git a/experiments/config_gno.yaml b/experiments/config_gno.yaml index cfdad66..46d81c4 100644 --- a/experiments/config_gno.yaml +++ b/experiments/config_gno.yaml @@ -44,12 +44,12 @@ model: increase_every: 10 increase_rate: 2 max_iters: 2000 - accumulation_iters: 320 + accumulation_iters: 160 data: class_path: ThermalSolver.graph_datamodule.GraphDataModule init_args: hf_repo: "SISSAmathLab/thermal-conduction" - split_name: "1000_40x30" + split_name: "1000_1_40x30" batch_size: 32 train_size: 0.8 test_size: 0.1 diff --git a/experiments/config_gno_inference.yaml b/experiments/config_gno_inference.yaml new file mode 100644 index 0000000..5272eb7 --- /dev/null +++ b/experiments/config_gno_inference.yaml @@ -0,0 +1,62 @@ +# lightning.pytorch==2.5.5 +seed_everything: 1999 +trainer: + accelerator: gpu + strategy: auto + devices: 1 + num_nodes: 1 + precision: null + logger: + - class_path: lightning.pytorch.loggers.TensorBoardLogger + init_args: + save_dir: logs_inference + name: "test" + version: null + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val/loss + mode: min + save_top_k: 1 + filename: best-checkpoint + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val/loss + mode: min + patience: 25 + verbose: false + max_epochs: 1000 + min_epochs: null + max_steps: -1 + min_steps: null + overfit_batches: 0.0 + log_every_n_steps: null + # inference_mode: true + default_root_dir: null + # accumulate_grad_batches: 2 + # gradient_clip_val: 1.0 +model: + class_path: ThermalSolver.graph_module.GraphSolver + init_args: + model_class_path: ThermalSolver.model.finite_difference.FiniteDifferenceStep + curriculum_learning: true + start_iters: 5 + increase_every: 10 + increase_rate: 2 + max_iters: 2000 + accumulation_iters: 320 +data: + class_path: ThermalSolver.graph_datamodule.GraphDataModule + init_args: + hf_repo: "SISSAmathLab/thermal-conduction" + split_name: "1000_3_40x30" + batch_size: 10 + train_size: 0.8 + test_size: 0.1 + test_size: 0.1 + build_radial_graph: True + radius: 1.2 + remove_boundary_edges: false +optimizer: null +lr_scheduler: null +# ckpt_path: logs/test/version_2/checkpoints/best-checkpoint.ckpt diff --git a/experiments/config_pointnet.yaml b/experiments/config_pointnet.yaml new file mode 100644 index 0000000..e6461f6 --- /dev/null +++ b/experiments/config_pointnet.yaml @@ -0,0 +1,56 @@ +# lightning.pytorch==2.5.5 +seed_everything: 1999 +trainer: + accelerator: gpu + strategy: auto + devices: 1 + num_nodes: 1 + precision: null + logger: + - class_path: lightning.pytorch.loggers.TensorBoardLogger + init_args: + save_dir: lightning_logs + name: "pointnet" + version: null + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val/loss + mode: min + save_top_k: 1 + filename: best-checkpoint + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val/loss + mode: min + patience: 10 + verbose: false + max_epochs: 200 + min_epochs: null + max_steps: -1 + min_steps: null + overfit_batches: 0.0 + log_every_n_steps: null + inference_mode: true + default_root_dir: null + accumulate_grad_batches: 2 + gradient_clip_val: 1.0 +model: + class_path: ThermalSolver.point_module.PointSolver + init_args: + model_class_path: ThermalSolver.model.point_net.PointNet + model_init_args: + input_dim: 4 + output_dim: 1 +data: + class_path: ThermalSolver.point_datamodule.PointDataModule + init_args: + hf_repo: "SISSAmathLab/thermal-conduction" + split_name: "2000" + batch_size: 10 + train_size: 0.8 + test_size: 0.1 + test_size: 0.1 +optimizer: null +lr_scheduler: null +# ckpt_path: lightning_logs/pointnet/version_0/checkpoints/best-checkpoint.ckpt