def apply_func(obj_index: tp.Index) -> tp.Index:
return index_fns.drop_duplicate_levels(obj_index, keep=keep)
def broadcast_index(args,
to_shape,
index_from=None,
axis=0,
ignore_single='default',
drop_duplicates='default',
keep='default'):
"""Produce a broadcasted index/columns.
Args:
*args (array_like): Array-like objects.
to_shape (tuple): Target shape.
index_from (None, int, str or array_like): Broadcasting rule for this index/these columns.
Accepts the following values:
* `'default'` - take the value from `vectorbt.defaults.broadcasting`
* `None` - use the original index/columns of the objects in `args`
* `int` - use the index/columns of the i-nth object in `args`
* `'strict'` - ensure that all pandas objects have the same index/columns
* `'stack'` - stack different indexes/columns using `vectorbt.base.index_fns.stack_indexes`
* everything else will be converted to `pd.Index`
axis (int): Set to 0 for index and 1 for columns.
ignore_single (bool): If `True`, won't consider index/columns with a single value.
drop_duplicates (bool): See `vectorbt.base.index_fns.drop_duplicate_levels`.
keep (bool): See `vectorbt.base.index_fns.drop_duplicate_levels`.
For defaults, see `vectorbt.defaults.broadcasting`.
"""
if ignore_single == 'default':
ignore_single = defaults.broadcasting['ignore_single']
if drop_duplicates == 'default':
drop_duplicates = defaults.broadcasting['drop_duplicates']
if keep == 'default':
keep = defaults.broadcasting['keep']
index_str = 'columns' if axis == 1 else 'index'
new_index = None
if axis == 1 and len(to_shape) == 1:
to_shape = (to_shape[0], 1)
maxlen = to_shape[1] if axis == 1 else to_shape[0]
if index_from is not None:
if isinstance(index_from, int):
# Take index/columns of the object indexed by index_from
if not checks.is_pandas(args[index_from]):
raise TypeError(
f"Argument under index {index_from} must be a pandas object"
)
new_index = index_fns.get_index(args[index_from], axis)
elif isinstance(index_from, str):
if index_from in ('stack', 'strict'):
# If pandas objects have different index/columns, stack them together
# maxlen stores the length of the longest index
for arg in args:
if checks.is_pandas(arg):
index = index_fns.get_index(arg, axis)
if pd.Index.equals(
index,
pd.RangeIndex(start=0, stop=len(index),
step=1)):
# ignore simple ranges without name
continue
if new_index is None:
new_index = index
else:
if index_from == 'strict':
# If pandas objects have different index/columns, raise an exception
if not pd.Index.equals(index, new_index):
raise ValueError(
f"Broadcasting {index_str} is not allowed for {index_str}_from=strict"
)
# Broadcasting index must follow the rules of a regular broadcasting operation
# https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html#general-broadcasting-rules
# 1. rule: if indexes are of the same length, they are simply stacked
# 2. rule: if index has one element, it gets repeated and then stacked
if pd.Index.equals(index, new_index):
continue
if len(index) != len(new_index):
if len(index) > 1 and len(new_index) > 1:
raise ValueError(
"Indexes could not be broadcast together"
)
if ignore_single:
# Columns of length 1 should be ignored
if len(index) > len(new_index):
new_index = index
continue
else:
if len(index) > len(new_index):
new_index = index_fns.repeat_index(
new_index, len(index))
elif len(index) < len(new_index):
index = index_fns.repeat_index(
index, len(new_index))
new_index = index_fns.stack_indexes(
new_index, index)
else:
raise ValueError(
f"Invalid value {index_from} for {'columns' if axis == 1 else 'index'}_from"
)
else:
new_index = index_from
if new_index is not None:
if maxlen > len(new_index):
if index_from == 'strict':
raise ValueError(
f"Broadcasting {index_str} is not allowed for {index_str}_from=strict"
)
# This happens only when some numpy object is longer than the new pandas index
# In this case, new pandas index (one element) should be repeated to match this length.
if maxlen > 1 and len(new_index) > 1:
raise ValueError("Indexes could not be broadcast together")
new_index = index_fns.repeat_index(new_index, maxlen)
if drop_duplicates:
new_index = index_fns.drop_duplicate_levels(new_index,
keep=keep)
return new_index