random changes

ThermalSolver/graph_module.py

@@ -2,6 +2,8 @@ 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
 
 
 def import_class(class_path: str):
@@ -11,6 +13,32 @@ def import_class(class_path: str):
     return cls
 
 
+def _plot_mesh(pos, y, y_pred, batch):
+
+    idx = batch == 0
+    y = y[idx].detach().cpu()
+    y_pred = y_pred[idx].detach().cpu()
+    pos = pos[idx].detach().cpu()
+
+    pos = pos.detach().cpu()
+    tria = Triangulation(pos[:, 0], pos[:, 1])
+    plt.figure(figsize=(18, 5))
+    plt.subplot(1, 3, 1)
+    plt.tricontourf(tria, y.squeeze().numpy(), levels=14)
+    plt.colorbar()
+    plt.title("True temperature")
+    plt.subplot(1, 3, 2)
+    plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=14)
+    plt.colorbar()
+    plt.title("Predicted temperature")
+    plt.subplot(1, 3, 3)
+    plt.tricontourf(tria, torch.abs(y_pred - y).squeeze().numpy(), levels=14)
+    plt.colorbar()
+    plt.title("Error")
+    plt.suptitle("GNO", fontsize=16)
+    plt.savefig("gno.png", dpi=300)
+
+
 class GraphSolver(LightningModule):
     def __init__(
         self,
@@ -32,14 +60,16 @@ class GraphSolver(LightningModule):
         edge_attr: torch.Tensor,
         unrolling_steps: int = None,
         boundary_mask: torch.Tensor = None,
+        boundary_values: torch.Tensor = None,
     ):
         return self.model(
-            x,
-            c,
-            edge_index,
-            edge_attr,
-            unrolling_steps,
+            x=x,
+            c=c,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            unrolling_steps=unrolling_steps,
             boundary_mask=boundary_mask,
+            boundary_values=boundary_values,
         )
 
     def _compute_loss(self, x, y):
@@ -61,52 +91,82 @@ class GraphSolver(LightningModule):
 
     def training_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred = self(
+        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]
         )
         self._log_loss(loss, batch, "train")
-        self._log_loss(boundary_loss, batch, "train_boundary")
+        # self._log_loss(boundary_loss, batch, "train_boundary")
+        self.log(
+            "train/iterations",
+            it,
+            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
 
     def validation_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred = self(
+        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]
         )
         self._log_loss(loss, batch, "val")
-        self._log_loss(boundary_loss, batch, "val_boundary")
+        self.log(
+            "val/iterations",
+            it,
+            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,
-            c,
+        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,
-            plot_results=True,
+            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