new plotting strategy

.gitignore

@@ -210,6 +210,7 @@ __marimo__/
 logs/
 *.log
 models/
+logs*
 
 # Images
 *.png

ThermalSolver/autoregressive_module.py

@@ -0,0 +1,340 @@
+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, cells, i, batch_idx):
+    # print(pos_.shape, y_.shape, y_pred_.shape, y_true_.shape)
+    for j in [0]:
+        idx = (batch == j).nonzero(as_tuple=True)[0]
+        y = y_[idx].detach().cpu()
+        y_pred = y_pred_[idx].detach().cpu()
+        pos = pos_[idx].detach().cpu()
+        # print(pos.shape, y.shape, y_pred.shape)
+        y_true = y_true_[idx].detach().cpu()
+        y_true = torch.clamp(y_true, min=0)
+        folder = f"{batch_idx:02d}_images"
+        if os.path.exists(folder) is False:
+            os.makedirs(folder)
+        triangles = torch.vstack([cells[:, [0, 1, 2]], cells[:, [0, 2, 3]]])
+        tria = Triangulation(pos[:, 0], pos[:, 1], triangles=triangles)
+        plt.figure(figsize=(24, 6))
+        # plt.subplot(1, 4, 1)
+        # plt.tricontourf(tria, y.squeeze().numpy(), levels=100)
+        # plt.colorbar()
+        # plt.title("Step t-1")
+        # plt.tripcolor(tria, y_pred.squeeze().numpy()
+        # plt.savefig("test_scatter_step_before.png", dpi=72)
+        # x = z
+        plt.subplot(1, 4, 1)
+        plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=100)
+        # plt.scatter(pos[:, 0], pos[:, 1], c=y_pred.squeeze().numpy(), s=20, cmap="viridis",)
+        plt.colorbar()
+        plt.title(f"Prediction at timestep {i:03d}")
+        plt.subplot(1, 4, 2)
+        plt.tricontourf(tria, y_true.squeeze().numpy(), levels=100)
+        # plt.scatter(pos[:, 0], pos[:, 1], c=y_true.squeeze().numpy(), s=20, cmap="viridis")
+        plt.colorbar()
+        plt.title("Ground Truth Steady State")
+        plt.subplot(1, 4, 3)
+        per_element_relative_error = torch.abs(y_pred - y_true) / (
+            y_true + 1e-6
+        )
+        per_element_relative_error = torch.clamp(
+            per_element_relative_error, max=1.0, min=0.0
+        )
+        plt.tricontourf(
+            tria,
+            per_element_relative_error.squeeze(),
+            levels=100,
+            vmin=0,
+            vmax=1.0,
+        )
+        # plt.scatter(pos[:, 0], pos[:, 1], c=per_element_relative_error.squeeze().numpy(), s=20, cmap="viridis", vmin=0, vmax=1.0)
+        plt.colorbar()
+        plt.title("Relative Error")
+        plt.subplot(1, 4, 4)
+        absolute_error = torch.abs(y_pred - y_true)
+        plt.tricontourf(tria, absolute_error.squeeze(), levels=100)
+        # plt.scatter(pos[:, 0], pos[:, 1], c=absolute_error.squeeze().numpy(), s=20, cmap="viridis")
+        plt.colorbar()
+        plt.title("Absolute Error")
+        plt.suptitle("GNO", fontsize=16)
+        name = f"{folder}/{j:04d}_graph_iter_{i:04d}.png"
+        plt.savefig(name, dpi=72)
+        plt.close()
+
+
+def _plot_losses(relative_errors, test_losses, relative_update, batch_idx):
+    # folder = f"{batch_idx:02d}_images"
+    plt.figure(figsize=(18, 6))
+    plt.subplot(1, 3, 1)
+    for i, losses in enumerate(test_losses):
+        plt.plot(losses)
+        if i == 3:
+            break
+    plt.yscale("log")
+    plt.xlabel("Iteration")
+    plt.ylabel("Test Loss")
+    plt.title("Test Loss over Iterations")
+    plt.grid(True)
+    plt.subplot(1, 3, 2)
+    for i, losses in enumerate(relative_errors):
+        plt.plot(losses)
+        if i == 3:
+            break
+    plt.yscale("log")
+    plt.xlabel("Iteration")
+    plt.ylabel("Relative Error")
+    plt.title("Relative error over Iterations")
+    plt.grid(True)
+    plt.subplot(1, 3, 3)
+    for i, updates in enumerate(relative_update):
+        plt.plot(updates)
+        if i == 3:
+            break
+    plt.yscale("log")
+    plt.xlabel("Iteration")
+    plt.ylabel("Relative Update")
+    plt.title("Relative update over Iterations")
+    plt.grid(True)
+    file_name = f"test_errors.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,
+        unrolling_steps: int = 1,
+    ):
+        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.loss = loss if loss is not None else torch.nn.MSELoss()
+        self.unrolling_steps = unrolling_steps
+        self.test_losses = []
+        self.test_relative_errors = []
+        self.test_relative_updates = []
+
+    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=False,
+            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,
+        conductivity,
+    ):
+        out = self.model(x, edge_index, edge_attr, conductivity)
+        out[boundary_mask] = boundary_values.unsqueeze(-1)
+        return out
+
+    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
+        conductivity = self._compute_c_ij(c, edge_index)
+        edge_attr = edge_attr * conductivity
+        return x, y, edge_index, edge_attr, conductivity
+
+    def training_step(self, batch: Batch):
+        x, y, edge_index, edge_attr, conductivity = self._preprocess_batch(
+            batch
+        )
+        losses = []
+        for i in range(self.unrolling_steps):
+            # print(f"Training step {i+1}/{self.unrolling_steps}")
+            out = self._compute_model_steps(
+                x,
+                edge_index,
+                edge_attr,
+                # deg,
+                batch.boundary_mask,
+                batch.boundary_values,
+                conductivity,
+            )
+            x = out
+            losses.append(self.loss(out.flatten(), y[:, i, :].flatten()))
+        loss = torch.stack(losses).mean()
+        self._log_loss(loss, batch, "train")
+        for i, layer in enumerate(self.model.layers):
+            self.log(
+                f"{i:03d}_alpha",
+                layer.alpha,
+                prog_bar=True,
+                on_epoch=True,
+                on_step=False,
+                batch_size=int(batch.num_graphs),
+            )
+        self.log(
+            "dt",
+            self.model.dt,
+            prog_bar=True,
+            on_epoch=True,
+            on_step=False,
+            batch_size=int(batch.num_graphs),
+        )
+        return loss
+
+    def validation_step(self, batch: Batch, batch_idx):
+        x, y, edge_index, edge_attr, conductivity = self._preprocess_batch(
+            batch
+        )
+        # deg = self._compute_deg(edge_index, edge_attr, x.size(0))
+        losses = []
+        pos = batch.pos
+        for i in range(self.unrolling_steps):
+            out = self._compute_model_steps(
+                # torch.cat([x,pos], dim=-1),
+                x,
+                edge_index,
+                edge_attr,
+                # deg,
+                batch.boundary_mask,
+                batch.boundary_values,
+                conductivity,
+            )
+            # if (
+            #     batch_idx == 0
+            #     and self.current_epoch % 10 == 0
+            #     and self.current_epoch > 0
+            # ):
+            #     _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 _check_convergence(self, y_new, y_old, tol=1e-4):
+        l2_norm = torch.norm(y_new - y_old, p=2)
+        y_old_norm = torch.norm(y_old, p=2)
+        rel_error = l2_norm / (y_old_norm)
+        return rel_error.item() < tol, rel_error.item()
+
+    def test_step(self, batch: Batch, batch_idx):
+        x, y, edge_index, edge_attr, conductivity = self._preprocess_batch(
+            batch
+        )
+        # deg = self._compute_deg(edge_index, edge_attr, x.size(0))
+        losses = []
+        all_losses = []
+        norms = []
+        s = []
+        relative_updates = []
+        sequence_length = y.size(1)
+        y = y[:, -1, :].unsqueeze(1)
+        _plot_mesh(
+            batch.pos, x, x, y[:, -1, :], batch.batch, batch.cells, 0, batch_idx
+        )
+        for i in range(200):
+            out = self._compute_model_steps(
+                # torch.cat([x,pos], dim=-1),
+                x,
+                edge_index,
+                edge_attr,
+                # deg,
+                batch.boundary_mask,
+                batch.boundary_values,
+                conductivity,
+            )
+            norms.append(torch.norm(out - x, p=2).item())
+            converged, relative_update = self._check_convergence(out, x)
+            relative_updates.append(relative_update)
+            if batch_idx <= 4:
+                print(f"Plotting iteration {i}, norm diff: {norms[-1]}")
+                _plot_mesh(
+                    batch.pos,
+                    x,
+                    out,
+                    y[:, -1, :],
+                    batch.batch,
+                    batch.cells,
+                    i + 1,
+                    batch_idx,
+                )
+            x = out
+            loss = self.loss(out, y[:, -1, :])
+            relative_error = torch.abs(out - y[:, -1, :]) / (
+                torch.abs(y[:, -1, :]) + 1e-6
+            )
+            mean_relative_error = relative_error.mean()
+            all_losses.append(mean_relative_error.item())
+            losses.append(loss)
+            if converged:
+                print(
+                    f"Test step converged at iteration {i} for batch {batch_idx}"
+                )
+                break
+        loss = torch.stack(losses).mean()
+        self.test_losses.append(losses)
+        self.test_relative_errors.append(all_losses)
+        self.test_relative_updates.append(relative_updates)
+        self._log_loss(loss, batch, "test")
+        return loss
+
+    def on_test_end(self):
+        if len(self.test_losses) > 0:
+            _plot_losses(
+                self.test_relative_errors,
+                self.test_losses,
+                self.test_relative_updates,
+                batch_idx=0,
+            )
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3)
+        return optimizer

