Synthetic diffusion MGN example.

[YohannDudouit]

Summary
This PR adds a pure-Torch MGN-like graph diffusion validation workflow for AMS homogeneous graph surrogates. It trains a small synthetic diffusion model, exports an AMS-compatible TorchScript model, writes runtime fixtures, and validates C++ AMS inference against Python TorchScript reference output.

What Changed

• Added generate_mgn_graph_diffusion.py with separate modes:
  • feasibility: eager/script/reload checks on dynamic graph sizes
  • train: trains the synthetic graph diffusion surrogate and saves checkpoint/metrics
  • fixtures: exports TorchScript and writes runtime fixture in fixtures.json
• Added test_graph_mgn_surrogate.cpp to load runtime fixtures and compare AMS node:delta_u output against Python TorchScript references for N=24 and N=73.
• Added MGN_DIFFUSION_* CTest workflow entries with dependency ordering:
  • feasibility → train → fixtures → AMS parity
• Added ENABLE_TRAINING_TESTS option, Off by default. Enabling ENABLE_TRAINING_TESTS currently enables the train and fixtures modes of the MGN diffusion tests.

Testing

python3 -m py_compile tests/AMSlib/models/generate_mgn_graph_diffusion.py
git diff --check
ctest --test-dir build -R MGN_DIFFUSION -j 8 --output-on-failure

Normal CTest can skip this training workflow with:

ctest --test-dir build -LE MGN_DIFFUSION --output-on-failure

Note: This PR is built on top of #194 and therefore include the graph construction and unit tests, until #194 is merged.

[github-actions]

Cpp-linter Review

Used clang-format v18.1.8

Click here for the full clang-format patch

diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp
index ed1b823..9e4cfeb 100644
--- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp
+++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp
@@ -176,4 +176,6 @@ static json loadManifest(const std::filesystem::path& fixture_dir)
-    CATCH_FAIL("Missing MGN graph diffusion fixtures at "
-               << manifest_path
-               << ". Regenerate and commit mgn_graph_diffusion.pt and "
-                  "fixtures.json.");
+    CATCH_FAIL("Missing MGN graph diffusion fixtures at " << manifest_path
+                                                          << ". Regenerate and "
+                                                             "commit "
+                                                             "mgn_graph_"
+                                                             "diffusion.pt and "
+                                                             "fixtures.json.");
@@ -248,8 +250,12 @@ static AMSHomogeneousGraph makeGraph(const json& graph_case)
-  auto node_features = readTensorValues<float>(
-      tensors.at("node_features"), "float32", {num_nodes, node_dim});
-  auto edge_index = readTensorValues<std::int64_t>(
-      tensors.at("edge_index"), "int64", {2, num_edges});
-  auto edge_features = readTensorValues<float>(
-      tensors.at("edge_features"), "float32", {num_edges, edge_dim});
-  auto global_features = readTensorValues<float>(
-      tensors.at("global_features"), "float32", {global_dim});
+  auto node_features = readTensorValues<float>(tensors.at("node_features"),
+                                               "float32",
+                                               {num_nodes, node_dim});
+  auto edge_index = readTensorValues<std::int64_t>(tensors.at("edge_index"),
+                                                   "int64",
+                                                   {2, num_edges});
+  auto edge_features = readTensorValues<float>(tensors.at("edge_features"),
+                                               "float32",
+                                               {num_edges, edge_dim});
+  auto global_features = readTensorValues<float>(tensors.at("global_features"),
+                                                 "float32",
+                                                 {global_dim});
@@ -273,4 +279,4 @@ static std::vector<float> loadReferenceDeltaU(const json& graph_case)
-  return readTensorValues<float>(
-      graph_case.at("tensors").at("reference_delta_u"),
-      "float32",
-      {num_nodes, output_dim});
+  return readTensorValues<float>(graph_case.at("tensors").at("reference_delta_"
+                                                             "u"),
+                                 "float32",
+                                 {num_nodes, output_dim});
@@ -383,2 +389 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate",
-      std::vector<float> reference_delta_u =
-          loadReferenceDeltaU(graph_case);
+      std::vector<float> reference_delta_u = loadReferenceDeltaU(graph_case);

Have any feedback or feature suggestions? Share it here.

[github-actions]

Cpp-linter Review

Used clang-format v18.1.8

Click here for the full clang-format patch

diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp
index ed1b823..9e4cfeb 100644
--- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp
+++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp
@@ -176,4 +176,6 @@ static json loadManifest(const std::filesystem::path& fixture_dir)
-    CATCH_FAIL("Missing MGN graph diffusion fixtures at "
-               << manifest_path
-               << ". Regenerate and commit mgn_graph_diffusion.pt and "
-                  "fixtures.json.");
+    CATCH_FAIL("Missing MGN graph diffusion fixtures at " << manifest_path
+                                                          << ". Regenerate and "
+                                                             "commit "
+                                                             "mgn_graph_"
+                                                             "diffusion.pt and "
+                                                             "fixtures.json.");
@@ -248,8 +250,12 @@ static AMSHomogeneousGraph makeGraph(const json& graph_case)
-  auto node_features = readTensorValues<float>(
-      tensors.at("node_features"), "float32", {num_nodes, node_dim});
-  auto edge_index = readTensorValues<std::int64_t>(
-      tensors.at("edge_index"), "int64", {2, num_edges});
-  auto edge_features = readTensorValues<float>(
-      tensors.at("edge_features"), "float32", {num_edges, edge_dim});
-  auto global_features = readTensorValues<float>(
-      tensors.at("global_features"), "float32", {global_dim});
+  auto node_features = readTensorValues<float>(tensors.at("node_features"),
+                                               "float32",
+                                               {num_nodes, node_dim});
+  auto edge_index = readTensorValues<std::int64_t>(tensors.at("edge_index"),
+                                                   "int64",
+                                                   {2, num_edges});
+  auto edge_features = readTensorValues<float>(tensors.at("edge_features"),
+                                               "float32",
+                                               {num_edges, edge_dim});
+  auto global_features = readTensorValues<float>(tensors.at("global_features"),
+                                                 "float32",
+                                                 {global_dim});
@@ -273,4 +279,4 @@ static std::vector<float> loadReferenceDeltaU(const json& graph_case)
-  return readTensorValues<float>(
-      graph_case.at("tensors").at("reference_delta_u"),
-      "float32",
-      {num_nodes, output_dim});
+  return readTensorValues<float>(graph_case.at("tensors").at("reference_delta_"
+                                                             "u"),
+                                 "float32",
+                                 {num_nodes, output_dim});
@@ -383,2 +389 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate",
-      std::vector<float> reference_delta_u =
-          loadReferenceDeltaU(graph_case);
+      std::vector<float> reference_delta_u = loadReferenceDeltaU(graph_case);

Have any feedback or feature suggestions? Share it here.

[lpottier]

LGTM overall, I've added a few comments. The only thing that is annoying in the current version is the dependency to torch. MGN_DIFFUSION_FEASIBILITY runs by default, so it makes a Python torch install a requirement at test time. Can we deactivate it by default or maybe check for torch like we do for hdf5 in tests/CMakeLists.txt and if torch is not there we mark these tests DISABLED?