def wrap_index(idx, dim):
    """
    Helper function that wraps index when greater than maximum dimension.

    Args:
        idx (int): Unwrapped index
        dim (int): Maximum dimension

    Returns:
        idx (int): idx if idx < dim or idx - dim
    """
    if idx < dim:
        return idx
    else:
        return idx - dim