implement new model

ThermalSolver/module.py

@@ -1,31 +1,26 @@
 import torch
 from lightning import LightningModule
 from torch_geometric.data import Batch
-from matplotlib.tri import Triangulation
+import importlib
 
 
-# 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()
+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
 
 
 class GraphSolver(LightningModule):
     def __init__(
         self,
-        model: torch.nn.Module,
+        model_class_path: str,
+        model_init_args: dict,
         loss: torch.nn.Module = None,
         unrolling_steps: int = 48,
     ):
         super().__init__()
-        self.model = model
+        self.model = import_class(model_class_path)(**model_init_args)
         self.loss = loss if loss is not None else torch.nn.MSELoss()
         self.unrolling_steps = unrolling_steps
 
@@ -57,7 +52,7 @@ class GraphSolver(LightningModule):
 
     def _log_loss(self, loss, batch, stage: str):
         self.log(
-            f"{stage}_loss",
+            f"{stage}/loss",
             loss,
             on_step=False,
             on_epoch=True,
@@ -68,8 +63,6 @@ class GraphSolver(LightningModule):
 
     def training_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        # x = self._impose_bc(x, batch)
-        # for _ in range(self.unrolling_steps):
         y_pred = self(
             x,
             c,
@@ -79,9 +72,12 @@ class GraphSolver(LightningModule):
             edge_attr=edge_attr,
             unrolling_steps=self.unrolling_steps,
         )
-        # x = self._impose_bc(x, batch)
         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")
         return loss
 
     def validation_step(self, batch: Batch, _):
@@ -96,12 +92,15 @@ class GraphSolver(LightningModule):
             unrolling_steps=self.unrolling_steps,
         )
         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")
         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):
         y_pred = self.model(
             x,
             c,
@@ -114,8 +113,6 @@ class GraphSolver(LightningModule):
             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