diff --git a/lumen/ai/agents/base_view.py b/lumen/ai/agents/base_view.py index c2b1ce8c6..be622ddcb 100644 --- a/lumen/ai/agents/base_view.py +++ b/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, + detect_gridded, get_root_exception, get_schema, 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 = 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, ) diff --git a/lumen/ai/agents/deck_gl.py b/lumen/ai/agents/deck_gl.py index 81243c26e..6b68ad048 100644 --- a/lumen/ai/agents/deck_gl.py +++ b/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, + detect_gridded, get_data, get_schema, 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 = detect_gridded(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 = detect_gridded(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, ) diff --git a/lumen/ai/agents/hvplot.py b/lumen/ai/agents/hvplot.py index 141a09af7..ca19e0d4f 100644 --- a/lumen/ai/agents/hvplot.py +++ b/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 diff --git a/lumen/ai/prompts/BaseViewAgent/main.jinja2 b/lumen/ai/prompts/BaseViewAgent/main.jinja2 index 7adc6fe69..a22f5703c 100644 --- a/lumen/ai/prompts/BaseViewAgent/main.jinja2 +++ b/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 diff --git a/lumen/ai/prompts/DeckGLAgent/main.jinja2 b/lumen/ai/prompts/DeckGLAgent/main.jinja2 index 33034fed6..32c89274c 100644 --- a/lumen/ai/prompts/DeckGLAgent/main.jinja2 +++ b/lumen/ai/prompts/DeckGLAgent/main.jinja2 @@ -45,6 +45,15 @@ 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**: The data is a long-form grid from an xarray dataset. +- dimensions: {{ gridded.dims | join(', ') }} +- data variables: {{ gridded.data_vars | join(', ') }} +Use HexagonLayer with `getPosition: "@@=[{{ gridded.dims[1] if gridded.dims | length > 1 else gridded.dims[0] }}, {{ gridded.dims[0] }}]"` for 3D aggregation. +Alternatively use ColumnLayer with `getElevation: "@@={{ gridded.data_vars[0] }}"` for per-cell elevation. +Large grids (>250,000 rows) require setting `limit` on the view to avoid browser OOM. +{%- endif %} {% endblock %} {% block examples %} @@ -125,4 +134,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 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 %} diff --git a/lumen/ai/prompts/DeckGLAgent/main_pydeck.jinja2 b/lumen/ai/prompts/DeckGLAgent/main_pydeck.jinja2 index 45257ab51..970737a4a 100644 --- a/lumen/ai/prompts/DeckGLAgent/main_pydeck.jinja2 +++ b/lumen/ai/prompts/DeckGLAgent/main_pydeck.jinja2 @@ -25,6 +25,13 @@ 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**: The data is a long-form grid from an xarray dataset. +- dimensions: {{ gridded.dims | join(', ') }} +- data 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] }}"`. +{%- endif %} # Map Styles (no API key needed) - `"https://basemaps.cartocdn.com/gl/dark-matter-gl-style/style.json"` (dark - best for 3D) diff --git a/lumen/ai/prompts/VegaLiteAgent/main.jinja2 b/lumen/ai/prompts/VegaLiteAgent/main.jinja2 index b0e368f30..f1ca7ade2 100644 --- a/lumen/ai/prompts/VegaLiteAgent/main.jinja2 +++ b/lumen/ai/prompts/VegaLiteAgent/main.jinja2 @@ -58,6 +58,14 @@ 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: `, `from: {data: {...}, key: , fields: [...]}` (key/fields inside from!) +{%- if gridded is defined and gridded %} + +**Gridded xarray data**: The data is a long-form grid from an xarray dataset. +- dimensions: {{ gridded.dims | join(', ') }} +- data variables: {{ gridded.data_vars | join(', ') }} +Use `mark: rect` with x/y mapped to coordinate dimensions and color mapped to the value column. +Do not use `longitude`/`latitude` encodings or geographic projections for this data. +{%- endif %} {% endblock %} {% block examples %} @@ -203,6 +211,26 @@ 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: +layer: + - 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: diff --git a/lumen/ai/prompts/hvPlotAgent/main.jinja2 b/lumen/ai/prompts/hvPlotAgent/main.jinja2 index a64ff6155..0f4cf17ab 100644 --- a/lumen/ai/prompts/hvPlotAgent/main.jinja2 +++ b/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 %} diff --git a/lumen/ai/utils.py b/lumen/ai/utils.py index 3e89bb768..93fd20748 100644 --- a/lumen/ai/utils.py +++ b/lumen/ai/utils.py @@ -40,8 +40,9 @@ 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 check_xarray_available, log from .config import ( PROMPTS_DIR, SOURCE_TABLE_SEPARATOR, UNRECOVERABLE_ERRORS, MissingContextError, RetriesExceededError, @@ -1409,3 +1410,23 @@ def result_to_dataframe(result) -> pd.DataFrame | None: return None return None + + +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), + } diff --git a/lumen/tests/ai/test_deckgl_gridded.py b/lumen/tests/ai/test_deckgl_gridded.py new file mode 100644 index 000000000..d5538e7ec --- /dev/null +++ b/lumen/tests/ai/test_deckgl_gridded.py @@ -0,0 +1,114 @@ +import datetime + +import pandas as pd +import pytest + +from lumen.ai.config import PROMPTS_DIR +from lumen.ai.utils import render_template +from lumen.pipeline import Pipeline +from lumen.sources.base import InMemorySource +from lumen.views.base import _DECKGL_ROW_LIMIT, DeckGLView + + +def _base_context(gridded=None): + 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", + "table": "air", + }) + return dict( + memory=memory, + doc="DeckGLView description", + gridded=gridded, + current_datetime=datetime.datetime(2026, 1, 1), + actor_name="DeckGLAgent", + ) + + +def test_deckgl_prompt_includes_gridded_block(): + gridded = { + "source_type": "xarray", + "dims": ["lat", "lon"], + "coords": {"lat": [3], "lon": [4]}, + "data_vars": ["air"], + } + rendered = render_template( + PROMPTS_DIR / "DeckGLAgent" / "main.jinja2", + **_base_context(gridded=gridded), + ) + assert "ColumnLayer" in rendered + assert "air" in rendered + assert "lat" in rendered + assert "lon" in rendered + + +def test_deckgl_prompt_no_gridded_block_for_tabular(): + rendered = render_template( + PROMPTS_DIR / "DeckGLAgent" / "main.jinja2", + **_base_context(gridded=None), + ) + assert "Gridded xarray data" not in rendered + + +def test_deckgl_pydeck_prompt_includes_gridded_block(): + gridded = { + "source_type": "xarray", + "dims": ["lat", "lon"], + "coords": {"lat": [3], "lon": [4]}, + "data_vars": ["air"], + } + rendered = render_template( + PROMPTS_DIR / "DeckGLAgent" / "main_pydeck.jinja2", + **_base_context(gridded=gridded), + ) + assert "ColumnLayer" in rendered + assert "air" in rendered + + +def test_deckglview_row_limit_raises(): + n = _DECKGL_ROW_LIMIT + 1 + df = pd.DataFrame({"lon": range(n), "lat": range(n), "value": range(n)}) + source = InMemorySource(tables={"big": df}) + pipeline = Pipeline(source=source, table="big") + + view = DeckGLView( + pipeline=pipeline, + spec={ + "initialViewState": {"latitude": 0, "longitude": 0, "zoom": 2}, + "layers": [{"@@type": "ScatterplotLayer", "getPosition": "@@=[lon, lat]"}], + "views": [{"@@type": "MapView", "controller": True}], + }, + ) + with pytest.raises(ValueError, match="250,000"): + view._get_params() + + +def test_deckglview_row_limit_passes_under_threshold(): + df = pd.DataFrame({"lon": [0.0, 1.0], "lat": [0.0, 1.0], "value": [1.0, 2.0]}) + source = InMemorySource(tables={"small": df}) + pipeline = Pipeline(source=source, table="small") + + view = DeckGLView( + pipeline=pipeline, + spec={ + "initialViewState": {"latitude": 0, "longitude": 0, "zoom": 2}, + "layers": [{"@@type": "ScatterplotLayer", "getPosition": "@@=[lon, lat]"}], + "views": [{"@@type": "MapView", "controller": True}], + }, + ) + params = view._get_params() + assert "object" in params + assert params["object"]["layers"][0]["data"] == [ + {"lon": 0.0, "lat": 0.0, "value": 1.0}, + {"lon": 1.0, "lat": 1.0, "value": 2.0}, + ] diff --git a/lumen/tests/ai/test_gridded.py b/lumen/tests/ai/test_gridded.py new file mode 100644 index 000000000..46889cf38 --- /dev/null +++ b/lumen/tests/ai/test_gridded.py @@ -0,0 +1,145 @@ +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.utils import detect_gridded, 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 diff --git a/lumen/tests/ai/test_vega_lite_gridded.py b/lumen/tests/ai/test_vega_lite_gridded.py new file mode 100644 index 000000000..ee6725bcf --- /dev/null +++ b/lumen/tests/ai/test_vega_lite_gridded.py @@ -0,0 +1,55 @@ +import datetime + +from lumen.ai.config import PROMPTS_DIR +from lumen.ai.utils import render_template + + +def _base_context(gridded=None): + 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", + "table": "air", + }) + return dict( + memory=memory, + doc="VegaLiteView description", + gridded=gridded, + current_datetime=datetime.datetime(2026, 1, 1), + actor_name="VegaLiteAgent", + ) + + +def test_vegalite_prompt_includes_gridded_block(): + gridded = { + "source_type": "xarray", + "dims": ["lat", "lon"], + "coords": {"lat": [3], "lon": [4]}, + "data_vars": ["air"], + } + rendered = render_template( + PROMPTS_DIR / "VegaLiteAgent" / "main.jinja2", + **_base_context(gridded=gridded), + ) + assert "mark: rect" in rendered + assert "viridis" in rendered + assert "lat" in rendered + assert "air" in rendered + + +def test_vegalite_prompt_no_gridded_block_for_tabular(): + rendered = render_template( + PROMPTS_DIR / "VegaLiteAgent" / "main.jinja2", + **_base_context(gridded=None), + ) + assert "Gridded xarray data" not in rendered + assert "Gridded heatmap" not in rendered diff --git a/lumen/tests/views/test_hvplot_gridded.py b/lumen/tests/views/test_hvplot_gridded.py new file mode 100644 index 000000000..19fd67afc --- /dev/null +++ b/lumen/tests/views/test_hvplot_gridded.py @@ -0,0 +1,194 @@ +from unittest.mock import MagicMock, patch + +import holoviews as hv +import numpy as np +import pandas as pd +import pytest + +from lumen.pipeline import Pipeline +from lumen.sources.base import InMemorySource +from lumen.views.base import hvPlotView + +# ---- Fixtures ---- + +@pytest.fixture +def gridded_df(): + """Long-form 2D grid: lat x lon with an 'air' value column.""" + lats = np.repeat([0.0, 1.0, 2.0], 4) + lons = np.tile([10.0, 20.0, 30.0, 40.0], 3) + air = np.arange(12, dtype=float) + return pd.DataFrame({"lat": lats, "lon": lons, "air": air}) + + +@pytest.fixture +def gridded_pipeline(gridded_df): + source = InMemorySource(tables={"grid": gridded_df}) + return Pipeline(source=source, table="grid") + + +@pytest.fixture +def tabular_df(): + return pd.DataFrame({"x": [1.0, 2.0, 3.0], "y": [4.0, 5.0, 6.0]}) + + +@pytest.fixture +def tabular_pipeline(tabular_df): + source = InMemorySource(tables={"data": tabular_df}) + return Pipeline(source=source, table="data") + + +# ---- Tests ---- + +def test_hvplot_z_param_accepted(): + """z param is a valid Selector on hvPlotBaseView (no AttributeError).""" + assert "z" in hvPlotView.param + + +def test_hvplot_quadmesh_in_kind_objects(): + """quadmesh is a recognised kind value.""" + assert "quadmesh" in hvPlotView.param.kind.objects + + +def test_hvplot_heatmap_from_longform_df(gridded_pipeline): + """kind=heatmap with x, y, z renders a 2D HeatMap from a long-form grid. + + The public z= param is mapped internally to hvPlot's C= for heatmap. + """ + view = hvPlotView( + pipeline=gridded_pipeline, + kind="heatmap", + x="lon", + y="lat", + z="air", + ) + panel = view.get_panel() + assert panel is not None + plot = view.get_plot(view.get_data()) + assert isinstance(plot, hv.HeatMap) + + +def test_hvplot_z_maps_to_C_only_for_heatmap(gridded_pipeline): + """z= is mapped to C= for heatmap, and passed as z= for other gridded kinds. + + heatmap stays on pandas (no pivot), so we can patch DataFrame.hvplot. + For contourf/quadmesh/image the view pivots to xarray, so we patch the + pivoted DataArray's hvplot accessor instead. + """ + pytest.importorskip("xarray") + heatmap_view = hvPlotView( + pipeline=gridded_pipeline, kind="heatmap", x="lon", y="lat", z="air", + ) + df = heatmap_view.get_data() + with patch.object(type(df), "hvplot", MagicMock()) as mock: + heatmap_view.get_plot(df) + call_kwargs = mock.call_args.kwargs + assert call_kwargs.get("C") == "air" + assert "z" not in call_kwargs + + contour_view = hvPlotView( + pipeline=gridded_pipeline, kind="contourf", x="lon", y="lat", z="air", + ) + df = contour_view.get_data() + gridded = contour_view._to_gridded(df) + with patch.object(type(gridded), "hvplot", MagicMock()) as mock: + contour_view.get_plot(df) + call_kwargs = mock.call_args.kwargs + assert call_kwargs.get("z") == "air" + assert "C" not in call_kwargs + + +def test_hvplot_z_none_does_not_break_scatter(tabular_pipeline): + """When z is None (default), scatter plots work unchanged.""" + view = hvPlotView( + pipeline=tabular_pipeline, + kind="scatter", + x="x", + y="y", + ) + panel = view.get_panel() + assert panel is not None + + +def test_hvplot_quadmesh_from_longform_df(gridded_pipeline): + """kind=quadmesh renders a QuadMesh from a long-form pandas DataFrame. + + Long-form pandas is the only data format XArraySQLSource produces, so the + view must pivot it to a 2D array before handing to hvPlot, otherwise + hvPlot raises NotImplementedError for quadmesh on pandas. Addresses the + review on #1823: quadmesh/image are semantically more correct than + heatmap for continuous-coordinate gridded data. + """ + pytest.importorskip("xarray") + view = hvPlotView( + pipeline=gridded_pipeline, + kind="quadmesh", + x="lon", + y="lat", + z="air", + ) + plot = view.get_plot(view.get_data()) + assert isinstance(plot, hv.QuadMesh) + + +def test_hvplot_image_from_longform_df(gridded_pipeline): + """kind=image renders an Image from a long-form pandas DataFrame via pivot.""" + pytest.importorskip("xarray") + view = hvPlotView( + pipeline=gridded_pipeline, + kind="image", + x="lon", + y="lat", + z="air", + ) + plot = view.get_plot(view.get_data()) + assert isinstance(plot, hv.Image) + + +def test_hvplot_gridded_size_guard_raises(gridded_pipeline): + """Pivoting a grid larger than _GRIDDED_MAX_CELLS raises ValueError.""" + pytest.importorskip("xarray") + view = hvPlotView( + pipeline=gridded_pipeline, + kind="quadmesh", + x="lon", + y="lat", + z="air", + ) + view._GRIDDED_MAX_CELLS = 4 # the fixture has 3 lats x 4 lons = 12 cells + with pytest.raises(ValueError, match="exceeding the 4 safety cap"): + view.get_plot(view.get_data()) + + +def test_hvplot_gridded_skips_pivot_if_xarray_unavailable(gridded_pipeline): + """If hvplot.xarray failed to import, _to_gridded returns the df unchanged. + + Build a view with a non-gridded kind so construction doesn't trigger the + quadmesh render path, then flip the flag and call _to_gridded directly. + """ + view = hvPlotView( + pipeline=gridded_pipeline, + kind="scatter", + x="lon", + y="lat", + z="air", + ) + view._hvplot_xarray_available = False + result = view._to_gridded(view.get_data()) + assert isinstance(result, pd.DataFrame), "expected df fallback when xarray unavailable" + + +def test_hvplot_gridded_passes_through_non_dataframe(gridded_pipeline): + """If the pipeline hands over an already-gridded xarray object, skip pivot.""" + xr = pytest.importorskip("xarray") + view = hvPlotView( + pipeline=gridded_pipeline, + kind="quadmesh", + x="lon", + y="lat", + z="air", + ) + df = view.get_data() + already_gridded = df.set_index(["lat", "lon"])["air"].to_xarray() + assert isinstance(already_gridded, xr.DataArray) + result = view._to_gridded(already_gridded) + assert result is already_gridded, "xarray input should pass through unchanged" diff --git a/lumen/views/base.py b/lumen/views/base.py index b45d5d3c4..f6296084c 100644 --- a/lumen/views/base.py +++ b/lumen/views/base.py @@ -49,6 +49,7 @@ from holoviews.selection import link_selections # type: ignore DOWNLOAD_FORMATS = ['csv', 'xlsx', 'json', 'parquet'] +_DECKGL_ROW_LIMIT = 250_000 class View(MultiTypeComponent, Viewer): @@ -926,7 +927,8 @@ class hvPlotBaseView(View): objects=[ 'area', 'bar', 'barh', 'bivariate', 'box', 'contour', 'contourf', 'errorbars', 'hist', 'image', 'kde', 'labels', - 'line', 'scatter', 'heatmap', 'hexbin', 'ohlc', 'points', 'step', 'violin' + 'line', 'scatter', 'heatmap', 'hexbin', 'ohlc', 'points', 'quadmesh', + 'step', 'violin' ] ) @@ -938,19 +940,31 @@ class hvPlotBaseView(View): groupby = param.ListSelector(doc="The column(s) to group by.") + z = param.Selector(doc=""" + Column of z-values for gridded plot kinds (image, quadmesh, heatmap, contourf). + Internally mapped to hvPlot's C= for kind='heatmap' and z= for other kinds.""") + geo = param.Boolean( default=False, doc="Toggle True if the plot is on a geographic map." ) - _field_params = ['x', 'y', 'by', 'groupby'] + _field_params = ['x', 'y', 'by', 'groupby', 'z'] __abstract = True + _hvplot_xarray_available = False + def __init__(self, **params): import hvplot.pandas # type: ignore + try: + import hvplot.xarray # type: ignore + except Exception: + type(self)._hvplot_xarray_available = False + else: + type(self)._hvplot_xarray_available = True if 'dask' in sys.modules: try: - import hvplot.dask # type: ignore # noqa + import hvplot.dask # type: ignore # noqa: F401 except Exception: pass if 'by' in params and isinstance(params['by'], str): @@ -1029,6 +1043,46 @@ def __init__(self, **params): self._linked_objs = [] super().__init__(**params) + _GRIDDED_KINDS = ('quadmesh', 'image', 'contourf') + + # Safety cap on pivoted grid size. 10M cells = ~80MB for float64. + # Anything larger is rejected with a clear error rather than risking OOM. + _GRIDDED_MAX_CELLS = 10_000_000 + + def _to_gridded(self, df): + """Pivot a long-form DataFrame (y, x, z columns) to a 2D xarray + DataArray that hvPlot's quadmesh/image/contourf kinds can consume. + + Returns the original df unchanged if x, y, or z are not set, if + required columns are missing, if duplicate (y, x) combinations are + present, or if the input is already an xarray object (falls back + to hvPlot's own path so it raises the informative error). + + Raises ValueError if the pivoted grid would exceed + _GRIDDED_MAX_CELLS to prevent OOM. + """ + import pandas as pd + if not isinstance(df, pd.DataFrame): + return df + if not self._hvplot_xarray_available: + return df + if not (self.x and self.y and self.z): + return df + required_cols = {self.x, self.y, self.z} + if not required_cols.issubset(df.columns): + return df + if df.duplicated(subset=[self.y, self.x]).any(): + return df + n_cells = df[self.y].nunique() * df[self.x].nunique() + if n_cells > self._GRIDDED_MAX_CELLS: + raise ValueError( + f"Pivoted grid would have {n_cells:,} cells, exceeding the " + f"{self._GRIDDED_MAX_CELLS:,} safety cap. Reduce the data " + f"(filter, resample, or aggregate) before using kind={self.kind!r}." + ) + indexed = df.set_index([self.y, self.x])[self.z] + return indexed.to_xarray() + def get_plot(self, df): processed = {} for k, v in self.kwargs.items(): @@ -1039,8 +1093,14 @@ def get_plot(self, df): processed[k] = v if self.streaming: processed['stream'] = self._data_stream + if self.z is not None: + processed['C' if self.kind == 'heatmap' else 'z'] = self.z + + plot_source = df + if self.kind in self._GRIDDED_KINDS: + plot_source = self._to_gridded(df) - plot = df.hvplot( + plot = plot_source.hvplot( kind=self.kind, x=self.x, y=self.y, by=self.by, groupby=self.groupby, **processed ) if self.operations: @@ -1347,6 +1407,11 @@ class DeckGLView(View): def _get_params(self) -> dict[str, Any]: df = self.get_data() + if len(df) > _DECKGL_ROW_LIMIT: + raise ValueError( + f"DeckGLView received {len(df):,} rows which may cause browser " + f"memory issues. Set limit={_DECKGL_ROW_LIMIT:,} or less on this view." + ) # Deep copy to avoid modifying self.spec when injecting data spec = copy.deepcopy(self.spec)