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)