ThermalSolver/graph_datamodule.py

@@ -18,6 +18,8 @@ class GraphDataModule(LightningDataModule):
         test_size: float = 0.1,
         batch_size: int = 32,
         remove_boundary_edges: bool = False,
+        build_radial_graph: bool = False,
+        radius: float = None,
     ):
         super().__init__()
         self.hf_repo = hf_repo
@@ -29,6 +31,8 @@ class GraphDataModule(LightningDataModule):
         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]
@@ -78,11 +82,12 @@ class GraphDataModule(LightningDataModule):
         conductivity = torch.tensor(
             snapshot["conductivity"], dtype=torch.float32
         )
-        temperature = torch.tensor(snapshot["temperature"], dtype=torch.float32)
-
-        edge_index = torch.tensor(geometry["edge_index"], dtype=torch.int64).T
+        temperature = torch.tensor(
+            snapshot["temperature"], dtype=torch.float32
+        )[:50]
 
         pos = torch.tensor(geometry["points"], dtype=torch.float32)[:, :2]
+
         bottom_ids = torch.tensor(
             geometry["bottom_boundary_ids"], dtype=torch.long
         )
@@ -92,20 +97,38 @@ class GraphDataModule(LightningDataModule):
             geometry["right_boundary_ids"], dtype=torch.long
         )
 
-        edge_index = to_undirected(edge_index, num_nodes=pos.size(0))
+        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, temperature
         )
+
         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 = pos[edge_index[0]] - pos[edge_index[1]]
-        edge_attr = torch.cat(
-            [edge_attr, torch.norm(edge_attr, dim=1).unsqueeze(-1)], dim=1
-        )
+        # edge_attr = pos[edge_index[0]] - pos[edge_index[1]]
+        # edge_attr = torch.cat(
+        #     [edge_attr, torch.norm(edge_attr, dim=1).unsqueeze(-1)], dim=1
+        # )
+        edge_attr = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
 
         x = torch.zeros_like(temperature, dtype=torch.float32).unsqueeze(-1)
         if self.remove_boundary_edges:

ThermalSolver/graph_datamodule_unsteady.py

@@ -0,0 +1,266 @@
+import torch
+from tqdm import tqdm
+from lightning import LightningDataModule
+from datasets import load_dataset, concatenate_datasets
+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 typing import List, Union
+
+
+class GraphDataModule(LightningDataModule):
+    def __init__(
+        self,
+        hf_repo: str,
+        split_name: Union[str, List[str]],
+        n_elements: int = None,
+        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,
+        unrolling_steps: int = 1,
+        aggregate_timesteps: int = 1,
+        min_normalized_diff: float = 1e-3,
+    ):
+        super().__init__()
+        self.hf_repo = hf_repo
+        self.split_name = split_name
+        self.n_elements = n_elements
+        self.dataset_dict = {}
+        self.train_dataset, self.val_dataset, self.test_dataset = (
+            None,
+            None,
+            None,
+        )
+        self.unrolling_steps = unrolling_steps
+        self.aggregate_timesteps = aggregate_timesteps
+        self.min_normalized_diff = min_normalized_diff
+
+        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):
+        if isinstance(self.split_name, list):
+            dataset_list = []
+            geometry_list = []
+            for split in self.split_name:
+                dataset_list.append(
+                    load_dataset(self.hf_repo, name="snapshots")[split]
+                )
+                geometry_list.append(
+                    load_dataset(self.hf_repo, name="geometry")[split]
+                )
+
+            dataset = concatenate_datasets(dataset_list)
+            geometry = concatenate_datasets(geometry_list)
+            idx = torch.randperm(len(dataset))
+            dataset = dataset.select(idx.tolist())
+            geometry = geometry.select(idx.tolist())
+        else:
+            dataset = load_dataset(self.hf_repo, name="snapshots")[
+                self.split_name
+            ]
+            geometry = load_dataset(self.hf_repo, name="geometry")[
+                self.split_name
+            ]
+
+        if self.n_elements is not None:
+            dataset = dataset.select(range(self.n_elements))
+            geometry = geometry.select(range(self.n_elements))
+
+        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 _build_dataset(
+        self,
+        snapshot: dict,
+        geometry: dict,
+        test: bool = False,
+    ) -> Data:
+        conductivity = torch.tensor(
+            geometry["conductivity"], dtype=torch.float32
+        )
+        temperatures = (
+            torch.tensor(snapshot["unsteady"], dtype=torch.float32)
+            if not test
+            else torch.stack(
+                [
+                    torch.tensor(snapshot["unsteady"], dtype=torch.float32)[
+                        0, ...
+                    ],
+                    torch.tensor(snapshot["steady"], dtype=torch.float32),
+                ],
+                dim=0,
+            )
+        )
+        if not test:
+            for t in range(1, temperatures.size(0)):
+                diff = temperatures[t, :] - temperatures[t - 1, :]
+                norm_diff = torch.norm(diff, p=2) / torch.norm(
+                    temperatures[t - 1], p=2
+                )
+                if norm_diff < self.min_normalized_diff:
+                    temperatures = temperatures[: t + 1, :]
+                    break
+        pos = torch.tensor(geometry["points"], dtype=torch.float32)[:, :2]
+        if self.build_radial_graph:
+            raise NotImplementedError(
+                "Radial graph building not implemented yet."
+            )
+        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 = torch.tensor(
+            geometry["constraints_mask"], dtype=torch.int64
+        )
+        boundary_values = temperatures[0, boundary_mask]
+
+        edge_attr = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
+        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 = edge_attr[edge_index_mask]
+
+        n_data = max(temperatures.size(0) - self.unrolling_steps, 1)
+        data = []
+
+        if test:
+            cells = geometry.get("cells", None)
+            if cells is not None:
+                cells = torch.tensor(cells, dtype=torch.int64)
+            data.append(
+                MeshData(
+                    x=temperatures[0, :].unsqueeze(-1),
+                    y=temperatures[1:2, :].unsqueeze(-1).permute(1, 0, 2),
+                    c=conductivity.unsqueeze(-1),
+                    edge_index=edge_index,
+                    pos=pos,
+                    edge_attr=edge_attr,
+                    boundary_mask=boundary_mask,
+                    boundary_values=boundary_values,
+                    cells=cells,
+                )
+            )
+            return data
+        for i in range(n_data):
+            x = temperatures[i, :].unsqueeze(-1)
+            y = (
+                temperatures[i + 1 : i + 1 + self.unrolling_steps, :]
+                .unsqueeze(-1)
+                .permute(1, 0, 2)
+            )
+            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, test=True)
+                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(
+            f"\nLoading training data, using {self.unrolling_steps} unrolling steps..."
+        )
+        return DataLoader(
+            ds,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=8,
+            pin_memory=False,
+        )
+
+    def val_dataloader(self):
+        print(
+            f"\nLoading validation data, using {self.unrolling_steps} unrolling steps..."
+        )
+        ds = [i for data in self.val_data for i in data]
+        return DataLoader(
+            ds,
+            batch_size=128,
+            shuffle=False,
+            num_workers=8,
+            pin_memory=False,
+        )
+
+    def test_dataloader(self):
+        ds = [i for data in self.test_data for i in data]
+        return DataLoader(
+            ds,
+            batch_size=1,
+            shuffle=False,
+            num_workers=8,
+            pin_memory=False,
+        )

