From fda6b9c251f5759a8049d191d604d2d6f49c86df Mon Sep 17 00:00:00 2001
From: Aarab Malik <178716638+aarab-malik@users.noreply.github.com>
Date: Wed, 11 Nov 2026 09:29:06 +0000
Subject: [PATCH] prefer geospatial_vertical_units when inferring length unit

Geodetic datasets were always assigned a length unit of 1.0 before
geospatial_vertical_units was checked, so code_length became 1 cm
even when the file recorded km. Check the attribute first.
---
 yt_xarray/accessor/accessor.py  |  6 +++---
 yt_xarray/tests/test_accesor.py | 27 +++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/yt_xarray/accessor/accessor.py b/yt_xarray/accessor/accessor.py
index 02b6dd3..480bcc7 100644
--- a/yt_xarray/accessor/accessor.py
+++ b/yt_xarray/accessor/accessor.py
@@ -400,11 +400,11 @@ def _field_data(grid: GridSpec) -> dict[str, Any]:
         return ds_yt
 
     def _infer_length_unit(self) -> str | float | None:
-        if self.geometry == "geodetic":
-            return 1.0
-        elif hasattr(self._obj, "geospatial_vertical_units"):
+        if hasattr(self._obj, "geospatial_vertical_units"):
             # some netcdf conventions have this!
             return str(self._obj.geospatial_vertical_units)
+        if self.geometry == "geodetic":
+            return 1.0
         return None
 
     _geometry = None
diff --git a/yt_xarray/tests/test_accesor.py b/yt_xarray/tests/test_accesor.py
index 751684e..0d762a3 100644
--- a/yt_xarray/tests/test_accesor.py
+++ b/yt_xarray/tests/test_accesor.py
@@ -234,3 +234,30 @@ def test_last_yt_cache(ds_xr):
     flds = "a_new_field_0"
     _ = ds_xr.yt.load_grid(flds)
     assert ds_xr.yt._yt_ds is not None
+
+
+def test_geospatial_vertical_units_on_geodetic():
+    # length unit should follow geospatial_vertical_units for lat/lon data,
+    # whether geometry is passed to load_grid or set beforehand.
+    def make_ds():
+        ds = xr.Dataset(
+            {"vs": (("depth", "latitude", "longitude"), np.random.rand(8, 7, 12))},
+            coords={
+                "depth": np.linspace(50.0, 2800.0, 8),
+                "latitude": np.linspace(10.0, 70.0, 7),
+                "longitude": np.linspace(-20.0, 90.0, 12),
+            },
+            attrs={"geospatial_vertical_units": "km"},
+        )
+        ds["depth"].attrs["units"] = "km"
+        return ds
+
+    ds = make_ds()
+    assert ds.yt.geometry == "geodetic"
+    yds = ds.yt.load_grid(geometry="internal_geographic")
+    assert np.isclose(yds.quan(1, "code_length").to("km").d, 1.0)
+
+    ds = make_ds()
+    ds.yt.set_geometry("internal_geographic")
+    yds = ds.yt.load_grid()
+    assert np.isclose(yds.quan(1, "code_length").to("km").d, 1.0)