DLESYM pdates

[daviddpruitt]

PhysicsNeMo Pull Request

Description

Updates DLESyM architectures and features.

• Updates support for diagnostic output variables in HEALPixRecUNet and HEALPixUNet
• Adds conditional layer norm to DLESyM UNetEncoder and DLESyM UNetDecoder
• Adds physical constraint support in HEALPixRecUNet and HEALPixUNet
  • NonnegativeConstraint added
• Adds/updates HEALPix blocks: SymmetricConvNeXtBlock, Multi_SymmetricConvNeXtBlock, DoubleConvNeXtBlock, ConvGRUBlock, ConvNeXtBlock, BasicConvBlock, TransposedConvUpsample, Interpolate
• Zarr backed datapipline
• Updated dlwp_healpix coupler
• Updated HEALPix loss funcitons: ConditionalWeightLoss, WeightedCRPSLoss, and WeightedCRPSLossSpectral
• Hydrostatic constraint support in losses

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

No new 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.

[copy-pr-bot]

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]

Greptile Summary

This PR updates the DLESyM HEALPix model stack and adds Zarr-backed data support. The main changes are:

• New Zarr time-series and coupled HEALPix datapipes.
• Updated HEALPix UNet and recurrent UNet support for diagnostics, constraints, residual prediction, and conditional layer norm.
• New and updated HEALPix blocks, normalization layers, constraints, and padding behavior.
• New CRPS and hydrostatic climate losses.
• Expanded tests for HEALPix datapipes, model layers, encoder/decoder behavior, and hydrostatic balance.

Important Files Changed

Filename	Overview
physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py	Adds the coupled Zarr dataset path; coupled-variable training noise currently indexes the time axis.
physicsnemo/metrics/climate/healpix_loss.py	Adds conditional and CRPS losses; several configuration and device paths can crash.
physicsnemo/metrics/climate/hydrostasy.py	Adds hydrostatic constraints and losses; the weighted hydrostatic path assumes humidity channels are present.
physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py	Adds recurrent UNet support for residual outputs, constraints, coupling layout, and CLN conditions.
physicsnemo/models/dlwp_healpix/HEALPixUNet.py	Adds UNet support for residual outputs, constraints, coupling layout, and CLN conditions.

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[negin513]

Great work!! Minor comments. Also, thanks for walking me through the changes.

[negin513]

Hey @daviddpruitt , most comments are with regards to coupler changes. other than that it looks good to me.

[negin513]

/codeowners-info

[coreyjadams]

/codeowners-info

[github-actions]

CODEOWNERS review map

Current for commit 66b66782f3cc. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @coreyjadams — 4 file(s)

• test/datapipes/test_healpix.py
• test/datapipes/test_healpix_couple.py
• test/datapipes/test_healpix_couplers.py
• test/datapipes/test_healpix_diagnostic_outputs.py

⏳ @loliverhennigh — 1 file(s)

• physicsnemo/nn/module/hpx/padding.py

⏳ @pzharrington — 12 file(s)

• physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py
• physicsnemo/datapipes/healpix/couplers.py
• physicsnemo/datapipes/healpix/timeseries_dataset.py
• physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py
• physicsnemo/models/dlwp_healpix/HEALPixUNet.py
• physicsnemo/models/dlwp_healpix/init.py
• physicsnemo/models/dlwp_healpix/layers/init.py
• physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py
• physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py
• physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py
• physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py
• physicsnemo/models/dlwp_healpix/layers/normalization.py

No CODEOWNER

• test/models/dlwp_healpix/_cln_reference.py
• test/models/dlwp_healpix/test_conditional_layer_norm.py
• test/models/dlwp_healpix/test_healpix_blocks.py
• test/models/dlwp_healpix/test_healpix_encoder_decoder.py
• test/models/dlwp_healpix/test_healpix_layers.py
• test/models/dlwp_healpix/test_healpix_recunet_model.py
• test/models/dlwp_healpix/test_healpix_unet_model.py
• test/models/dlwp_healpix/test_layers_backward_compat.py
• test/nn/module/healpix_helpers.py
• test/nn/module/test_healpix.py

Comment /codeowners-info to refresh.

[coreyjadams]

I have left a few comments that are very minor but need to be addressed for downstream users who might not be working with pandas and xarray, but want to run tests. Currently, this will break their development. Please address them before merge?

[loliverhennigh]

The module stuff looks fine to me

[pzharrington]

/ok to test 442a354