ThermalSolver/graph_module.py

@@ -4,6 +4,8 @@ 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):
@@ -13,13 +15,15 @@ def import_class(class_path: str):
     return cls
 
 
-def _plot_mesh(pos, y, y_pred, batch):
+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))
@@ -36,48 +40,62 @@ def _plot_mesh(pos, y, y_pred, batch):
     plt.colorbar()
     plt.title("Error")
     plt.suptitle("GNO", fontsize=16)
-    plt.savefig("gno.png", dpi=300)
+    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,
+        model_init_args: dict = {},
         loss: torch.nn.Module = None,
-        unrolling_steps: int = 48,
+        curriculum_learning: bool = False,
+        start_iters: int = 10,
+        increase_every: int = 100,
+        increase_rate: float = 1.1,
+        max_iters: int = 1000,
+        accumulation_iters: int = None,
     ):
         super().__init__()
         self.model = import_class(model_class_path)(**model_init_args)
+        self.fd_net = FiniteDifferenceStep()
         self.loss = loss if loss is not None else torch.nn.MSELoss()
-        self.unrolling_steps = unrolling_steps
+        self.curriculum_learning = curriculum_learning
+        self.start_iters = start_iters
+        self.increase_every = increase_every
+        self.increase_rate = increase_rate
+        self.max_iters = max_iters
+        self.current_iters = start_iters
+        self.accumulation_iters = accumulation_iters
+        self.automatic_optimization = False
+        self.threshold = 1e-5
 
-    def forward(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor,
-        edge_index: torch.Tensor,
-        edge_attr: torch.Tensor,
-        unrolling_steps: int = None,
-        boundary_mask: torch.Tensor = None,
-        boundary_values: torch.Tensor = None,
-    ):
-        return self.model(
-            x=x,
-            c=c,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            unrolling_steps=unrolling_steps,
-            boundary_mask=boundary_mask,
-            boundary_values=boundary_values,
-        )
+        self.alpha = torch.nn.Parameter(torch.tensor(0.1))
+
+    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 _preprocess_batch(self, batch: Batch):
-        return batch.x, batch.y, batch.c, batch.edge_index, batch.edge_attr
-
     def _log_loss(self, loss, batch, stage: str):
         self.log(
             f"{stage}/loss",
@@ -89,89 +107,177 @@ class GraphSolver(LightningModule):
         )
         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, deg, boundary_mask, boundary_values
+    ):
+        # with torch.no_grad():
+        # out = self.fd_net(x, edge_index, edge_attr, deg)
+        # out[boundary_mask] = boundary_values.unsqueeze(-1)
+        # diff = out - x
+        # out = self.model(out, edge_index, edge_attr, deg)
+        # out = out + self.alpha * correction
+        # out[boundary_mask] = boundary_values.unsqueeze(-1)
+        out = self.model(x, edge_index, edge_attr, deg)
+        out[boundary_mask] = boundary_values.unsqueeze(-1)
+        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 accumulate_gradients(self, losses):
+        loss_ = torch.stack(losses, dim=0).mean()
+        self.manual_backward(loss_, retain_graph=True)
+        return loss_.item()
+
+    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
+        return x, y, edge_index, edge_attr
+
     def training_step(self, batch: Batch, _):
-        x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred, it = self(
-            x,
-            c,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            unrolling_steps=self.unrolling_steps,
-            boundary_mask=batch.boundary_mask,
-            boundary_values=batch.boundary_values,
+        optim = self.optimizers()
+        optim.zero_grad()
+        x, y, edge_index, edge_attr = self._preprocess_batch(batch)
+
+        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
+        losses = []
+        acc_loss, acc_it = 0, 0
+
+        for i in range(self.current_iters):
+            out = self._compute_model_steps(
+                x,
+                edge_index,
+                edge_attr.unsqueeze(-1),
+                deg,
+                batch.boundary_mask,
+                batch.boundary_values,
+            )
+            losses.append(self.loss(out, y))
+
+            # Accumulate gradients if reached accumulation iters
+            if (
+                self.accumulation_iters is not None
+                and (i + 1) % self.accumulation_iters == 0
+            ):
+                loss = self.accumulate_gradients(losses)
+                losses = []
+                acc_it += 1
+                out = out.detach()
+                acc_loss = acc_loss + loss
+
+            # Check for convergence and break if converged (with final accumulation)
+            converged = self._check_convergence(out, x)
+            if converged:
+                if losses:
+                    loss = self.accumulate_gradients(losses)
+                    acc_it += 1
+                    acc_loss = acc_loss + loss
+                break
+
+            # Final accumulation if we are at the last iteration
+            if i == self.current_iters - 1:
+                if losses:
+                    loss = self.accumulate_gradients(losses)
+                    acc_it += 1
+                    acc_loss = acc_loss + loss
+
+            x = out
+
+        loss = self.loss(out, y)
+        for param in self.model.parameters():
+            if param.grad is not None:
+                param.grad /= acc_it
+        optim.step()
+        optim.zero_grad()
+
+        self.log(
+            "train/accumulated_loss",
+            (acc_loss / acc_it if acc_it > 0 else acc_loss),
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            batch_size=int(batch.num_graphs),
         )
-        loss = self.loss(y_pred, y)
-        boundary_loss = self.loss(
-            y_pred[batch.boundary_mask], y[batch.boundary_mask]
-        )
-        self._log_loss(loss, batch, "train")
-        # self._log_loss(boundary_loss, batch, "train_boundary")
         self.log(
             "train/iterations",
-            it,
+            i + 1,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
             batch_size=int(batch.num_graphs),
         )
-        self.log(
-            "train/param_p",
-            self.model.fd_step.p,
-            on_step=False,
-            on_epoch=True,
-            prog_bar=True,
-            batch_size=int(batch.num_graphs),
-        )
-        # self.log("train/param_a", self.model.fd_step.a, on_step=False, on_epoch=True, prog_bar=True, batch_size=int(batch.num_graphs))
-        return loss
+        if hasattr(self.model, "p"):
+            self.log(
+                "train/p",
+                self.model.p,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                batch_size=int(batch.num_graphs),
+            )
+
+    def on_train_epoch_end(self):
+        if self.curriculum_learning:
+            if (self.current_iters < self.max_iters) and (
+                self.current_epoch % self.increase_every == 0
+            ):
+                self.current_iters = min(
+                    int(self.current_iters * self.increase_rate), self.max_iters
+                )
+        return super().on_train_epoch_end()
 
     def validation_step(self, batch: Batch, _):
-        x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred, it = self(
-            x,
-            c,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            unrolling_steps=self.unrolling_steps,
-            boundary_mask=batch.boundary_mask,
-            boundary_values=batch.boundary_values,
-        )
-        loss = self.loss(y_pred, y)
-        boundary_loss = self.loss(
-            y_pred[batch.boundary_mask], y[batch.boundary_mask]
-        )
+        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,
+                edge_index,
+                edge_attr.unsqueeze(-1),
+                deg,
+                batch.boundary_mask,
+                batch.boundary_values,
+            )
+            converged = self._check_convergence(out, x)
+
+            if converged:
+                break
+            x = out
+        loss = self.loss(out, y)
         self._log_loss(loss, batch, "val")
         self.log(
             "val/iterations",
-            it,
+            i + 1,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
             batch_size=int(batch.num_graphs),
         )
