Add urban_flow_diffusion_3d example: 3D diffusion prior for unconditional sampling in a simplified turbulent urban environment by AbVishwas · Pull Request #1823 · NVIDIA/physicsnemo

AbVishwas · 2026-07-09T23:31:43Z

#Add urban_flow_diffusion_3d example: 3D diffusion prior for unconditional sampling in a simplified turbulent urban environment

Description

Ports a research training pipeline (previously developed against an older
physicsnemo fork with SongUNet3D/EDMPrecond3D) onto the current
physicsnemo.diffusion stack: DiffusionUNet3D + EDMPreconditioner +
EDMNoiseScheduler + MSEDSMLoss. The denoising backbone architecture is
based on Diff-SPORT (arXiv:2506.00214); this example implements only the
unconditional generative-prior training and sampling stage of that
pipeline -- not Diff-SPORT's sensor-placement optimization or its
conditional/posterior reconstruction from sparse observations. Includes
Hydra-based train/generate scripts, dataset loaders supporting both
combined and separate U/V/W HDF5 layouts, and Reynolds-stress evaluation
plotting.

Dataset used for development/smoke-testing:
https://huggingface.co/datasets/abvish/UrbanFlow-oneObstacle-NoTrip

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

None beyond what's already in examples/cfd/urban_flow_diffusion_3d/requirements.txt (h5py, hydra-core, matplotlib, omegaconf, psutil, tqdm, wandb).

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.

…onal sampling in a simplified turbulent urban environment Ports a research training pipeline (previously developed against an older physicsnemo fork with SongUNet3D/EDMPrecond3D) onto the current physicsnemo.diffusion stack: DiffusionUNet3D + EDMPreconditioner + EDMNoiseScheduler + MSEDSMLoss. The denoising backbone architecture is based on Diff-SPORT (arXiv:2506.00214); this example implements only the unconditional generative-prior training and sampling stage of that pipeline, not sensor-placement optimization or conditional/posterior reconstruction. Includes Hydra-based train/generate scripts, dataset loaders supporting both combined and separate U/V/W HDF5 layouts, and Reynolds-stress evaluation plotting. Signed-off-by: Abhijeet Vishwasrao <abvish@lh-login2.arc-ts.umich.edu>

copy-pr-bot · 2026-07-09T23:31:47Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-07-09T23:36:52Z

Greptile Summary

This PR adds a new urban_flow_diffusion_3d example that ports a 3D diffusion-based urban-flow generative prior onto the current physicsnemo.diffusion stack (DiffusionUNet3D + EDMPreconditioner + EDMNoiseScheduler + MSEDSMLoss), with Hydra configs, HDF5 dataloaders, and Reynolds-stress evaluation plots.

Several concrete bugs were found:

train.py: DistributedSampler.set_epoch uses the frozen checkpoint-start epoch instead of the current loop variable ep, causing every epoch to use the same shuffle permutation.
dataset_utils.py: resize_spatial_tensor_yz calls F.pad but torch.nn.functional is never imported as F; padding will crash at runtime.
dataset_utils.py: lazy_load_h5 returns a raw HDF5 file handle that no caller closes, leaking the file descriptor.
src/utils.py: plot_two_comps references plot_config which is not in its signature or scope — every call raises a NameError.

Important Files Changed

Filename	Overview
examples/cfd/urban_flow_diffusion_3d/train.py	Main training loop — `set_epoch(epoch)` uses the frozen start epoch instead of the current `ep`, causing identical shuffle permutations every epoch in distributed training.
examples/cfd/urban_flow_diffusion_3d/src/dataloaders/dataset_utils.py	Dataset utilities — `resize_spatial_tensor_yz` calls undefined `F.pad` (NameError), `lazy_load_h5` leaks HDF5 file handles, and debug emoji prints clutter logs.
examples/cfd/urban_flow_diffusion_3d/src/utils.py	`plot_two_comps` references undefined `plot_config` — every call raises NameError.
examples/cfd/urban_flow_diffusion_3d/src/dataloaders/uflow_dataset.py	Dataset classes — `@cache` on instance methods creates GC-preventing reference cycles; `OldUflowDataset3D` is legacy code.
examples/cfd/urban_flow_diffusion_3d/src/gen_utils/gen_helpers.py	`rescale_and_crop_ds4` is unconditionally applied for all dataset types with a hardcoded 3-channel assumption.
examples/cfd/urban_flow_diffusion_3d/src/dataloaders/dataset_builder.py	Dataset/dataloader builder — logic is correct; `datatype` kwarg is passed but effectively ignored by `UflowDataset3D`.
examples/cfd/urban_flow_diffusion_3d/src/train_utils/train_helpers.py	Training utility helpers — logic looks correct; `configure_cuda_for_consistent_precision` is called redundantly in `train.py` as well (harmless).
examples/cfd/urban_flow_diffusion_3d/generate.py	Inference entry point — straightforward; inherits risks from `gen_helpers.py`.
examples/cfd/urban_flow_diffusion_3d/src/metrics/statistics.py	Reynolds-stress evaluation and plotting module — no critical logic bugs identified.

Comments Outside Diff (4)

examples/cfd/urban_flow_diffusion_3d/train.py, line 826 (link)

Sampler epoch not advancing in training loop

set_epoch is called with the frozen epoch variable (the checkpoint-loaded starting epoch) instead of the loop variable ep. This means every epoch will use the exact same shuffle permutation, breaking the distributed sampler's designed variance between epochs.
examples/cfd/urban_flow_diffusion_3d/src/dataloaders/dataset_utils.py, line 293-320 (link)

lazy_load_h5 leaks the HDF5 file handle

hf = h5py.File(file_path, "r") is opened inside a plain function and returned to the caller. None of the callers in this PR call hf.close(), so the file descriptor stays open for the entire process lifetime. There is also no else/raise after the three if/elif branches — an invalid load_components value silently returns None.
examples/cfd/urban_flow_diffusion_3d/src/utils.py, line 875-915 (link)

plot_two_comps references undefined plot_config

The function body references plot_config.figure.figsize and plot_config.axes.fontsize, but plot_config is not in the function signature and not in scope. Every call raises a NameError.
examples/cfd/urban_flow_diffusion_3d/src/dataloaders/uflow_dataset.py, line 455-461 (link)

@cache on instance methods causes a memory-cycle / GC issue

functools.cache holds a strong reference back to the instance via self in the cache key, preventing garbage collection. A simple instance attribute set during __init__ would be cleaner and avoids the cycle.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "Add urban_flow_diffusion_3d example: 3D ..." | Re-trigger Greptile}

greptile-apps · 2026-07-09T23:36:56Z

+import numpy as np
+import torch
+import h5py
+from typing import Optional, Union


torch.nn.functional is never imported as F in this module, so calling F.pad(...) will raise a NameError at runtime whenever the padding branch is taken.

Suggested change

import numpy as np

import torch

import h5py

from typing import Optional, Union

import numpy as np

import torch

import torch.nn.functional as F

import h5py

from typing import Optional, Union

greptile-apps · 2026-07-09T23:37:00Z

+        preds_np = rescale_and_crop_ds4(preds, min_scaler, max_scaler, sample_shape[2:])
+
+        save_uncond_samples(filename_h5, preds_np, batch_seeds.cpu().numpy())
+
+    if dist.world_size > 1:
+        torch.distributed.barrier()
+    logger0.info("Generation complete.")


rescale_and_crop_ds4 unconditionally applied regardless of dataset type

The function is commented as "only used for ds4" but is called for all inference runs. The hardcoded (1,3,1,1,1) min/max reshape will silently produce wrong values for any non-3-channel dataset.

greptile-apps · 2026-07-09T23:37:04Z

+            return F.pad(x, pad=pad, mode="constant", value=pad_value)
+        else:
+            pad_width = [(0, 0)] * x.ndim
+            pad_width[-2] = (pad_top, pad_bottom)


Debug print with emoji in library-level utility

These prints fire on every dataset construction and will clutter distributed-training logs across all ranks. Use the logger infrastructure already present in the training scripts, or remove them.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

CharlelieLrt · 2026-07-10T02:41:54Z

This file seems to contain numerical values specific to the dataset. I think it should not be there and instead should be part of the dataset hosted on HuggingFace.

These numbers are used by the code to normalise the batches frequently each epoch. I can add a copy of this file on HF, but we will still need them here, right?

For all the numerical values, image resolutions, etc... I think they can be packaged directly with the dataset. The only thing to retain in the example yaml is the path of the dataset (where it is downloaded from HF).

Yes, It was just done this way to be flexible during development, but now it makes sense to add these values in h5 directly.

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

…rban_flow_diffusion_3d Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

…p32 inputs, default fp32 Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

…_flow_diffusion_3d Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

…fusion_3d Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

…s on-device, drop redundant barriers Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

greptile-apps Bot reviewed Jul 9, 2026

View reviewed changes

CharlelieLrt reviewed Jul 10, 2026

View reviewed changes

CharlelieLrt added 6 commits July 10, 2026 11:56

Made 3D diffusion UNet AMP-compatible

994f05c

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Add dataset download instructions and huggingface_hub dependency to u…

87643ea

…rban_flow_diffusion_3d Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Follow standard torch.autocast AMP in urban_flow_diffusion_3d: keep f…

c2e956a

…p32 inputs, default fp32 Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Replace hardcoded evaluate pred_dir with a run-relative path in urban…

1592932

…_flow_diffusion_3d Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Replace wandb logging with PhysicsNeMo PythonLogger in urban_flow_dif…

423b519

…fusion_3d Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Speed up urban_flow_diffusion_3d: stop disabling TF32, accumulate los…

3643b0f

…s on-device, drop redundant barriers Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

CharlelieLrt force-pushed the example/uflow-diffusion-3d branch from 62123f0 to 3643b0f Compare July 11, 2026 02:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add urban_flow_diffusion_3d example: 3D diffusion prior for unconditional sampling in a simplified turbulent urban environment#1823

Add urban_flow_diffusion_3d example: 3D diffusion prior for unconditional sampling in a simplified turbulent urban environment#1823
AbVishwas wants to merge 7 commits into
NVIDIA:mainfrom
AbVishwas:example/uflow-diffusion-3d

AbVishwas commented Jul 9, 2026

Uh oh!

copy-pr-bot Bot commented Jul 9, 2026

Uh oh!

greptile-apps Bot commented Jul 9, 2026 •

edited

Loading

Comments Outside Diff (4)

Uh oh!

greptile-apps Bot Jul 9, 2026

Uh oh!

greptile-apps Bot Jul 9, 2026

Uh oh!

greptile-apps Bot Jul 9, 2026

Uh oh!

CharlelieLrt Jul 10, 2026

Uh oh!

AbVishwas Jul 10, 2026

Uh oh!

CharlelieLrt Jul 10, 2026

Uh oh!

AbVishwas Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

AbVishwas commented Jul 9, 2026

Description

Checklist

Dependencies

Review Process

Uh oh!

copy-pr-bot Bot commented Jul 9, 2026

Uh oh!

greptile-apps Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Important Files Changed

Comments Outside Diff (4)

Uh oh!

greptile-apps Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

CharlelieLrt Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

AbVishwas Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

CharlelieLrt Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

AbVishwas Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

greptile-apps Bot commented Jul 9, 2026 •

edited

Loading