52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
import torch
|
|
from torch_geometric.data import Data
|
|
|
|
D_IN_KEYS = "x"
|
|
D_ATTR_KEYS = ["c", "edge_attr"]
|
|
D_OUT_KEY = "y"
|
|
D_KEYS = [D_IN_KEYS] + [D_OUT_KEY] + D_ATTR_KEYS
|
|
D_BOUNDS_KEYS = "boundary_temperatures"
|
|
|
|
|
|
class Normalizer:
|
|
def __init__(self, data):
|
|
self.mean, self.std = self._compute_stats(data)
|
|
|
|
def _compute_stats(self, data: list[Data]) -> tuple[dict, dict]:
|
|
mean = {}
|
|
std = {}
|
|
for key in D_KEYS:
|
|
tmp = torch.empty(0)
|
|
for d in data:
|
|
if not hasattr(d, key):
|
|
raise AttributeError(f"Manca '{key}' in uno dei Data.")
|
|
if tmp.numel() == 0:
|
|
tmp = d[key]
|
|
else:
|
|
tmp = torch.cat([tmp, d[key]], dim=0)
|
|
mean[key] = tmp.mean(dim=0, keepdim=True)
|
|
std[key] = tmp.std(dim=0, keepdim=True) + 1e-6
|
|
return mean, std
|
|
|
|
def normalize(self, data):
|
|
for d in data:
|
|
for key in D_KEYS:
|
|
if not hasattr(d, key):
|
|
raise AttributeError(f"Manca '{key}' in uno dei Data.")
|
|
d[key] = (d[key] - self.mean[key]) / self.std[key]
|
|
self._recompute_boundary_temperatures(data)
|
|
|
|
def _recompute_boundary_temperatures(self, data):
|
|
for d in data:
|
|
bottom_bc = d.y[d.bottom_boundary_ids].median()
|
|
top_bc = d.y[d.top_boundary_ids].median()
|
|
left_bc = d.y[d.left_boundary_ids].median()
|
|
right_bc = d.y[d.right_boundary_ids].median()
|
|
boundaries_temperatures = torch.tensor(
|
|
[bottom_bc, right_bc, top_bc, left_bc], dtype=torch.float32
|
|
)
|
|
d.boundary_temperatures = boundaries_temperatures.unsqueeze(0)
|
|
|
|
def denormalize(self, y: torch.tensor):
|
|
return y * self.std[D_OUT_KEY] + self.mean[D_OUT_KEY]
|