Compiled pipelines drop coordinate units on the pytree round trip

[d-chambers]

Every patch that goes through a compiled pipeline comes back with its coordinate units stripped. Downstream anything reading coord.units sees None, and DASCore's plots quietly lose their axis labels — a waterfall's x axis goes from Distance [m] to Distance.

The round trip is where it happens, so this is not specific to any operation:

import dascore as dc
from dasjax import JaxPatchPipeline

patch = dc.get_example_patch("random_das").set_units(distance="m", time="s")
out = JaxPatchPipeline().abs().compile()(patch)

show = lambda p: {n: str(p.coords.coord_map[n].units) for n in p.coords.coord_map}
print(show(patch))   # {'distance': '1 m', 'time': '1 s'}
print(show(out))     # {'distance': 'None', 'time': '1 s'}

time survives only because it is a datetime coordinate whose units are re-derived; anything carrying real units does not.

Cause

patch_from_leaves rebuilds each coordinate with update_data, which takes units as a keyword and defaults it to None rather than to what the coordinate already had:

https://github.com/DASDAE/dasjax/blob/main/src/dasjax/pytree.py#L77

coord = meta["coord"].update_data(data=values)

DASCore's BaseCoord.update_data ends in get_coord(data=data, units=kwargs.get("units")), so omitting the keyword is an instruction to drop them.

patch_to_leaves already stashes the original coordinate in coord_meta, so the information is right there.

Suggested fix

coord = meta["coord"].update_data(data=values, units=meta["coord"].units)

Verified: with that line the pipeline above returns {'distance': '1 m', 'time': '1 s'}.

Why the parity suite does not catch it

test_method_equivalence._assert_patch_data_close ends with assert left.coords == right.coords, and DASCore's coordinate comparison returns True across this difference:

print(patch.coords == out.coords)   # True, despite distance being 1 m vs None

So a units assertion has to be explicit — comparing {name: coord.units} before and after, say — or the equality on its own will keep passing. (Whether Coords.__eq__ should ignore units is a DASCore question; I have not raised it there.)

Tested against dasjax main (4df967b) and 0.0.2, with DASCore dev.

Found while writing the DASJax notebook for the EGU Galileo tutorial, where a piped patch's waterfall came out with an unlabelled distance axis.