Feat(datapipes): Zarr mesh schema and readers — a zarr storage path for mesh-native CAE pipelines

[negin513]

Description

Adds a chunked, compressed, cloud-capable zarr storage path for mesh-native CAE pipelines, bridging the previously disjoint physicsnemo.mesh memmap (.pmsh/.pdmsh) and zarr data stacks with a documented, versioned schema and drop-in readers. No changes to models, transforms, collate, or training loops — the readers honor the existing (Mesh, metadata) / (DomainMesh, metadata) contracts.

What's included

Library (physicsnemo/datapipes/readers/zarr_mesh.py)

• save_mesh_to_zarr / save_domain_mesh_to_zarr: reference writers for the physicsnemo-mesh-zarr / physicsnemo-domainmesh-zarr schemas (spec: physicsnemo/datapipes/MESH_ZARR_SCHEMA.md). DomainMesh groups mirror the .pdmsh tree with case-level global_data single-sourced at the root. Writers accept paths or zarr StoreLike objects.
• ZarrMeshReader: (Mesh, metadata) per sample; subpath= navigation into DomainMesh groups; merge_global_data_from= read-time metadata merge (collision-checked); contiguous-block subsampling with chunk-aligned I/O; zarr-python and tensorstore backends.
• ZarrDomainMeshReader: (DomainMesh, metadata) per case for volume pipelines; per-submesh subsample selection; full_resolution_boundaries= for exact-geometry consumers (SDF).
• to_cell_soup + machine-verified layout attr: opt-in per-cell vertex-soup curation for cell-centric training; soup groups skip the cells read entirely and issue all array reads concurrently.
• validate_mesh_zarr: schema conformance checker; readers reject stores written by a newer schema_version.

Recipe (examples/.../unified_external_aero_recipe)

