fix model and datamodule

This commit is contained in:
Filippo Olivo
2025-12-15 09:08:21 +01:00
parent 3cc1d230e4
commit a9d56a3ed9
2 changed files with 69 additions and 50 deletions

View File

@@ -17,7 +17,7 @@ def import_class(class_path: str):
def _plot_mesh(pos_, y_, y_pred_, y_true_, batch, i, batch_idx):
# print(pos_.shape, y_.shape, y_pred_.shape, y_true_.shape)
for j in [0]:
for j in [0, 5, 10, 20]:
idx = (batch == j).nonzero(as_tuple=True)[0]
y = y_[idx].detach().cpu()
y_pred = y_pred_[idx].detach().cpu()
@@ -38,39 +38,37 @@ def _plot_mesh(pos_, y_, y_pred_, y_true_, batch, i, batch_idx):
# plt.savefig("test_scatter_step_before.png", dpi=72)
# x = z
plt.subplot(1, 3, 1)
# plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=100)
plt.scatter(
pos[:, 0],
pos[:, 1],
c=y_pred.squeeze().numpy(),
s=20,
cmap="viridis",
)
plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=100)
# plt.scatter(
# pos[:, 0],
# pos[:, 1],
# c=y_pred.squeeze().numpy(),
# s=20,
# cmap="viridis",
# )
plt.colorbar()
plt.title("Step t Predicted")
plt.subplot(1, 3, 2)
# plt.tricontourf(tria, y_true.squeeze().numpy(), levels=100)
plt.scatter(
pos[:, 0],
pos[:, 1],
c=y_true.squeeze().numpy(),
s=20,
cmap="viridis",
)
plt.tricontourf(tria, y_true.squeeze().numpy(), levels=100)
# plt.scatter(
# pos[:, 0],
# pos[:, 1],
# c=y_true.squeeze().numpy(),
# s=20,
# cmap="viridis",
# )
plt.colorbar()
plt.title("t True")
plt.subplot(1, 3, 3)
per_element_relative_error = torch.abs(y_pred - y_true) / torch.clamp(
torch.abs(y_true), min=1e-6
)
# plt.tricontourf(tria, per_element_relative_error.squeeze(), levels=100)
plt.scatter(
pos[:, 0],
pos[:, 1],
c=per_element_relative_error.squeeze().numpy(),
s=20,
cmap="viridis",
)
per_element_relative_error = torch.abs(y_pred - y_true)
plt.tricontourf(tria, per_element_relative_error.squeeze(), levels=100)
# plt.scatter(
# pos[:, 0],
# pos[:, 1],
# c=per_element_relative_error.squeeze().numpy(),
# s=20,
# cmap="viridis",
# )
plt.colorbar()
plt.title("Relative Error")
plt.suptitle("GNO", fontsize=16)
@@ -216,20 +214,20 @@ class GraphSolver(LightningModule):
batch.boundary_values,
conductivity,
)
if (
batch_idx == 0
and self.current_epoch % 10 == 0
and self.current_epoch > 0
):
_plot_mesh(
batch.pos,
x,
out,
y[:, i, :],
batch.batch,
i,
self.current_epoch,
)
# if (
# batch_idx == 0
# and self.current_epoch % 10 == 0
# and self.current_epoch > 0
# ):
# _plot_mesh(
# batch.pos,
# x,
# out,
# y[:, i, :],
# batch.batch,
# i,
# self.current_epoch,
# )
x = out
losses.append(self.loss(out, y[:, i, :]))