improve unrolling

ThermalSolver/module.py

@@ -1,6 +1,20 @@
 import torch
 from lightning import LightningModule
 from torch_geometric.data import Batch
+from matplotlib.tri import Triangulation
+
+
+# def plot_results(x, pos, step, i, batch):
+#     x = x[batch == 0]
+#     pos = pos[batch == 0]
+#     tria = Triangulation(pos[:, 0].cpu(), pos[:, 1].cpu())
+#     import matplotlib.pyplot as plt
+
+#     plt.tricontourf(tria, x[:, 0].cpu(), levels=14)
+#     plt.colorbar()
+#     plt.savefig(f"{step:03d}_out_{i:03d}.png")
+#     plt.axis("equal")
+#     plt.close()
 
 
 class GraphSolver(LightningModule):
@@ -8,7 +22,7 @@ class GraphSolver(LightningModule):
         self,
         model: torch.nn.Module,
         loss: torch.nn.Module = None,
-        unrolling_steps: int = 10,
+        unrolling_steps: int = 48,
     ):
         super().__init__()
         self.model = model
@@ -19,13 +33,21 @@ class GraphSolver(LightningModule):
         self,
         x: torch.Tensor,
         c: torch.Tensor,
+        boundary: torch.Tensor,
+        boundary_mask: torch.Tensor,
         edge_index: torch.Tensor,
         edge_attr: torch.Tensor,
+        unrolling_steps: int = None,
     ):
-        return self.model(x, c, edge_index, edge_attr)
-
-    def _compute_loss_train(self, x, x_prev, y):
-        return self.loss(x, y) + self.loss(x, x_prev)
+        return self.model(
+            x,
+            c,
+            boundary,
+            boundary_mask,
+            edge_index,
+            edge_attr,
+            unrolling_steps,
+        )
 
     def _compute_loss(self, x, y):
         return self.loss(x, y)
@@ -46,35 +68,55 @@ class GraphSolver(LightningModule):
 
     def training_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        loss = 0.0
-        for _ in range(self.unrolling_steps):
-            x_prev = x.detach()
-            x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
-            actual_loss = self.loss(x, y)
-            loss += actual_loss
-            print(f"Train step loss: {actual_loss.item()}")
-
+        # x = self._impose_bc(x, batch)
+        # for _ in range(self.unrolling_steps):
+        y_pred = self(
+            x,
+            c,
+            batch.boundary_values,
+            batch.boundary_mask,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            unrolling_steps=self.unrolling_steps,
+        )
+        # x = self._impose_bc(x, batch)
+        loss = self.loss(y_pred, y)
         self._log_loss(loss, batch, "train")
         return loss
 
     def validation_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        for _ in range(self.unrolling_steps):
-            x_prev = x.detach()
-            x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
-            loss = self.loss(x, x_prev)
-            if loss < 1e-5:
-                break
-        loss = self._compute_loss(x, y)
+        y_pred = self(
+            x,
+            c,
+            batch.boundary_values,
+            batch.boundary_mask,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            unrolling_steps=self.unrolling_steps,
+        )
+        loss = self.loss(y_pred, y)
         self._log_loss(loss, batch, "val")
         return loss
 
     def test_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        for _ in range(self.unrolling_steps):
-            x_prev = x.detach()
-            x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
-        loss = self._compute_loss(x, y)
+        # for _ in range(self.unrolling_steps):
+        y_pred = self.model(
+            x,
+            c,
+            batch.boundary_values,
+            batch.boundary_mask,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            unrolling_steps=self.unrolling_steps,
+            plot_results=True,
+            batch=batch.batch,
+            pos=batch.pos,
+        )
+        # x = self._impose_bc(x, batch)
+        # plot_results(x, batch.pos, self.global_step, _, batch.batch)
+        loss = self._compute_loss(y_pred, y)
         self._log_loss(loss, batch, "test")
         return loss
 
@@ -82,6 +124,6 @@ class GraphSolver(LightningModule):
         optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
         return optimizer
 
-    def scale_bc(self, data: Batch, y: torch.Tensor):
-        t = data.boundary_temperature[data.batch]
-        return y * t
+    def _impose_bc(self, x: torch.Tensor, data: Batch):
+        x[data.boundary_mask] = data.boundary_values
+        return x