def broadcast_to(arg1,
arg2,
index_from=1,
columns_from=1,
writeable=False,
copy_kwargs={},
raw=False,
**kwargs):
"""Bring first argument to the shape of second argument.
Closely resembles the other broadcast function."""
if not checks.is_array_like(arg1):
arg1 = np.asarray(arg1)
if not checks.is_array_like(arg2):
arg2 = np.asarray(arg2)
is_2d = arg1.ndim > 1 or arg2.ndim > 1
is_pd = checks.is_pandas(arg1) or checks.is_pandas(arg2)
if is_pd:
if is_2d:
if checks.is_series(arg1):
arg1 = arg1.to_frame()
if checks.is_series(arg2):
arg2 = arg2.to_frame()
new_index = broadcast_index(arg1,
arg2,
index_from=index_from,
axis=0,
is_2d=is_2d,
**kwargs)
new_columns = broadcast_index(arg1,
arg2,
index_from=columns_from,
axis=1,
is_2d=is_2d,
**kwargs)
else:
new_index, new_columns = None, None
if is_broadcasting_needed(arg1, arg2):
arg1_new = np.broadcast_to(arg1, arg2.shape, subok=True)
arg1_new = np.array(arg1_new, copy=writeable, **copy_kwargs)
else:
arg1_new = np.array(arg1, copy=False, **copy_kwargs)
return wrap_broadcasted(arg1,
arg1_new,
is_pd=is_pd,
new_index=new_index,
new_columns=new_columns)
def is_broadcasting_needed(*args):
"""Broadcasting may be expensive, do we really need it?"""
args = list(args)
shapes = []
for i in range(len(args)):
if not checks.is_array_like(args[i]):
args[i] = np.asarray(args[i])
shapes.append(args[i].shape)
return len(set(shapes)) > 1
def soft_broadcast_to_ndim(arg, ndim):
"""Try to softly bring the argument to the specified number of dimensions (max 2)."""
if not checks.is_array_like(arg):
arg = np.asarray(arg)
if ndim == 1:
if arg.ndim == 2:
if arg.shape[1] == 1:
if checks.is_pandas(arg):
return arg.iloc[:, 0]
return arg[:, 0] # downgrade
if ndim == 2:
if arg.ndim == 1:
if checks.is_pandas(arg):
return arg.to_frame()
return arg[:, None] # upgrade
return arg # do nothing
def repeat(arg, n, along_axis=1):
"""Repeat array n times along specified axis."""
if not checks.is_array_like(arg):
arg = np.asarray(arg)
if along_axis == 0:
if checks.is_pandas(arg):
return arg.vbt.wrap_array(np.repeat(arg.values, n, axis=0),
index=index_fns.repeat(arg.index, n))
return np.repeat(arg, n, axis=0)
elif along_axis == 1:
arg = to_2d(arg)
if checks.is_pandas(arg):
return arg.vbt.wrap_array(np.repeat(arg.values, n, axis=1),
columns=index_fns.repeat(arg.columns, n))
return np.repeat(arg, n, axis=1)
else:
raise ValueError("Only axis 0 and 1 are supported")
def to_1d(arg, raw=False):
"""Reshape argument to one dimension."""
if raw:
arg = np.asarray(arg)
if not checks.is_array_like(arg):
arg = np.asarray(arg)
if arg.ndim == 2:
if arg.shape[1] == 1:
if checks.is_frame(arg):
return arg.iloc[:, 0]
return arg[:, 0]
if arg.ndim == 1:
return arg
elif arg.ndim == 0:
return arg.reshape((1, ))
raise ValueError(
f"Cannot reshape a {arg.ndim}-dimensional array to 1 dimension")
def to_2d(arg, raw=False, expand_axis=1):
"""Reshape argument to two dimensions."""
if raw:
arg = np.asarray(arg)
if not checks.is_array_like(arg):
arg = np.asarray(arg)
if arg.ndim == 2:
return arg
elif arg.ndim == 1:
if checks.is_series(arg):
if expand_axis == 0:
return pd.DataFrame(arg.values[None, :], columns=arg.index)
elif expand_axis == 1:
return arg.to_frame()
return np.expand_dims(arg, expand_axis)
elif arg.ndim == 0:
return arg.reshape((1, 1))
raise ValueError(
f"Cannot reshape a {arg.ndim}-dimensional array to 2 dimensions")
def tile(arg, n, along_axis=1):
"""Tile array n times along specified axis."""
if not checks.is_array_like(arg):
arg = np.asarray(arg)
if along_axis == 0:
if arg.ndim == 1:
if checks.is_pandas(arg):
return arg.vbt.wrap_array(np.tile(arg.values, n),
index=index_fns.tile(arg.index, n))
return np.tile(arg, n)
if arg.ndim == 2:
if checks.is_pandas(arg):
return arg.vbt.wrap_array(np.tile(arg.values, (n, 1)),
index=index_fns.tile(arg.index, n))
return np.tile(arg, (n, 1))
elif along_axis == 1:
arg = to_2d(arg)
if checks.is_pandas(arg):
return arg.vbt.wrap_array(np.tile(arg.values, (1, n)),
columns=index_fns.tile(arg.columns, n))
return np.tile(arg, (1, n))
else:
raise ValueError("Only axis 0 and 1 are supported")
def broadcast(*args,
index_from='default',
columns_from='default',
writeable=False,
copy_kwargs={},
**kwargs):
"""Bring multiple arguments to the same shape."""
is_pd = False
is_2d = False
args = list(args)
# Convert to np.ndarray object if not numpy or pandas
for i in range(len(args)):
if not checks.is_array_like(args[i]):
args[i] = np.asarray(args[i])
if args[i].ndim > 1:
is_2d = True
if checks.is_pandas(args[i]):
is_pd = True
if is_pd:
# Convert all pd.Series objects to pd.DataFrame
if is_2d:
for i in range(len(args)):
if checks.is_series(args[i]):
args[i] = args[i].to_frame()
# Decide on index and columns
if index_from == 'default':
index_from = defaults.broadcast['index_from']
if columns_from == 'default':
columns_from = defaults.broadcast['columns_from']
new_index = broadcast_index(*args,
index_from=index_from,
axis=0,
is_2d=is_2d,
**kwargs)
new_columns = broadcast_index(*args,
index_from=columns_from,
axis=1,
is_2d=is_2d,
**kwargs)
else:
new_index, new_columns = None, None
# Perform broadcasting operation if needed
if is_broadcasting_needed(*args):
new_args = np.broadcast_arrays(*args, subok=True)
# The problem is that broadcasting creates readonly objects and numba requires writable ones.
# So we have to copy all of them, which is ok for small-sized arrays and not ok for large ones.
# copy kwarg is only applied when broadcasting was done to avoid deprecation warnings
# NOTE: If copy=False, then the resulting arrays will be readonly in the future!
new_args = list(
map(lambda x: np.array(x, copy=writeable, **copy_kwargs),
new_args))
else:
# No copy here, just pandas -> numpy and any order to contiguous
new_args = list(
map(lambda x: np.array(x, copy=False, **copy_kwargs), args))
# Bring arrays to their old types (e.g. array -> pandas)
for i in range(len(new_args)):
new_args[i] = wrap_broadcasted(args[i],
new_args[i],
is_pd=is_pd,
new_index=new_index,
new_columns=new_columns)
return tuple(new_args)