-        return loss
 
-    def test_step(self, batch: Batch, _):
-        x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred, _ = self.model(
-            x=x,
-            c=c,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            unrolling_steps=self.unrolling_steps,
-            batch=batch.batch,
-            pos=batch.pos,
-            boundary_mask=batch.boundary_mask,
-            boundary_values=batch.boundary_values,
-            plot_results=False,
-        )
-        loss = self._compute_loss(y_pred, y)
-        _plot_mesh(batch.pos, y, y_pred, batch.batch)
-        self._log_loss(loss, batch, "test")
-        return loss
+    def test_step(self, batch: Batch, batch_idx):
+        pass
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3)
         return optimizer
 
     def _impose_bc(self, x: torch.Tensor, data: Batch):

ThermalSolver/model/__init__.py

@@ -1,13 +1,13 @@
 __all__ = [
-    "GraphFiniteDifference",
+    # "GraphFiniteDifference",
     "GatingGNO",
-    "LearnableGraphFiniteDifference",
+    # "LearnableGraphFiniteDifference",
     "PointNet",
 ]
 
-from .learnable_finite_difference import (
-    GraphFiniteDifference as LearnableGraphFiniteDifference,
-)
-from .finite_difference import GraphFiniteDifference as GraphFiniteDifference
+# from .learnable_finite_difference import (
+# GraphFiniteDifference as LearnableGraphFiniteDifference,
+# )
+# from .finite_difference import GraphFiniteDifference as GraphFiniteDifference
 from .local_gno import GatingGNO
 from .point_net import PointNet

ThermalSolver/model/diffusion_net.py

@@ -0,0 +1,142 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import MessagePassing
+from torch.nn.utils import spectral_norm
+
+
+class LogPhysEncoder(nn.Module):
+    """
+    Processes 1/dx in log-space to handle multiple scales of geometry
+    (from micro-meshes to macro-meshes) without numerical instability.
+    """
+
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            spectral_norm(nn.Linear(1, hidden_dim)),
+            nn.GELU(),
+            spectral_norm(nn.Linear(hidden_dim, 1)),
+            nn.Softplus(),  # Physical conductance must be positive
+        )
+
+    def forward(self, inv_dx):
+        # We use log(1/dx) to linearize the scale of different geometries
+        log_inv_dx = torch.log(inv_dx + 1e-9)
+        return self.mlp(log_inv_dx)
+
+
+class DiffusionLayer(MessagePassing):
+    """
+    Modella: T_new = T_old + dt * Divergenza(Flusso)
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+
+        self.conductivity_net = nn.Sequential(
+            spectral_norm(nn.Linear(channels, channels, bias=False)),
+            nn.GELU(),
+            spectral_norm(nn.Linear(channels, channels, bias=False)),
+        )
+
+        self.phys_encoder = LogPhysEncoder(hidden_dim=channels)
+
+        self.alpha_param = nn.Parameter(torch.tensor(1e-2))
+
+    @property
+    def alpha(self):
+        return torch.clamp(self.alpha_param, min=1e-7, max=1.0)
+
+    def forward(self, x, edge_index, edge_weight, conductivity):
+        edge_weight = edge_weight.unsqueeze(-1)
+        conductance = self.phys_encoder(edge_weight)
+        net_flux = self.propagate(edge_index, x=x, conductance=conductance)
+        return x + self.alpha * net_flux
+
+    def message(self, x_i, x_j, conductance):
+        delta = x_j - x_i
+        flux = delta * conductance
+        flux = flux + self.conductivity_net(flux)
+        return flux
+
+
+class DiffusionNet(nn.Module):
+    def __init__(
+        self,
+        input_dim=1,
+        output_dim=1,
+        hidden_dim=8,
+        n_layers=4,
+        shared_weights=False,
+    ):
+        super().__init__()
+
+        # Encoder: Projects input temperature to hidden feature space
+        self.enc = nn.Sequential(
+            spectral_norm(nn.Linear(input_dim, hidden_dim, bias=True)),
+            nn.GELU(),
+            spectral_norm(nn.Linear(hidden_dim, hidden_dim, bias=True)),
+        )
+
+        self.scale_x = nn.Parameter(torch.zeros(hidden_dim))
+
+        # Scale parameters for conditioning
+        self.scale_edge_attr = nn.Parameter(torch.zeros(1))
+
+        # If shared_weights is True, use the same DiffusionLayer multiple times
+        if shared_weights:
+            diffusion_layer = DiffusionLayer(hidden_dim)
+            self.layers = torch.nn.ModuleList(
+                [diffusion_layer for _ in range(n_layers)]
+            )
+        # If shared_weights is False, use separate DiffusionLayers
+        else:
+            # Stack of Diffusion Layers
+            self.layers = torch.nn.ModuleList(
+                [DiffusionLayer(hidden_dim) for _ in range(n_layers)]
+            )
+
+        # Decoder: Projects hidden features back to Temperature space
+        self.dec = nn.Sequential(
+            spectral_norm(nn.Linear(hidden_dim, hidden_dim, bias=True)),
+            nn.GELU(),
+            spectral_norm(nn.Linear(hidden_dim, output_dim, bias=True)),
+            nn.Softplus(),  # Ensure positive temperature output
+        )
+
+        self.func = torch.nn.GELU()
+
+        self.dt_param = nn.Parameter(torch.tensor(1e-2))
+
+    @property
+    def dt(self):
+        return torch.clamp(self.dt_param, min=1e-5, max=0.5)
+
+    def forward(self, x, edge_index, edge_attr, conductivity):
+        # 1. Global Residual Connection setup
+        # We save the input to add it back at the very end.
+        # The network learns the correction (Delta T), not the absolute T.
+        x_input = x
+        # 2. Encode
+        h = self.enc(x) * torch.exp(self.scale_x)
+
+        # Scale edge attributes (learnable gating of physical conductivity)
+        w = edge_attr * torch.exp(self.scale_edge_attr)
+
+        # 4. Message Passing (Diffusion Steps)
+        for layer in self.layers:
+            # h is updated internally via residual connection in DiffusionLayer
+            h = layer(h, edge_index, w, conductivity)
+            h = self.func(h)
+
+        # 5. Decode
+        delta_x = self.dec(h)
+
+        # 6. Final Update (Explicit Euler Step)
+        # T_new = T_old + Correction
+        return delta_x + x_input * self.dt
+        # return delta_x

ThermalSolver/model/finite_difference.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from torch_geometric.nn import MessagePassing
+from torch.nn.utils import spectral_norm
 
 
 class FiniteDifferenceStep(MessagePassing):
@@ -8,14 +9,8 @@ class FiniteDifferenceStep(MessagePassing):
     TODO: add docstring.
     """
 
-    def __init__(self, aggr: str = "add", root_weight: float = 1.0):
-        super().__init__(aggr=aggr)
-        assert (
-            aggr == "add"
-        ), "Per somme pesate, l'aggregazione deve essere 'add'."
-        # self.root_weight = float(root_weight)
-        self.p = torch.nn.Parameter(torch.tensor(0.8))
-        self.a = root_weight
+    def __init__(self):
+        super().__init__(aggr="add")
 
     def forward(self, x, edge_index, edge_attr, deg):
         """
@@ -28,8 +23,13 @@ class FiniteDifferenceStep(MessagePassing):
         """
         TODO: add docstring.
         """
-        p = torch.clamp(self.p, 0.0, 1.0)
-        return p * edge_attr.view(-1, 1) * x_j
+        return x_j * edge_attr
+
+    def update(self, aggr_out, _):
+        """
+        TODO: add docstring.
+        """
+        return aggr_out
 
     def aggregate(self, inputs, index, deg):
         """
@@ -38,84 +38,3 @@ class FiniteDifferenceStep(MessagePassing):
         out = super().aggregate(inputs, index)
         deg = deg + 1e-7
         return out / deg.view(-1, 1)
-
-    def update(self, aggr_out, x):
-        """
-        TODO: add docstring.
-        """
-        a = torch.clamp(self.a, 0.0, 1.0)
-        return a * aggr_out + (1 - a) * x
-        # return self.a * aggr_out + (1 - self.a) * x
-
-
-class GraphFiniteDifference(nn.Module):
-    """
-    TODO: add docstring.
-    """
-
-    def __init__(self, max_iters: int = 5000, threshold: float = 1e-4):
-        """
-        TODO: add docstring.
-        """
-        super().__init__()
-        self.max_iters = max_iters
-        self.threshold = threshold
-        self.fd_step = FiniteDifferenceStep(aggr="add", root_weight=1.0)
-
-    @staticmethod
-    def _compute_deg(edge_index, edge_attr, num_nodes):
-        """
-        TODO: add docstring.
-        """
-        deg = torch.zeros(num_nodes, device=edge_index.device)
-        deg = deg.scatter_add(0, edge_index[1], edge_attr)
-        return deg + 1e-7
-
-    @staticmethod
-    def _compute_c_ij(c, edge_index):
-        """
-        TODO: add docstring.
-        """
-        return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
-
-    def forward(
-        self,
-        x,
-        edge_index,
-        edge_attr,
-        c,
-        boundary_mask,
-        boundary_values,
-        **kwargs,
-    ):
-        """
-        TODO: add docstring.
-        """
-        edge_attr = 1 / edge_attr[:, -1]
-        c_ij = self._compute_c_ij(c, edge_index)
-        edge_attr = edge_attr * c_ij
-        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-
-        # Calcola la soglia staccando x dal grafo
-        conv_thres = self.threshold * torch.norm(x.detach())
-
-        for _i in range(self.max_iters):
-            out = self.fd_step(x, edge_index, edge_attr, deg)
-            out[boundary_mask] = boundary_values.unsqueeze(-1)
-
-            # Controllo convergenza senza tracciamento gradienti
-            with torch.no_grad():
-                residual_norm = torch.norm(out - x)
-
-            if residual_norm < conv_thres:
-                break
-
-            # --- OTTIMIZZAZIONE CHIAVE ---
-            # Stacca 'out' dal grafo prima della prossima iterazione
-            # per evitare BPTT e risparmiare memoria.
-            x = out.detach()
-
-        # Il 'out' finale restituito mantiene i gradienti
-        # dell'ULTIMA chiamata a fd_step, permettendo al modello
-        # di apprendere correttamente.
-        return out, _i + 1

ThermalSolver/model/learnable_finite_difference.py

@@ -2,128 +2,211 @@ 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, SAGEConv, GatedGraphConv, GraphConv
+
+# class GCNConvLayer(MessagePassing):
+#     def __init__(
+#         self,
+#         in_channels,
+#         out_channels,
+#         aggr: str = 'mean',
+#         bias: bool = True,
+#         **kwargs,
+#     ):
+#         super().__init__(aggr=aggr, **kwargs)
+
+#         self.in_channels = in_channels
+#         self.out_channels = out_channels
+
+#         if isinstance(in_channels, int):
+#             in_channels = (in_channels, in_channels)
+
+#         self.lin_rel = nn.Linear(in_channels[0], out_channels, bias=bias)
+#         self.lin_root = nn.Linear(in_channels[1], out_channels, bias=False)
+
+#         self.reset_parameters()
+
+#     def reset_parameters(self):
+#         super().reset_parameters()
+#         self.lin_rel.reset_parameters()
+#         self.lin_root.reset_parameters()
 
 
-class FiniteDifferenceStep(MessagePassing):
+#     def forward(self, x, edge_index,
+#                 edge_weight = None, size = None):
+
+#         edge_weight = self.normalize(edge_weight, edge_index, x.size(0), dtype=x.dtype)
+#         out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
+#                              size=size)
+#         out = self.lin_rel(out)
+#         out = out + self.lin_root(x)
+#         return out
+
+#     def message(self, x_j, edge_weight):
+#         return x_j * edge_weight.view(-1, 1)
+
+#     @staticmethod
+#     def normalize(edge_weights, edge_index, num_nodes, dtype=None):
+#         """Symmetrically normalize edge weights."""
+#         if dtype is None:
+#             dtype = edge_weights.dtype
+#         device = edge_index.device
+
+#         row, col = edge_index
+#         deg = torch.zeros(num_nodes, device=device, dtype=dtype)
+#         deg = deg.scatter_add(0, row, edge_weights)
+#         deg_inv_sqrt = deg.pow(-0.5)
+#         deg_inv_sqrt[deg_inv_sqrt == float("inf")] = 0
+#         return deg_inv_sqrt[row] * edge_weights * deg_inv_sqrt[col]
+
+
+# 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=True),
+
+#         self.scale_x = nn.Parameter(torch.zeros(hidden_dim))
+#         self.scale_edge_attr = nn.Parameter(torch.zeros(1))
+#         self.layers = torch.nn.ModuleList(
+#             [GCNConv(hidden_dim, hidden_dim, aggr="mean") for _ in range(n_layers)]
+#         )
+#         self.dec = nn.Linear(hidden_dim, output_dim, bias=True),
+#         self.func = torch.nn.GELU()
+
+#     def forward(self, x, edge_index, edge_attr,):
+#         h = self.enc(x) # * torch.exp(self.scale_x)
+#         edge_attr = edge_attr # * torch.exp(self.scale_edge_attr)
+#         h = self.func(h)
+#         for l in self.layers:
+#             h = l(h, edge_index, edge_attr)
+#             h = self.func(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)
+
+
+# import torch
+# import torch.nn as nn
+# from torch_geometric.nn import MessagePassing
+
+# import torch
+# import torch.nn as nn
+# from torch_geometric.nn import MessagePassing
+
+
+class DiffusionLayer(MessagePassing):
     """
-    TODO: add docstring.
+    Modella: T_new = T_old + dt * Divergenza(Flusso)
     """
 
-    def __init__(self, aggr: str = "add", root_weight: float = 1.0):
-        super().__init__(aggr=aggr)
-        assert (
-            aggr == "add"
-        ), "Per somme pesate, l'aggregazione deve essere 'add'."
-
-        self.correction_net = nn.Sequential(
-            nn.Linear(2, 6),
-            nn.Tanh(),
-            nn.Linear(6, 1),
-            nn.Tanh(),
-        )
-        self.update_net = nn.Sequential(
-            spectral_norm(nn.Linear(1, 6)),
-            nn.Softplus(),
-            spectral_norm(nn.Linear(6, 1)),
-            nn.Softplus(),
-        )
-
-        self.message_net = nn.Sequential(
-            spectral_norm(nn.Linear(1, 6)),
-            nn.Softplus(),
-            spectral_norm(nn.Linear(6, 1)),
-            nn.Softplus(),
-        )
-        self.p = torch.nn.Parameter(torch.tensor(0.5))
-        # self.a = torch.nn.Parameter(torch.tensor(root_weight))
-
-    def forward(self, x, edge_index, edge_attr, deg):
-        """
-        TODO: add docstring.
-        """
-        out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
-        return out
-
-    def message(self, x_j, edge_attr):
-        """
-        TODO: add docstring.
-        """
-        # x_in = torch.cat([x_j, edge_attr.view(-1, 1)], dim=-1)
-        # correction = self.correction_net(x_in)
-        # p = torch.sigmoid(self.p)
-        # return (p * edge_attr.view(-1, 1) + (1 - p) * correction) * x_j
-        return edge_attr.view(-1, 1) * x_j
-
-    def aggregate(self, inputs, index, deg):
-        """
-        TODO: add docstring.
-        """
-        out = super().aggregate(inputs, index)
-        deg = deg + 1e-7
-        return out / deg.view(-1, 1)
-
-    def update(self, aggr_out, x):
-        """
-        TODO: add docstring.
-        """
-        return self.update_net(aggr_out)
-
-
-class GraphFiniteDifference(nn.Module):
-    """
-    TODO: add docstring.
-    """
-
-    def __init__(self, max_iters: int = 5000, threshold: float = 1e-4):
-        """
-        TODO: add docstring.
-        """
-        super().__init__()
-        self.max_iters = max_iters
-        self.threshold = threshold
-        self.fd_step = FiniteDifferenceStep(aggr="add", root_weight=1.0)
-
-    @staticmethod
-    def _compute_deg(edge_index, edge_attr, num_nodes):
-        """
-        TODO: add docstring.
-        """
-        deg = torch.zeros(num_nodes, device=edge_index.device)
-        deg = deg.scatter_add(0, edge_index[1], edge_attr)
-        return deg + 1e-7
-
-    @staticmethod
-    def _compute_c_ij(c, edge_index):
-        """
-        TODO: add docstring.
-        """
-        return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
-
-    def forward(
+    def __init__(
         self,
-        x,
-        edge_index,
-        edge_attr,
-        c,
-        boundary_mask,
-        boundary_values,
+        channels: int,
         **kwargs,
     ):
-        """
-        TODO: add docstring.
-        """
-        edge_attr = 1 / edge_attr[:, -1]
-        c_ij = self._compute_c_ij(c, edge_index)
-        edge_attr = edge_attr * c_ij
-        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-        conv_thres = self.threshold * torch.norm(x.detach())
 
-        for _i in range(self.max_iters):
-            out = self.fd_step(x, edge_index, edge_attr, deg)
-            out[boundary_mask] = boundary_values.unsqueeze(-1)
-            with torch.no_grad():
-                residual_norm = torch.norm(out - x)
-            if residual_norm < conv_thres:
-                break
-            x = out.detach()
-        return out, _i + 1
+        super().__init__(aggr="add", **kwargs)
+
+        self.dt = nn.Parameter(torch.tensor(1e-4))
+        self.conductivity_net = nn.Sequential(
+            nn.Linear(channels, channels, bias=False),
+            nn.GELU(),
+            nn.Linear(channels, channels, bias=False),
+        )
+
+        self.phys_encoder = nn.Sequential(
+            nn.Linear(1, 8, bias=False),
+            nn.Tanh(),
+            nn.Linear(8, 1, bias=False),
+            nn.Softplus(),
+        )
+
+    def forward(self, x, edge_index, edge_weight):
+        edge_weight = edge_weight.unsqueeze(-1)
+        conductance = self.phys_encoder(edge_weight)
+        net_flux = self.propagate(edge_index, x=x, conductance=conductance)
+        return x + (net_flux * self.dt)
+
+    def message(self, x_i, x_j, conductance):
+        delta = x_j - x_i
+        flux = delta * conductance
+        flux = flux + self.conductivity_net(flux)
+        return flux
+
+
+class CorrectionNet(nn.Module):
+    def __init__(self, input_dim=1, output_dim=1, hidden_dim=32, n_layers=4):
+        super().__init__()
+
+        # Encoder: Projects input temperature to hidden feature space
+        self.enc = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim, bias=True),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim, bias=True),
+            nn.GELU(),
+        )
+
+        self.scale_x = nn.Parameter(torch.zeros(hidden_dim))
+
+        # Scale parameters for conditioning
+        self.scale_edge_attr = nn.Parameter(torch.zeros(1))
+
+        # Stack of Diffusion Layers
+        self.layers = torch.nn.ModuleList(
+            [DiffusionLayer(hidden_dim) for _ in range(n_layers)]
+        )
+
+        # Decoder: Projects hidden features back to Temperature space
+        self.dec = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim, bias=True),
+            nn.GELU(),
+            nn.Linear(hidden_dim, output_dim, bias=True),
+            nn.Softplus(),  # Ensure positive temperature output
+        )
+
+        self.func = torch.nn.GELU()
+
+    def forward(self, x, edge_index, edge_attr):
+        # 1. Global Residual Connection setup
+        # We save the input to add it back at the very end.
+        # The network learns the correction (Delta T), not the absolute T.
+        x_input = x
+
+        # 2. Encode
+        h = self.enc(x) * torch.exp(self.scale_x)
+
+        # Scale edge attributes (learnable gating of physical conductivity)
+        w = edge_attr * torch.exp(self.scale_edge_attr)
+
+        # 4. Message Passing (Diffusion Steps)
+        for layer in self.layers:
+            # h is updated internally via residual connection in DiffusionLayer
+            h = layer(h, edge_index, w)
+            h = self.func(h)
+
+        # 5. Decode
+        delta_x = self.dec(h)
+
+        # 6. Final Update (Explicit Euler Step)
+        # T_new = T_old + Correction
+        # return x_input + delta_x
+        return delta_x

ThermalSolver/switch_dataloader_callback.py

@@ -0,0 +1,104 @@
+import torch
+from lightning.pytorch.callbacks import Callback
+import os
+
+
+class SwitchDataLoaderCallback(Callback):
+    def __init__(
+        self,
+        ckpt_path,
+        increase_unrolling_steps_by,
+        increase_unrolling_steps_every,
+        max_unrolling_steps=10,
+        patience=None,
+        last_patience=None,
+        metric="val/loss",
+    ):
+        super().__init__()
+        self.ckpt_path = ckpt_path
+        if os.path.exists(ckpt_path) is False:
+            os.makedirs(ckpt_path)
+        self.increase_unrolling_steps_by = increase_unrolling_steps_by
+        self.increase_unrolling_steps_every = increase_unrolling_steps_every
+        self.max_unrolling_steps = max_unrolling_steps
+        self.metric = metric
+        self.actual_loss = torch.inf
+        if patience is not None:
+            self.patience = patience
+        if last_patience is not None:
+            self.last_patience = last_patience
+        self.no_improvement_epochs = 0
+        self.last_step_reached = False
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        self._metric_tracker(trainer, pl_module)
+        if self.last_step_reached is False:
+            self._unrolling_steps_handler(pl_module, trainer)
+        else:
+            if self.no_improvement_epochs >= self.last_patience:
+                trainer.should_stop = True
+
+    def _metric_tracker(self, trainer, pl_module):
+        if trainer.callback_metrics.get(self.metric) < self.actual_loss:
+            self.actual_loss = trainer.callback_metrics.get(self.metric)
+            self._save_model(pl_module, trainer)
+            self.no_improvement_epochs = 0
+            print(f"\nNew best {self.metric}: {self.actual_loss:.4f}")
+        else:
+            self.no_improvement_epochs += 1
+            print(
+                f"\nNo improvement in {self.metric} for {self.no_improvement_epochs} epochs."
+            )
+
+    def _should_reload_dataloader(self, trainer):
+        if self.patience is not None:
+            print(
+                f"Checking patience: {self.no_improvement_epochs} / {self.patience}"
+            )
+            if self.no_improvement_epochs >= self.patience:
+                return True
+        elif (
+            trainer.current_epoch + 1 % self.increase_unrolling_steps_every == 0
+        ):
+            print("Reached scheduled epoch for increasing unrolling steps.")
+            return True
+        return False
+
+    def _unrolling_steps_handler(self, pl_module, trainer):
+        if self._should_reload_dataloader(trainer):
+            self._load_model(pl_module)
+            if pl_module.unrolling_steps >= self.max_unrolling_steps:
+                return
+            pl_module.unrolling_steps += self.increase_unrolling_steps_by
+            trainer.datamodule.unrolling_steps = pl_module.unrolling_steps
+            print(f"Incremented unrolling steps to {pl_module.unrolling_steps}")
+            trainer.datamodule.setup(stage="fit")
+            trainer.manual_dataloader_reload()
+            self.actual_loss = torch.inf
+            if pl_module.unrolling_steps >= self.max_unrolling_steps:
+                print(
+                    "Reached max unrolling steps. Stopping further increments."
+                )
+                self.last_step_reached = True
+
+    def _save_model(self, pl_module, trainer):
+        pt_path = os.path.join(
+            self.ckpt_path,
+            f"{pl_module.unrolling_steps}_unrolling_best_model.pt",
+        )
+        torch.save(pl_module.state_dict(), pt_path)  # <--- CHANGED THIS
+        ckpt_path = os.path.join(
+            self.ckpt_path,
+            f"{pl_module.unrolling_steps}_unrolling_best_checkpoint.ckpt",
+        )
+        trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+    def _load_model(self, pl_module):
+        pt_path = os.path.join(
+            self.ckpt_path,
+            f"{pl_module.unrolling_steps}_unrolling_best_model.pt",
+        )
+        pl_module.load_state_dict(torch.load(pt_path, weights_only=True))
+        print(
+            f"Loaded model weights from {pt_path} for unrolling steps = {pl_module.unrolling_steps}"
+        )

experiments/10_steps/config_16_layer_16_hidden_adaptive_refined.yaml

@@ -0,0 +1,72 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.adaptive_refined"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 15
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.adaptive_refined/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "3_stripes.basic.1_adaptive_refined"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+optimizer: null
+lr_scheduler: null

experiments/10_steps/config_16_layer_16_hidden_adaptive_refined_combined.yaml

@@ -0,0 +1,63 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.adaptive_refined.combined"
+  callbacks:
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 10
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.adaptive_refined.combined/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: 
+      - "4_stripes.basic.1_adaptive_refined"
+      - "3_stripes.basic.1_adaptive_refined"
+      - "2_stripes.basic.1_adaptive_refined"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+    
+optimizer: null
+lr_scheduler: null

experiments/10_steps/config_16_layer_16_hidden_refined.yaml

@@ -0,0 +1,72 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.refined"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 15
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.refined/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "3_stripes.basic.refined"
+    n_elements: 50
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+optimizer: null
+lr_scheduler: null

experiments/10_steps/config_16_layer_16_hidden_refined_combined.yaml

@@ -0,0 +1,77 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.refined.combined"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 10
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.refined.combined/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 10
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: 
+      - "4_stripes.basic.refined"
+      - "3_stripes.basic.refined"
+      - "2_stripes.basic.refined"
+    n_elements: 75
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 10
+    min_normalized_diff: 1e-4
+    
+optimizer: null
+lr_scheduler: null
+ckpt_path: logs.autoregressive.wandb/10_steps/basic.refined.combined/16_layer_16_hidden/6_unrolling_best_checkpoint.ckpt

experiments/10_steps/config_16_layer_16_hidden_refined_star.yaml

@@ -0,0 +1,72 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.refined.star"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 15
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.refined.star/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "3_stripes.star"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+optimizer: null
+lr_scheduler: null

experiments/10_steps/config_16_layer_16_hidden_star.yaml

@@ -0,0 +1,72 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.star"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 15
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.star/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "3_stripes.star.refined"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+optimizer: null
+lr_scheduler: null

experiments/10_steps/config_16_layer_16_hidden_star_combined.yaml

@@ -0,0 +1,75 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "16_layer_16_hidden.star.combined"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 15
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.star.combined/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: 
+      - "4_stripes.star"
+      - "3_stripes.star"
+      - "2_stripes.star"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+optimizer: null
+lr_scheduler: null

experiments/10_steps/config_8_layer_16_hidden_star.yaml

@@ -0,0 +1,72 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-10.steps"
+        name: "8_layer_16_hidden.star"
+  callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+      init_args:
+        increase_unrolling_steps_by: 4
+        patience: 5
+        last_patience: 15
+        max_unrolling_steps: 10
+        ckpt_path: logs.autoregressive.wandb/10_steps/basic.star/8_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 8
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "3_stripes.star.refined"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+    min_normalized_diff: 1e-4
+optimizer: null
+lr_scheduler: null

experiments/10_steps/test_config_16_layer_16_hidden_star.yaml

@@ -0,0 +1,76 @@
+# lightning.pytorch==2.5.5
+seed_everything: 1999
+trainer:
+  accelerator: cpu
+  strategy: auto
+  devices: 1
+  num_nodes: 1
+  precision: null
+  # logger:
+  #   - class_path: lightning.pytorch.loggers.WandbLogger
+  #     init_args:
+  #       save_dir: logs.autoregressive.wandb
+  #       project: "thermal-conduction-unsteady-10.steps"
+  #       name: "16_layer_16_hidden.adaptive_refined.combined"
+  # callbacks:
+    # - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+    #   init_args:
+    #     dirpath: logs.autoregressive.wandb/16_refined.10_steps/checkpoints
+    #     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: 30
+    #     verbose: false
+    # - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+    #   init_args:
+    #     increase_unrolling_steps_by: 4
+    #     patience: 15
+    #     last_patience: 20
+    #     max_unrolling_steps: 10
+    #     ckpt_path: logs.autoregressive.wandb/10_steps/basic.adaptive_refined.combined/16_layer_16_hidden/
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 2
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: 
+      # - "2_stripes.basic.refined"
+      # - "3_stripes.basic.refined"
+      # - "4_stripes.basic.1_adaptive_refined"
+      - "3_stripes.star"
+    n_elements: 50
+    batch_size: 32
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 2
+optimizer: null
+lr_scheduler: null
+ckpt_path: /home/folivo/storage/thermal-conduction-ml/logs.autoregressive.wandb/10_steps/basic.star.combined/16_layer_16_hidden/10_unrolling_best_checkpoint.ckpt

experiments/5_steps/config_16_layer_16_hidden_refined.yaml

@@ -0,0 +1,71 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-5.steps"
+        name: "16_layer_16_hidden"
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        dirpath: logs.autoregressive.wandb/16_refined/checkpoints
+        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: 30
+        verbose: false
+    # - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+    #   init_args:
+    #     increase_unrolling_steps_by: 5
+    #     patience: 10
+    #     last_patience: 15
+    #     max_unrolling_steps: 20
+    #     ckpt_path: logs.autoregressive.wandb/16_16_refined/checkpoints
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 16
+    unrolling_steps: 5
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "basic.refined"
+    n_elements: 100
+    batch_size: 32
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 5
+optimizer: null
+lr_scheduler: null

experiments/5_steps/config_32_layer_16_hidden_refined.yaml

@@ -0,0 +1,71 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-5.steps"
+        name: "32_layer_16_hidden"
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        dirpath: logs.autoregressive.wandb/32_refined/checkpoints
+        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: 30
+        verbose: false
+    # - class_path: ThermalSolver.switch_dataloader_callback.SwitchDataLoaderCallback
+    #   init_args:
+    #     increase_unrolling_steps_by: 5
+    #     patience: 10
+    #     last_patience: 15
+    #     max_unrolling_steps: 20
+    #     ckpt_path: logs.autoregressive.wandb/16_16_refined/checkpoints
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: 0
+  accumulate_grad_batches: 1
+  default_root_dir: null
+  gradient_clip_val: 1.0
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 32
+    unrolling_steps: 5
+    
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "basic.refined"
+    n_elements: 100
+    batch_size: 24
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 5
+optimizer: null
+lr_scheduler: null

experiments/5_steps/config_8_layer_16_hidden_refined.yaml

@@ -0,0 +1,63 @@
+# 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.WandbLogger
+      init_args:
+        save_dir: logs.autoregressive.wandb
+        project: "thermal-conduction-unsteady-5.steps"
+        name: "8_layer_16_hidden"
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        dirpath: logs.autoregressive.wandb/8_refined/checkpoints
+        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: 30
+        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
+  default_root_dir: null
+
+model: 
+  class_path: ThermalSolver.autoregressive_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.diffusion_net.DiffusionNet
+    model_init_args: 
+      input_dim: 1
+      hidden_dim: 16
+      output_dim: 1
+      n_layers: 8
+    unrolling_steps: 5
+
+data: 
+  class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
+    split_name: "basic.refined"
+    n_elements: 100
+    batch_size: 32
+    train_size: 0.7
+    val_size: 0.2
+    test_size: 0.1
+    build_radial_graph: false
+    remove_boundary_edges: true
+    unrolling_steps: 5
+optimizer: null
+lr_scheduler: null

experiments/config.yaml

@@ -1,56 +0,0 @@
-# lightning.pytorch==2.5.5
-seed_everything: 1999
-trainer:
-  accelerator: gpu
-  strategy: auto
-  devices: 1
-  num_nodes: 1
-  precision: null
-  logger: 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.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 1
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 64
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 10
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-ckpt_path: null

experiments/config_01.yaml

@@ -1,61 +0,0 @@
-# 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: "01"
-        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: 6
-  # gradient_clip_val: 1.0
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 1
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 64
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 4
-    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

experiments/config_02.yaml

@@ -1,62 +0,0 @@
-# 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: "02"
-        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: 6
-  gradient_clip_val: 1.0
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 2
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 32
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 4
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-# ckpt_path: lightning_logs/version_15/checkpoints/best-checkpoint.ckpt
-ckpt_path: null

experiments/config_04.yaml

@@ -1,62 +0,0 @@
-# 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: "04"
-        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: 6
-  gradient_clip_val: 1.0
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 4
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 16
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 4
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-# ckpt_path: lightning_logs/version_15/checkpoints/best-checkpoint.ckpt
-ckpt_path: null

experiments/config_08.yaml

@@ -1,62 +0,0 @@
-# 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: "08"
-        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: 6
-  gradient_clip_val: 1.0
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 8
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 8
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 4
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-# ckpt_path: lightning_logs/version_15/checkpoints/best-checkpoint.ckpt
-ckpt_path: null

experiments/config_16.yaml

@@ -1,62 +0,0 @@
-# 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: "16"
-        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: 6
-  gradient_clip_val: 1.0
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 16
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 4
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 4
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-# ckpt_path: lightning_logs/version_15/checkpoints/best-checkpoint.ckpt
-ckpt_path: null

experiments/config_32.yaml

@@ -1,62 +0,0 @@
-# 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: "32"
-        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: 6
-  gradient_clip_val: 1.0
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 32
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 2
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000_ref_1"
-    batch_size: 4
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-# ckpt_path: lightning_logs/version_15/checkpoints/best-checkpoint.ckpt
-ckpt_path: null

experiments/config_64.yaml

@@ -1,63 +0,0 @@
-# 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: "64"
-        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: 6
-  gradient_clip_val: 1.0
-
-model: 
-  class_path: ThermalSolver.module.GraphSolver
-  init_args:
-    model_class_path: ThermalSolver.model.local_gno.GatingGNO
-    model_init_args:
-      x_ch_node: 1 
-      f_ch_node: 1 
-      hidden: 16
-      layers: 64
-      edge_ch: 3 
-      out_ch: 1
-    unrolling_steps: 1
-data: 
-  class_path: ThermalSolver.data_module.GraphDataModule
-  init_args:
-    hf_repo: "SISSAmathLab/thermal-conduction"
-    split_name: "2000"
-    batch_size: 4
-    train_size: 0.8
-    test_size: 0.1
-    test_size: 0.1
-optimizer: null
-lr_scheduler: null
-# ckpt_path: lightning_logs/64/version_0/checkpoints/best-checkpoint.ckpt
-ckpt_path: null

run.py

@@ -5,7 +5,11 @@ torch.set_float32_matmul_precision("medium")
 
 
 def main():
-    LightningCLI(subclass_mode_data=True, subclass_mode_model=True)
+    LightningCLI(
+        subclass_mode_data=True,
+        subclass_mode_model=True,
+        save_config_kwargs={"overwrite": True},
+    )
 
 
 if __name__ == "__main__":

submit.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+export CUDA_VISIBLE_DEVICES=1
+python run.py fit --config experiments/5_steps/config_16_layer_16_hidden_refined.yaml
+python run.py fit --config experiments/5_steps/config_32_layer_16_hidden_refined.yaml
+python run.py fit --config experiments/5_steps/config_8_layer_16_hidden_refined.yaml