⚡️ Speed up method Cube._add_unique_aux_coord by 9%

[codeflash-ai]

📄 9% (0.09x) speedup for Cube._add_unique_aux_coord in lib/iris/cube.py

⏱️ Runtime : 3.38 milliseconds → 3.09 milliseconds (best of 48 runs)

📝 Explanation and details

The optimized code improves performance by reducing type checks and redundant operations in the _check_multi_dim_metadata method, which is called frequently during cube coordinate initialization.

Key optimizations:

1. Direct type checking optimization: Instead of using isinstance(data_dims, Iterable) which requires a method call and attribute lookup, the code now uses type(data_dims) and direct type comparisons (is int, is tuple, is list). This avoids expensive isinstance checks and method resolution.

2. Early specialization for common types: The code now handles the most common cases (int, tuple, list) with fast paths before falling back to the generic Iterable check, reducing unnecessary type checking overhead.

3. Reduced attribute lookups: The optimization caches self.shape and metadata.shape in local variables (shape, coord_shape) within the validation loop, avoiding repeated attribute access during dimension checking.

4. Streamlined mesh coordinate handling: In _add_unique_aux_coord, the code replaces the inner function definition and TypeGuard check with a direct hasattr(coord, "mesh") check, eliminating function call overhead.

5. Variable renaming consistency: Uses data_dims_tuple instead of reusing data_dims to avoid variable reassignment overhead.

Performance impact: The optimizations show 18-27% speedup for most test cases, with particularly strong improvements for:

• Adding multiple coordinates (21-27% faster)
• Large-scale operations with many coordinates (22-25% faster)
• List/tuple dimension handling (20-24% faster)

These optimizations are most effective for workloads involving frequent coordinate additions during cube construction, which is a common pattern in scientific data processing with Iris.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from iris.cube import Cube

# --- Minimal stubs for required Iris classes and exceptions ---

class CannotAddError(Exception):
    pass

class DummyCoord:
    """
    Minimal stand-in for iris.coords.AuxCoord or DimCoord.
    """
    def __init__(self, points, shape=None, ndim=None, standard_name=None, long_name=None, mesh=None, location=None):
        self.points = points
        self.shape = shape if shape is not None else (len(points),)
        self.ndim = ndim if ndim is not None else len(self.shape)
        self.standard_name = standard_name
        self.long_name = long_name
        self.mesh = mesh
        self.location = location
        self.metadata = (self.standard_name, self.long_name)
    def name(self):
        return self.standard_name or self.long_name or "unknown"
    def __eq__(self, other):
        return (
            isinstance(other, DummyCoord)
            and self.points == other.points
            and self.shape == other.shape
            and self.ndim == other.ndim
            and self.standard_name == other.standard_name
            and self.long_name == other.long_name
            and self.mesh == other.mesh
            and self.location == other.location
        )
    def __repr__(self):
        return f"DummyCoord({self.points}, shape={self.shape}, std={self.standard_name}, long={self.long_name}, mesh={self.mesh}, location={self.location})"

class DummyMesh:
    def __init__(self, id):
        self.id = id
    def __eq__(self, other):
        return isinstance(other, DummyMesh) and self.id == other.id
    def __repr__(self):
        return f"DummyMesh({self.id})"

def _add_unique_aux_coord(coord, data_dims, cube):
    """
    Add an auxiliary coordinate to the cube, checking dimension compatibility.
    If the coord is a mesh coord, check mesh/location/dims.
    """
    data_dims = _check_multi_dim_metadata(coord, data_dims, cube['shape'])

    def is_mesh_coord(anycoord):
        return hasattr(anycoord, "mesh") and anycoord.mesh is not None

    if is_mesh_coord(coord):
        mesh = cube.get('mesh')
        if mesh:
            msg = (
                "{item} of Meshcoord {coord!r} is "
                "{thisval!r}, which does not match existing "
                "cube {item} of {ownval!r}."
            )
            if coord.mesh != mesh:
                raise CannotAddError(
                    msg.format(
                        item="mesh",
                        coord=coord,
                        thisval=coord.mesh,
                        ownval=mesh,
                    )
                )
            location = cube.get('location')
            if coord.location != location:
                raise CannotAddError(
                    msg.format(
                        item="location",
                        coord=coord,
                        thisval=coord.location,
                        ownval=location,
                    )
                )
            mesh_dims = (cube['mesh_dim'],)
            if data_dims != mesh_dims:
                raise CannotAddError(
                    msg.format(
                        item="mesh dimension",
                        coord=coord,
                        thisval=data_dims,
                        ownval=mesh_dims,
                    )
                )
    cube['aux_coords_and_dims'].append((coord, data_dims))

# --- Pytest fixtures and helpers ---

@pytest.fixture
def basic_cube():
    # Cube with shape (4, 5)
    return {
        'shape': (4, 5),
        'aux_coords_and_dims': [],
        'mesh': None,
        'location': None,
        'mesh_dim': None,
    }

@pytest.fixture
def mesh_cube():
    # Cube with mesh
    return {
        'shape': (10, 20),
        'aux_coords_and_dims': [],
        'mesh': DummyMesh('m1'),
        'location': 'face',
        'mesh_dim': 1,
    }

# --- Unit tests ---

# 1. Basic Test Cases


















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


# Minimal mock classes to simulate Iris coordinate objects and cube
class AuxCoord:
    def __init__(self, points, standard_name=None, long_name=None, shape=None):
        self.points = points
        self.standard_name = standard_name
        self.long_name = long_name
        self.shape = shape if shape is not None else (len(points),)
        self.ndim = len(self.shape)
        self.metadata = (self.standard_name, self.long_name)
    def name(self):
        return self.standard_name or self.long_name or "unknown"
    def __eq__(self, other):
        return (
            isinstance(other, AuxCoord)
            and self.points == other.points
            and self.standard_name == other.standard_name
            and self.long_name == other.long_name
        )

class MeshCoord(AuxCoord):
    def __init__(self, points, mesh, location, standard_name=None, long_name=None, shape=None):
        super().__init__(points, standard_name, long_name, shape)
        self.mesh = mesh
        self.location = location

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

# ----------- BASIC TEST CASES -----------


def test_add_1d_aux_coord():
    # Adding a 1D aux coord mapped to one dimension
    cube = Cube(shape=(5, 2))
    coord = AuxCoord(points=[1, 2, 3, 4, 5], standard_name="bar")
    cube._add_unique_aux_coord(coord, 0) # 3.70μs -> 3.02μs (22.3% faster)

def test_add_2d_aux_coord():
    # Adding a 2D aux coord mapped to two dimensions
    cube = Cube(shape=(2, 3))
    coord = AuxCoord(points=[[1,2,3],[4,5,6]], standard_name="baz", shape=(2,3))
    cube._add_unique_aux_coord(coord, (0,1)) # 4.72μs -> 3.98μs (18.7% faster)

def test_add_multiple_aux_coords():
    # Adding several aux coords with different dims
    cube = Cube(shape=(2, 2))
    coord1 = AuxCoord(points=[7,8], standard_name="a")
    coord2 = AuxCoord(points=[9,10], standard_name="b")
    cube._add_unique_aux_coord(coord1, 0) # 3.20μs -> 2.63μs (21.7% faster)
    cube._add_unique_aux_coord(coord2, 1) # 1.44μs -> 1.14μs (26.8% faster)

# ----------- EDGE TEST CASES -----------

def test_missing_dims_for_multi_valued_coord_raises():
    # Adding a multi-valued coord without dims should fail
    cube = Cube(shape=(4, 4))
    coord = AuxCoord(points=[1,2,3,4], standard_name="fail")
    with pytest.raises(ValueError) as excinfo:
        cube._add_unique_aux_coord(coord, None) # 4.62μs -> 4.52μs (2.06% faster)

def test_incorrect_dim_length_raises():
    # Adding a coord with mismatched length for the dimension
    cube = Cube(shape=(4, 5))
    coord = AuxCoord(points=[1,2,3], standard_name="fail")
    with pytest.raises(ValueError) as excinfo:
        cube._add_unique_aux_coord(coord, 0) # 5.86μs -> 5.28μs (11.0% faster)

def test_incorrect_number_of_dims_raises():
    # Adding a 2D coord but passing only one dimension
    cube = Cube(shape=(3, 3))
    coord = AuxCoord(points=[[1,2,3],[4,5,6],[7,8,9]], standard_name="fail", shape=(3,3))
    with pytest.raises(ValueError) as excinfo:
        cube._add_unique_aux_coord(coord, 0) # 4.00μs -> 3.76μs (6.44% faster)

def test_scalar_coord_with_dims_raises():
    # Adding a scalar coord but passing dims should fail
    cube = Cube(shape=(2, 2))
    coord = AuxCoord(points=[99], standard_name="scalar", shape=(1,))
    with pytest.raises(ValueError) as excinfo:
        cube._add_unique_aux_coord(coord, 0) # 5.31μs -> 5.07μs (4.79% faster)




def test_add_coord_with_non_tuple_dims():
    # Dims as list should work
    cube = Cube(shape=(2,3))
    coord = AuxCoord(points=[[1,2,3],[4,5,6]], standard_name="listdims", shape=(2,3))
    cube._add_unique_aux_coord(coord, [0,1]) # 5.50μs -> 4.43μs (24.2% faster)

def test_add_coord_with_dims_as_int():
    # Dims as int for 1D coord should work
    cube = Cube(shape=(5,2))
    coord = AuxCoord(points=[1,2,3,4,5], standard_name="intdims")
    cube._add_unique_aux_coord(coord, 0) # 3.38μs -> 2.70μs (25.6% faster)

# ----------- LARGE SCALE TEST CASES -----------

def test_add_many_aux_coords():
    # Add up to 1000 aux coords to test scalability
    cube = Cube(shape=(1000,))
    coords = [AuxCoord(points=[i], standard_name=f"coord_{i}", shape=(1,)) for i in range(1000)]
    for i, coord in enumerate(coords):
        cube._add_unique_aux_coord(coord, None) # 433μs -> 353μs (22.6% faster)

def test_add_large_2d_aux_coord():
    # Add a large 2D aux coord
    shape = (100, 10)
    cube = Cube(shape=shape)
    points = [[i*j for j in range(shape[1])] for i in range(shape[0])]
    coord = AuxCoord(points=points, standard_name="large2d", shape=shape)
    cube._add_unique_aux_coord(coord, (0,1)) # 5.53μs -> 4.57μs (21.0% faster)

def test_add_large_1d_aux_coord():
    # Add a large 1D aux coord
    shape = (999,)
    cube = Cube(shape=shape)
    points = list(range(999))
    coord = AuxCoord(points=points, standard_name="large1d", shape=(999,))
    cube._add_unique_aux_coord(coord, 0) # 3.34μs -> 2.67μs (25.2% faster)

def test_add_large_scalar_aux_coord():
    # Add a large number of scalar aux coords
    cube = Cube(shape=(1,))
    for i in range(500):
        coord = AuxCoord(points=[i], standard_name=f"scalar_{i}", shape=(1,))
        cube._add_unique_aux_coord(coord, None) # 218μs -> 178μs (22.3% faster)

def test_add_large_2d_aux_coord_with_list_dims():
    # Add a large 2D aux coord with dims as list
    shape = (50, 20)
    cube = Cube(shape=shape)
    points = [[i+j for j in range(shape[1])] for i in range(shape[0])]
    coord = AuxCoord(points=points, standard_name="large2dlist", shape=shape)
    cube._add_unique_aux_coord(coord, [0,1]) # 5.36μs -> 4.43μs (20.8% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

⏪ 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__add_unique_aux_coord	2.67ms	2.51ms	6.44%✅

To edit these changes git checkout codeflash/optimize-Cube._add_unique_aux_coord-mh51r6bo and push.