⚡️ Speed up method Cube.coord_dims by 55%

[codeflash-ai]

📄 55% (0.55x) speedup for Cube.coord_dims in lib/iris/cube.py

⏱️ Runtime : 315 microseconds → 203 microseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 55% speedup by replacing expensive dictionary construction with direct iteration in the coord_dims method.

Key optimizations:

1. Eliminated dictionary creation overhead: The original code built dims_by_id dictionaries on every call using comprehensions ({id(c): (d,) for c, d in self._dim_coords_and_dims}), which creates temporary objects and performs hash operations. The optimized version uses simple for loops that check object identity directly and return immediately upon finding a match.

2. Reduced memory allocations: Instead of creating dictionaries that may never be used (when the coordinate is found quickly), the optimized code performs direct iteration with early termination, avoiding unnecessary memory allocation.

3. Streamlined control flow in coord: The condition checking was simplified from nested if-elif to a single if n_coords != 1 with early branching, reducing the number of comparisons in the common case.

Performance characteristics: These optimizations are particularly effective for test cases with small numbers of coordinates (like the annotated tests showing 57-66% improvements), where the coordinate being searched for is likely to be found early in the iteration. The direct iteration approach scales better than dictionary construction, especially when coordinates are found quickly, which is the common case in typical Iris cube operations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 280 Passed
🌀 Generated Regression Tests	✅ 8 Passed
⏪ Replay Tests	✅ 255 Passed
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	84.0%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_cdm.py::TestBasicCubeConstruction.test_immutable_dimcoord_dims	3.66μs	2.32μs	57.4%✅
unit/mesh/components/test_MeshCoord.py::Test_cube_containment.test_cube_dims	2.72μs	2.13μs	27.6%✅

🌀 Generated Regression Tests and Runtime

import pytest
from iris.cube import Cube


# Minimal stub classes to simulate Iris coordinate objects for testing.
class DimCoord:
    def __init__(self, points, standard_name=None, long_name=None):
        self.points = points
        self.standard_name = standard_name
        self.long_name = long_name
        self.metadata = (standard_name, long_name)
        self.shape = (len(points),)

    def name(self):
        return self.standard_name or self.long_name or "unknown"

class AuxCoord:
    def __init__(self, points, standard_name=None, long_name=None):
        self.points = points
        self.standard_name = standard_name
        self.long_name = long_name
        self.metadata = (standard_name, long_name)
        self.shape = (len(points),)

    def name(self):
        return self.standard_name or self.long_name or "unknown"

class AuxCoordFactory:
    def __init__(self, metadata, dims):
        self.metadata = metadata
        self._dims = dims

    def derived_dims(self, coord_dims_func):
        # For testing, just return the dims specified
        return self._dims

# Exception for not found coordinates.
class CoordinateNotFoundError(Exception):
    pass

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

# 1. Basic Test Cases









def test_edge_dimcoord_multi_dim():
    # DimCoord mapped to a tuple (should be single int)
    lat = DimCoord([1, 2, 3], standard_name="latitude")
    cube = Cube(shape=(3,), dim_coords_and_dims=[(lat, (0,))])
    codeflash_output = cube.coord_dims(lat) # 2.23μs -> 1.41μs (57.4% faster)






def test_edge_coord_dims_return_type():
    # Ensure coord_dims always returns a tuple
    lat = DimCoord([1, 2, 3], standard_name="latitude")
    cube = Cube(shape=(3,), dim_coords_and_dims=[(lat, 0)])
    codeflash_output = cube.coord_dims(lat); result = codeflash_output # 2.25μs -> 1.35μs (65.7% faster)

# 3. Large Scale Test Cases






#------------------------------------------------
import pytest
from iris.cube import Cube


# Minimal mock classes to simulate Iris coordinate and cube behavior for testing coord_dims
class DimCoord:
    def __init__(self, name, points=None, long_name=None, standard_name=None):
        self._name = name
        self.points = points if points is not None else []
        self.long_name = long_name
        self.standard_name = standard_name
        self.metadata = (self.standard_name, self.long_name)
    def name(self):
        return self._name

class AuxCoord:
    def __init__(self, name, points=None, long_name=None, standard_name=None):
        self._name = name
        self.points = points if points is not None else []
        self.long_name = long_name
        self.standard_name = standard_name
        self.metadata = (self.standard_name, self.long_name)
    def name(self):
        return self._name

class DummyAuxFactory:
    def __init__(self, name, dims, metadata=None):
        self._name = name
        self._dims = dims
        self.metadata = metadata if metadata is not None else name
    def name(self):
        return self._name
    def derived_dims(self, coord_dims_func):
        return self._dims

class CoordinateNotFoundError(Exception):
    pass

# ------------------- UNIT TESTS START HERE -------------------

# 1. Basic Test Cases

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_libiristestsintegrationtest_netcdf__loadsaveattrs_py_libiristestsunitlazy_datatest_non_lazy_p__replay_test_0.py::test_iris_cube_Cube_coord_dims	304μs	196μs	55.2%✅

To edit these changes git checkout codeflash/optimize-Cube.coord_dims-mh52xrqq and push.