Merge pull request #203 from CU-ESIIL/codex/fix-prism-loader-return-type-and-viewer-interactivity

Fix PRISM single-variable return and restore cube viewer rotation interactivity

Author: ttuff

docs/changelog.md

@@ -7,6 +7,8 @@
 - Documented the `Pipe` helper and new operations reference structure.
 - `fire_plot` now requests daily gridMET/PRISM data by default and propagates provenance metadata on returned cubes.
 - Added an `allow_synthetic` safety switch to gridMET/PRISM loaders with clearer empty-time/all-NaN error messages.
+- `load_prism_cube` now returns a DataArray when a single variable is requested, matching docs examples.
+- Fixed cube viewer rotation so drag/zoom updates the cube as well as axis labels.
 
 ## Earlier work
 

src/cubedynamics/data/prism.py

@@ -36,7 +36,7 @@ def load_prism_cube(
     prefer_streaming: bool = True,
     show_progress: bool = True,
     allow_synthetic: bool = False,
-) -> xr.Dataset:
+) -> xr.Dataset | xr.DataArray:
     """Load a PRISM-like cube.
 
     Parameters
@@ -54,7 +54,8 @@ def load_prism_cube(
         Temporal extent for the request.
     variable : str or sequence of str, default "ppt"
         PRISM variable(s) to request. ``variables`` may also be used for
-        clarity when passing multiple entries.
+        clarity when passing multiple entries. When a single variable is
+        requested through ``variable``, the loader returns an ``xarray.DataArray``.
     time_res, freq : str, default "ME"
         Temporal frequency code. ``freq`` overrides ``time_res`` when set.
     chunks : mapping, optional
@@ -114,7 +115,7 @@ def load_prism_cube(
 
     aoi = _coerce_aoi(lat=lat, lon=lon, bbox=bbox, aoi_geojson=aoi_geojson)
 
-    return _load_prism_cube_impl(
+    ds = _load_prism_cube_impl(
         normalized_variables,
         start_ts.isoformat(),
         end_ts.isoformat(),
@@ -125,6 +126,12 @@ def load_prism_cube(
         show_progress,
         allow_synthetic,
     )
+    if variables is None and len(normalized_variables) == 1:
+        da = ds[normalized_variables[0]]
+        da = da.copy()
+        da.attrs.update(ds.attrs)
+        return da
+    return ds
 
 
 def _load_prism_cube_impl(

src/cubedynamics/fire_time_hull.py

@@ -992,9 +992,12 @@ def load_climate_cube_for_event(
             allow_synthetic=allow_synth,
         )
 
-    target_var = variable if variable in ds.data_vars else next(iter(ds.data_vars))
-    cube_da = ds[target_var]
-    cube_da.attrs.update(ds.attrs)
+    if isinstance(ds, xr.DataArray):
+        cube_da = ds
+    else:
+        target_var = variable if variable in ds.data_vars else next(iter(ds.data_vars))
+        cube_da = ds[target_var]
+        cube_da.attrs.update(ds.attrs)
     log(verbose, f"{source.upper()} source: {cube_da.attrs.get('source')}")
     return ClimateCube(da=cube_da)
 

src/cubedynamics/plotting/cube_viewer.py

@@ -702,7 +702,7 @@ def _render_cube_html(
         document.getElementById("cube-drag-" + viewerId)
         || sceneDragSurface
         || cubeWrapper;
-      const rotationTarget = scene || cubeWrapper;
+      const rotationTarget = cubeWrapper || scene;
       const jsWarning = document.getElementById("cube-js-warning-" + viewerId);
       const jsWarningText = jsWarning ? jsWarning.querySelector(".cube-warning-text") : null;
       let updateAxisRigBillboard = null;
@@ -734,9 +734,21 @@ def _render_cube_html(
         // Maintain a universal reference frame via CSS variables so scenes with
         // multiple cubes can share rotation/zoom state.
         function applyCubeRotation() {{
-          rotationTarget.style.setProperty("--rot-x", rotationX + "rad");
-          rotationTarget.style.setProperty("--rot-y", rotationY + "rad");
-          rotationTarget.style.setProperty("--zoom", zoom);
+          if (cubeWrapper) {{
+            cubeWrapper.style.setProperty("--rot-x", rotationX + "rad");
+            cubeWrapper.style.setProperty("--rot-y", rotationY + "rad");
+            cubeWrapper.style.setProperty("--zoom", zoom);
+          }}
+          if (scene && scene !== cubeWrapper) {{
+            scene.style.setProperty("--rot-x", rotationX + "rad");
+            scene.style.setProperty("--rot-y", rotationY + "rad");
+            scene.style.setProperty("--zoom", zoom);
+          }}
+          if (!cubeWrapper && rotationTarget) {{
+            rotationTarget.style.setProperty("--rot-x", rotationX + "rad");
+            rotationTarget.style.setProperty("--rot-y", rotationY + "rad");
+            rotationTarget.style.setProperty("--zoom", zoom);
+          }}
           if (updateAxisRigBillboard) {{
             updateAxisRigBillboard();
           }}

src/cubedynamics/variables.py

@@ -128,8 +128,11 @@ def _load_temperature(
             **kwargs,
         )
 
-    da = ds[var_name]
-    da = da.copy()
+    if isinstance(ds, xr.DataArray):
+        da = ds
+    else:
+        da = ds[var_name]
+        da = da.copy()
     da.attrs.setdefault("variable", var_name)
     da.attrs.setdefault("source", source)
     return da

tests/test_cube_viewer_interactivity.py

@@ -45,6 +45,33 @@ def test_cube_viewer_emits_interactive_markup(tmp_path):
     assert (tmp_path / "viewer.html").exists()
 
 
+def test_cube_viewer_rotation_writes_to_wrapper(tmp_path):
+    data = xr.DataArray(
+        np.arange(2 * 3 * 4, dtype=float).reshape(2, 3, 4),
+        dims=("time", "y", "x"),
+        coords={
+            "time": pd.date_range("2023-01-01", periods=2, freq="D"),
+            "y": np.arange(3),
+            "x": np.arange(4),
+        },
+        name="demo",
+    )
+
+    html = cube_from_dataarray(
+        data,
+        out_html=str(tmp_path / "viewer.html"),
+        return_html=True,
+        show_progress=False,
+        thin_time_factor=1,
+    )
+
+    assert "const rotationTarget = cubeWrapper || scene" in html
+    assert "const rotationTarget = scene || cubeWrapper" not in html
+    assert 'cubeWrapper.style.setProperty("--rot-x"' in html
+    assert 'cubeWrapper.style.setProperty("--rot-y"' in html
+    assert 'cubeWrapper.style.setProperty("--zoom"' in html
+
+
 def test_cube_viewer_wraps_html_in_iframe(tmp_path):
     data = xr.DataArray(
         np.arange(4 * 4 * 4, dtype=float).reshape(4, 4, 4),

tests/test_prism_loader.py

@@ -52,29 +52,31 @@ def fake_download(*args, **kwargs):  # pragma: no cover - fallback not used
 
 
 def test_load_prism_cube_with_point_aoi(stub_prism_backends):
-    ds = prism.load_prism_cube(
+    da = prism.load_prism_cube(
         lat=40.0,
         lon=-105.25,
         start="2000-01-01",
         end="2000-12-31",
         variable="ppt",
     )
 
-    assert "ppt" in ds.data_vars
+    assert isinstance(da, xr.DataArray)
+    assert da.name == "ppt"
     aoi = stub_prism_backends["stream"]["aoi"]
     assert aoi["min_lat"] < 40.0 < aoi["max_lat"]
     assert aoi["min_lon"] < -105.25 < aoi["max_lon"]
 
 
 def test_load_prism_cube_with_bbox(stub_prism_backends):
-    ds = prism.load_prism_cube(
+    da = prism.load_prism_cube(
         bbox=[-105.4, 40.0, -105.2, 40.2],
         start="2000-01-01",
         end="2000-12-31",
         variable="ppt",
     )
 
-    assert "ppt" in ds.data_vars
+    assert isinstance(da, xr.DataArray)
+    assert da.name == "ppt"
     assert stub_prism_backends["stream"]["aoi"] == {
         "min_lon": -105.4,
         "min_lat": 40.0,
@@ -99,14 +101,15 @@ def test_load_prism_cube_with_geojson(stub_prism_backends):
         "properties": {"name": "Boulder"},
     }
 
-    ds = prism.load_prism_cube(
+    da = prism.load_prism_cube(
         aoi_geojson=boulder,
         start="2000-01-01",
         end="2000-12-31",
         variable="ppt",
     )
 
-    assert "ppt" in ds.data_vars
+    assert isinstance(da, xr.DataArray)
+    assert da.name == "ppt"
     aoi = stub_prism_backends["stream"]["aoi"]
     assert aoi["min_lon"] == pytest.approx(-105.35)
     assert aoi["max_lon"] == pytest.approx(-105.20)
@@ -145,3 +148,16 @@ def test_load_prism_cube_legacy_positional(stub_prism_backends):
 
     assert "ppt" in ds.data_vars
     assert stub_prism_backends["stream"]["aoi"] == aoi
+
+
+def test_load_prism_cube_variable_returns_dataarray(stub_prism_backends):
+    da = prism.load_prism_cube(
+        lat=40.0,
+        lon=-105.25,
+        start="2000-01-01",
+        end="2000-12-31",
+        variable="ppt",
+    )
+
+    assert isinstance(da, xr.DataArray)
+    assert da.name == "ppt"