⚡️ Speed up function _infer_xy_labels_3d by 23%

[codeflash-ai]

📄 23% (0.23x) speedup for _infer_xy_labels_3d in xarray/plot/utils.py

⏱️ Runtime : 352 microseconds → 285 microseconds (best of 5 runs)

📝 Explanation and details

The optimization achieves a 23% speedup through several targeted micro-optimizations in two critical xarray functions used for plotting and dataset indexing:

Key Optimizations Applied

1. _infer_xy_labels_3d Function (xarray/plot/utils.py)

• Cached darray.dims: Stored frequently accessed darray.dims in a local variable to eliminate repeated attribute lookups
• Streamlined validation logic: Replaced list comprehension + set operations with a more direct validation loop that builds the seen set incrementally
• Optimized color dimension detection: Changed from list comprehension with in checks to a direct loop that avoids redundant membership tests
• Eliminated unnecessary intermediate data structures: Removed not_none list creation and merged validation steps

2. Dataset.isel Method (xarray/core/dataset.py)

• Cached indexers.values(): Stored the values in a local variable to avoid repeated method calls in the is_fancy_indexer check
• Optimized variable indexer construction: Instead of always building a dictionary with dict comprehension, first check if any shared dimensions exist before creating var_indexers
• Used list comprehension for shared dimensions: Replaced set intersection with a more direct list comprehension for finding shared keys

Performance Impact Analysis

Based on the annotated test results, these optimizations are particularly effective for:

• Error cases: 20-65% speedup on validation failures (dimension not found, wrong RGB size)
• Standard inference cases: 10-15% speedup on typical RGB dimension detection
• Large arrays: 13-16% speedup scales well with data size

Hot Path Benefits

The _infer_xy_labels_3d function is called from _infer_xy_labels during 3D plotting operations (imshow), making these optimizations valuable for:

• Interactive plotting workflows where these functions are called repeatedly
• Large-scale data visualization pipelines
• Applications that frequently create RGB/RGBA visualizations from xarray DataArrays

The Dataset.isel optimization benefits any indexing operation on xarray Datasets, which is fundamental to data selection and manipulation workflows throughout the xarray ecosystem.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 18 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import warnings

# imports
import pytest
from xarray.plot.utils import _infer_xy_labels_3d


# Minimal DataArray and Dataset mocks for testing
class DataArray:
    def __init__(self, data, dims, coords=None):
        self.data = data
        self.dims = tuple(dims)
        self.coords = coords or {}
        self.ndim = len(self.dims)
        self.shape = tuple(
            len(data) if isinstance(data, list) else 1 for _ in self.dims
        )
        # Simulate .size for each dimension
        self._dim_sizes = {}
        for i, dim in enumerate(self.dims):
            if isinstance(data, list) and len(data) > 0 and isinstance(data[0], list):
                if i == 0:
                    self._dim_sizes[dim] = len(data)
                elif i == 1:
                    self._dim_sizes[dim] = len(data[0])
                elif i == 2:
                    self._dim_sizes[dim] = len(data[0][0])
            elif isinstance(data, list):
                self._dim_sizes[dim] = len(data)
            else:
                self._dim_sizes[dim] = 1
        # For .isel
        self._isel_called = False

    def __getitem__(self, key):
        # For .size: darray[dim].size
        class DummyDim:
            def __init__(self, size):
                self.size = size

        if key in self._dim_sizes:
            return DummyDim(self._dim_sizes[key])
        raise KeyError(f"Dimension {key} not found")

    def isel(self, indexers):
        # Return a new DataArray with one less dimension
        dims = [dim for dim in self.dims if dim not in indexers]
        # For testing, just return a new DataArray with dims reduced
        # Data shape not important for label inference, so use dummy data
        return DataArray([0], dims)


from xarray.plot.utils import _infer_xy_labels_3d

# ---- TESTS ----

# Basic Test Cases


def test_basic_rgb_inference_single_color_dim():
    # 3D array with dims ('row', 'col', 'band'), band size 3
    arr = DataArray(
        data=[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]],  # shape (2,2,3)
        dims=("row", "col", "band"),
    )
    # Should infer 'band' as rgb, and return ('row', 'col')
    x, y = _infer_xy_labels_3d(
        arr, x=None, y=None, rgb=None
    )  # 36.5μs -> 32.4μs (12.8% faster)


def test_basic_rgb_explicit_x_y_none():
    # Explicit rgb, x/y None
    arr = DataArray(
        data=[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]],
        dims=("lat", "lon", "rgb"),
    )
    x, y = _infer_xy_labels_3d(
        arr, x=None, y=None, rgb="rgb"
    )  # 33.7μs -> 32.5μs (3.66% faster)


