⚡️ Speed up function _chunked_aware_interpnd by 27%

[codeflash-ai]

📄 27% (0.27x) speedup for _chunked_aware_interpnd in xarray/core/missing.py

⏱️ Runtime : 4.51 milliseconds → 3.55 milliseconds (best of 9 runs)

📝 Explanation and details

The optimized code achieves a 26% speedup by focusing on the most expensive operations identified in the line profiler. The key optimizations target the _localize function, which accounts for 85% of the total runtime in _chunked_aware_interpnd.

Primary optimization in _localize:

• Batched index lookups: Instead of calling index.get_indexer() twice separately for minval and maxval (41.2% + 27.2% = 68.4% of original runtime), the optimized version combines them into a single call using np.array([minval, maxval]). This reduces expensive pandas Index operations from two calls to one.
• Cached .values access: The new_x.values is cached once rather than accessed twice, eliminating redundant attribute lookups.
• Variable unpacking optimization: Using for dim, pair in indexes_coords.items() with x, new_x = pair is more efficient than the list unpacking in the loop header.

Secondary optimizations in other functions:

• Tuple-based transpositions: Using *tuple(range(...)) instead of list(range(...)) for numpy transpose operations is more efficient as tuples are faster to construct and unpack.
• Reduced list comprehension overhead: Pre-computing dimension name lists (const_dims, var_dims) avoids repeated string formatting in loops.

Impact on workloads:
The optimizations are particularly effective for test cases involving:

• Datetime coordinates (24.9% faster): Benefits from reduced index operations
• Multidimensional interpolation targets (36.7% faster): Gains from optimized dimension handling
• Large-scale data (25.6% faster): Batched operations scale better with data size
• Localized interpolation (30-38% faster): Direct beneficiary of _localize optimizations

The optimizations maintain identical functionality while reducing computational overhead in the critical interpolation preprocessing path, making them especially valuable for applications performing repeated spatial/temporal interpolations.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np

# imports
import pytest  # used for our unit tests
from xarray.core.missing import _chunked_aware_interpnd


# Minimal stubs for required xarray internals
class Variable:
    # Simulate xarray Variable for our test purposes
    def __init__(self, dims, data):
        self.dims = dims
        self.data = np.array(data)
        self.shape = self.data.shape

    def isel(self, **indexers):
        # Only supports slices for this stub
        idx = []
        for dim in self.dims:
            if dim in indexers:
                slc = indexers[dim]
                if isinstance(slc, slice):
                    idx.append(slc)
                else:  # int or array
                    idx.append(slc)
            else:
                idx.append(slice(None))
        # Use tuple for numpy indexing
        return Variable(self.dims, self.data[tuple(idx)])

    @property
    def values(self):
        return self.data

    def transpose(self, axes):
        # axes: list of axis indices
        return Variable(self.dims, self.data.transpose(axes))

    def to_index(self):
        # For 1D coordinate arrays, treat as pandas Index
        return DummyIndex(self.data)

    def __getitem__(self, key):
        return Variable(self.dims, self.data[key])


class DummyIndex:
    # Simulate pandas Index for nearest lookup
    def __init__(self, arr):
        self.arr = np.array(arr)

    def get_indexer(self, values, method=None):
        # Only supports method='nearest'
        idxs = []
        for v in values:
            idxs.append(np.abs(self.arr - v).argmin())
        return np.array(idxs)


def reshape(array, shape):
    return np.reshape(array, shape)


# Interpolator stub (linear)
def linear_interp(x, y, assume_sorted=True, **kwargs):
    def interp(new_x):
        return np.interp(new_x, x, y)

    return interp


# Interpolator stub (nearest)
def nearest_interp(x, y, assume_sorted=True, **kwargs):
    def interp(new_x):
        idxs = np.abs(x[:, None] - new_x).argmin(axis=0)
        return y[idxs]

    return interp


from xarray.core.missing import _chunked_aware_interpnd

# -------------------- UNIT TESTS --------------------

# 1. Basic Test Cases


def test_basic_1d_linear():
    # Interpolate a simple 1D array linearly
    x = np.array([0, 1, 2, 3])
    y = np.array([0, 10, 20, 30])
    new_x = np.array([0.5, 1.5, 2.5])
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 505μs -> 372μs (35.7% faster)
    # Should interpolate linearly
    expected = np.array([5, 15, 25])


def test_basic_1d_no_localize():
    # Interpolate without localization
    x = np.array([0, 1, 2, 3])
    y = np.array([0, 10, 20, 30])
    new_x = np.array([0.5, 1.5, 2.5])
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=False
    )
    result = codeflash_output  # 50.3μs -> 49.0μs (2.50% faster)
    expected = np.array([5, 15, 25])


# 2. Edge Test Cases


def test_edge_empty_input():
    # Interpolate with empty input arrays
    x = np.array([])
    y = np.array([])
    new_x = np.array([0.5, 1.5])
    with pytest.raises(ValueError):
        _chunked_aware_interpnd(
            y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
        )  # 312μs -> 257μs (21.5% faster)


def test_edge_out_of_bounds():
    # Interpolate outside the bounds of x
    x = np.array([0, 1, 2])
    y = np.array([0, 10, 20])
    new_x = np.array([-1, 3])
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 382μs -> 294μs (30.0% faster)
    # np.interp returns edge values for out-of-bounds
    expected = np.array([0, 20])


def test_edge_nan_in_x():
    # Interpolate with NaN in x
    x = np.array([0, np.nan, 2])
    y = np.array([0, 10, 20])
    new_x = np.array([1])
    # Should raise error due to NaN in x
    with pytest.raises(Exception):
        _chunked_aware_interpnd(
            y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
        )  # 282μs -> 266μs (6.03% faster)


def test_edge_nan_in_new_x():
    # Interpolate with NaN in new_x
    x = np.array([0, 1, 2])
    y = np.array([0, 10, 20])
    new_x = np.array([np.nan, 1])
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 486μs -> 352μs (38.2% faster)


def test_edge_datetime64():
    # Interpolate with datetime64 coordinates
    x = np.array(["2020-01-01", "2020-01-02", "2020-01-03"], dtype="datetime64[D]")
    y = np.array([0, 10, 20])
    new_x = np.array(["2020-01-01T12", "2020-01-02T12"], dtype="datetime64[h]")
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 1.21ms -> 972μs (24.9% faster)
    # Should interpolate halfway between days
    expected = np.array([5, 15])


def test_edge_multidim_new_x():
    # Interpolate with multidimensional new_x
    x = np.array([0, 1, 2])
    y = np.array([0, 10, 20])
    new_x = np.array([[0.5, 1.5], [1.0, 2.0]])
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 498μs -> 364μs (36.7% faster)
    expected = np.array([[5, 15], [10, 20]])


# 3. Large Scale Test Cases


def test_large_1d_linear():
    # Interpolate a large 1D array
    x = np.linspace(0, 999, 1000)
    y = np.linspace(0, 9990, 1000)
    new_x = np.linspace(0, 999, 1000)
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 378μs -> 301μs (25.6% faster)


def test_large_multidim_new_x():
    # Interpolate with large multidimensional new_x
    x = np.linspace(0, 999, 1000)
    y = np.linspace(0, 9990, 1000)
    new_x = np.linspace(0, 999, 1000).reshape(10, 100)
    codeflash_output = _chunked_aware_interpnd(
        y, x, new_x, interp_func=linear_interp, interp_kwargs={}, localize=True
    )
    result = codeflash_output  # 396μs -> 319μs (24.0% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_chunked_aware_interpnd-miijyp0r and push.