def check_and_set_idx(ids, idx, prefix):
""" Reconciles passed-in IDs and indices and returns indices, as well as unique IDs
in the order specified by the indices. If only IDs supplied, returns the sort-arg
as the index. If only indices supplied, returns None for IDs. If both supplied,
checks that the correspondence is unique and returns unique IDs in the sort order of
the associated index.
:param np.ndarray ids: array of IDs
:param np.ndarray[int] idx: array of indices
:param str prefix: variable name (for error logging)
:return: unique IDs and indices (passed in or derived from the IDs)
:rtype: np.ndarray, np.ndarray
"""
if ids is None and idx is None:
raise ValueError('Both {}_ids and {}_idx cannot be None'.format(prefix, prefix))
if ids is None:
return None, np.asarray_chkfinite(idx)
if idx is None:
return np.unique(ids, return_inverse=True)
else:
ids = np.asarray(ids)
idx = np.asarray_chkfinite(idx)
if len(idx) != len(ids):
raise ValueError('{}_ids ({}) and {}_idx ({}) must have the same length'.format(
prefix, len(ids), prefix, len(idx)))
uniq_idx, idx_sort_index = np.unique(idx, return_index=True)
# make sure each unique index corresponds to a unique id
if not all(len(set(ids[idx == i])) == 1 for i in uniq_idx):
raise ValueError("Each index must correspond to a unique {}_id".format(prefix))
return ids[idx_sort_index], idx
评论列表
文章目录