Diffusion fwi refactor

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

Closes #1536 .

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

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.

[greptile-apps]

Greptile Summary

This PR refactors the diffusion FWI example to replace several local utility classes (DiffusionAdapter, ModelBasedGuidance, EDMLoss, edm_precond, EDMStochasticSampler) with the canonical implementations from physicsnemo.diffusion, and moves previously hard-coded EDM noise-schedule parameters into the config files. As part of this, DiffusionFWINet.forward was renamed from (x, y, sigma) to (x, t, condition).

• test_diffusion_fwi_net.py was not updated to match the new signature: all three test functions call model(x, y, sigma), which now passes the seismic-observation tensor as t and the scalar noise label as condition, causing an immediate ValueError on the shape guard.

Important Files Changed

Filename	Overview
examples/geophysics/diffusion_fwi/tests/test_diffusion_fwi_net.py	Not modified in this PR, but broken by it: three test functions call model(x, y, sigma) using the old argument order, which now maps seismic observations to the noise-label slot and vice versa, causing an immediate ValueError on shape checks.
examples/geophysics/diffusion_fwi/utils/nn.py	Forward signature changed from (x, y, sigma) to (x, t, condition=None) to align with EDMPreconditioner interface; adds jaxtyping annotations and guards validation checks with not torch.compiler.is_compiling(). The signature change breaks the existing test file which was not updated.
examples/geophysics/diffusion_fwi/generate.py	Replaces local EDMStochasticSampler/ModelBasedGuidance/generate with physicsnemo's EDMNoiseScheduler, DPSScorePredictor, ModelConsistencyDPSGuidance, and sample(); ClippedGuidance is re-defined every data iteration; torch.cat on a potentially-empty list is an edge-case crash.
examples/geophysics/diffusion_fwi/train.py	Replaces local DiffusionAdapter/edm_precond/EDMLoss with physicsnemo's EDMPreconditioner, MSEDSMLoss, and EDMNoiseScheduler; noise schedule parameters moved to a new config section; clean overall refactor with no logic issues found.
examples/geophysics/diffusion_fwi/conf/config_train.yaml	Adds a new noise_schedule section (P_mean, P_std, sigma_data) to expose previously hard-coded EDM parameters.
examples/geophysics/diffusion_fwi/conf/config_generate.yaml	Adds noise_schedule.sigma_data; removes several guidance parameters (scale, power, norm_ord, magnitude_scaling) replaced by the simpler ModelConsistencyDPSGuidance interface; renames std to std_y.

Comments Outside Diff (1)

1. examples/geophysics/diffusion_fwi/tests/test_diffusion_fwi_net.py, line 127-128 (link)

   Test broken by forward-signature change

   DiffusionFWINet.forward was renamed from (x, y, sigma) to (x, t, condition), but this test file was not updated. The call model(x, y, sigma) now maps y (shape (B, 6, 50, 32)) to t and sigma (shape (B,)) to condition, which is the exact opposite of the new signature. The validator if t.shape != (B,) will raise a ValueError immediately, causing all three test functions in this file to fail.

   The same incorrect call pattern appears at lines 157 and 211 in the same file.

_{Reviews (1): Last reviewed commit: "Merge branch 'main' into diffusion-fwi-r..." | Re-trigger Greptile}

[CharlelieLrt]

@Dibyajyoti-Chakraborty could you please check the modifications on the wave operator I made in this PR?
I changed a few things:

• Re-interpolation to original 70x70 resolution before the wave solver: can you confirm that this is aligned with what we discussed before?
  ⚠️ Now we are essentially doing 70x70 -> 80x80 -> 70x70 on the velocity model that is used as input for the wave prediction. I verified, and this double-interpolation is not reversible, so it does introduce a slight smoothing of the velocity model. This means that the y generated from the wave_operator in the guidance will have missing high frequencies.
• Problem in the smooth_clamp function, which was disabling the guidance form most of the velocity range: might have affected the experiments you ran
• The output of the wave operator was not normalized so the guidance was comparing normalized y vs. non-normalized y: can you confirm that you agree with this change? I am not sure how it is handled in the version of the code you used for your experiments, but it might have created some problems

[pzharrington]

Adoption of physicsnemo.diffusion APIs looks good to me

[CharlelieLrt]

/blossom-ci