Field GP - #1887
Merged
Merged
Conversation
Adds a pointwise independent multitask variational GP head to physicsnemo.experimental.uq. It is the field sibling of VariationalGPHead: where that head pools a geometry to one embedding and predicts a scalar, this one keeps the point dimension and returns a Gaussian posterior per point per channel. A single forward pass therefore yields the field prediction, the total predictive variance and the epistemic-only variance, with no ensembling or MC-Dropout sampling. The head is backbone-agnostic by construction: it consumes only a (..., input_dim) feature tensor, so any point-wise encoder can drive it. To make that usable, GeoTransolver gains a `return_point_features` flag that exposes the per-point latents computed just before its output projection (default False, so existing callers are unaffected). Also documents VariationalGPHead as the scalar counterpart, since its name does not say so, and keeps `FieldGPHead` / `FieldGPPrediction` as aliases for back-compatibility. Note the aliases do not preserve checkpoint filenames: save_checkpoint derives the file stem from type(model).__name__, so new runs write FieldVariationalGPHead.0.<tag>.pt.
Three changes to the shared transformer-model training scripts, all needed to run a field-GP recipe alongside the existing baselines on equal terms: - gp_utils: compute_drag_target_from_batch assumed the full-mesh `surface_normals` / `surface_areas` keys. With point subsampling (data.resolution below the full mesh) the datapipe emits subsampled, field-aligned `*_sub` arrays instead and omits the full ones, so the drag integral was computed against mismatched geometry. Fall back to the subsampled keys when the full arrays are absent. - train: add optional gradient accumulation to train_epoch, so the deterministic baseline can match the GP run's effective batch size. Defaults to 1, which is bit-identical to the previous behaviour. - train_gp_combined: log the per-epoch drag-ranking Spearman across all validation geometries, gathered across ranks. This is the active-learning objective, and logging it the same way the field-GP recipe does makes the scalar and field heads directly comparable.
Trains GeoTransolver with FieldVariationalGPHead replacing its readout, so the GP posterior mean is the surface field (pressure + 3 wall-shear-stress components) and the posterior variance is the per-point uncertainty. The config defaults are the settled recipe, so the run reproduces with no hyperparameter overrides -- only data paths and a run id. Three of those defaults are load-bearing and documented in the README as such: the l2_radial feature normalisation (fixes the GP-input feature scale while keeping the radial out-of-distribution cue that a plain unit-sphere projection erases), the heteroscedastic noise MLP (a constant per-channel noise makes the total predictive std rank points identically to the epistemic std), and the noise std floor (the heteroscedastic ELBO weights points by 1/sigma^2(x), so too low a floor lets one collapsing point destabilise training). DrivAerStar surface fields are in raw physical units, so this recipe gets its own normalisation stats under src/normalization/drivaerstar/ rather than overwriting the nondimensional stats in src/, which belong to the DrivAer AWS runs and are shared with the other recipes. Checkpointing is made crash-consistent: Slurm preemption during the (non-atomic, three-file) end-of-epoch save could otherwise leave a torn or partial set, and physicsnemo's loader globs the latest index of each file independently, so a resume could pair a backbone from epoch N with an optimizer from N-1. Saves now stage and fsync before an atomic rename, committing the training-state file last as a marker, and resume selects only a complete epoch set.
Adds a README section spelling out that the total predictive variance is the GP posterior variance plus an input-dependent observation noise, which of the two to use for which job, and how to refer to the second term. The naming matters and the previous wording was loose. In the loss, sigma^2(x) is just the variance of the Gaussian likelihood, so "input-dependent observation noise" is the right methods-level name. But DrivAerStar targets are deterministic steady-RANS solutions, so there is no observational scatter to learn and reading it as measurement noise or "inherent flow variability" is wrong. What it absorbs is the residual the GP mean cannot represent, making it a learned model-discrepancy variance in the Kennedy & O'Hagan sense -- which is also what makes the per-point noise field worth plotting. The word "aleatoric" is dropped except to note that we borrow Kendall & Gal's mechanism (an amortized network on the kernel's own features) rather than their interpretation. One in-code comment claimed a wake has "higher irreducible variance than the hood", which is exactly the reading to avoid; it now says discrepancy. Also records that the noise model is a point estimate: the classical variational heteroscedastic GP puts a second GP on the log-noise and gets a KL term that resists the noise running to extremes, whereas amortizing with a network has no prior, so noise_std_range and gradient clipping are the only guard. That is why those two settings are load-bearing rather than cosmetic. References now cover the whole composition rather than a subset -- the sparse variational bound (Titsias; Hensman, Fusi & Lawrence, whose normalisation _hetero_neg_elbo mirrors), deep kernel learning and its documented failure modes (Wilson et al.; Ober et al.), and the heteroscedastic-noise line (Goldberg et al.; Lazaro-Gredilla & Titsias; Liu et al. for the sparse heteroscedastic ELBO closest to ours). Docs only; predict() output remains bit-identical on the reference checkpoint.
Contributor
CODEOWNERS review mapCurrent for commit ⏳ @coreyjadams — 11 file(s)
No CODEOWNER
Comment |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
Greptile SummaryAdds a field-level variational GP implementation and a GeoTransolver training and evaluation recipe.
Important Files Changed
Reviews (1): Last reviewed commit: "Document the variance decomposition and ..." | Re-trigger Greptile |
The head is now a hydra `_target_` block under `gp_head`, so training and inference build it from one source of truth instead of each assembling flat `gp_*` keys by hand. That is what lets an eval run load a training checkpoint: the two call sites can no longer drift into structurally different heads. Removes what physicsnemo-cfd already measures. The training-time drag UQ stats were only ever logged, and plot_field_gp.py's calibration, error-vs-std and ID/OOD diagnostics are covered by the benchmark's uncertainty_error_spearman, calibration_zrms and per-class rows. The normalized std arrays in the VTK output are dropped too, since nothing read them and they are recoverable from the physical arrays. Documents how the head learning rates are set relative to the backbone's, and returns gpytorch to the uq-extras optional dependency so the deterministic models and Concrete Dropout stay installable without it, matching how the scalar GP head handles the same dependency.
The datapipe samples geometry points without replacement, so requesting more points than a mesh has is now an error rather than silently returning duplicate indices, following the upstream switch from poisson_sample_indices_fixed to weighted_multinomial. The smallest class_F meshes hold ~265k points and 13% of both splits sit below the previous 300k request, so training died on the first epoch: the ranks holding those geometries raised, and the rest hung in an allreduce until the NCCL watchdog fired. 200k clears every geometry in the train and val splits.
Nothing imports the pre-rename names any more: the physicsnemo-cfd wrapper now builds FieldVariationalGPHead directly, and the only remaining reference was the test asserting the aliases resolved to the new classes. Checkpoints are state dicts rather than pickled classes, so files written under the old name still load from an explicit path with the aliases gone. The resume scan looks for a single head-checkpoint stem as a result, which lets _ckpt_indices drop the tuple-of-alternatives handling that existed only to match both spellings. Each entry's trailing suffix element goes too; nothing read it. Resuming a checkpoint directory written before the rename now needs the files renamed, which only affects the superseded runs. On the docs: the head-comparison table was malformed, its rows overflowing the column rulers, so it is rewritten with B, D and num_tasks spelled out. n_train is the number of training points and is easy to read as a count of geometries, so it now says 10 geometries of N points is 10 * N. The heteroscedastic-noise rationale loses its DrivAerStar-specific variance figures, since this is the general API, and points at _hetero_neg_elbo for the second-GP trade-off.
Adds a figure of the per-point epistemic std of surface pressure for a head trained on Fastback alone and evaluated on Notchback and Estateback, which is the clearest evidence that the uncertainty is smooth and tracks how far a geometry sits from the training set. The scalar GP's KDE figure never rendered: this README is four directories deep but reached for ../../../docs/img, which resolves inside examples/ rather than the repository root. Every other README under external_aerodynamics uses four levels, so both figures now do as well.
The head is not DDP-wrapped, and only its inducing points were broadcast. Torch seeds itself per process, so each rank built a different random DKL and noise MLP and then kept that offset: averaging gradients shares the update, not the parameters. The documented eight-GPU recipe was therefore training eight slightly different heads and checkpointing rank 0's. Broadcasting the whole state dict after construction is what DDP does for the backbone. The feature-norm BatchNorm's running stats still follow each rank's own shard afterwards, which costs nothing: train mode normalises by batch statistics, and eval and the checkpoint both read rank 0's.
The DrivAerStar surface statistics are Pa-scale — pressure mean and std are -138.7 and 287.6 — so unstandardizing already lands in physical units, and the datapipe's own re-dimensionalization is commented out for the same reason. Multiplying by rho * u^2 on top of that inflated every mean and std in the VTK output by 1084.5x with the documented defaults. physicsnemo-cfd never did this (REDIMENSIONALIZE_OUTPUTS is False for these checkpoints), so the benchmark numbers were never affected. Also fails when no checkpoint is loaded, rather than writing confident-looking predictions from random weights. Checking load_checkpoint's return value is not sufficient on its own: it resolves each model separately and a missing file is only a warning, so a directory holding a backbone but no head returns a nonzero epoch and leaves the head randomly initialized. The files are checked up front instead, with the return value guarded behind that. Both raise CheckpointError, which is re-raised past the per-run handler that would otherwise log the failure and move on to the next run.
The Gamma lengthscale and outputscale priors only reach the objective through GPyTorch's VariationalELBO, which sums registered prior log-probabilities into the loss. Setting noise_mlp_hidden routes training through the hand-built heteroscedastic ELBO instead, which has no prior term, so the priors are inert and the hard lengthscale_range constraint is what shapes the kernel. The surface recipe enables the noise MLP while also setting both priors, which reads as though they are doing something. Documented rather than changed: adding the terms would alter the objective and invalidate the settled recipe.
Collaborator
|
I've been bogged down with toolkit-ops reviews and whatnot, but I'll review first thing next week. |
melo-gonzo
approved these changes
Aug 4, 2026
melo-gonzo
left a comment
melo-gonzo
left a comment
Collaborator
There was a problem hiding this comment.
Overall lgtm. Added some consolidation between train/inference and the utils files + a few other minor comments.
laserkelvin
requested changes
Aug 5, 2026
laserkelvin
left a comment
laserkelvin
left a comment
Collaborator
There was a problem hiding this comment.
Various comments; one universal theme is please use jaxtyping for tensor annotations. I think a lot of them have known shapes but they're just hinted as torch.Tensor
Review asked for jaxtyping on tensor annotations rather than bare torch.Tensor, so the head's arguments and returns now carry their shapes: points, tasks, the feature width and the kernel-input width, which differ and were previously indistinguishable at the signature. n_train becomes a required keyword argument. It was `int | None = None` with a ValueError in the body, which advertised a default that does not exist -- the ELBO's normalisation constant is the caller's dataset size and nothing can stand in for it. Everything after input_dim is keyword-only now: with fifteen parameters, all of them set from a config block, positional construction only invites a silent mis-ordering. The Matern order was hardcoded at 2.5 with no statement of what it meant. It is an assumption about the field -- how many times the sample paths are differentiable -- so it is now a documented matern_nu parameter, defaulting to the same 2.5. feature_norm and GeoTransolver's attention_type become Literals, which is what their two- and three-valued domains were already enforcing at runtime. On the docs: the subpackage claimed to produce "calibrated uncertainty estimates", which overstates it -- these are posterior variances, and whether their scale tracks observed error depends on the recipe and the train/test shift, so both docstrings now say that and point at held-out validation. The subpackage docstring also read as though the two GP heads were the whole UQ API; the heads move under their own section. The head's module docstring gains the recipe the architecture does not carry: inducing points seeded from real features, the mean anchor, the KL ramp, and the noise floor with gradient clipping. Those live in the training script because they are not properties of the model, but a caller who skips them gets a collapsed variance rather than a merely worse one, so the docstring now says which they are and where the worked example is. It also explains why the float64 cast and the public transform_features exist, both of which read as unexplained machinery.
…es from Deduplication. _probe_feature_dim existed twice, once per script, differing only in whether it truncated the point dimension and whether it logged; a full boundary mesh is far larger than the training batches, which is why inference grew its own copy. It is now one probe_feature_dim in field_gp_utils with an optional max_points, restoring the backbone's previous mode before returning rather than leaving it in eval. It stays with the other field-GP helpers rather than moving to gp_utils, since it needs return_point_features, which only the field head uses. beta_ramp_weight was gp_ramp_weight with an extra guard for an empty warmup window, so the guard moves into gp_ramp_weight and the duplicate goes. collect_inducing_features no longer round-trips through host memory: the caller sends the result straight back to the device, and the tensor is at most n_inducing x D. field_gp_predict_full_mesh drops its unused device argument. Determinism. Everything stochastic in this recipe -- the initialisation, the datapipe's point sampling, the inducing-point draw, the per-step GP subsample, the distance penalty's pairs -- draws from torch's global RNG, so threading explicit generators would have to reach into the datapipe to cover it. A cfg.seed knob seeds that RNG instead, offset per rank so ranks keep drawing different point subsets, and is passed to DistributedSampler, whose own default seed of 0 otherwise repeats the same epoch order in every run. Left null the behaviour is unchanged, which is what the reference results were produced with. The functions that draw say so in their docstrings. Typing and validation. jaxtyping shapes on the tensor arguments and returns, the head annotated as FieldVariationalGPHead under a TYPE_CHECKING guard so its "install gpytorch" message is not traded for an ImportError on the name, and precision as a Precision literal checked by validate_precision -- cast_precisions silently no-ops on an unrecognised value, so a typo would otherwise run in float32. The VTP path is checked before pyvista reads it, and the unstandardizing is done as tensor arithmetic. Strict config access. getattr(cfg, key, default) turns a mistyped YAML key or override into a silent fallback that reads exactly like a set value in the logs. The twenty knobs this recipe needs are now listed and checked once at startup, naming all the missing ones at once, and read as plain attributes thereafter. cfg.get survives only where absent genuinely means "derive it". Also vectorises the distance penalty's per-geometry loop into one batched pair draw, comments why the OmegaConf types have to be registered as safe globals for weights_only unpickling, records where the freestream defaults come from, and keeps the head's pre-clip gradient norm as a tensor so the device sync happens when the log line is formatted rather than every step. The config's acronyms are spelled out on first use.
The DOI given for Liu, Cai & Ong's "Remarks on multi-output Gaussian process regression" resolved to a different Knowledge-Based Systems article from the same year, Goyal & Ferrara's graph-embedding survey. The paper is volume 144, pages 102-121, doi 10.1016/j.knosys.2017.12.034; the authors are named now rather than abbreviated, since "Liu et al." appears twice in this list for different papers.
The review landed on the field head and the field-GP scripts, but four of the comments describe patterns that recur in the other files this PR touches, where no reviewer happened to look. VariationalGPHead had the same two constructor problems as its field sibling: nu was hardcoded at 2.5, and n_train was annotated optional while the body raised on None. It gains a matern_nu argument and takes n_train as a required keyword-only argument, matching FieldVariationalGPHead's signature; every call site already passed it by keyword, so nothing has to change around it. Its claim to produce calibrated uncertainty is relaxed the same way the subpackage docstring's was -- calibration is a property of the recipe, and the head cannot promise it. Tests now pin the constructor contract, which was previously untested. GeoTransolver's state_mixing_mode is the same shape as attention_type: two values, validated in the GALE layer, annotated str. It becomes a Literal too. reinitialize_inducing_points staged its collected embeddings on the host and copied them straight back to the device eight lines later, the same round trip the field-GP collector was pulled up on. The function is under no_grad, so dropping it holds nothing extra. train.py carries the same OmegaConf safe-globals block as the inference script and had no explanation of why weights_only unpickling needs it, only a note that it should eventually go away. Not touched: train_gp_combined.py's getattr(cfg, ...) defaults and train.py's own precision and normalization_dir fallbacks are the same silent-fallback pattern, but they sit on the settled scalar-GP and deterministic paths; they get the same pass in a follow-up rather than a behaviour change buried in this one.
Reverts the Literal on state_mixing_mode from the previous commit. The review asked for it on attention_type, and this PR has no other reason to be on that line; "the same pattern is next door" is not enough to widen the diff of a model every other recipe in the repository instantiates. Nothing was wrong with the change itself -- annotations are not enforced at runtime here, and checkpoints store constructor values rather than types -- so it stands on its own whenever someone wants it.
Making n_train required and keyword-only changes an existing public class, so it belongs under Changed rather than only in the commit that did it. Also notes matern_nu, and corrects the field head's entry now that its kernel order is configurable rather than fixed at 5/2.
Main promoted GeoTransolver to physicsnemo.models.geotransolver and left the experimental module as a checkpoint shim, so the field-GP additions move with it: return_point_features and the attention_type Literal now live on the mainline model, with the GALE cross-reference pointing at physicsnemo.nn.
The head docstring says a working recipe needs a mean anchor, two ramps, a noise floor and gradient clipping, and points at the README for the reasoning. The README only covered the noise floor: the ramp windows, the clipping norm and the distance penalty appeared as names in the loss formula and nowhere else. Add a table under Training giving each term its config key, default and role, and extend the porting notes so a new backbone knows which carry over unchanged and which need retuning against its own target scale and epoch budget.
The docstrings, comments and README added for the two GP heads had drifted into British spellings — normalised, initialisation, behaviour, optimiser. Convert them, along with the handful that predate this branch in the same files, so a reader does not meet both conventions in one docstring. Prose only.
Collaborator
Author
|
/ok to test f0e5178 |
coreyjadams
approved these changes
Aug 7, 2026
coreyjadams
left a comment
coreyjadams
left a comment
Collaborator
There was a problem hiding this comment.
Hi @ktangsali I reviewed the GeoTransolver changes with just a few small requests. No show stoppers so I'm approving the PR.
return_embedding_states and return_point_features share a return signature, so a positional True at the call site does not say which was meant. The * goes after time, keeping the four positional arguments positional: the shared test helpers call the model with a positional tuple of those. Every caller across the repos already passed the flags by keyword, so none needed updating. Cover the paths while here. One test runs a model through all four combinations and checks the prediction is identical in each, since the risk with two flags sharing a shape is that one path quietly diverges; the rest pin the positional call raising, the point-feature width under local features, and the multi-stream branch.
laserkelvin
reviewed
Aug 7, 2026
laserkelvin
reviewed
Aug 7, 2026
laserkelvin
reviewed
Aug 7, 2026
"Entropy-seeded" was jargon for unseeded: the config comment now says that with seed null torch keeps the nondeterministic seed it takes at startup, so runs differ. The inducing-feature helper's jaxtyping axes read "inducing dim", where the first looks like a flag rather than a count and the second is the head's input_dim, not its gp_dim. Name both for what they are, and make the head's inducing annotations agree.
Collaborator
|
@ktangsali @mnabian Where do you envision drawing a line between |
Collaborator
Author
I think both of these are serving different aspects. I like to think of UQ and Guardrails as a Venn diagram - yes, there is an intersection where both can be used for OOD detection, but their use cases and the way they operate differ. Roughly: a guardrail looks at what you feed the model and warns you when it falls outside what the model was trained on; UQ looks at what came out and puts an error bar on it. And a model can be uncertain on a perfectly in-distribution sample, which is something a guardrail has nothing to say about. I don't think we need to settle this in a PR though. It's worth doing properly as part of the namespace design when this comes out of experimental. |
Collaborator
I agree - guardrails generally can't do UQ, and not all of our UQ tools can do OOD detection. We should clarify the overlaps and differences in the user dcumentation. |
Collaborator
|
They are two sides of a venn diagram, as long as that venn diagram is narrowly scoped to "ways to ask if my model is right?" :) I don't think we want to have too many high level folders that can appear similar to users, I'm imagining a generalized uq/guardrails folder outside of experimental where these might be siblings. @laserkelvin do you think this API discussion in experimental should hold the PR? If there are other concerns, that's different. This isn't moving out of experimental today, I think we should put it as a discussion soon for how this eventual API could look, including CFD, alchemi, and e2s use cases. But that's not really this PR. |
laserkelvin
approved these changes
Aug 10, 2026
laserkelvin
left a comment
laserkelvin
left a comment
Collaborator
There was a problem hiding this comment.
I'll approve to unblock the PR, but when I have a second I'll post replies to the discussion
Collaborator
Author
|
/ok to test d37a8d9 |
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
5 participants
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.
PhysicsNeMo Pull Request
Description
Adds GPs for fields
Checklist
Dependencies
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.
AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.