Support plotting xarray grids (hvPlot + VegaLite + DeckGL)

lumen/ai/agents/base_view.py

@@ -8,6 +8,7 @@
 from ...pipeline import Pipeline
 from ..config import PROMPTS_DIR, SOURCE_TABLE_SEPARATOR, MissingContextError
 from ..context import ContextModel, TContext
+from ..gridded import detect_gridded
 from ..llm import Message
 from ..schemas import Metaset
 from ..utils import (
@@ -88,12 +89,14 @@ async def _generate_yaml_spec(
     ) -> dict[str, Any]:
         errors_context = self._build_errors_context(pipeline, context, errors)
         doc = self.view_type.__doc__.split("\n\n")[0] if self.view_type.__doc__ else self.view_type.__name__
+        gridded = detect_gridded(pipeline)
         response = await self._stream_prompt(
             "main",
             messages,
             context,
             table=pipeline.table,
             doc=doc,
+            gridded=gridded,
             model_kwargs=dict(schema=schema),
             **errors_context,
         )

lumen/ai/agents/hvplot.py

@@ -64,6 +64,9 @@ async def _extract_spec(self, context: TContext, spec: dict[str, Any]):
 
         # Add defaults
         spec["responsive"] = True
+        if spec.get("kind") == "heatmap":
+            spec.pop("by", None)
+            spec.pop("groupby", None)
         data = await get_data(pipeline)
         if len(data) > 20000 and spec["kind"] in ("line", "scatter", "points"):
             spec["rasterize"] = True

lumen/ai/gridded.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ..pipeline import Pipeline
+from ..sources.xarray_sql import XArraySQLSource
+from ..util import check_xarray_available
+
+
+def detect_gridded(pipeline: Pipeline) -> dict[str, Any] | None:
+    """Return gridded metadata if pipeline.source is an XArraySQLSource, else None.
+
+    Returns a dict with keys ``dims``, ``coords``, ``data_vars``, ``source_type``,
+    or ``None`` for tabular pipelines.
+    """
+    source = pipeline.source
+    if not isinstance(source, XArraySQLSource):
+        return None
+    if not check_xarray_available():
+        return None
+    ds = source.dataset
+    return {
+        'source_type': 'xarray',
+        'dims': list(ds.dims),
+        'coords': {k: list(ds.coords[k].shape) for k in ds.coords},
+        'data_vars': list(ds.data_vars),
+    }

lumen/ai/prompts/BaseViewAgent/main.jinja2

@@ -9,6 +9,16 @@ The current SQL table name: {{ memory['pipeline'].table }}.
 Here is the current dataset to use:
 {{ memory['data'] }}
 
+{%- if gridded is defined and gridded %}
+
+Data originates from an xarray dataset.
+- dimensions: {{ gridded.dims | join(', ') }}
+- data variables: {{ gridded.data_vars | join(', ') }}
+Use a 2D gridded encoding with axes on the coordinate dimensions and
+colour/value on the data variable, using whichever syntax this agent
+accepts (see agent-specific instructions below).
+{%- endif %}
+
 {%- if "view" in memory and memory["view"].get("agent") and actor_name.startswith(memory["view"].get("agent")) %}
 Adapt from the previous view specification below as needed:
 ```yaml

lumen/ai/prompts/hvPlotAgent/main.jinja2

@@ -5,4 +5,16 @@ Generate the plot the user requested. Note that `x`, `y`, `by` and `groupby` fie
 no repeated columns are allowed. Do not arbitrarily set `groupby` and `by` fields unless explicitly requested. If a
 histogram is requested, use `y` instead of `x`. If x is categorical or strings, prefer barh over bar, and use `y` for
 the values.
+{%- if gridded is defined and gridded %}
+
+For gridded xarray data:
+- Prefer kind='quadmesh'. It preserves continuous coordinate axes and is
+  the right default for xarray-derived grids. Use 'image' only for strictly
+  regular grids, 'contourf' for filled contours, or 'heatmap' if you
+  specifically want ordinal/categorical axes.
+- Set x and y to coordinate names (e.g. x='lon', y='lat').
+- Set z to the value column (e.g. z='air'). The view pivots long-form
+  pandas to xarray internally and maps z= to hvPlot's C= for heatmap only.
+- Do not set by or groupby.
+{%- endif %}
 {% endblock %}

lumen/tests/ai/test_gridded.py

@@ -0,0 +1,146 @@
+import datetime
+
+import numpy as np
+import pandas as pd
+import pytest
+
+xr = pytest.importorskip("xarray")
+pytest.importorskip("xarray_sql")
+
+from lumen.ai.config import PROMPTS_DIR
+from lumen.ai.gridded import detect_gridded
+from lumen.ai.utils import render_template
+from lumen.pipeline import Pipeline
+from lumen.sources.base import InMemorySource
+from lumen.sources.xarray_sql import XArraySQLSource
+from lumen.util import check_xarray_available
+
+pytestmark = pytest.mark.skipif(
+    not check_xarray_available(), reason="xarray not installed"
+)
+
+
+# ---- Fixtures ----
+
+@pytest.fixture
+def simple_dataset():
+    lats = np.array([0.0, 1.0, 2.0])
+    lons = np.array([10.0, 20.0, 30.0, 40.0])
+    air = np.arange(12, dtype=float).reshape(3, 4)
+    return xr.Dataset(
+        {"air": (["lat", "lon"], air)},
+        coords={"lat": lats, "lon": lons},
+    )
+
+
+@pytest.fixture
+def xarray_source(simple_dataset):
+    return XArraySQLSource(_dataset=simple_dataset)
+
+
+@pytest.fixture
+def xarray_pipeline(xarray_source):
+    return Pipeline(source=xarray_source, table="air")
+
+
+@pytest.fixture
+def tabular_pipeline():
+    df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
+    source = InMemorySource(tables={"data": df})
+    return Pipeline(source=source, table="data")
+
+
+# ---- detect_gridded tests ----
+
+def test_detect_gridded_xarray_source(xarray_pipeline):
+    result = detect_gridded(xarray_pipeline)
+    assert result is not None
+    assert result["source_type"] == "xarray"
+    assert "lat" in result["dims"]
+    assert "lon" in result["dims"]
+    assert "air" in result["data_vars"]
+
+
+def test_detect_gridded_pandas_source(tabular_pipeline):
+    result = detect_gridded(tabular_pipeline)
+    assert result is None
+
+
+# ---- Prompt rendering tests ----
+
+def _base_context(gridded=None):
+    """Minimal context required to render BaseViewAgent/main.jinja2."""
+    class _FakeSource:
+        dialect = None
+
+    class _FakePipeline:
+        table = "air"
+        source = _FakeSource()
+
+    class _FakeMemory(dict):
+        pass
+
+    memory = _FakeMemory({
+        "pipeline": _FakePipeline(),
+        "data": "lat,lon,air\n0,10,1.0",
+    })
+    return dict(
+        memory=memory,
+        doc="hvPlotUIView description",
+        gridded=gridded,
+        current_datetime=datetime.datetime(2026, 1, 1),
+        actor_name="hvPlotAgent",
+    )
+
+
+def test_baseview_prompt_includes_gridded_block():
+    gridded = {
+        "source_type": "xarray",
+        "dims": ["lat", "lon"],
+        "coords": {"lat": [3], "lon": [4]},
+        "data_vars": ["air"],
+    }
+    rendered = render_template(
+        PROMPTS_DIR / "BaseViewAgent" / "main.jinja2",
+        **_base_context(gridded=gridded),
+    )
+    assert "xarray dataset" in rendered
+    assert "lat, lon" in rendered
+    assert "air" in rendered
+    # The base prompt must stay agent-agnostic: no hvPlot-specific kind
+    # values, since VegaLite and DeckGL inherit this template too.
+    assert "kind=" not in rendered
+    assert "quadmesh" not in rendered
+    assert "heatmap" not in rendered
+
+
+def test_baseview_prompt_no_gridded_block_for_tabular():
+    rendered = render_template(
+        PROMPTS_DIR / "BaseViewAgent" / "main.jinja2",
+        **_base_context(gridded=None),
+    )
+    assert "xarray dataset" not in rendered
+    assert "gridded plot kinds" not in rendered
+
+
+def test_hvplot_prompt_includes_gridded_rules():
+    gridded = {
+        "source_type": "xarray",
+        "dims": ["lat", "lon"],
+        "coords": {"lat": [3], "lon": [4]},
+        "data_vars": ["air"],
+    }
+    rendered = render_template(
+        PROMPTS_DIR / "hvPlotAgent" / "main.jinja2",
+        **_base_context(gridded=gridded),
+    )
+    assert "quadmesh" in rendered
+    assert "z" in rendered
+
+
+def test_hvplot_prompt_no_gridded_rules_for_tabular():
+    rendered = render_template(
+        PROMPTS_DIR / "hvPlotAgent" / "main.jinja2",
+        **_base_context(gridded=None),
+    )
+    assert "For gridded xarray data" not in rendered