def test_error_dim_not_in_array():
    arr = DataArray(
        data=[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]], dims=("x", "y", "rgb")
    )
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="foo", y="y", rgb="rgb"
        )  # 4.07μs -> 3.04μs (33.6% faster)


def test_error_no_color_dim():
    arr = DataArray(
        data=[[[0, 1], [3, 4]], [[6, 7], [9, 10]]],  # shape (2,2,2)
        dims=("x", "y", "z"),
    )
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x=None, y=None, rgb=None
        )  # 24.5μs -> 24.3μs (0.719% faster)


def test_error_rgb_dim_wrong_size():
    arr = DataArray(
        data=[
            [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]],
            [[10, 11, 12, 13, 14], [15, 16, 17, 18, 19]],
        ],
        dims=("x", "y", "foo"),
    )
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="x", y="y", rgb="foo"
        )  # 29.6μs -> 20.0μs (48.0% faster)


def test_multiple_color_dims_warns_and_chooses_last():
    arr = DataArray(
        data=[
            [
                [[0, 1, 2], [3, 4, 5], [6, 7, 8]],
                [[9, 10, 11], [12, 13, 14], [15, 16, 17]],
            ]
        ],
        dims=("a", "b", "c", "d"),
    )
    # Set c and d both size 3 (simulate by overriding .size)
    arr._dim_sizes["c"] = 3
    arr._dim_sizes["d"] = 3
    arr.dims = ("a", "b", "c", "d")
    arr.ndim = 4
    # Reduce to 3D for test (simulate by slicing)
    arr3d = DataArray(
        data=[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]], dims=("a", "c", "d")
    )
    arr3d._dim_sizes["c"] = 3
    arr3d._dim_sizes["d"] = 3
    arr3d.ndim = 3
    # Should warn and pick 'd'
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        x, y = _infer_xy_labels_3d(
            arr3d, x=None, y=None, rgb=None
        )  # 37.7μs -> 33.5μs (12.6% faster)


def test_assert_ndim_is_3():
    arr = DataArray(data=[[0, 1], [2, 3]], dims=("x", "y"))
    arr.ndim = 2
    with pytest.raises(AssertionError):
        _infer_xy_labels_3d(
            arr, x=None, y=None, rgb=None
        )  # 1.73μs -> 1.57μs (9.98% faster)


def test_assert_rgb_not_x_or_y():
    arr = DataArray(
        data=[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]], dims=("rgb", "y", "z")
    )
    with pytest.raises(AssertionError):
        _infer_xy_labels_3d(
            arr, x="rgb", y="y", rgb="rgb"
        )  # 1.47μs -> 1.71μs (14.4% slower)
    with pytest.raises(AssertionError):
        _infer_xy_labels_3d(
            arr, x="z", y="rgb", rgb="rgb"
        )  # 930ns -> 730ns (27.4% faster)


def test_x_equals_y_error():
    arr = DataArray(
        data=[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]], dims=("x", "y", "rgb")
    )
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="x", y="x", rgb="rgb"
        )  # 3.81μs -> 3.21μs (18.7% faster)


# Large Scale Test Cases


def test_large_array_rgb_inference():
    # 3D array, shape (100, 200, 3)
    arr = DataArray(
        data=[[[i + j + k for k in range(3)] for j in range(200)] for i in range(100)],
        dims=("row", "col", "rgb"),
    )
    x, y = _infer_xy_labels_3d(
        arr, x=None, y=None, rgb=None
    )  # 37.1μs -> 32.1μs (15.6% faster)


def test_large_multiple_color_dims_warns_and_chooses_last():
    # 3D array, dims ('x', 'y', 'rgb', 'rgba'), both color dims
    arr = DataArray(
        data=[
            [[[i + j + k + l for l in range(4)] for k in range(3)] for j in range(10)]
            for i in range(20)
        ],
        dims=("x", "y", "rgb", "rgba"),
    )
    arr._dim_sizes["rgb"] = 3
    arr._dim_sizes["rgba"] = 4
    arr.dims = ("x", "y", "rgb", "rgba")
    arr.ndim = 4
    # Reduce to 3D for test
    arr3d = DataArray(
        data=[[[i + j + k for k in range(4)] for j in range(10)] for i in range(20)],
        dims=("x", "y", "rgba"),
    )
    arr3d._dim_sizes["rgba"] = 4
    arr3d.ndim = 3
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        x, y = _infer_xy_labels_3d(
            arr3d, x=None, y=None, rgb=None
        )  # 32.2μs -> 28.4μs (13.3% faster)

