diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index 17ee7034..38d89399 100644 --- a/tests/AMSlib/ams_interface/CMakeLists.txt +++ b/tests/AMSlib/ams_interface/CMakeLists.txt @@ -33,3 +33,21 @@ ADD_AMS_UNIT_TEST(AMS_GRAPH_FALLBACK ams_graph_fallback) BUILD_UNIT_TEST(ams_graph_surrogate test_graph_surrogate.cpp) ADD_AMS_UNIT_TEST(AMS_GRAPH_SURROGATE ams_graph_surrogate) +BUILD_UNIT_TEST(ams_graph_mgn_surrogate test_graph_mgn_surrogate.cpp) +target_link_libraries(ams_graph_mgn_surrogate PRIVATE nlohmann_json::nlohmann_json) +target_compile_definitions(ams_graph_mgn_surrogate + PRIVATE + AMS_MGN_DIFFUSION_FIXTURE_DIR="${AMS_TEST_ROOT}/models/mgn_graph_diffusion" +) + +# Final runtime check for the MGN diffusion workflow: fixtures mode must have +# produced the TorchScript model and binary graph cases before AMS parity runs. +add_test( + NAME MGN_DIFFUSION_AMS_PARITY + COMMAND $ -s --reporter console +) +set_tests_properties(MGN_DIFFUSION_AMS_PARITY PROPERTIES + DEPENDS MGN_DIFFUSION_FIXTURES + LABELS "AMS_INTERFACE;MGN_DIFFUSION" + WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" +) diff --git a/tests/AMSlib/ams_interface/test_graph_fallback.cpp b/tests/AMSlib/ams_interface/test_graph_fallback.cpp index 16e6e75f..5b25ef8f 100644 --- a/tests/AMSlib/ams_interface/test_graph_fallback.cpp +++ b/tests/AMSlib/ams_interface/test_graph_fallback.cpp @@ -124,21 +124,9 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", AMSInit(); CATCH_REQUIRE_NOTHROW(makeValidGraph()); - { - AMSHomogeneousGraph graph(makeNodeFeatures(), - makeEdgeIndex32(), - makeEdgeFeatures()); - CATCH_REQUIRE(graph.global_features.shape().size() == 1); - CATCH_REQUIRE(graph.global_features.shape()[0] == 0); - } - CATCH_REQUIRE_NOTHROW(AMSHomogeneousGraph(makeNodeFeatures(), - makeEdgeIndex64(), - makeEdgeFeatures(), - makeTensor({0}))); CATCH_REQUIRE_NOTHROW(AMSHomogeneousGraph(makeNodeFeatures(), - makeEdgeIndex64(), - makeEdgeFeatures(), - makeGlobalFeatures())); + makeEdgeIndex32(), + makeEdgeFeatures())); CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeTensor({3, 2}), makeEdgeIndex64(), @@ -172,11 +160,6 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", makeEdgeIndex64(), makeTensor({2, 1})), std::runtime_error); - CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), - makeEdgeIndex64(), - makeEdgeFeatures(), - makeTensor({1, 2})), - std::runtime_error); CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), makeEdgeIndex64(), makeEdgeFeatures(), diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp new file mode 100644 index 00000000..20cb7c76 --- /dev/null +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -0,0 +1,506 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "AMS.h" +#include "AMSGraph.hpp" +#include "AMSTensor.hpp" + +// This test is the C++ half of the MGN diffusion example. +// +// The Python generator does the learned-model work: +// +// 1. create synthetic graph diffusion data; +// 2. train a small pure-Torch MeshGraphNet-like model; +// 3. export that model as an AMS-wrapped TorchScript file; +// 4. write "fixtures": fixed graph inputs plus Python TorchScript reference +// outputs. +// +// This C++ test does not know the diffusion formula and does not try to judge +// model quality. Its job is deployment parity: load the exact TorchScript model +// and graph tensors from the generated fixture directory, run them through AMS, +// and verify that the AMS output field node:delta_u matches the Python +// TorchScript reference output. +// +// Said differently: Python proves "the model learned something"; this test +// proves "AMS runs the exported model the same way Python does." + +using namespace ams; +using json = nlohmann::json; + +using Dim = AMSTensor::IntDimType; + +// The MGN diffusion workflow writes fixtures into the build tree at CTest +// runtime. CMake compiles that build-tree directory into this focused parity +// test so the test can be launched like any other Catch executable while still +// keeping generated tensors out of the source tree. +#ifndef AMS_MGN_DIFFUSION_FIXTURE_DIR +#define AMS_MGN_DIFFUSION_FIXTURE_DIR "" +#endif + +namespace +{ + +constexpr int kFixtureFormatVersion = 1; +constexpr std::int64_t kNodeFeatureDim = 4; +constexpr std::int64_t kEdgeFeatureDim = 4; +constexpr std::int64_t kGlobalFeatureDim = 1; +constexpr std::int64_t kReferenceOutputDim = 1; +constexpr double kUnusedGraphSurrogateThreshold = 1.0; +constexpr int kSingleRankProcessId = 0; +constexpr int kSingleRankWorldSize = 1; +constexpr bool kStoreTrainingData = false; +constexpr const char* kDomainName = "test_mgn_graph_diffusion_surrogate"; + +// Glossary for readers new to this path: +// +// manifest: fixtures.json, the table of contents for generated fixture files. +// fixture directory: build-tree directory containing fixtures.json, the +// TorchScript model, and raw tensor binaries. +// raw tensor binary: contiguous bytes for one tensor; shape/dtype live in JSON. +// AMSHomogeneousGraph: AMS C++ container for node_features, edge_index, +// edge_features, and optional global_features. +// node:delta_u: TorchScript output key. AMS parses this into the node field +// named "delta_u". +// fallback callback: application callback AMS calls only when surrogate +// inference is unavailable or rejected. +// parity: equality of AMS/LibTorch inference and Python TorchScript inference. + +struct TensorMetadata { + std::filesystem::path path; + std::string dtype; + std::vector shape; + std::string endianness; + std::uintmax_t byte_size; +}; + +static std::vector contiguousStrides(const std::vector& shape) +{ + // AMSTensor stores both shape and strides. These fixtures are written as + // simple row-major contiguous arrays, so the last dimension has stride 1 and + // each earlier stride is the product of dimensions to its right. + std::vector strides(shape.size(), 1); + Dim stride = 1; + for (std::size_t i = shape.size(); i-- > 0;) { + strides[i] = stride; + stride *= shape[i]; + } + return strides; +} + +template +static AMSTensor makeTensor(std::vector shape, + const std::vector& values) +{ + // Convert ordinary C++ vectors read from fixture binaries into owned + // AMSTensors. From this point on, the graph looks like application-provided + // AMS input rather than test-specific storage. + std::vector strides = contiguousStrides(shape); + auto tensor = AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); + CATCH_REQUIRE(values.size() == static_cast(tensor.elements())); + std::copy(values.begin(), values.end(), tensor.template data()); + return tensor; +} + +static Dim toDim(std::int64_t value) +{ + CATCH_REQUIRE(value >= 0); + CATCH_REQUIRE(value <= std::numeric_limits::max()); + return static_cast(value); +} + +static bool hostIsLittleEndian() +{ + // Fixture binaries are written little-endian. AMS currently reads them by + // copying bytes directly into native scalar arrays, so running this test on a + // big-endian host should fail explicitly instead of silently byte-swapping + // wrong values. + const std::uint16_t value = 1; + return *reinterpret_cast(&value) == 1; +} + +static std::uintmax_t dtypeByteWidth(const std::string& dtype) +{ + // The generator only writes float32 tensors for real-valued graph data and + // int64 tensors for edge_index. Supporting more dtypes would require extending + // both the manifest contract and the AMSTensor construction below. + if (dtype == "float32") { + return sizeof(float); + } + if (dtype == "int64") { + return sizeof(std::int64_t); + } + CATCH_FAIL( + "Unsupported MGN graph diffusion tensor dtype in manifest: " << dtype); + return 0; +} + +static std::uintmax_t shapeElementCount(const std::vector& shape) +{ + std::uintmax_t count = 1; + for (std::int64_t dim : shape) { + CATCH_REQUIRE(dim >= 0); + count *= static_cast(dim); + } + return count; +} + +static TensorMetadata parseTensorMetadata(const json& tensor) +{ + // Pull the small per-tensor manifest record into a typed C++ struct. This is + // intentionally separate from validation so missing fields, unsupported + // dtypes, and wrong shapes fail at the most helpful point in the test. + CATCH_REQUIRE(tensor.contains("path")); + CATCH_REQUIRE(tensor.contains("dtype")); + CATCH_REQUIRE(tensor.contains("shape")); + CATCH_REQUIRE(tensor.contains("endianness")); + CATCH_REQUIRE(tensor.contains("byte_size")); + + TensorMetadata metadata; + metadata.path = tensor.at("path").get(); + metadata.dtype = tensor.at("dtype").get(); + metadata.shape = tensor.at("shape").get>(); + metadata.endianness = tensor.at("endianness").get(); + metadata.byte_size = tensor.at("byte_size").get(); + return metadata; +} + +static void validateTensorMetadata( + const TensorMetadata& metadata, + const std::string& expected_dtype, + const std::vector& expected_shape) +{ + // Validate the manifest before allocating AMSTensors. + // + // The binary files deliberately contain no headers: no magic number, no shape, + // no dtype tag. That keeps them easy for C++ to read without Python or NumPy, + // but it means fixtures.json is the source of truth. If the manifest says a + // tensor is [N, 4] float32, this check makes sure the file metadata agrees + // with exactly the tensor shape this test is about to construct. + CATCH_REQUIRE_FALSE(metadata.path.empty()); + CATCH_REQUIRE_FALSE(metadata.path.is_absolute()); + for (const auto& part : metadata.path) { + CATCH_REQUIRE(part.string() != ".."); + } + CATCH_REQUIRE(metadata.dtype == expected_dtype); + CATCH_REQUIRE(metadata.shape == expected_shape); + + // The generator writes little-endian bytes. The host-endianness check makes + // the limitation visible if this test ever runs on a different architecture. + CATCH_REQUIRE(metadata.endianness == "little"); + CATCH_REQUIRE(hostIsLittleEndian()); + + const std::uintmax_t expected_size = + dtypeByteWidth(metadata.dtype) * shapeElementCount(metadata.shape); + CATCH_REQUIRE(metadata.byte_size == expected_size); +} + +template +static std::vector readTensorBinary( + const std::filesystem::path& manifest_dir, + const json& tensor, + const std::string& expected_dtype, + const std::vector& expected_shape) +{ + TensorMetadata metadata = parseTensorMetadata(tensor); + validateTensorMetadata(metadata, expected_dtype, expected_shape); + + // All tensor paths are relative to fixtures.json. That keeps the generated + // fixture directory relocatable: moving the build directory as a unit does not + // invalidate absolute paths baked into the manifest. + const std::filesystem::path tensor_path = manifest_dir / metadata.path; + CATCH_REQUIRE(std::filesystem::exists(tensor_path)); + CATCH_REQUIRE(std::filesystem::is_regular_file(tensor_path)); + CATCH_REQUIRE(std::filesystem::file_size(tensor_path) == metadata.byte_size); + + const std::uintmax_t element_count = shapeElementCount(metadata.shape); + CATCH_REQUIRE(element_count <= std::numeric_limits::max()); + std::vector values(static_cast(element_count)); + + std::ifstream input(tensor_path, std::ios::binary); + CATCH_REQUIRE(input); + if (metadata.byte_size > 0) { + CATCH_REQUIRE(metadata.byte_size <= + static_cast( + std::numeric_limits::max())); + input.read(reinterpret_cast(values.data()), + static_cast(metadata.byte_size)); + // gcount verifies we actually read the promised number of bytes. Combined + // with file_size above, this catches truncated or stale fixture files before + // they can become misleading numerical comparisons. + CATCH_REQUIRE(static_cast(input.gcount()) == + metadata.byte_size); + } + return values; +} + +static json loadManifest(const std::filesystem::path& fixture_dir) +{ + // fixtures.json is generated by the Python fixtures mode. If this file is + // missing, the C++ test cannot make progress because it does not know which + // model or tensors to load. + const std::filesystem::path manifest_path = fixture_dir / "fixtures.json"; + if (!std::filesystem::exists(manifest_path)) { + CATCH_FAIL("Missing MGN graph diffusion fixtures at " + << manifest_path + << ". Run `ctest -R MGN_DIFFUSION_FIXTURES` or the full " + "`ctest -R MGN_DIFFUSION` chain."); + } + + std::ifstream input(manifest_path); + CATCH_REQUIRE(input); + json manifest = json::parse(input); + + // Keep the manifest version check close to parsing so future fixture format + // changes fail loudly instead of being interpreted as the current raw-binary + // contract. + CATCH_REQUIRE(manifest.contains("format_version")); + CATCH_REQUIRE(manifest.contains("endianness")); + CATCH_REQUIRE(manifest.at("format_version").get() == + kFixtureFormatVersion); + CATCH_REQUIRE(manifest.at("endianness").get() == "little"); + CATCH_REQUIRE(manifest.contains("model")); + CATCH_REQUIRE(manifest.contains("cases")); + CATCH_REQUIRE(manifest.contains("comparison")); + CATCH_REQUIRE(manifest.at("cases").is_array()); + return manifest; +} + +static std::filesystem::path resolveManifestRelativePath( + const std::filesystem::path& manifest_dir, + const json& object) +{ + // Model and tensor paths in the manifest are relative by design. Avoiding + // absolute paths makes fixtures reusable if the whole build directory moves. + const std::filesystem::path relative_path = + object.at("path").get(); + CATCH_REQUIRE_FALSE(relative_path.empty()); + CATCH_REQUIRE_FALSE(relative_path.is_absolute()); + for (const auto& part : relative_path) { + CATCH_REQUIRE(part.string() != ".."); + } + return manifest_dir / relative_path; +} + +static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, + const json& graph_case) +{ + // Runtime fixture binaries become the same AMSTensor-backed homogeneous graph + // that an application would pass to AMS: + // + // node_features [N, 4] -> x, y, u, kappa + // edge_index [2, E] -> source row, destination row + // edge_features [E, 4] -> dx, dy, distance, message + // global_features [1] -> dt + // + // Keeping this conversion here, instead of compiling arrays into the test, + // makes the test closer to a real AMS deployment path: model and tensors are + // runtime files. + const std::int64_t num_nodes = graph_case.at("num_nodes").get(); + const std::int64_t num_edges = graph_case.at("num_edges").get(); + const std::int64_t node_dim = + graph_case.at("node_feature_dim").get(); + const std::int64_t edge_dim = + graph_case.at("edge_feature_dim").get(); + const std::int64_t global_dim = + graph_case.at("global_feature_dim").get(); + + // These feature dimensions are part of this particular learned example. If + // the generator changes the graph contract, the C++ parity test should fail + // and be updated deliberately. + CATCH_REQUIRE(node_dim == kNodeFeatureDim); + CATCH_REQUIRE(edge_dim == kEdgeFeatureDim); + CATCH_REQUIRE(global_dim == kGlobalFeatureDim); + + const json& tensors = graph_case.at("tensors"); + + // edge_index is the only integer tensor. It must remain int64 because the + // TorchScript model canonicalizes and indexes with 64-bit node ids. + auto node_features = readTensorBinary(manifest_dir, + tensors.at("node_features"), + "float32", + {num_nodes, node_dim}); + auto edge_index = readTensorBinary(manifest_dir, + tensors.at("edge_index"), + "int64", + {2, num_edges}); + auto edge_features = readTensorBinary(manifest_dir, + tensors.at("edge_features"), + "float32", + {num_edges, edge_dim}); + auto global_features = readTensorBinary(manifest_dir, + tensors.at("global_features"), + "float32", + {global_dim}); + + return AMSHomogeneousGraph( + makeTensor({toDim(num_nodes), toDim(node_dim)}, node_features), + makeTensor({2, toDim(num_edges)}, edge_index), + makeTensor({toDim(num_edges), toDim(edge_dim)}, edge_features), + makeTensor({toDim(global_dim)}, global_features)); +} + +static std::vector loadReferenceDeltaU( + const std::filesystem::path& manifest_dir, + const json& graph_case) +{ + // Reference output is one scalar delta_u per node. It was produced in Python + // by reloading the exported TorchScript model, so it is the closest available + // representation of what LibTorch should compute in AMS. + const std::int64_t num_nodes = graph_case.at("num_nodes").get(); + const std::int64_t output_dim = + graph_case.at("reference_output_dim").get(); + CATCH_REQUIRE(output_dim == kReferenceOutputDim); + return readTensorBinary(manifest_dir, + graph_case.at("tensors").at("reference_delta_" + "u"), + "float32", + {num_nodes, output_dim}); +} + +static void verifyDeltaU(const json& graph_case, + const std::vector& reference_delta_u, + const AMSHomogeneousGraphFields& outputs, + double rtol, + double atol) +{ + // The reference is the Python TorchScript output, not the exact synthetic + // diffusion target. Training mode already checks whether the model learned + // the synthetic formula. This test is about AMS/LibTorch deployment parity. + const std::int64_t num_nodes = graph_case.at("num_nodes").get(); + const std::int64_t output_dim = + graph_case.at("reference_output_dim").get(); + CATCH_REQUIRE(output_dim == kReferenceOutputDim); + + CATCH_REQUIRE(outputs.node_fields.contains("delta_u")); + + // The Python model returns "node:delta_u". AMS strips the "node:" namespace + // and stores the tensor under outputs.node_fields.at("delta_u"). + const auto& delta_u = outputs.node_fields.at("delta_u"); + CATCH_REQUIRE(delta_u.shape()[0] == num_nodes); + CATCH_REQUIRE(delta_u.shape()[1] == output_dim); + CATCH_REQUIRE(reference_delta_u.size() == + static_cast(num_nodes * output_dim)); + + const float* actual = delta_u.data(); + const std::int64_t count = num_nodes * output_dim; + for (std::int64_t i = 0; i < count; ++i) { + CATCH_REQUIRE( + actual[i] == + Catch::Approx(reference_delta_u[i]).epsilon(rtol).margin(atol)); + } +} + +} // namespace + +CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", + "[wf][graph][surrogate][mgn]") +{ + // CMake sets AMS_MGN_DIFFUSION_FIXTURE_DIR to the generated build-tree + // fixture directory. The manifest inside that directory tells us where the + // model and per-case tensor files live. + const std::filesystem::path fixture_dir = AMS_MGN_DIFFUSION_FIXTURE_DIR; + CATCH_REQUIRE_FALSE(fixture_dir.empty()); + + const json manifest = loadManifest(fixture_dir); + const std::filesystem::path manifest_dir = fixture_dir; + + // The model path also comes from the manifest rather than being hard-coded in + // the test. That keeps the C++ loader tied to the same artifact table of + // contents as the tensor files. + const std::filesystem::path model_path = + resolveManifestRelativePath(manifest_dir, manifest.at("model")); + CATCH_REQUIRE(std::filesystem::exists(model_path)); + CATCH_REQUIRE(std::filesystem::is_regular_file(model_path)); + + const double rtol = manifest.at("comparison").at("rtol").get(); + const double atol = manifest.at("comparison").at("atol").get(); + CATCH_REQUIRE(manifest.at("cases").size() == 2); + + // Register the exported TorchScript model with AMS and create an executor + // exactly as an application would before calling AMSExecute. + // + // This test runs as a single-process executor, so process_id/world_size are + // named constants rather than bare 0/1 literals. + // + // The threshold argument has no acceptance/rejection meaning for homogeneous + // graph surrogate execution today: the graph path either runs the model or + // falls back if the model cannot be used. AMSRegisterAbstractModel still + // requires a threshold because the API is shared with pointwise surrogates. + // Use 1.0 here as a readable "accept everything" value for that unused + // graph-path parameter. + // + // The store_data argument is also inherited from the shared pointwise + // workflow API. It controls whether an AMSWorkflow opens a training-data DB, + // but this graph parity test is inference-only and requires the fallback + // callback to remain unused. Current graph fallback data storage is also not + // implemented, so false keeps the test from opening unrelated DB state. + AMSInit(); + const std::string model_path_string = model_path.string(); + auto model = AMSRegisterAbstractModel(kDomainName, + kUnusedGraphSurrogateThreshold, + model_path_string.c_str(), + kStoreTrainingData); + AMSExecutor executor = + AMSCreateExecutor(model, kSingleRankProcessId, kSingleRankWorldSize); + + // The two fixed fixture sizes intentionally have different N and E values. + // Running both cases catches accidental static-shape assumptions in the + // exported TorchScript model or in AMS graph tensor handling. + // + // N is checked explicitly here because this test is meant to prove the model + // handles both fixture graph sizes, not just "whatever happened to be in the + // manifest". + const std::vector expected_node_counts = {24, 73}; + std::size_t case_index = 0; + for (const json& graph_case : manifest.at("cases")) { + CATCH_REQUIRE(graph_case.at("num_nodes").get() == + expected_node_counts.at(case_index)); + CATCH_DYNAMIC_SECTION("fixture " + << graph_case.at("name").get()) + { + // Load one graph case from runtime files, then run it through AMS. Each + // case has its own N, E, input tensors, and reference output. + AMSHomogeneousGraph graph = makeGraph(manifest_dir, graph_case); + std::vector reference_delta_u = + loadReferenceDeltaU(manifest_dir, graph_case); + + bool callback_invoked = false; + // If the surrogate path fails, AMS would call the domain fallback. For + // this parity test that would hide a deployment failure, so the callback + // records whether it was invoked and supplies the reference only as a + // valid fallback value. + // + // The assertion below requires callback_invoked == false. That means AMS + // accepted the graph surrogate output instead of falling back to this + // lambda. + HomogeneousGraphDomainFn callback = + [&](const AMSHomogeneousGraph&, AMSHomogeneousGraphFields& outputs) { + callback_invoked = true; + const std::int64_t num_nodes = + graph_case.at("num_nodes").get(); + const std::int64_t output_dim = + graph_case.at("reference_output_dim").get(); + outputs.node_fields.set("delta_u", + makeTensor({toDim(num_nodes), + toDim(output_dim)}, + reference_delta_u)); + }; + + AMSHomogeneousGraphFields outputs; + AMSExecute(executor, callback, graph, outputs); + + CATCH_REQUIRE_FALSE(callback_invoked); + verifyDeltaU(graph_case, reference_delta_u, outputs, rtol, atol); + } + ++case_index; + } +} diff --git a/tests/AMSlib/models/CMakeLists.txt b/tests/AMSlib/models/CMakeLists.txt index ee1e5847..ec40888d 100644 --- a/tests/AMSlib/models/CMakeLists.txt +++ b/tests/AMSlib/models/CMakeLists.txt @@ -4,6 +4,7 @@ set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/generate_linear_model.py" "${CMAKE_CURRENT_SOURCE_DIR}/generate_base_models.py" "${CMAKE_CURRENT_SOURCE_DIR}/generate_graph_models.py" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" ) # Where to drop the .pt files @@ -114,3 +115,122 @@ configure_file( ${CMAKE_CURRENT_BINARY_DIR}/ams_test_linear_models.hpp COPYONLY ) + +find_program(AMS_PYTHON_EXECUTABLE NAMES python3 python REQUIRED) + +# MGN diffusion artifacts are generated in the build tree. The custom target is +# available for explicit regeneration, but normal C++ builds do not depend on +# training/export; the CTest workflow below owns the runtime validation path. +set(MGN_GRAPH_DIFFUSION_DIR + "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion" +) +set(MGN_GRAPH_DIFFUSION_CHECKPOINT + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_graph_diffusion_checkpoint.pt" +) +set(MGN_GRAPH_DIFFUSION_MODEL + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_graph_diffusion.pt" +) +set(MGN_GRAPH_DIFFUSION_FEASIBILITY_MODEL + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_graph_diffusion_feasibility.pt" +) +set(MGN_GRAPH_DIFFUSION_MANIFEST + "${MGN_GRAPH_DIFFUSION_DIR}/fixtures.json" +) +set(MGN_GRAPH_DIFFUSION_METRICS + "${MGN_GRAPH_DIFFUSION_DIR}/training_metrics.json" +) +set(MGN_GRAPH_DIFFUSION_TENSOR_FILES + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/node_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/edge_index.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/edge_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/global_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/reference_delta_u.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/node_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/edge_index.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/edge_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/global_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/reference_delta_u.bin" +) + +add_custom_command( + OUTPUT + "${MGN_GRAPH_DIFFUSION_CHECKPOINT}" + "${MGN_GRAPH_DIFFUSION_MODEL}" + "${MGN_GRAPH_DIFFUSION_MANIFEST}" + ${MGN_GRAPH_DIFFUSION_TENSOR_FILES} + BYPRODUCTS + "${MGN_GRAPH_DIFFUSION_FEASIBILITY_MODEL}" + "${MGN_GRAPH_DIFFUSION_METRICS}" + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode feasibility + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode train + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode fixtures + DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + "${CMAKE_CURRENT_SOURCE_DIR}/ams_model.py" + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" + COMMENT "Generating pure-Torch MGN graph diffusion model and fixtures" + VERBATIM +) + +add_custom_target(generate_mgn_graph_diffusion + DEPENDS + "${MGN_GRAPH_DIFFUSION_CHECKPOINT}" + "${MGN_GRAPH_DIFFUSION_MODEL}" + "${MGN_GRAPH_DIFFUSION_MANIFEST}" + ${MGN_GRAPH_DIFFUSION_TENSOR_FILES} +) + +# Runtime workflow tests. The dependency chain keeps the expensive learned-model +# path ordered even under parallel CTest: feasibility -> train -> fixtures. +add_test( + NAME MGN_DIFFUSION_FEASIBILITY + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode feasibility +) +set_tests_properties(MGN_DIFFUSION_FEASIBILITY PROPERTIES + LABELS MGN_DIFFUSION + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" +) + +add_test( + NAME MGN_DIFFUSION_TRAIN + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode train +) +set_tests_properties(MGN_DIFFUSION_TRAIN PROPERTIES + DEPENDS MGN_DIFFUSION_FEASIBILITY + LABELS MGN_DIFFUSION + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" +) + +add_test( + NAME MGN_DIFFUSION_FIXTURES + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode fixtures +) +set_tests_properties(MGN_DIFFUSION_FIXTURES PROPERTIES + DEPENDS MGN_DIFFUSION_TRAIN + LABELS MGN_DIFFUSION + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" +) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py new file mode 100644 index 00000000..2c9c9952 --- /dev/null +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -0,0 +1,781 @@ +#!/usr/bin/env python3 +"""Train/export a tiny pure-Torch graph diffusion surrogate. + +This file is meant to be readable as a small end-to-end example, even if you +have not worked with AI surrogate models before. + +In this example, a "surrogate" is a learned replacement for a known +calculation. We first create synthetic graph data where the correct answer is +known exactly. Then we train a small neural network to imitate that answer. +Finally, we export the trained network and check that AMS/C++ gets the same +answer as Python/TorchScript for fixed test graphs. + +The data is a graph: + +* Nodes are sample points in a 2D square. Each node has features such as its + x/y position, a scalar state value u, and a conductivity-like value kappa. +* Edges connect nearby nodes. Each edge has a source node and a destination + node, so a message can travel "from source to destination". +* Global features describe the whole graph. Here the only global feature is + dt, the time-step size. + +The synthetic target is a diffusion update: + + delta_u_i = dt * sum_{src -> i} weight * (u_src - u_i) + +Intuitively, each node is pulled toward its neighbors. If a neighboring source +node has a larger u value, it sends a positive contribution. If it has a +smaller u value, it sends a negative contribution. This is close in spirit to +heat diffusion, but cheap and fully deterministic. + +The model is "MGN-like" (MeshGraphNet-like) because it repeats the same graph +communication pattern: + +1. build/update an edge message from source node, destination node, edge data, + and global data; +2. add incoming edge messages onto each destination node; +3. update each node from its old state and the aggregated incoming messages. + +The workflow has three explicit modes: + +* feasibility: prove that the model can run eagerly, be scripted, be reloaded, + and still accept different graph sizes. This checks deployability before + paying any training cost. +* train: train the tiny model on deterministic synthetic graphs, then save a + checkpoint containing learned weights and metrics. +* fixtures: load the checkpoint, export the AMS-wrapped TorchScript model, and + save fixed graph inputs plus Python TorchScript reference outputs. + +The C++ parity test compares AMS output to Python TorchScript output, not to +the exact synthetic formula. Training already checks that the model learned the +synthetic target. The C++ test asks a different question: "when AMS deploys this +TorchScript graph model, does it reproduce Python inference?" +""" + +import argparse +import copy +import json +import os +import struct +import sys +from pathlib import Path +from typing import Dict, Iterable, List, Tuple + +try: + import torch + import torch.nn as nn + from torch import Tensor +except ModuleNotFoundError as exc: + raise SystemExit( + "PyTorch is required for MGN graph diffusion generation. " + "Install PyTorch or run this script on a Torch-enabled machine." + ) from exc + + +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(CURRENT_DIR) +from ams_model import create_ams_homogeneous_graph_model + + +NODE_FEATURE_DIM = 4 +EDGE_FEATURE_DIM = 4 +GLOBAL_FEATURE_DIM = 1 +REFERENCE_OUTPUT_DIM = 1 + +FIXTURE_GRAPH_SIZES = (24, 73) +FEASIBILITY_SEEDS = (70024, 70073) +FIXTURE_SEEDS = (90024, 90073) + +MODEL_SEED = 314159 +TRAIN_DATA_SEED = 271828 +VAL_DATA_SEED = 161803 + +LATENT_DIM = 64 +NUM_PROCESSOR_BLOCKS = 2 +K_NEIGHBORS = 6 +TRAIN_STEPS = 3000 +VAL_GRAPHS = 32 +LEARNING_RATE = 1.0e-3 +WEIGHT_DECAY = 1.0e-6 +PRIMARY_BASELINE_FACTOR = 0.1 +SECONDARY_ABSOLUTE_MSE = 1.0e-4 + +CHECKPOINT_NAME = "mgn_graph_diffusion_checkpoint.pt" +MODEL_NAME = "mgn_graph_diffusion.pt" +FIXTURE_MANIFEST_NAME = "fixtures.json" +TRAINING_METRICS_NAME = "training_metrics.json" +FIXTURE_FORMAT_VERSION = 1 +FIXTURE_ENDIANNESS = "little" +COMPARISON_RTOL = 2.0e-5 +COMPARISON_ATOL = 2.0e-5 + +# Glossary for readers new to graph surrogates: +# +# N: number of nodes in one graph. +# E: number of directed edges in one graph. +# node_features: table with one row per node. Here [x, y, u, kappa]. +# edge_index: integer connectivity table with shape [2, E]. Row 0 is source +# node id, row 1 is destination node id. +# edge_features: table with one row per edge. Here [dx, dy, distance, message]. +# global_features: graph-wide values. Here just [dt]. +# delta_u: the node-wise output we want the model to predict. +# dt: time-step size in the synthetic diffusion update. +# kappa: conductivity-like node value used to define edge weights. +# checkpoint: saved training result containing model weights and metadata. +# fixture: fixed test input plus expected output used by C++ for repeatability. +# TorchScript: PyTorch's serialized model format that LibTorch/C++ can load. + + +class MLP(nn.Module): + """A small trainable function used as the edge/node update rule. + + An MLP, or multilayer perceptron, is just a few linear layers with a + nonlinearity. Here it learns simple formulas from examples instead of us + hand-writing the formulas in C++. + """ + + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, output_dim), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.layers(x) + + +class ProcessorBlock(nn.Module): + """One round of graph communication. + + A MeshGraphNet-style model alternates between updating edges and updating + nodes. Think of this block as one "conversation round": edges look at their + source and destination nodes, then nodes collect all incoming edge messages. + """ + + def __init__(self, latent_dim: int, global_dim: int): + super().__init__() + self.edge_mlp = MLP(latent_dim * 3 + global_dim, latent_dim, latent_dim) + self.node_mlp = MLP(latent_dim * 2 + global_dim, latent_dim, latent_dim) + + def forward( + self, + node_latent: Tensor, + edge_latent: Tensor, + edge_index: Tensor, + global_features: Tensor, + ) -> Tuple[Tensor, Tensor]: + # AMS homogeneous graphs provide global_features as one graph-level + # vector with shape [F_global]. For concatenation, this block needs one + # copy per edge and one copy per node, so the broadcast happens here at + # the point where graph-wide data is turned into per-entity inputs. + global_row = global_features.unsqueeze(0) + + # The canonical AMS edge_index uses row 0 as source and row 1 as + # destination. index_select gathers per-edge source/destination node + # states without any graph-library dependency. + src = edge_index[0] + dst = edge_index[1] + + # Edge update: each edge sees its current latent state, source node, + # destination node, and the graph-level dt feature. The residual update + # keeps this tiny model easy to train. + global_edges = global_row.expand(edge_latent.shape[0], global_row.shape[1]) + edge_input = torch.cat( + ( + edge_latent, + node_latent.index_select(0, src), + node_latent.index_select(0, dst), + global_edges, + ), + dim=1, + ) + edge_latent = edge_latent + self.edge_mlp(edge_input) + + # Node aggregation: incoming edge states are summed onto destination + # nodes with index_add, the same primitive TorchScript and LibTorch will + # execute after export. + aggregated = torch.zeros( + (node_latent.shape[0], edge_latent.shape[1]), + dtype=node_latent.dtype, + device=node_latent.device, + ) + aggregated = aggregated.index_add(0, dst, edge_latent) + + # Node update: each node sees its previous latent state, the aggregated + # incoming message, and the global dt feature. + global_nodes = global_row.expand(node_latent.shape[0], global_row.shape[1]) + node_input = torch.cat((node_latent, aggregated, global_nodes), dim=1) + node_latent = node_latent + self.node_mlp(node_input) + return node_latent, edge_latent + + +class TinyGraphDiffusionMGN(nn.Module): + """Small MeshGraphNet-like model for the AMS graph contract. + + The structure is: + + * encoder: turn raw node/edge numbers into latent vectors the model can + learn with; + * processor blocks: let neighboring nodes exchange information through + edges; + * decoder: turn the final node latent vectors into node:delta_u. + + The model is intentionally tiny and fixed to two processor blocks. The goal + is to validate the learned-model loop through AMS, not to build a large or + configurable production architecture. + """ + + def __init__( + self, + node_dim: int = NODE_FEATURE_DIM, + edge_dim: int = EDGE_FEATURE_DIM, + global_dim: int = GLOBAL_FEATURE_DIM, + latent_dim: int = LATENT_DIM, + ): + super().__init__() + self.node_encoder = MLP(node_dim, latent_dim, latent_dim) + self.edge_encoder = MLP(edge_dim, latent_dim, latent_dim) + self.processor1 = ProcessorBlock(latent_dim, global_dim) + self.processor2 = ProcessorBlock(latent_dim, global_dim) + self.node_decoder = MLP(latent_dim, latent_dim, REFERENCE_OUTPUT_DIM) + + def forward(self, graph: Dict[str, Tensor]) -> Dict[str, Tensor]: + # Keep this signature and return type aligned with the AMS homogeneous + # graph surrogate contract. The output key becomes outputs.node_fields + # entry "delta_u" in C++. + node_features = graph["node_features"] + edge_index = graph["edge_index"].to(torch.int64) + edge_features = graph["edge_features"] + global_features = graph["global_features"] + + # Raw features such as x/y/u/kappa are low-dimensional physical-looking + # values. Encoders map them into a latent space where the trainable + # processor blocks can represent richer relationships. + node_latent = self.node_encoder(node_features) + edge_latent = self.edge_encoder(edge_features) + + # Two graph communication rounds are enough for this toy diffusion + # problem: edges gather source/destination context, then destination + # nodes receive aggregated incoming messages. + node_latent, edge_latent = self.processor1( + node_latent, edge_latent, edge_index, global_features + ) + node_latent, edge_latent = self.processor2( + node_latent, edge_latent, edge_index, global_features + ) + + # Decode one scalar per node. The "node:" prefix tells AMS this belongs + # in outputs.node_fields rather than edge or global fields. + delta_u = self.node_decoder(node_latent) + return {"node:delta_u": delta_u} + + +def model_config() -> Dict[str, int]: + return { + "node_feature_dim": NODE_FEATURE_DIM, + "edge_feature_dim": EDGE_FEATURE_DIM, + "global_feature_dim": GLOBAL_FEATURE_DIM, + "reference_output_dim": REFERENCE_OUTPUT_DIM, + "latent_dim": LATENT_DIM, + "num_processor_blocks": NUM_PROCESSOR_BLOCKS, + "k_neighbors": K_NEIGHBORS, + } + + +def make_model() -> TinyGraphDiffusionMGN: + torch.manual_seed(MODEL_SEED) + model = TinyGraphDiffusionMGN() + return model.to(device=torch.device("cpu"), dtype=torch.float32) + + +def make_generator(seed: int) -> torch.Generator: + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + return generator + + +def _unique_directed_edges(knn: Tensor) -> Tensor: + # kNN is computed from each node's perspective. Adding both directions makes + # the fixture exercise directed source/destination semantics while still + # representing an undirected diffusion neighborhood. + pairs = set() + num_nodes = int(knn.shape[0]) + for dst in range(num_nodes): + for n in range(int(knn.shape[1])): + src = int(knn[dst, n].item()) + if src == dst: + continue + pairs.add((src, dst)) + pairs.add((dst, src)) + ordered = sorted(pairs) + return torch.tensor(ordered, dtype=torch.int64).t().contiguous() + + +def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor]: + """Create one deterministic synthetic graph and its exact diffusion target. + + This function is the "data generator" for the example. It replaces a real + simulation code with a cheap formula that still has the graph ingredients we + care about: node values, edge directions, edge weights, aggregation onto + destination nodes, and a node-wise output field. + """ + + generator = make_generator(seed) + + # 1. Sample node locations in a unit square. In a mesh-based application + # these would come from mesh vertices or degrees of freedom. + positions = torch.rand((num_nodes, 2), generator=generator, dtype=torch.float32) + + # 2. Assign each node a smooth scalar state u. Smooth values are easier to + # relate to heat diffusion than independent random noise: neighboring points + # tend to have related values, but still differ enough to produce diffusion. + phases = torch.rand((1, 4), generator=generator, dtype=torch.float32) * (2.0 * torch.pi) + amps = torch.rand((1, 4), generator=generator, dtype=torch.float32) * 0.5 + 0.5 + x = positions[:, 0:1] + y = positions[:, 1:2] + u_raw = ( + amps[:, 0:1] * torch.sin(2.0 * torch.pi * x + phases[:, 0:1]) + + amps[:, 1:2] * torch.cos(2.0 * torch.pi * y + phases[:, 1:2]) + + amps[:, 2:3] * torch.sin(2.0 * torch.pi * (x + y) + phases[:, 2:3]) + + amps[:, 3:4] * torch.cos(2.0 * torch.pi * (x - y) + phases[:, 3:4]) + ) + u = torch.tanh(0.5 * u_raw) + + # 3. Add a conductivity-like scalar kappa and a graph-wide time step dt. + # kappa changes how strongly values diffuse across edges; dt scales the + # final update for every node in the graph. + kappa = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) + 0.5 + dt = torch.rand((1,), generator=generator, dtype=torch.float32) * 0.06 + 0.02 + + # 4. Connect nearby nodes with a k-nearest-neighbor graph. The model sees + # directed edges, so edge_index[0, e] is the source node and edge_index[1, e] + # is the destination node for edge e. E varies with N because duplicate + # edges are removed after adding the reverse direction; that helps test + # dynamic node and edge counts in TorchScript and C++. + distances = torch.cdist(positions, positions) + masked = distances + torch.eye(num_nodes, dtype=torch.float32) * 1.0e6 + knn = torch.topk(masked, k=K_NEIGHBORS, largest=False, dim=1).indices + edge_index = _unique_directed_edges(knn) + src = edge_index[0] + dst = edge_index[1] + + delta_pos = positions.index_select(0, src) - positions.index_select(0, dst) + distance = delta_pos.norm(dim=1, keepdim=True).clamp_min(1.0e-6) + conductivity = 0.5 * (kappa.index_select(0, src) + kappa.index_select(0, dst)) + raw_weight = conductivity / (distance + 0.05) + + # 5. Turn geometric distance and kappa into a diffusion weight. The incoming + # weights for each destination node are normalized so every node receives a + # bounded weighted average from its neighbors. This is only edge-weight + # normalization in the synthetic formula, not model feature normalization. + incoming_sum = torch.zeros((num_nodes, 1), dtype=torch.float32) + incoming_sum = incoming_sum.index_add(0, dst, raw_weight) + weight = raw_weight / incoming_sum.index_select(0, dst).clamp_min(1.0e-8) + + # 6. Build the tensors the model will receive. The fourth edge feature is + # the weighted scalar message weight * (u_src - u_dst). Including it keeps + # the learned problem small while still requiring source/destination indexing + # and destination aggregation to recover node:delta_u. + node_features = torch.cat((positions, u, kappa), dim=1).contiguous() + messages = weight * (u.index_select(0, src) - u.index_select(0, dst)) + edge_features = torch.cat((delta_pos, distance, messages), dim=1).contiguous() + + # 7. Compute the exact supervised answer. For each edge, the source node + # pushes the destination node toward u_src. index_add sums all incoming edge + # messages for each destination node, then dt scales the update. + target = torch.zeros((num_nodes, 1), dtype=torch.float32) + target = target.index_add(0, dst, messages) + target = dt * target + + graph = { + "node_features": node_features, + "edge_index": edge_index, + "edge_features": edge_features, + "global_features": dt.contiguous(), + } + return graph, target.contiguous() + + +def graph_to_model_input(graph: Dict[str, Tensor]) -> Dict[str, Tensor]: + """Canonicalize dtypes before eager, scripted, or reloaded model calls.""" + + return { + "node_features": graph["node_features"].to(dtype=torch.float32), + "edge_index": graph["edge_index"].to(dtype=torch.int64), + "edge_features": graph["edge_features"].to(dtype=torch.float32), + "global_features": graph["global_features"].to(dtype=torch.float32), + } + + +def assert_close(name: str, actual: Tensor, expected: Tensor, atol: float = 1.0e-6, rtol: float = 1.0e-6) -> None: + if not torch.allclose(actual, expected, atol=atol, rtol=rtol): + max_diff = (actual - expected).abs().max().item() + raise RuntimeError(f"{name} mismatch: max abs diff={max_diff}") + + +def run_feasibility(out_dir: Path) -> None: + # Mode 1: feasibility. + # + # Before training anything, prove that this model is deployable in the form + # AMS needs. We run the same randomly initialized model four ways: + # + # 1. normal eager PyTorch; + # 2. torch.jit.script output; + # 3. the scripted model after saving and loading it; + # 4. the AMS graph wrapper around the model. + # + # If these disagree, training would only hide a deployment problem. This + # mode intentionally does not write fixtures or a trained checkpoint. + out_dir.mkdir(parents=True, exist_ok=True) + print("[info] MGN graph diffusion feasibility gate") + print(f"[info] model_seed={MODEL_SEED}, fixture_sizes={FIXTURE_GRAPH_SIZES}, config={model_config()}") + + model = make_model().eval() + graphs = [graph_to_model_input(generate_graph(n, seed)[0]) for n, seed in zip(FIXTURE_GRAPH_SIZES, FEASIBILITY_SEEDS)] + + with torch.no_grad(): + eager = [model(graph)["node:delta_u"].detach() for graph in graphs] + + scripted = torch.jit.script(model) + with torch.no_grad(): + scripted_out = [scripted(graph)["node:delta_u"].detach() for graph in graphs] + for i, (actual, expected) in enumerate(zip(scripted_out, eager)): + assert_close(f"scripted graph {i}", actual, expected) + + script_path = out_dir / "mgn_graph_diffusion_feasibility.pt" + scripted.save(str(script_path)) + reloaded = torch.jit.load(str(script_path)).eval() + with torch.no_grad(): + reloaded_out = [reloaded(graph)["node:delta_u"].detach() for graph in graphs] + for i, (actual, expected) in enumerate(zip(reloaded_out, eager)): + assert_close(f"reloaded graph {i}", actual, expected) + + wrapped = create_ams_homogeneous_graph_model(model, torch.device("cpu"), torch.float32) + with torch.no_grad(): + wrapped_out = [wrapped(graph)["node:delta_u"].detach() for graph in graphs] + for i, (actual, expected) in enumerate(zip(wrapped_out, eager)): + assert_close(f"AMS-wrapped graph {i}", actual, expected) + + print(f"[info] Feasibility passed for dynamic graph sizes {FIXTURE_GRAPH_SIZES}") + + +def validation_graphs() -> List[Tuple[Dict[str, Tensor], Tensor]]: + # Fixed validation seeds make the primary "10x better than zero baseline" + # criterion reproducible across regeneration runs. + graphs = [] + for i in range(VAL_GRAPHS): + num_nodes = 16 + ((i * 37) % 113) + graphs.append(generate_graph(num_nodes, VAL_DATA_SEED + i)) + return graphs + + +def evaluate(model: nn.Module, cases: Iterable[Tuple[Dict[str, Tensor], Tensor]]) -> Tuple[float, float, Tensor]: + """Return model MSE, zero-prediction baseline MSE, and all targets.""" + + total_loss = 0.0 + total_baseline = 0.0 + total_nodes = 0 + targets = [] + model.eval() + with torch.no_grad(): + for graph, target in cases: + pred = model(graph_to_model_input(graph))["node:delta_u"] + loss_sum = (pred - target).pow(2).sum().item() + baseline_sum = target.pow(2).sum().item() + total_loss += loss_sum + total_baseline += baseline_sum + total_nodes += int(target.shape[0]) + targets.append(target) + return total_loss / total_nodes, total_baseline / total_nodes, torch.cat(targets, dim=0) + + +def run_train(out_dir: Path) -> None: + # Mode 2: train. + # + # The model sees many freshly generated synthetic graphs. For each graph we + # know the exact target_delta_u, so training is ordinary supervised learning: + # predict delta_u, measure mean-squared error, and update the network + # weights. This mode writes learned weights and metrics only. It does not + # export the AMS model or write C++ fixture tensors. + out_dir.mkdir(parents=True, exist_ok=True) + print("[info] MGN graph diffusion training") + print(f"[info] model_seed={MODEL_SEED}, train_data_seed={TRAIN_DATA_SEED}, val_data_seed={VAL_DATA_SEED}") + print(f"[info] config={model_config()}") + print( + "[info] training " + f"steps={TRAIN_STEPS}, lr={LEARNING_RATE}, weight_decay={WEIGHT_DECAY}, " + f"primary_pass=validation_mse <= {PRIMARY_BASELINE_FACTOR} * zero_baseline_mse" + ) + + model = make_model().train() + optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) + scheduler = torch.optim.lr_scheduler.MultiStepLR( + optimizer, + milestones=(TRAIN_STEPS // 2, (TRAIN_STEPS * 5) // 6), + gamma=0.3, + ) + val_cases = validation_graphs() + best_val_mse = float("inf") + best_zero_baseline_mse = float("inf") + best_step = 0 + best_state_dict = copy.deepcopy(model.state_dict()) + + for step in range(1, TRAIN_STEPS + 1): + # Generate one graph per optimization step. Avoiding batching keeps the + # example small and makes dynamic graph sizes part of normal training. + num_nodes = 16 + ((step * 53) % 113) + graph, target = generate_graph(num_nodes, TRAIN_DATA_SEED + step) + pred = model(graph_to_model_input(graph))["node:delta_u"] + loss = torch.nn.functional.mse_loss(pred, target) + + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + scheduler.step() + + if step == 1 or step % 200 == 0 or step == TRAIN_STEPS: + val_mse, zero_baseline_mse, _ = evaluate(model, val_cases) + if val_mse < best_val_mse: + best_val_mse = val_mse + best_zero_baseline_mse = zero_baseline_mse + best_step = step + best_state_dict = copy.deepcopy(model.state_dict()) + print( + f"[info] step={step:04d} train_mse={loss.item():.8e} " + f"val_mse={val_mse:.8e} zero_baseline_mse={zero_baseline_mse:.8e} " + f"best_step={best_step:04d} best_val_mse={best_val_mse:.8e}" + ) + + model.load_state_dict(best_state_dict) + val_mse, zero_baseline_mse, val_targets = evaluate(model, val_cases) + target_mean = val_targets.mean().item() + target_max_abs = val_targets.abs().max().item() + target_std = val_targets.std(unbiased=False).item() + primary_threshold = PRIMARY_BASELINE_FACTOR * zero_baseline_mse + + print( + "[info] target_delta_u stats " + f"mean={target_mean:.8e} max_abs={target_max_abs:.8e} std={target_std:.8e}" + ) + print( + f"[info] selected best_step={best_step}, validation_mse={val_mse:.8e}, " + f"zero_baseline_mse={zero_baseline_mse:.8e}, primary_threshold={primary_threshold:.8e}" + ) + # The primary success check is relative to a trivial model that always + # predicts zero. That is easier to explain and more robust to target scale + # changes than an absolute MSE alone. + if val_mse > SECONDARY_ABSOLUTE_MSE: + print( + f"[warn] validation_mse={val_mse:.8e} is above secondary diagnostic " + f"absolute MSE {SECONDARY_ABSOLUTE_MSE:.8e}" + ) + if val_mse > primary_threshold: + raise RuntimeError( + "Training failed primary MGN graph diffusion criterion: " + f"validation_mse={val_mse:.8e} > {primary_threshold:.8e}" + ) + + metadata = { + "model_config": model_config(), + "model_seed": MODEL_SEED, + "train_data_seed": TRAIN_DATA_SEED, + "val_data_seed": VAL_DATA_SEED, + "fixture_graph_sizes": list(FIXTURE_GRAPH_SIZES), + "fixture_seeds": list(FIXTURE_SEEDS), + "train_steps": TRAIN_STEPS, + "learning_rate": LEARNING_RATE, + "weight_decay": WEIGHT_DECAY, + "validation_mse": val_mse, + "best_step": best_step, + "best_validation_mse": best_val_mse, + "best_zero_baseline_mse": best_zero_baseline_mse, + "zero_baseline_mse": zero_baseline_mse, + "primary_baseline_factor": PRIMARY_BASELINE_FACTOR, + "secondary_absolute_mse": SECONDARY_ABSOLUTE_MSE, + "target_mean": target_mean, + "target_max_abs": target_max_abs, + "target_std": target_std, + "torch_version": str(torch.__version__), + } + checkpoint = { + "model_state_dict": model.state_dict(), + "metadata": metadata, + } + checkpoint_path = out_dir / CHECKPOINT_NAME + torch.save(checkpoint, checkpoint_path) + print(f"[info] Saved checkpoint: {checkpoint_path}") + + metrics_path = out_dir / TRAINING_METRICS_NAME + metrics_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"[info] Saved training metrics: {metrics_path}") + + +def load_checkpoint(path: Path) -> Dict[str, object]: + # PyTorch 2.6 changed torch.load's default weights_only behavior. The + # checkpoint is generated by this local script and contains metadata, so + # loading with weights_only=False is intentional here. + try: + return torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + return torch.load(path, map_location="cpu") + + +def write_tensor_binary(out_dir: Path, relative_path: str, tensor: Tensor, dtype: str) -> Dict[str, object]: + """Write one fixture tensor and return the matching manifest entry. + + The C++ test deliberately avoids NumPy, pickle, or torch file readers for + graph inputs. Those formats would make the C++ side depend on Python data + tooling. Instead each tensor is raw contiguous row-major little-endian bytes, + and fixtures.json records the dtype, shape, endianness, and byte count + needed to validate it before use. + """ + + rel_path = Path(relative_path) + if rel_path.is_absolute() or ".." in rel_path.parts: + raise ValueError(f"Fixture tensor path must stay relative: {relative_path}") + + if dtype == "float32": + packed_tensor = tensor.detach().cpu().to(dtype=torch.float32).contiguous() + values = [float(x) for x in packed_tensor.view(-1).tolist()] + data = struct.pack("<" + "f" * len(values), *values) if values else b"" + elif dtype == "int64": + packed_tensor = tensor.detach().cpu().to(dtype=torch.int64).contiguous() + values = [int(x) for x in packed_tensor.view(-1).tolist()] + data = struct.pack("<" + "q" * len(values), *values) if values else b"" + else: + raise ValueError(f"Unsupported fixture tensor dtype: {dtype}") + + tensor_path = out_dir / rel_path + tensor_path.parent.mkdir(parents=True, exist_ok=True) + tensor_path.write_bytes(data) + return { + "path": rel_path.as_posix(), + "dtype": dtype, + "shape": [int(dim) for dim in packed_tensor.shape], + "endianness": FIXTURE_ENDIANNESS, + "byte_size": len(data), + } + + +def run_fixtures(out_dir: Path) -> None: + # Mode 3: fixtures. + # + # A fixture is fixed test data: inputs plus the expected output. This mode + # loads the trained checkpoint, exports the AMS-wrapped TorchScript model, + # creates two fixed graph inputs, runs the exported model in Python, and + # saves those model outputs as the C++ reference. Fixtures mode never + # silently trains because parity tests should use an explicit checkpoint. + out_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = out_dir / CHECKPOINT_NAME + if not checkpoint_path.exists(): + raise RuntimeError( + f"Missing checkpoint {checkpoint_path}. Run --mode train, " + "or `ctest -R MGN_DIFFUSION_TRAIN`, before --mode fixtures." + ) + + checkpoint = load_checkpoint(checkpoint_path) + metadata = checkpoint["metadata"] + print(f"[info] Loading checkpoint: {checkpoint_path}") + print(f"[info] metadata={metadata}") + + model = make_model().eval() + model.load_state_dict(checkpoint["model_state_dict"]) + wrapped = create_ams_homogeneous_graph_model(model, torch.device("cpu"), torch.float32) + + model_path = out_dir / MODEL_NAME + wrapped.save(str(model_path)) + print(f"[info] Saved AMS TorchScript model: {model_path}") + + # Reload the saved artifact before generating references. That catches + # export-time issues and makes the Python reference match what C++ loads. + reloaded = torch.jit.load(str(model_path)).eval() + cases: List[Dict[str, object]] = [] + with torch.no_grad(): + for num_nodes, seed in zip(FIXTURE_GRAPH_SIZES, FIXTURE_SEEDS): + graph, _target = generate_graph(num_nodes, seed) + graph = graph_to_model_input(graph) + output = reloaded(graph)["node:delta_u"].detach().cpu().contiguous() + name = f"mgn_diffusion_n{num_nodes}" + case_dir = name + # Each tensor file is accompanied by manifest metadata. C++ uses the + # metadata to validate shape/type/size before constructing + # AMSTensor objects. + tensors = { + "node_features": write_tensor_binary( + out_dir, f"{case_dir}/node_features.bin", graph["node_features"], "float32" + ), + "edge_index": write_tensor_binary( + out_dir, f"{case_dir}/edge_index.bin", graph["edge_index"], "int64" + ), + "edge_features": write_tensor_binary( + out_dir, f"{case_dir}/edge_features.bin", graph["edge_features"], "float32" + ), + "global_features": write_tensor_binary( + out_dir, f"{case_dir}/global_features.bin", graph["global_features"], "float32" + ), + "reference_delta_u": write_tensor_binary( + out_dir, f"{case_dir}/reference_delta_u.bin", output, "float32" + ), + } + case = { + "name": name, + "seed": int(seed), + "num_nodes": int(num_nodes), + "num_edges": int(graph["edge_index"].shape[1]), + "node_feature_dim": NODE_FEATURE_DIM, + "edge_feature_dim": EDGE_FEATURE_DIM, + "global_feature_dim": GLOBAL_FEATURE_DIM, + "reference_output_dim": REFERENCE_OUTPUT_DIM, + "tensors": tensors, + } + print( + f"[info] fixture {case['name']} N={case['num_nodes']} " + f"E={case['num_edges']} reference_shape={tuple(output.shape)}" + ) + cases.append(case) + + # The manifest is the table of contents for the fixture directory. Paths are + # relative to fixtures.json so the generated build-tree directory can be + # moved as a unit. + manifest = { + "format_version": FIXTURE_FORMAT_VERSION, + "endianness": FIXTURE_ENDIANNESS, + "model": { + "path": MODEL_NAME, + }, + "training_metrics_path": TRAINING_METRICS_NAME, + "metadata": metadata, + "comparison": { + "rtol": COMPARISON_RTOL, + "atol": COMPARISON_ATOL, + }, + "cases": cases, + } + manifest_path = out_dir / FIXTURE_MANIFEST_NAME + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"[info] Saved fixture manifest: {manifest_path}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-dir", type=Path, required=True, help="Output directory for MGN graph diffusion artifacts") + parser.add_argument( + "--mode", + choices=("feasibility", "train", "fixtures", "all"), + required=True, + help="Separated MGN graph diffusion generation mode", + ) + args = parser.parse_args() + + if args.mode in ("feasibility", "all"): + run_feasibility(args.out_dir) + if args.mode in ("train", "all"): + run_train(args.out_dir) + if args.mode in ("fixtures", "all"): + run_fixtures(args.out_dir) + + +if __name__ == "__main__": + main()