• drivaer_ml_surface_zarr.yaml / drivaer_ml_volume_zarr.yaml: differ from their .pdmsh twins only in the reader: block; the full transform stack and all five recipe models run unchanged.
• src/convert_pdmsh_to_zarr.py: transitional CLI for datasets already curated to .pdmsh (new curation should emit the schema directly from PhysicsNeMo-Curator; companion PR: feat(mesh): DomainMeshZarrSink emitting the physicsnemo mesh-zarr schema physicsnemo-curator#64).

Benchmarks (benchmarks/physicsnemo/datapipes/): the reader-level, recipe-pipeline, and per-stage profiling scripts behind the numbers below (happy to split into a follow-up PR if preferred).

Performance (DrivAerML, 8 runs of 15-20M cells, 200k-cell subsample, Lustre)

Level	.pdmsh/memmap	zarr (this PR)
Reader, 8 workers (cold / warm samples/s)	17.6 / 45.2	115.5 / 139.1
Recipe DataLoader (benchmark_io, cold / warm batches/s)	3.7 / 7.6	9.9 / 18.7
GeoTransolver training (warm s/step, train / val)	0.56 / 0.24	0.475 / 0.15

At 1 GPU, prefetch hides most of the loader gap (~17% faster training epochs, ~40% faster validation); the full loader ratio re-emerges for cold epochs, larger-than-cache datasets, and multi-GPU nodes. Stores are ~6% smaller than .pdmsh (zstd) while additionally embedding the STL geometry in-group. Per-stage profiling shows the entire format difference is in the read stage (230 ms vs 58 ms warm); GPU-side transforms are unaffected.

Correctness

• Bitwise equivalence vs .pdmsh on real data: per-cell geometry points[cells], all cell fields, and merged global_data exact; 166M-point volume interior verified by shape + random 1M-point probes.
• Matched training A/Bs (seed 42): surface — final train loss 0.1266 (.pdmsh) vs 0.1265 (zarr); volume — 0.2648 vs 0.2680 train / 0.0998 vs 0.0987 val at epoch 3; trajectories track epoch-for-epoch within subsample-RNG noise.
• End-to-end: GeoTransolver surface and volume both train from zarr through the recipe's train.py (volume incl. on-the-fly SDF against the in-group full-resolution STL boundary).

Scope notes

The soup layout is an opt-in curation choice for cell-centric consumers, recorded as a verified layout attr; connectivity-dependent consumers (GNNs) should curate layout: indexed groups, which the schema and readers support. Version-1 schema requires fixed-width cells; an offsets-based ragged/mixed-element encoding is explicitly reserved for v2 in the spec (readers reject newer versions rather than misread).

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes (26 tests in test/datapipes/readers/test_zarr_mesh.py; 262 datapipes core+readers tests pass).
• The documentation is up to date with these changes (readers API .rst, MESH_ZARR_SCHEMA.md).
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

None new: zarr was already optional (OptionalImport); tensorstore optional with automatic zarr-python fallback.

[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

Adds a complete zarr storage path for mesh-native CAE pipelines: ZarrMeshReader / ZarrDomainMeshReader readers (with zarr-python and tensorstore backends), reference writers (save_mesh_to_zarr / save_domain_mesh_to_zarr), a versioned schema spec (MESH_ZARR_SCHEMA.md), schema validator, and a conversion CLI from .pdmsh — all without touching models, transforms, or training loops.

• Readers honor the existing (Mesh, metadata) / (DomainMesh, metadata) contract with a soup fast-path (synthesized cell indices, concurrent array reads) and chunk-aligned contiguous-block subsampling; set_epoch / set_generator mirror the MeshReader epoch-control API.
• Schema and validation include a machine-verified layout attr, schema-version rejection for forward-incompatible stores, and validate_mesh_zarr for conformance checking across mesh and domainmesh groups.
• Tests (26 cases) cover round-trips, both backends, soup detection, merge-collision, domain reader subsample/full-res boundary, and pipeline integration, but miss the sequential set_epoch path that would expose a base-seed drift bug in the current implementation.

Important Files Changed

Filename	Overview
physicsnemo/datapipes/readers/zarr_mesh.py	Core new reader module — well-structured with lazy backends and soup fast-path; contains a P1 set_epoch reproducibility bug (drifting base seed on sequential calls) and a P1 propagation gap in ZarrDomainMeshReader.set_epoch, plus minor P2 issues: silent truncation in _block, hardcoded zarr3 tensorstore driver, and incomplete DomainMesh schema validation for root-level global_data.
test/datapipes/readers/test_zarr_mesh.py	Good coverage of round-trips, subsampling, soup fast-path, merge collision, schema version rejection, and pipeline integration; missing a test for sequential set_epoch calls that would have caught the base-seed accumulation bug.
examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/convert_pdmsh_to_zarr.py	Thin conversion CLI; fragile int(d.name.split("_")[1]) sort key would crash on non-numeric or multi-component run directory names with an unguarded ValueError.
physicsnemo/datapipes/init.py	Adds ZarrMeshReader, ZarrDomainMeshReader, save_mesh_to_zarr, save_domain_mesh_to_zarr, to_cell_soup, and validate_mesh_zarr to the public namespace and __all__; no issues.
physicsnemo/datapipes/readers/init.py	Re-exports new zarr mesh symbols from zarr_mesh.py; no issues.
physicsnemo/datapipes/MESH_ZARR_SCHEMA.md	Schema specification for the new zarr format; clearly documents v1 requirements, opt-in soup layout, and v2 reservation for ragged/mixed-element encoding.
benchmarks/physicsnemo/datapipes/benchmark_mesh_vs_zarr.py	Benchmark comparing memmap vs zarr reader throughput; self-contained profiling script with no issues.

_{Reviews (1): Last reviewed commit: "feat(datapipes): mesh-zarr schema spec, ..." | Re-trigger Greptile}

[negin513]

Re-verified the dataloader numbers on exactly the code in this PR (the table in the description was measured on an earlier iteration of the branch). Same protocol: recipe's benchmark_io=true, 8 DrivAerML cases, 3 epochs, page cache evicted before each run, single node, Lustre. Raw logs: https://gist.github.com/negin513/72b651f343d52fecf07f804a6bfab69f

benchmark_io (train batches/s)	epoch 1 (cold)	epochs 2–3 (warm, best)	warm median latency
dataset=drivaer_ml_surface (.pdmsh)	3.66	7.96	117 ms
dataset=drivaer_ml_surface_zarr (this PR)	9.23	23.64	42 ms

Slightly better than the description's table (18.7 → 23.6 warm): the final reader includes the soup fast path (verified layout attr ⇒ cells read skipped, all array reads issued concurrently), which the earlier measurement predated. Cold ratio unchanged (~2.5×), warm now ~3×. The zarr store here is the full-DomainMesh schema (all 5 boundaries + volume interior + STL in-group) read via subpath="boundaries/vehicle" + read-time global_data merge — i.e., the exact configuration the recipe configs in this PR ship.

[peterdsharpe]

Deep review done locally (worktree checkout of c7ca4787), with empirical probes run against this branch rather than desk-checking alone. Overall: the code quality is high, and several design choices are exactly right — the schema versioning policy (reject-newer rather than best-effort), the machine-verified layout attr, the collision-checked read-time global_data merge (a big cleanup over globbing into _tensordict/ internals), and store-object writer support. Tests are well-designed for what they cover.

Two blocking themes, detailed in the inline comments:

1. Nested-TensorDict and key-name semantics (the blocker). Mesh's data containers are TensorDicts, where any value can be another TensorDict — and the .pmsh memmap path round-trips that today. The v1 schema/writer/reader assume exactly-one-level-flat dicts of tensors. Probed on this branch: nested TensorDicts crash the writer; a field key containing / (plausible in CFD exports, handled fine by .pmsh) writes successfully, reads back silently missing, and passes validate_mesh_zarr; NonTensorData crashes the writer. The clean fix is small: map the TensorDict tree to a zarr group tree recursively (zarr is hierarchical, like HDF5) — no flattening, so the whole key-collision class vanishes by construction, and the special-cased per-subgroup loops plus prefix-string plumbing collapse into one recursion. If v1 is to stay flat instead, it needs write-time validation with clear errors, reader/validator detection of unexpected subgroups, and the restriction stated in MESH_ZARR_SCHEMA.md (the spec is what PhysicsNeMo-Curator's independent writer implements against).

2. The performance story needs the controls the harness already has. The benchmark thoughtfully builds a pmsh_soup control store precisely to separate layout effects from format/reader effects — but its numbers appear nowhere (PR description, follow-up comment, gist). The claimed speedup conflates four mechanisms (soup layout; skipping unique() compaction; tensorstore concurrency; compression + buffered-vs-mmap cold I/O), and only the last is inherently "zarr". Running this PR's own synthetic benchmark (layout-controlled: soup on both sides) on local NVMe: memmap wins warm single-worker vs the shipped default backend (33.1 vs 25.8 samples/s, zarr-python+zstd), while zarr wins cold, multi-worker, and via tensorstore. That's consistent with the Lustre numbers being real — but the honest pitch is "zarr wins cold reads, parallel loaders, and network/cloud storage", and the pmsh_soup numbers on DrivAerML/Lustre are the evidence that decides how much a pure re-curation of .pdmsh would have captured. A sentence positioning zarr against TensorDict's own HDF5 path (to_h5/PersistentTensorDict: nested-native, single file per case) would also preempt the obvious question.

On size: the library footprint is ~1,150 lines, not 2,900 — benchmarks are ~1,000 and tests 410. I'd take up the offer to split the benchmarks into a follow-up PR. Worth a maintainer conversation: this adds the repo's third and fourth zarr access implementations (ZarrReader, TensorStoreZarrReader exist); the private backend pair here is small and justified by the dependent-read pattern, but a shared backend abstraction would trim ~100 lines and unify tensorstore spec handling. If the recursive rewrite happens, save_mesh_to_zarr becomes "walk the tensorclass fields", which argues for it living next to Mesh.save/Mesh.load in the mesh io layer, with only the readers in datapipes.

Smaller items inline: set_epoch seed drift (inherited from MeshReader — verified sequential-vs-resume divergence), the indexed (non-soup) subsample path has zero test coverage (I probed it on both backends — it is correct, but it's the trickiest code in the file), ZarrMeshReader never checks the format attr, first-case-only boundary_names, tensorstore backend hardcodes zarr v3 while discovery accepts v2 stores, and readers/validator are filesystem-only while writers accept StoreLike.

---

Negin — per our conversation, I'll push fixes for the review items directly to this branch, starting with the nested-TensorDict/key-semantics blocker (recursive TensorDict-tree ↔ zarr-group-tree mapping, plus the spec updates in MESH_ZARR_SCHEMA.md), then the smaller items (set_epoch drift, format checks, tests for the indexed path). The inline comments document the rationale; happy to discuss any of them before or after the commits land.

[peterdsharpe]

This control is exactly the right experimental design — and its numbers are reported nowhere (PR description, the re-verified follow-up comment, or the gist). Please publish the pmsh soup variant's results from the DrivAerML/Lustre runs. They're the datum that decides how much of the memmap-vs-zarr gap is the storage format versus the soup re-layout that the zarr conversion also performs.

For calibration, I ran this script's synthetic mode (layout-controlled: soup on both sides) on local NVMe, 8×500k-cell samples: memmap wins warm single-worker against the shipped default backend (33.1 vs 25.8 samples/s, zarr-python+zstd); zarr wins cold (19.6-39.2 vs 8.0), multi-worker (memmap degrades to 20.1 at 8 threads), and via tensorstore (up to 116). So the honest headline is "zarr wins cold reads, parallel loaders, and network storage" — not a flat 3-6×. The four mechanisms in the reported speedup (soup layout, skipping unique() compaction, tensorstore concurrency, compression + buffered-vs-mmap cold I/O) deserve to be separated in the writeup, since the first two are available to the memmap path with no format change.

Related: the PR should say a sentence about why zarr over TensorDict's own HDF5 path (to_h5/PersistentTensorDict) — which handles nested TensorDicts natively, stores one file per case (friendlier to Lustre metadata than many chunk files, absent v3 sharding), and does chunked+compressed slice reads. The genuine zarr advantages (object-store-native access, curator-side parallel writes, tensorstore async reads) are worth stating explicitly so the format choice is on the record.

[peterdsharpe]

Pushed the first batch of review fixes directly to the branch (as discussed):

• ec56ca95 — nested-TensorDict semantics (the blocker from my review): the writer/reader now map the TensorDict tree to a zarr group tree recursively, /-in-field-names and NonTensorData are rejected at write time with clear errors, the validator checks leaves at any depth, and MESH_ZARR_SCHEMA.md now specifies field-name/nesting rules and pins zarr format 3. Two intentional behavior tightenings to be aware of: zarr format-2 stores now fail with a clear error on both backends (previously they half-worked on zarr-python and failed opaquely on tensorstore), and ZarrDomainMeshReader now raises on heterogeneous boundary_names instead of silently dropping boundaries a later case adds. No schema_version bump since v1 has not shipped in any release — this PR is its first appearance.
• 79d25079 — set_epoch resume-reproducibility fix in the pre-existing .pmsh readers (same fix as applied to the zarr readers in the commit above). Note for later: the same idiom exists in transforms/base.py, transforms/mesh/base.py, and the numpy/zarr/tensorstore readers — I'll propose a small separate PR to fix those, so don't consider this one exhaustive.
• ea1b481f — benchmark docstring staleness + converter run-dir sort robustness.

All 36 test_zarr_mesh.py tests pass (both backends), plus the full test/datapipes/readers/ suite. @negin513 the schema changes affect the Curator sink in NVIDIA/physicsnemo-curator#64 — it should mirror the write-time rejections (/ in names, non-tensor leaves) and emit zarr format 3; running validate_mesh_zarr in its tests would keep the two implementations honest.

[peterdsharpe]

One scope proposal to close out my review (the two open threads above both point here):

1. Pull the benchmarks (~1,000 lines) into a separate PR — as you already offered in the description; +1 on doing it.

2. Pull the unified-aero-recipe changes (drivaer_ml_*_zarr.yaml, convert_pdmsh_to_zarr.py, the dataset_paths entry, ~315 lines) into a separate PR as well. Right now we're doing early experiments with the zarr-backed Mesh path — until we decide whether to invest in migrating our datasets, it's probably best kept as an open item rather than wired into the flagship recipe. Once it's in the recipe, we're inviting "which format should I use?" questions from end-users, and we may not have firm answers until we've put it through its paces (the pmsh_soup control numbers, the HDF5 comparison, and the Curator sink landing all feed that decision). The split is mechanically free: the recipe files are purely additive (the only pre-existing file touched is dataset_paths.yaml, +4 lines) and nothing in the library depends on them. The converter travels with the configs since it's currently the only way to produce data they can read.

That would leave this PR as the clean core: library + schema + validator + tests + docs (~1,600 lines). Two small consequences if you take this: the CHANGELOG line mentioning the recipe configs would trim down, and the description's training A/B evidence should note the configs it references are coming in the follow-up (the evidence itself stands). Happy to help with the extraction if useful.

[negin513]

"One scope proposal to close out my review… 1. Pull the benchmarks (~1,000 lines) into a separate PR… 2. Pull the
unified-aero-recipe changes… into a separate PR as well. …That would leave this PR as the clean core: library + schema + validator + tests + docs (~1,600 lines)."

@peterdsharpe done!

[negin513]

The .pmsh soup control, from the same July 1 run as the published reader numbers (8 real DrivAerML cases, Lustre, 200k-cell subsample; script now in #1829 for reference... ):

Variant	Store (GB)	1w cold	1w warm	8w cold	8w warm (samples/s)
.pmsh indexed (current pipeline)	7.2	3.9	6.9	17.6	45.2
.pmsh soup (control)	11.2	4.6	6.2	22.3	31.6
Zarr + Zstd (zarr-python)	7.0	11.6	8.7	55.6	52.0
Zarr + Zstd (tensorstore)	7.0	25.4	37.9	115.5	139.1
Zarr raw (zarr-python)	11.3	8.0	8.0	47.0	29.8
Zarr raw (tensorstore)	11.3	20.5	11.8	140.1	34.8

The relayout alone is not the win. Soup on memmap is slower when warm than what ships today (45.2 → 31.6 samples/s) because it stores 56% more data from duplicated vertices, with no compression to recover that overhead.

The win comes from the format:

• Zstd compression: Without it, warm TensorStore throughput drops from 139.1 to 34.8 samples/s. It also reduces the soup representation to 7.0 GB, smaller than the 7.2 GB indexed original.
• TensorStore concurrent reads: Approximately 2.7× faster than zarr-python for the warm 8-worker case.
• Chunk sizing: Chunks are aligned with the subsample size, so each field read touches at most two chunks.

Re-curating .pdmsh as soup is therefore not a shortcut to these results. We can erchunk the pdmsh though?? @peterdsharpe

The NVMe result reflects the same mechanics in reverse: the page cache leaves compression little I/O to save, while a single worker leaves concurrency little opportunity to help.

[megnvidia]

Suggested change

	Mesh-zarr schema utilities
	Mesh-Zarr Schema Utilities

[megnvidia]

Suggested change

	## Mesh group (`format = "physicsnemo-mesh-zarr"`)
	## Mesh Group (`format = "physicsnemo-mesh-zarr"`)

[megnvidia]

Suggested change

	### Field names and nesting
	### Field Names and Nesting

[megnvidia]

Suggested change

	Data fields mirror the `TensorDict` tree: each `<f>` under `point_data` /
	Data fields mirror the `TensorDict` tree. Each `<f>` under `point_data`,

[megnvidia]

Suggested change

	`cell_data` / `global_data` is either an **array** (a tensor leaf) or a
	`cell_data`, or `global_data` is an *array* (a tensor leaf) or a

[peterdsharpe]

I notice that both the DomainMesh save utility and load utility are in physicsnemo/datapipes/readers/zarr_mesh.py.

@coreyjadams will probably have stronger opinions (as codeowner for datapipes), but my impression was that readers/ was only for load utilities. One API design option: you could put this in physicsnemo/mesh/io/ instead? Perhaps a more natural fit - there we have from_pyvista() and to_pyvista() already (both from VTK-family formats, basically), so this would slot right in. What do you think @negin513?

[coreyjadams]

Hi @negin513 - thanks for working on this PR. In general I am not opposed to this backend for mesh. But I think we have some design changes needed here before this can merge.

First, I think the logic for mesh-zarr backend needs to deploy at mesh/io/zarr.py. It's important if we have zarr supported that it be supported with or without datapipes, here, from a software design. Peter mentioned that the writing utilities shouldn't be in the datapipes but I'd actually think we should put the reading and writing utilites as much as reasonably possible into mesh/io/zarr.py and the MeshReader and DomainMesh reader can tap into that as a backend instead. Then, swapping backends at user recipe time becomes a very simple config change even.

So, abstractly, datapipes/readers/mesh.py might get an update of Mesh.load(mesh_path, backend="zarr") (and other associated settings at constructor time) and similarly for domain mesh reader.

Overall, I think the integration logic of a zarr backend becomes a lot cleaner to datapipes and every other consumer too if we contain it to just mesh.

For testing, I haven't actually looked closely at the tests, but we should consider several components:

• If we ship multiple IO backends, the save and load methods should default to memmap and auto correspondingly, and we should test that those deliver as promised.
• The tests should validate that `load(...., backend="auto") successfully identifies and deserializes .pmsh and .pdmsh files of all supported backends.
• We should have really good "round trip" testing: make a synthetic file, write it to both formats, load it with "auto", check that it's identical when restored. Do this across a suite of dimensions / .pmsh / .pdmsh though no need to go crazy of course?
• @peterdsharpe what other tests can tighten up the Mesh IO to help prevent user facing churn?

Overall, my top priority is making sure we don't introduce churn into the mesh dataset ecosystem, even though its admittedly very young.

[negin513]

Thanks Corey! I'm actually fully on board with where you want this to end up —mesh/io/zarr.py owning the format, Mesh.save(backend="memmap"|"zarr"), load with backend="auto", and the round-trip tests you described. That's what I proposed to you earlier.

Where I'd push back a little is on sequencing. I'd rather merge this PR as-is and do the mesh/io move as the immediate next PR, and here's why I think that's actually the safer path...

FWIW, this is also how we've been operating on this PR already — when Peter asked, we split the converter, recipe, and benchmarks out to #1828/#1829. So "follow-up PR" isn't a promise that goes to die in the backlog; it's the pattern this work has actually followed....

That said, if you feel strongly it has to happen in this PR, I'll do it back again.

I will defer @coreyjadams and @peterdsharpe to finish this PR per your design of choice.

[peterdsharpe]

@negin513 I think that moving this interface to mesh.io (rather than in datapipes) is probably a critical enough contract to get right that should happen in this PR. I'm happy to push commits for this if you want, though it should be pretty straightforward - more or less just grafting the one file over.

[peterdsharpe]

Let's see what happens here: pytorch/tensordict#1754

[negin513]

On placement, this structure follows the existing codebase convention rather than departing from it.

Every on-disk format reader in PhysicsNeMo currently lives under datapipes/readers/: vtk.py reads STL, VTP, and VTU files directly through PyVista, while hdf5.py, numpy.py, zarr.py, and tensorstore_zarr.py each own their format-specific I/O.

By contrast, mesh/io/ has not historically handled file I/O. io_pyvista.py converts an already-loaded, in-memory PyVista object through from_pyvista(pv_mesh); it does not read files itself.

For that reason, zarr_mesh.py is currently placed alongside the other file readers, including the VTK reader it most closely parallels. Moving it into mesh/io/ would not follow the existing pattern; it would establish a new one. This is more inline with arguments around software design. (cc @ram-cherukuri ).

I’m open to that direction, but I don’t think this PR should be gated on a broader package reorganization.

My preference would be to land this consistently with the current layout, then address the mesh/io/ structure in a separate PR where we can deliberately define the convention and software design, including whether vtk.py should move as well...

[negin513]

@peterdsharpe and @coreyjadams I trust your judgment on this PR and will defer this to you. Please feel free to push the proposed structure to this PR.

@peterdsharpe thanks for the tensordict issue. This would be really cool.