import warnings

# imports
import pytest
from xarray.plot.utils import _infer_xy_labels_3d


# Minimal DataArray and Dataset mocks for testing
class DataArray:
    def __init__(self, data, dims, coords=None):
        self.data = data
        self.dims = tuple(dims)
        self.coords = coords or {}
        self.shape = tuple(len(data) if isinstance(data, list) else 1 for _ in dims)
        self.ndim = len(dims)
        # For simplicity, dims are always strings
        self._indexes = {}
        # For size queries
        self._dim_sizes = {}
        for i, dim in enumerate(dims):
            if isinstance(data, list):
                size = (
                    len(data)
                    if i == 0
                    else len(data[0]) if i == 1 else len(data[0][0]) if i == 2 else 1
                )
            else:
                size = 1
            self._dim_sizes[dim] = size

    def __getitem__(self, key):
        # For rgb size checking
        if key in self._dim_sizes:

            class Dummy:
                def __init__(self, size):
                    self.size = size

            return Dummy(self._dim_sizes[key])
        raise KeyError(key)

    def isel(self, indexer):
        # Simulate selecting a slice along one dimension
        dims = [d for d in self.dims if d not in indexer]
        # Remove the selected dimension
        return DataArray(data=[0], dims=dims, coords=self.coords)


class Dataset(DataArray):
    pass


from xarray.plot.utils import _infer_xy_labels_3d

# ------------------- Unit Tests -------------------

# ----------- Basic Test Cases -----------


def test_edge_invalid_ndim():
    # Should raise for ndim != 3
    arr = DataArray(data=[[1, 2], [3, 4]], dims=("y", "x"))
    arr._dim_sizes = {"y": 2, "x": 2}
    arr.ndim = 2
    arr.dims = ("y", "x")
    with pytest.raises(AssertionError):
        _infer_xy_labels_3d(
            arr, x="x", y="y", rgb=None
        )  # 1.84μs -> 1.65μs (12.0% faster)


def test_edge_duplicate_dim_names():
    # x and rgb are the same
    arr = DataArray(
        data=[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dims=("y", "x", "rgb")
    )
    arr._dim_sizes = {"y": 2, "x": 2, "rgb": 3}
    arr.ndim = 3
    arr.dims = ("y", "x", "rgb")
    with pytest.raises(AssertionError):
        _infer_xy_labels_3d(
            arr, x="rgb", y="y", rgb="rgb"
        )  # 1.63μs -> 1.49μs (9.59% faster)


def test_edge_nonexistent_dim():
    # x is not in dims
    arr = DataArray(
        data=[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dims=("y", "rgb", "z")
    )
    arr._dim_sizes = {"y": 2, "rgb": 3, "z": 2}
    arr.ndim = 3
    arr.dims = ("y", "rgb", "z")
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="x", y="y", rgb=None
        )  # 3.75μs -> 2.83μs (32.3% faster)


def test_edge_no_color_dim():
    # No dimension of size 3 or 4
    arr = DataArray(data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dims=("y", "x", "z"))
    arr._dim_sizes = {"y": 2, "x": 2, "z": 2}
    arr.ndim = 3
    arr.dims = ("y", "x", "z")
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="x", y="y", rgb=None
        )  # 24.8μs -> 15.1μs (64.0% faster)


def test_edge_invalid_rgb_size():
    # rgb specified but wrong size
    arr = DataArray(data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dims=("y", "x", "z"))
    arr._dim_sizes = {"y": 2, "x": 2, "z": 2}
    arr.ndim = 3
    arr.dims = ("y", "x", "z")
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="x", y="y", rgb="z"
        )  # 30.2μs -> 20.8μs (44.8% faster)


def test_edge_x_equals_y():
    # x and y are the same
    arr = DataArray(
        data=[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dims=("y", "x", "rgb")
    )
    arr._dim_sizes = {"y": 2, "x": 2, "rgb": 3}
    arr.ndim = 3
    arr.dims = ("y", "x", "rgb")
    with pytest.raises(ValueError):
        _infer_xy_labels_3d(
            arr, x="x", y="x", rgb=None
        )  # 3.84μs -> 3.20μs (20.2% faster)


# ----------- Large Scale Test Cases -----------

Timer unit: 1e-09 s

To edit these changes git checkout codeflash/optimize-_infer_xy_labels_3d-mio81wke and push.