Support plotting xarray grids (hvPlot + VegaLite + DeckGL)

lumen/ai/agents/base_view.py

@@ -11,7 +11,8 @@
 from ..llm import Message
 from ..schemas import Metaset
 from ..utils import (
-    get_root_exception, get_schema, log_debug, report_error, retry_llm_output,
+    get_root_exception, get_schema, gridded_metadata, log_debug, report_error,
+    retry_llm_output,
 )
 from .base_lumen import BaseLumenAgent
 
@@ -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 = gridded_metadata(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/deck_gl.py

@@ -21,7 +21,8 @@
 from ..editors import DeckGLEditor
 from ..models import BaseModel, EscapeBaseModel, RetrySpec
 from ..utils import (
-    get_data, get_schema, retry_llm_output, sanitize_column_names,
+    get_data, get_schema, gridded_metadata, retry_llm_output,
+    sanitize_column_names,
 )
 from .base_code import BaseCodeAgent
 
@@ -152,6 +153,7 @@ async def _generate_yaml_spec(
     ) -> dict[str, Any]:
         """Generate DeckGL spec via YAML/JSON (declarative mode)."""
         errors_context = self._build_errors_context(pipeline, context, errors)
+        gridded = gridded_metadata(pipeline)
 
         with self._add_step(title="Generating DeckGL specification", steps_layout=self._steps_layout) as step:
             response = self._stream_prompt(
@@ -160,6 +162,7 @@ async def _generate_yaml_spec(
                 context,
                 table=pipeline.table,
                 doc=doc,
+                gridded=gridded,
                 **errors_context,
             )
             async for output in response:
@@ -182,6 +185,7 @@ async def _generate_code_spec(
     ) -> dict[str, Any] | None:
         """Generate spec via PyDeck code execution."""
         errors_context = self._build_errors_context(pipeline, context, errors)
+        gridded = gridded_metadata(pipeline)
 
         with self._add_step(title="Generating PyDeck code", steps_layout=self._steps_layout) as step:
             response = self._stream_prompt(
@@ -190,6 +194,7 @@ async def _generate_code_spec(
                 context,
                 table=pipeline.table,
                 doc=doc,
+                gridded=gridded,
                 **errors_context,
             )
 

lumen/ai/agents/vega_lite.py

@@ -24,8 +24,8 @@
 from ..llm import Message, OpenAI
 from ..models import EscapeBaseModel, RetrySpec
 from ..utils import (
-    get_data, get_schema, load_json, log_debug, normalize_vegalite_spec,
-    retry_llm_output,
+    get_data, get_schema, gridded_metadata, load_json, log_debug,
+    normalize_vegalite_spec, retry_llm_output,
 )
 from ..vector_store import DuckDBVectorStore
 from .base_code import BaseCodeAgent
@@ -415,6 +415,7 @@ async def _generate_yaml_spec(
     ) -> dict[str, Any]:
         """Generate VegaLite spec via YAML (declarative mode)."""
         errors_context = self._build_errors_context(pipeline, context, errors)
+        gridded = gridded_metadata(pipeline)
         with self._add_step(title="Creating basic plot structure", steps_layout=self._steps_layout) as step:
             response = self._stream_prompt(
                 "main",
@@ -423,6 +424,7 @@ async def _generate_yaml_spec(
                 table=pipeline.table,
                 doc=doc,
                 doc_examples=doc_examples,
+                gridded=gridded,
                 **errors_context,
             )
             async for output in response:
@@ -444,6 +446,7 @@ async def _generate_code_spec(
     ) -> dict[str, Any] | None:
         """Generate spec via Altair code execution."""
         errors_context = self._build_errors_context(pipeline, context, errors)
+        gridded = gridded_metadata(pipeline)
 
         with self._add_step(title="Generating Altair code", steps_layout=self._steps_layout) as step:
             response = self._stream_prompt(
@@ -453,6 +456,7 @@ async def _generate_code_spec(
                 table=pipeline.table,
                 doc=doc,
                 doc_examples=doc_examples,
+                gridded=gridded,
                 **errors_context,
             )
             async for output in response:

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/DeckGLAgent/main.jinja2

@@ -45,6 +45,14 @@ Column names: spaces→underscores, non-alphanumeric removed (e.g., "Plant latit
 - **Elevation/Radius:** Use simple math only: `"@@=value"` or `"@@=value * 100"` (NO Math.min, Math.max, etc.)
 - Zoom levels: 1-3 continent, 4-6 country, 7-10 city, 11+ street
 - Pitch: 45-60 for good 3D perspective
+{%- if gridded is defined and gridded %}
+
+**Gridded xarray data**: This table is a long-form grid derived from an xarray dataset.
+- coordinate dimensions: {{ gridded.dims | join(', ') }}
+- data variables: {{ gridded.data_vars | join(', ') }}
+This only reads well on a map when the coordinate dimensions are geographic (longitude/latitude).
+Use `ColumnLayer` with `getPosition: "@@=[{{ gridded.dims[1] if gridded.dims | length > 1 else gridded.dims[0] }}, {{ gridded.dims[0] }}]"` and `getElevation: "@@={{ gridded.data_vars[0] }}"` for per-cell elevation, or `HexagonLayer` for spatial aggregation.
+{%- endif %}
 {% endblock %}
 
 {% block examples %}
@@ -125,4 +133,27 @@ ArcLayer (connections):
   "mapStyle": "https://basemaps.cartocdn.com/gl/dark-matter-gl-style/style.json"
 }
 ```
+{%- if gridded is defined and gridded %}
+
+Gridded data, ColumnLayer (per-cell elevation from a value field):
+```json
+{
+  "initialViewState": {"latitude": 0, "longitude": 30, "zoom": 3, "pitch": 50, "bearing": 0},
+  "layers": [{
+    "@@type": "ColumnLayer",
+    "id": "grid-values",
+    "getPosition": "@@=[{{ gridded.dims[1] if gridded.dims | length > 1 else gridded.dims[0] }}, {{ gridded.dims[0] }}]",
+    "getElevation": "@@={{ gridded.data_vars[0] }}",
+    "getFillColor": [65, 182, 196, 220],
+    "diskResolution": 6,
+    "radius": 50000,
+    "elevationScale": 10,
+    "extruded": true,
+    "pickable": true
+  }],
+  "views": [{"@@type": "MapView", "controller": true}],
+  "mapStyle": "https://basemaps.cartocdn.com/gl/dark-matter-gl-style/style.json"
+}
+```
+{%- endif %}
 {% endblock %}

lumen/ai/prompts/DeckGLAgent/main_pydeck.jinja2

@@ -25,6 +25,12 @@ Generate PyDeck code. Data is `df`, assign to `deck`. NO I/O methods (.to_html()
 - `bearing`: 0 or slight angle (-15 to 15) for visual interest
 - `elevation_scale`: Controls the height of 3D elements
 - `extruded`: Set to True for 3D extrusion
+{%- if gridded is defined and gridded %}
+
+## Gridded xarray data
+`df` is a long-form grid from an xarray dataset (dimensions: {{ gridded.dims | join(', ') }}; variables: {{ gridded.data_vars | join(', ') }}).
+Use `ColumnLayer` with `get_position=["{{ gridded.dims[1] if gridded.dims | length > 1 else gridded.dims[0] }}", "{{ gridded.dims[0] }}"]` and `get_elevation="{{ gridded.data_vars[0] }}"`. This reads well only when the coordinate dimensions are geographic.
+{%- endif %}
 
 ## Map Styles (no API key needed)
 - `"https://basemaps.cartocdn.com/gl/dark-matter-gl-style/style.json"` (dark - best for 3D)

lumen/ai/prompts/VegaLiteAgent/main.jinja2

@@ -58,6 +58,15 @@ Legends appear when you map data to `color`, `size`, or `shape` in encoding.
 **Geographic maps**: Use `longitude`/`latitude` encodings (NOT `x`/`y`) with `projection: {type: mercator}` at top level. Base map added automatically.
 
 **Choropleth maps**: Join data to map boundaries - `lookup: <map_field>`, `from: {data: {...}, key: <your_field>, fields: [...]}` (key/fields inside from!)
+{%- if gridded is defined and gridded %}
+
+**Gridded xarray data**: This table is a long-form grid derived from an xarray dataset.
+- coordinate dimensions: {{ gridded.dims | join(', ') }}
+- data variables: {{ gridded.data_vars | join(', ') }}
+Use `mark: rect` with `x`/`y` mapped to the coordinate dimensions (type `ordinal`) and
+`color` mapped to a data variable (type `quantitative`, `scale: {scheme: viridis}`).
+Do NOT use `longitude`/`latitude` encodings or geographic projections for this data.
+{%- endif %}
 {% endblock %}
 
 {% block examples %}
@@ -203,6 +212,24 @@ layer:
       color: {field: category, type: nominal, legend: {title: "Category", orient: "right"}}
 ```
 
+{%- if gridded is defined and gridded %}
+Gridded heatmap (xarray-derived long-form data):
+```yaml
+data:
+  name: <TABLE_NAME>
+mark: rect
+encoding:
+  x: {field: {{ gridded.dims[1] if gridded.dims | length > 1 else gridded.dims[0] }}, type: ordinal}
+  y: {field: {{ gridded.dims[0] }}, type: ordinal}
+  color:
+    field: {{ gridded.data_vars[0] }}
+    type: quantitative
+    scale: {scheme: viridis}
+    legend: {title: "{{ gridded.data_vars[0] }}"}
+config:
+  view: {stroke: null}
+```
+{%- endif %}
 Choropleth map:
 ```yaml
 data:

lumen/ai/prompts/VegaLiteAgent/main_altair.jinja2

@@ -16,6 +16,12 @@ Generate Altair code. Data is `df`, assign to `chart`, end with `.interactive()`
 - When the n_rows > 5000:
   - You **MUST** include `alt.data_transformers.disable_max_rows()` and `alt.renderers.set_embed_options(renderer="canvas")`
   - Prefer aggregated and binned plots
+{%- if gridded is defined and gridded %}
+
+## Gridded xarray data
+`df` is a long-form grid from an xarray dataset (dimensions: {{ gridded.dims | join(', ') }}; variables: {{ gridded.data_vars | join(', ') }}).
+Use `mark_rect()` for a heatmap: encode `x`/`y` as ordinal coordinate dimensions and `color` as the quantitative value with `alt.Scale(scheme='viridis')`. Do not use geographic projections.
+{%- endif %}
 {% endblock %}
 
 {% block examples %}

lumen/ai/prompts/hvPlotAgent/main.jinja2

@@ -5,4 +5,25 @@ 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:
+- Set x and y to coordinate names from: {{ gridded.dims | join(', ') }}.
+- Set z to one of the data variables: {{ gridded.data_vars | join(', ') }}.
+- Pick `kind` using this rule:
+  {%- if gridded.regular %}
+  - This dataset has regularly-spaced coordinates, so prefer kind='image'
+    (fastest, best for evenly-gridded data). Use 'quadmesh' if you need to
+    show coordinate cell boundaries explicitly, or 'contourf' for filled
+    contours.
+  {%- else %}
+  - This dataset has irregularly-spaced coordinates, so use kind='quadmesh'
+    (handles non-uniform cells correctly). Do NOT use 'image' here; it
+    assumes uniform spacing and will misrepresent the data.
+  {%- endif %}
+  - Use 'heatmap' only when axes are ordinal or categorical, not continuous.
+- 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/ai/utils.py

@@ -23,6 +23,7 @@
 from typing import TYPE_CHECKING, Any, Literal
 from urllib.parse import parse_qs
 
+import numpy as np
 import pandas as pd
 import param
 import sqlglot
@@ -41,8 +42,9 @@
 from ..config import dump_yaml
 from ..pipeline import Pipeline
 from ..sources.base import Source
+from ..sources.xarray_sql import XArraySQLSource
 from ..transforms import SQLRemoveSourceSeparator
-from ..util import log
+from ..util import log, try_import_xarray
 from .config import (
     PROMPTS_DIR, SOURCE_TABLE_SEPARATOR, UNRECOVERABLE_ERRORS, VEGA_MAP_LAYER,
     VEGA_ZOOMABLE_MAP_ITEMS, MissingContextError, RetriesExceededError,
@@ -1489,6 +1491,42 @@ def result_to_dataframe(result) -> pd.DataFrame | None:
     return None
 
 
+def gridded_metadata(pipeline: Pipeline) -> dict[str, Any] | None:
+    """Return gridded metadata for an xarray-backed pipeline, else None.
+
+    Returns a dict with keys ``source_type``, ``dims``, ``coords``,
+    ``data_vars``, and ``regular`` (True iff every coordinate dimension is
+    uniformly spaced). Coord arrays are 1-D and held in memory by xarray
+    even for large lazy datasets, so the regularity check is cheap; do not
+    call on a hot path.
+    """
+    if not isinstance(pipeline.source, XArraySQLSource) or try_import_xarray() is None:
+        return None
+
+    ds = pipeline.source.dataset
+
+    def _is_regular(values) -> bool:
+        arr = np.asarray(values)
+        if arr.ndim != 1 or arr.size < 2:
+            return False
+        # datetime64 coords need int64 conversion before allclose
+        if np.issubdtype(arr.dtype, np.datetime64):
+            arr = arr.astype('int64')
+        diffs = np.diff(arr)
+        return bool(np.allclose(diffs, diffs[0], rtol=1e-6))
+
+    regular = all(
+        _is_regular(ds.coords[d].values) for d in ds.dims if d in ds.coords
+    )
+    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),
+        'regular': regular,
+    }
+
+
 def normalize_vegalite_spec(
     vega_spec: dict[str, Any], *, editor_type: type[VegaLiteEditor] | None = None
 ) -> dict[str, Any]: