Function create_transform_da

properties/test_coordinate_transform.py:42–59

Create a DataArray with IdentityTransform CoordinateTransformIndex.

(sizes: dict[str, int])

def create_transform_da(sizes: dict[str, int]) -> xr.DataArray:
    """Create a DataArray with IdentityTransform CoordinateTransformIndex."""
    dims = list(sizes.keys())
    shape = tuple(sizes.values())
    data = np.arange(np.prod(shape)).reshape(shape)

    # Create dataset with transform index for each dimension
    ds = xr.Dataset({DATA_VAR_NAME: (dims, data)})
    indexes = [
        xr.Coordinates.from_xindex(
            CoordinateTransformIndex(
                IdentityTransform((dim,), {dim: size}, dtype=np.dtype(np.int64))
            )
        )
        for dim, size in sizes.items()
    ]
    coords = functools.reduce(operator.or_, indexes)
    return ds.assign_coords(coords).get(DATA_VAR_NAME)


def create_pandas_da(sizes: dict[str, int]) -> xr.DataArray: