Adds the PhysicsNeMo-Mesh changes required for GLOBE 3D

[peterdsharpe]

PhysicsNeMo Pull Request

Description

Extracts out the PhysicsNeMo-Mesh-related changes required to support GLOBE-3D.

A subset of changes in:
#1481

Nothing too spicy in here - mostly just minor improvements for torch.compile-friendliness and bugfixes for visualizing low-precision outputs (NumPy doesn't play nice with bf16, for example).

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.

[peterdsharpe]

/blossom-ci

[greptile-apps]

Greptile Summary

This PR extracts standalone compute_cell_areas and compute_cell_normals functions into physicsnemo/mesh/geometry/ with dimension-specific closed-form implementations (Lagrange identity, scalar triple product, Hodge star) that are faster and support reduced-precision dtypes, and refactors Mesh.cell_areas / Mesh.cell_normals to delegate to these new functions. It also adds _cache propagation to pad_to_size, fixes edgecolors=None → "none" in matplotlib, and adds .float() conversions throughout visualization and I/O code to handle bfloat16 tensors.

Key changes:

• New _cell_areas.py and _cell_normals.py with well-structured dimension-dispatched implementations and proper degenerate-case handling.
• mesh.py: Clean delegation to new geometry functions; safe _cache propagation in pad_to_size (the "cell" and "point" sub-TensorDicts are always present after __post_init__).
• io_pyvista.py: The .float() calls fix bfloat16 compatibility but also silently downcast float64 → float32 for all tensors on export, which is a precision regression for double-precision meshes. Integer tensors in global_data would be similarly affected.
• New dedicated unit tests for _cell_areas and _cell_normals with thorough branch, degenerate-case, and cross-branch consistency coverage replace the equivalent tests removed from test_optimizations.py.
• The _gram_det_volumes general fallback (n≥4) documents that it disables autocast to stay in native dtype, but doesn't explicitly upcast bfloat16/float16 inputs, which can cause cryptic failures if reduced-precision data reaches the n≥4 code path.

Important Files Changed

Filename	Overview
physicsnemo/mesh/io/io_pyvista.py	Adds .float() before .cpu().numpy() to handle bfloat16 tensors, but this unconditionally downcasts float64 → float32, silently losing precision for double-precision meshes.
physicsnemo/mesh/geometry/_cell_areas.py	New file implementing dimension-specific closed-form simplex volume computations. Logic is mathematically sound with correct Gram determinant, Lagrange identity, and scalar triple product branches. Proper clamping before sqrt for degenerate cases.
physicsnemo/mesh/geometry/_cell_normals.py	New file implementing unit normal computation for codimension-1 simplices with correct dimension-specific branches (2D rotation, 3D cross product, general Hodge star). Sign convention is mathematically correct.
physicsnemo/mesh/mesh.py	Refactors cell_areas and cell_normals to delegate to the new standalone functions, and adds _cache propagation to pad_to_size. Cache is always initialized with "cell" and "point" sub-TensorDicts, so the new pad_to_size code is safe.
test/mesh/geometry/test_cell_areas.py	Comprehensive tests covering all five computation branches with known analytic results, degenerate cases, cross-branch consistency, and device parametrization. Well-structured with float64 precision for correctness.
test/mesh/geometry/test_cell_normals.py	Thorough tests for all three normal computation branches, including degenerate cells, cross-branch consistency verification, and unit-length validation. Good coverage of edge cases.

Comments Outside Diff (1)

1. physicsnemo/mesh/geometry/_cell_areas.py, line 212-219 (link)

   _gram_det_volumes doesn't guard against reduced-precision input tensors

   The module docstring states the closed-form branches support bfloat16/float16 natively, while this general fallback path relies on torch.det (which dispatches to cuBLAS LU and does not support reduced-precision types). Disabling torch.autocast prevents automatic promotion of float32 inputs during an autocast region, but it doesn't protect against a caller explicitly passing a bfloat16 or float16 tensor directly to compute_cell_areas with n_manifold_dims >= 4.

   Consider adding an explicit upcast for the general path to fail loudly or handle gracefully:

   with torch.autocast(device_type=relative_vectors.device.type, enabled=False):
       # Upcast to at least float32 so torch.det can dispatch correctly
       rv = relative_vectors if relative_vectors.dtype not in (torch.bfloat16, torch.float16) else relative_vectors.float()
       gram_matrix = torch.matmul(rv, rv.transpose(-2, -1))
       ...

   This would make the n>=4 path consistent with the documented bfloat16 support and prevent a cryptic CUDA kernel error.

_{Last reviewed commit: 6344bfc}

[peterdsharpe]

/blossom-ci

[laserkelvin]

Nice refactor! Generally LGTM, and I just added some stylistic suggestions

[laserkelvin]

This is minor style, but match/case looks better IMO for >= 3 cases

[peterdsharpe]

Haha you're right, that's what the cool kids do :D I'll update

[laserkelvin]

Nitpicking, but does this really need to be in its own function? I guess it's nice for consistency, but functionally... I'll defer to you to decide how you want to proceed

[laserkelvin]

Looking at these, I wonder if it might be a bit faster (I don't know if it's critical, or helpful) if we did the tensordict.float() (maybe even with cpu()), and then within the loop work on the NumPy handles?

[peterdsharpe]

Oh yeah, good catch! Likely slightly faster, and also slightly better style. Will do.