random changes

ThermalSolver/model/finite_difference.py

@@ -9,32 +9,27 @@ class FiniteDifferenceStep(MessagePassing):
     TODO: add docstring.
     """
 
-    def __init__(
-        self,
-        aggr: str = "add",
-        normalize: bool = True,
-        root_weight: float = 1.0,
-    ):
+    def __init__(self, aggr: str = "add", root_weight: float = 1.0):
         super().__init__(aggr=aggr)
-
-        self.normalize = normalize
         assert (
             aggr == "add"
         ), "Per somme pesate, l'aggregazione deve essere 'add'."
         self.root_weight = float(root_weight)
 
-    def forward(self, x, edge_index, edge_weight, deg):
+    def forward(self, x, edge_index, edge_attr, deg, weight=1.0):
         """
         TODO: add docstring.
         """
-        out = self.propagate(edge_index, x=x, edge_weight=edge_weight, deg=deg)
+        out = self.propagate(
+            edge_index, x=x, edge_attr=edge_attr, deg=deg, weight=weight
+        )
         return out
 
-    def message(self, x_j, edge_weight):
+    def message(self, x_j, edge_attr):
         """
         TODO: add docstring.
         """
-        return edge_weight.view(-1, 1) * x_j
+        return edge_attr.view(-1, 1) * x_j
 
     def aggregate(self, inputs, index, deg):
         """
@@ -44,11 +39,12 @@ class FiniteDifferenceStep(MessagePassing):
         deg = deg + 1e-7
         return out / deg.view(-1, 1)
 
-    def update(self, aggr_out, x):
+    def update(self, aggr_out, x, weight):
         """
         TODO: add docstring.
         """
-        return self.root_weight * aggr_out + (1 - self.root_weight) * x
+        print(weight)
+        return weight * aggr_out + (1 - weight) * x
 
 
 class GraphFiniteDifference(nn.Module):
@@ -56,24 +52,22 @@ class GraphFiniteDifference(nn.Module):
     TODO: add docstring.
     """
 
-    def __init__(self, max_iters: int = 1000, threshold: float = 1e-4):
+    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", normalize=True, root_weight=1.0
-        )
+        self.fd_step = FiniteDifferenceStep(aggr="add", root_weight=1.0)
 
     @staticmethod
-    def _compute_deg(edge_index, edge_weight, num_nodes):
+    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_weight)
+        deg = deg.scatter_add(0, edge_index[1], edge_attr)
         return deg + 1e-7
 
     @staticmethod
@@ -84,19 +78,29 @@ class GraphFiniteDifference(nn.Module):
         return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
 
     def forward(
-        self, x, edge_index, edge_weight, c, boundary_mask, boundary_values
+        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_weight = edge_weight * c_ij
-        deg = self._compute_deg(edge_index, edge_weight, x.size(0))
+        edge_attr = edge_attr * c_ij
+        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
         conv_thres = self.threshold * torch.norm(x)
-        for _i in tqdm(range(self.max_iters)):
-            out = self.fd_step(x, edge_index, edge_weight, deg)
+        weight = 1.0
+        for _i in range(self.max_iters):
+            out = self.fd_step(x, edge_index, edge_attr, deg, weight=weight)
+            weight = weight * 0.9999
             out[boundary_mask] = boundary_values.unsqueeze(-1)
             if torch.norm(out - x) < conv_thres:
                 break
             x = out
-        return out
+        return out, _i + 1