Add hooks for PyTorch velocity emulator - #186
Draft
trhille wants to merge 26 commits into
Draft
Conversation
Add hooks for PyTorch velocity emulator using FTorch.
trhille
marked this pull request as draft
June 17, 2026 02:55
Add `use mpi` to make MPI constants available in GNN velocity hooks
Add build options for when FTORCH=true. User needs to define FTORCH_ROOT and LIBTORCH_ROOT environment variables.
Add batch tensor to pass to GNN model. For now, this is all zeros because all nodes in the graph belong to the same batch.
Re-order input features to match definition in GNN. Make units consistent with GNN, using km, yr, and °C instead of m, s, and K.
Remove garbage boundary cells, which caused an out-of-bounds error when calling the GNN.
Build a subgraph of ice-filled cells to send to GNN, to be consistent with the training approach. Non-ice cells get zero velocity. We should make sure that this does cause issues for code that was based on Albany's extended cell approach.
Pass single-precision arrays (real(4)) to GNN to match PyTorch's Float weights.
Disable FPE traps around LibTorch call. LibTorch has internal operations that trigger FPE but are handled internally. Disabling those traps around the LibTorch call enables us to compile in debug mode without getting FPE at runtime.
Add layout argument to torch_tensor_from_array to specify the order of indices.
Gather indexToCellID alongside features, then reorder global_features into global cell ID order on rank 0. Similarly, reorder velocity back to gathered order before scattering.
Assign a normalVelocity value at marginal edges that allow advection to produce ice advance. Without this, ice at the margins will thicken unphysically because the normalVelocity is an average of the non-zero upstream cell velocity and the zero downstream cell velocity. This was causing the seaward margins of the Filchner-Ronne and Ross ice shelves to become grounded in test simulations!
Update halos on uReconstructX and uReconstructY after populating them from the GNN prediction on rank 0.
Include ice-free cells that are adjacent to ice-filled cells to give the model a sense of where boundary conditions are applied.
Add FTORCH_CUDA compile flag to make for flexible build with either CPU or GPU.
This reverts commit f66b276.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This merge adds hooks for PyTorch-based velocity emulator using FTorch.