diff --git a/tests/test_figures.py b/tests/test_figures.py index b9124d6..8166200 100644 --- a/tests/test_figures.py +++ b/tests/test_figures.py @@ -797,6 +797,151 @@ def test_bar_zero_baseline(self) -> None: lo, _hi = combined.layout.yaxis.range assert lo <= 0, f"Bar y-axis range should include 0, got lo={lo}" + def test_x_axis_stays_on_autorange_when_frames_share_extent(self) -> None: + """x-axis must not be pinned when every frame spans the same x values. + + Pinning it froze a numeric range that became a garbage category-index + range if the user later switched the axis to type='category'. + """ + bar_values = np.full((4, 3, 2), 10.0) + bar_values[:, :, 1] = 50.0 # second frame exceeds the first + da = xr.DataArray( + bar_values, + dims=["period", "tech", "case"], + coords={"period": [2025, 2030, 2035, 2040]}, + name="heat", + ) + line_da = xr.DataArray( + np.full((4, 2), 20.0), + dims=["period", "case"], + coords={"period": [2025, 2030, 2035, 2040]}, + name="demand", + ) + combined = overlay( + xpx(da).bar(x="period", barmode="relative", animation_frame="case"), + xpx(line_da).line(animation_frame="case"), + ) + + # Same periods in every frame: initial autorange stays valid + assert combined.layout.xaxis.range is None + # y differs per frame, so it should be pinned + assert combined.layout.yaxis.range is not None + + def test_no_axis_pinned_when_frames_identical(self) -> None: + """Nothing is pinned when all frames have identical data extents.""" + da = xr.DataArray( + np.array([[1.0, 5.0], [1.0, 5.0], [1.0, 5.0]]).T, + dims=["x", "frame"], + name="val", + ) + combined = overlay( + xpx(da).line(animation_frame="frame"), + xpx(da).scatter(animation_frame="frame"), + ) + + assert combined.layout.xaxis.range is None + assert combined.layout.yaxis.range is None + + def test_y_axis_pinned_when_later_frame_exceeds_initial(self) -> None: + """y-axis is pinned to cover a later frame with larger values.""" + da = xr.DataArray( + np.array([[1.0, 100.0], [2.0, 200.0], [3.0, 300.0]]), + dims=["x", "frame"], + name="val", + ) + combined = overlay( + xpx(da).line(animation_frame="frame"), + xpx(da).scatter(animation_frame="frame"), + ) + + lo, hi = combined.layout.yaxis.range + assert lo <= 1.0 + assert hi >= 300.0 + + def test_stacked_bars_range_covers_summed_height(self) -> None: + """With barmode='relative', the y-range must cover stacked totals.""" + # In the second frame, 3 techs of 100 each stack to 300 per period — + # beyond the first frame's extent, so the y-axis gets pinned and must + # use stacked sums, not individual segment values (max segment: 100). + values = np.full((2, 3, 2), 100.0) + values[:, :, 0] = 50.0 + da = xr.DataArray( + values, + dims=["period", "tech", "case"], + name="heat", + ) + combined = overlay( + xpx(da).bar(x="period", barmode="relative", animation_frame="case"), + xpx(da.isel(tech=0)).line(animation_frame="case"), + ) + + _lo, hi = combined.layout.yaxis.range + assert hi >= 300.0, f"Stacked bars reach 300 but range top is {hi}" + + def test_stacked_bars_negative_values(self) -> None: + """Negative segments stack downward; the range must cover their sum.""" + values = np.full((2, 3, 2), -40.0) + values[:, :, 0] = -10.0 + da = xr.DataArray( + values, + dims=["period", "tech", "case"], + name="losses", + ) + combined = overlay( + xpx(da).bar(x="period", barmode="relative", animation_frame="case"), + xpx(da.isel(tech=0)).line(animation_frame="case"), + ) + + lo, hi = combined.layout.yaxis.range + assert lo <= -120.0, f"Stacked bars reach -120 but range bottom is {lo}" + assert hi >= 0.0, "Bar axis should include the zero baseline" + + def test_existing_explicit_range_respected(self) -> None: + """A user-set range on the base figure must not be overwritten.""" + da = xr.DataArray( + np.array([[1.0, 100.0], [2.0, 200.0]]), + dims=["x", "frame"], + name="val", + ) + fig1 = xpx(da).line(animation_frame="frame") + fig1.update_yaxes(range=[0, 42]) + combined = overlay(fig1, xpx(da).scatter(animation_frame="frame")) + + assert list(combined.layout.yaxis.range) == [0, 42] + + def test_log_axis_not_pinned(self) -> None: + """Log axes use exponent coordinates; a data-value range would corrupt them.""" + da = xr.DataArray( + np.array([[1.0, 100.0], [2.0, 200.0]]), + dims=["x", "frame"], + name="val", + ) + fig1 = xpx(da).line(animation_frame="frame") + fig1.update_yaxes(type="log") + combined = overlay(fig1, xpx(da).scatter(animation_frame="frame")) + + assert combined.layout.yaxis.range is None + + def test_category_conversion_after_overlay_works(self) -> None: + """The reported bug: type='category' after overlay showed an empty plot.""" + da = xr.DataArray( + np.random.rand(4, 3, 2), + dims=["period", "tech", "case"], + coords={"period": [2025, 2030, 2035, 2040]}, + name="heat", + ) + combined = overlay( + xpx(da).bar(x="period", barmode="relative", animation_frame="case"), + xpx(da.isel(tech=0)).line(animation_frame="case"), + ).update_xaxes( + type="category", + categoryorder="array", + categoryarray=[2025, 2030, 2035, 2040], + ) + + # No stale numeric range interpreted as category indices + assert combined.layout.xaxis.range is None + class TestSubplotsBasic: """Basic tests for subplots function.""" diff --git a/xarray_plotly/figures.py b/xarray_plotly/figures.py index 0c9f9d7..0c68053 100644 --- a/xarray_plotly/figures.py +++ b/xarray_plotly/figures.py @@ -6,11 +6,14 @@ import copy import re +from collections import defaultdict from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from collections.abc import Iterator + import numpy as np + import numpy.typing as npt import plotly.graph_objects as go from xarray_plotly.common import FacetTitlesMode @@ -54,8 +57,6 @@ def _ensure_legend_visibility( source_figs: The original source figures, in trace order. trace_slices: Slices into ``combined.data`` for each source figure. """ - from collections import defaultdict - # --- Step 1: label unnamed traces from source y-axis titles ----------- labels = [_get_yaxis_title(f) for f in source_figs] @@ -114,88 +115,135 @@ def _ensure_legend_visibility( setattr(frame_trace, attr, src_val) -def _fix_animation_axis_ranges(fig: go.Figure) -> None: - """Set axis ranges to encompass data across all animation frames. +def _numeric_values(vals: Any) -> npt.NDArray[np.float64] | None: + """Convert trace data to a 1-D float array, or None if not numeric. - Plotly.js computes autorange from ``fig.data`` only and does not - recalculate during animation. When different frames have very different - data ranges (e.g. population of Brazil vs China), values can go off-screen. - This function computes the global min/max for each axis across all frames - and sets explicit ranges on the layout. + Datetime/timedelta and categorical (string) data return None — those + axes are left on autorange. + """ + import numpy as np + + if vals is None: + return None + arr = np.atleast_1d(np.asarray(vals)) + if np.issubdtype(arr.dtype, np.datetime64) or np.issubdtype(arr.dtype, np.timedelta64): + return None + try: + return arr.astype(float) + except (ValueError, TypeError): + return None # Non-numeric (categorical) + + +def _collect_axis_extents(traces: Any, stacked: bool) -> dict[tuple[str, str], list[float]]: + """Collect candidate axis values for one trace collection. - Only numeric axes are handled; categorical/date axes are left to autorange. + A "trace collection" is either ``fig.data`` or a single frame's data — + stacking only happens between bars shown at the same time, so each + collection must be aggregated independently. Args: - fig: A Plotly figure with animation frames (mutated in place). + traces: Iterable of traces (from ``fig.data`` or ``frame.data``). + stacked: Whether the layout barmode stacks bars. When True, bar + values on the value axis contribute per-category positive and + negative sums instead of raw segment values. + + Returns: + Mapping of ``(axis_letter, axis_ref)`` (e.g. ``("y", "y2")``) to the + list of candidate values on that axis. """ import numpy as np - if not fig.frames: - return + values: dict[tuple[str, str], list[float]] = defaultdict(list) + # (letter, axis_ref) -> category value -> [positive_sum, negative_sum] + stack_sums: dict[tuple[str, str], dict[Any, list[float]]] = defaultdict( + lambda: defaultdict(lambda: [0.0, 0.0]) + ) - from collections import defaultdict + for trace in traces: + is_bar = getattr(trace, "type", None) == "bar" + value_letter = "x" if (getattr(trace, "orientation", None) or "v") == "h" else "y" - # Collect numeric y-values per axis across all traces (fig.data + frames) - y_by_axis: dict[str, list[float]] = defaultdict(list) - x_by_axis: dict[str, list[float]] = defaultdict(list) + for letter in ("x", "y"): + ref = getattr(trace, f"{letter}axis", None) or letter + arr = _numeric_values(getattr(trace, letter, None)) + if arr is None: + continue + if is_bar and letter == value_letter: + # Bars grow from a zero baseline, so 0 is part of the extent + values[(letter, ref)].append(0.0) + categories = getattr(trace, "y" if letter == "x" else "x", None) + if stacked and categories is not None: + sums = stack_sums[(letter, ref)] + cat_list = np.atleast_1d(np.asarray(categories, dtype=object)).tolist() + for cat, val in zip(cat_list, arr.tolist(), strict=False): + if np.isfinite(val): + sums[cat][0 if val >= 0 else 1] += val + continue + finite = arr[np.isfinite(arr)] + if len(finite): + values[(letter, ref)].extend(finite.tolist()) + + # Stacked sums are the candidate extremes for the value axis + for key, groups in stack_sums.items(): + for pos_sum, neg_sum in groups.values(): + values[key].extend((pos_sum, neg_sum)) + + return values - # Track which axes have bar traces (for zero-baseline clamping) - y_has_vbar: set[str] = set() # vertical bars → y-axis includes 0 - x_has_hbar: set[str] = set() # horizontal bars → x-axis includes 0 - for trace in _iter_all_traces(fig): - yaxis = getattr(trace, "yaxis", None) or "y" - xaxis = getattr(trace, "xaxis", None) or "x" +def _fix_animation_axis_ranges(fig: go.Figure) -> None: + """Pin axis ranges where animation frames exceed the initial view. - # Track bar orientations - if getattr(trace, "type", None) == "bar": - orientation = getattr(trace, "orientation", None) or "v" - if orientation == "h": - x_has_hbar.add(xaxis) - else: - y_has_vbar.add(yaxis) - - for data_attr, axis_ref, by_axis in [ - ("y", yaxis, y_by_axis), - ("x", xaxis, x_by_axis), - ]: - vals = getattr(trace, data_attr, None) - if vals is None: - continue - arr = np.asarray(vals) - # Skip datetime/timedelta — leave those axes on autorange - if np.issubdtype(arr.dtype, np.datetime64) or np.issubdtype(arr.dtype, np.timedelta64): - continue - try: - arr = arr.astype(float) - finite = arr[np.isfinite(arr)] - if len(finite): - by_axis[axis_ref].extend(finite.tolist()) - except (ValueError, TypeError): - pass # Non-numeric (categorical) — skip - - # Apply ranges to layout - for axis_ref, values in y_by_axis.items(): - if not values: + Plotly.js computes autorange from ``fig.data`` only and does not + recalculate during animation. When a later frame has data outside the + initial extent (e.g. population of Brazil vs China), values go + off-screen. This function computes the global min/max for each axis + across all frames and sets an explicit range on the layout — but only + for axes where that is actually necessary. + + Axes whose initial autorange already covers every frame are left on + live autorange, so plotly's native padding is preserved and downstream + changes (like switching the axis to ``type='category'``) keep working. + Axes with a pre-set explicit range, non-linear axes (log, category, + date), and non-numeric data are never touched. + + For ``barmode='stack'`` / ``'relative'``, bar extents are computed from + per-category stacked sums rather than individual segment values, so + tall stacks are not clipped. + + Args: + fig: A Plotly figure with animation frames (mutated in place). + """ + if not fig.frames: + return + + stacked = fig.layout.barmode in ("stack", "relative") + base_extents = _collect_axis_extents(fig.data, stacked) + frame_extents = [_collect_axis_extents(frame.data, stacked) for frame in fig.frames] + + all_keys = set(base_extents) | {key for fe in frame_extents for key in fe} + for key in sorted(all_keys): + _letter, ref = key + axis = fig.layout[_axis_layout_key(ref)] + # Respect explicit ranges; log/date/category range coordinates are + # not plain data values, so a computed range would corrupt the view. + if axis.range is not None or axis.type in ("log", "date", "category", "multicategory"): continue - lo, hi = min(values), max(values) - if axis_ref in y_has_vbar: - lo = min(lo, 0.0) - hi = max(hi, 0.0) - pad = (hi - lo) * 0.05 or 1 # 5% padding - layout_prop = "yaxis" if axis_ref == "y" else f"yaxis{axis_ref[1:]}" - fig.layout[layout_prop].range = [lo - pad, hi + pad] - for axis_ref, values in x_by_axis.items(): - if not values: + base_vals = base_extents.get(key, []) + all_vals = base_vals + [v for fe in frame_extents for v in fe.get(key, [])] + if not all_vals: continue - lo, hi = min(values), max(values) - if axis_ref in x_has_hbar: - lo = min(lo, 0.0) - hi = max(hi, 0.0) - pad = (hi - lo) * 0.05 or 1 - layout_prop = "xaxis" if axis_ref == "x" else f"xaxis{axis_ref[1:]}" - fig.layout[layout_prop].range = [lo - pad, hi + pad] + lo, hi = min(all_vals), max(all_vals) + + # Pin only when some frame exceeds the initial (fig.data) extent — + # otherwise the autorange computed at first render stays valid for + # the whole animation. + if base_vals and min(base_vals) <= lo and max(base_vals) >= hi: + continue + + pad = (hi - lo) * 0.05 or 1 # 5% padding + axis.range = [lo - pad, hi + pad] def _iter_all_traces(fig: go.Figure) -> Iterator[Any]: