Moving HealPix Ops into module folder

[NickGeneva]

PhysicsNeMo Pull Request

Description

Moving healpix layers, pooling, padding, tokenizers, etc for DLESym and HealDA

With Earth2Grid installed tests pass:

local-ngeneva@ipp2-1759:~/physicsnemo$ uv run --no-sync pytest test/nn/module/test_healpix.py
================================================= test session starts ==================================================
platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0
rootdir: /localhome/local-ngeneva/physicsnemo
configfile: pyproject.toml
plugins: testmon-2.2.0, timeout-2.4.0, anyio-4.12.1, dependency-0.6.0, zarr-3.1.5, hydra-core-1.3.2, jaxtyping-0.3.5
collected 45 items                                                                                                     

test/nn/module/test_healpix.py .............................................                                     [100%]
=== Test durations by file ===
test/nn/module/test_healpix.py: 0.60 seconds


================================================== 45 passed in 2.99s ==================================================

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 Overview

Greptile Summary

This PR relocates HEALPix-related ops into a new physicsnemo/nn/module/hpx/ package, splitting the former physicsnemo/nn/healpix.py implementation into layers.py, padding.py, pool.py, and tokenizer.py, and adding hpx/__init__.py re-exports.

Merge blockers were found: layers.py currently self-imports padding symbols and will fail at import time; padding.py regresses shape validation by constructing ValueError objects without raising them in multiple places; the new tokenizer classes are not nn.Module subclasses despite defining parameters/layers; tokenizer.py adds an unconditional earth2grid import (breaking environments where it’s optional); and physicsnemo/nn/__init__.py still imports from a deleted physicsnemo/nn/healpix.py, breaking the existing public API imports.

Important Files Changed

Filename	Overview
physicsnemo/nn/module/hpx/init.py	Adds new hpx submodule exports; no internal logic, but package integration depends on updating upstream imports.
physicsnemo/nn/module/hpx/layers.py	Introduces HEALPixLayer but contains a self-import (from .layers import ...) that will raise ImportError at import time, breaking any users of this module.
physicsnemo/nn/module/hpx/padding.py	Moves HEALPix padding utilities but regresses validation: several ValueError(...) are constructed but never raised, so invalid shapes will now crash later with unpack/reshape errors.
physicsnemo/nn/module/hpx/pool.py	Adds pooling wrappers around HEALPixLayer; likely fine but depends on broken layers.py import.
physicsnemo/nn/module/hpx/tokenizer.py	Adds HEALPix tokenizers but introduces hard earth2grid import (no optional gating) and tokenizer classes don’t subclass nn.Module despite calling super().__init__() and using Parameters.

[greptile-apps]

_{5 files reviewed, 5 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

physicsnemo/nn/__init__.py
Public imports now broken

physicsnemo/nn/__init__.py still re-exports HEALPix symbols from from .healpix import (...), but physicsnemo/nn/healpix.py is deleted in this PR. Any import physicsnemo.nn as nn; nn.HEALPixPadding (or similar) will fail after merge unless __init__.py is updated to import from the new physicsnemo.nn.module.hpx location (or a compatibility shim module is kept).

[pzharrington]

Approving to unblock

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci