def apply(self, *args, apply_func=None, to_2d=False, **kwargs):
"""Apply a function `apply_func`.
Arguments `*args` and `**kwargs` will be directly passed to `apply_func`.
If `to_2d` is True, 2-dimensional NumPy arrays will be passed, otherwise as is.
!!! note
The resulted array must have the same shape as the original array.
## Example
```python-repl
>>> import vectorbt as vbt
>>> import pandas as pd
>>> sr = pd.Series([1, 2], index=['x', 'y'])
>>> sr2.vbt.apply(apply_func=lambda x: x ** 2)
i2
x2 1
y2 4
z2 9
Name: a2, dtype: int64
```
"""
checks.assert_not_none(apply_func)
# Optionally cast to 2d array
if to_2d:
obj = reshape_fns.to_2d(self._obj, raw=True)
else:
obj = np.asarray(self._obj)
result = apply_func(obj, *args, **kwargs)
return self.wrapper.wrap(result, group_by=False)
def test_assert_not_none(self):
checks.assert_not_none(0)
try:
checks.assert_not_none(None)
raise Exception
except:
pass
def map_reduce_partitions(self, map_func_nb=None, map_args=None,
reduce_func_nb=None, reduce_args=None):
"""See `vectorbt.signals.nb.map_reduce_partitions_nb`.
## Example
Get average length of each partition in `sig`:
```python-repl
>>> distance_map_nb = njit(lambda from_i, to_i, col: to_i - from_i)
>>> mean_reduce_nb = njit(lambda col, a: np.nanmean(a))
>>> sig.vbt.signals.map_reduce_partitions(
... map_func_nb=distance_map_nb,
... reduce_func_nb=mean_reduce_nb)
a 1.0
b 1.0
c 3.0
dtype: float64
```
"""
checks.assert_not_none(map_func_nb)
checks.assert_not_none(reduce_func_nb)
checks.assert_numba_func(map_func_nb)
checks.assert_numba_func(reduce_func_nb)
if map_args is None:
map_args = ()
if reduce_args is None:
reduce_args = ()
result = nb.map_reduce_partitions_nb(
self.to_2d_array(),
map_func_nb, map_args,
reduce_func_nb, reduce_args
)
return self.wrapper.wrap_reduced(result)
def map_reduce_between(self,
other=None,
map_func_nb=None,
map_args=None,
reduce_func_nb=None,
reduce_args=None,
broadcast_kwargs=None):
"""See `vectorbt.signals.nb.map_reduce_between_nb`.
If `other` specified, see `vectorbt.signals.nb.map_reduce_between_two_nb`.
Both will be broadcast using `vectorbt.base.reshape_fns.broadcast`
with `broadcast_kwargs`.
Note that `map_args` and `reduce_args` won't be broadcast.
## Example
Get average distance between signals in `sig`:
```python-repl
>>> distance_map_nb = njit(lambda from_i, to_i, col: to_i - from_i)
>>> mean_reduce_nb = njit(lambda col, a: np.nanmean(a))
>>> sig.vbt.signals.map_reduce_between(
... map_func_nb=distance_map_nb,
... reduce_func_nb=mean_reduce_nb)
a NaN
b 2.0
c 1.0
dtype: float64
```
"""
if broadcast_kwargs is None:
broadcast_kwargs = {}
checks.assert_not_none(map_func_nb)
checks.assert_not_none(reduce_func_nb)
checks.assert_numba_func(map_func_nb)
checks.assert_numba_func(reduce_func_nb)
if map_args is None:
map_args = ()
if reduce_args is None:
reduce_args = ()
if other is None:
# One input array
result = nb.map_reduce_between_nb(self.to_2d_array(), map_func_nb,
map_args, reduce_func_nb,
reduce_args)
if isinstance(self._obj, pd.Series):
return result[0]
return pd.Series(result, index=self.wrapper.columns)
else:
# Two input arrays
obj, other = reshape_fns.broadcast(self._obj, other,
**broadcast_kwargs)
checks.assert_dtype(other, np.bool)
result = nb.map_reduce_between_two_nb(obj.vbt.to_2d_array(),
other.vbt.to_2d_array(),
map_func_nb, map_args,
reduce_func_nb, reduce_args)
return obj.vbt.wrapper.wrap_reduced(result)
def apply_and_concat(self,
ntimes,
*args,
apply_func=None,
as_columns=None,
**kwargs):
"""Apply a function n times and concatenate results into a single dataframe."""
checks.assert_not_none(apply_func)
if checks.is_numba_func(apply_func):
# NOTE: your apply_func must a numba-compiled function and arguments must be numba-compatible
# Also NOTE: outputs of apply_func must always be 2-dimensional
result = combine_fns.apply_and_concat_nb(np.asarray(self._obj),
ntimes, apply_func, *args,
**kwargs)
else:
result = combine_fns.apply_and_concat(np.asarray(self._obj),
ntimes, apply_func, *args,
**kwargs)
# Build column hierarchy
if as_columns is not None:
new_columns = index_fns.combine(
as_columns,
reshape_fns.to_2d(self._obj).columns)
else:
new_columns = index_fns.tile(
reshape_fns.to_2d(self._obj).columns, ntimes)
return self.wrap_array(result, columns=new_columns)
def __init__(self,
wrapper: ArrayWrapper,
records_arr: tp.RecordArray,
col_mapper: tp.Optional[ColumnMapper] = None,
**kwargs) -> None:
Wrapping.__init__(
self,
wrapper,
records_arr=records_arr,
col_mapper=col_mapper,
**kwargs
)
StatsBuilderMixin.__init__(self)
# Check fields
records_arr = np.asarray(records_arr)
checks.assert_not_none(records_arr.dtype.fields)
field_names = {
dct.get('name', field_name)
for field_name, dct in self.field_config.get('settings', {}).items()
}
dtype = self.field_config.get('dtype', None)
if dtype is not None:
for field in dtype.names:
if field not in records_arr.dtype.names:
if field not in field_names:
raise TypeError(f"Field '{field}' from {dtype} cannot be found in records or config")
self._records_arr = records_arr
if col_mapper is None:
col_mapper = ColumnMapper(wrapper, self.col_arr)
self._col_mapper = col_mapper
def wrap_reduced(self,
a,
index=None,
columns=None,
time_units=False,
dtype=None,
group_by=None):
"""Wrap result of reduction.
`index` can be set when reducing to an array of values (vs. one value) per column.
`columns` can be set to override object's default columns.
If `time_units` is set, calls `to_time_units`."""
checks.assert_not_none(self.ndim)
_self = self.resolve(group_by=group_by)
if columns is None:
columns = _self.columns
a = np.asarray(a)
if dtype is not None:
try:
a = a.astype(dtype)
except Exception as e:
warnings.warn(repr(e), stacklevel=2)
if time_units:
a = _self.to_time_units(a)
if a.ndim == 0:
# Scalar per Series/DataFrame
if time_units:
return pd.to_timedelta(a.item())
return a.item()
if a.ndim == 1:
if _self.ndim == 1:
if a.shape[0] == 1:
# Scalar per Series/DataFrame with one column
if time_units:
return pd.to_timedelta(a[0])
return a[0]
# Array per Series
name = columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a, index=index, name=name, dtype=dtype)
# Scalar per column in a DataFrame
if index is None:
index = columns
return pd.Series(a, index=index, dtype=dtype)
if a.ndim == 2:
if a.shape[1] == 1 and _self.ndim == 1:
# Array per Series
name = columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a[:, 0], index=index, name=name, dtype=dtype)
# Array per column in a DataFrame
return pd.DataFrame(a, index=index, columns=columns, dtype=dtype)
raise ValueError(f"{a.ndim}-d input is not supported")
def apply_and_concat(self,
ntimes,
*args,
apply_func=None,
to_2d=False,
keys=None,
wrap_kwargs=None,
**kwargs):
"""Apply `apply_func` `ntimes` times and concatenate the results along columns.
See `vectorbt.base.combine_fns.apply_and_concat_one`.
Arguments `*args` and `**kwargs` will be directly passed to `apply_func`.
If `to_2d` is True, 2-dimensional NumPy arrays will be passed, otherwise as is.
Use `keys` as the outermost level.
!!! note
The resulted arrays to be concatenated must have the same shape as broadcast input arrays.
## Example
```python-repl
>>> import vectorbt as vbt
>>> import pandas as pd
>>> df = pd.DataFrame([[3, 4], [5, 6]], index=['x', 'y'], columns=['a', 'b'])
>>> df.vbt.apply_and_concat(3, [1, 2, 3],
... apply_func=lambda i, a, b: a * b[i], keys=['c', 'd', 'e'])
c d e
a b a b a b
x 3 4 6 8 9 12
y 5 6 10 12 15 18
```
"""
checks.assert_not_none(apply_func)
# Optionally cast to 2d array
if to_2d:
obj_arr = reshape_fns.to_2d(self._obj, raw=True)
else:
obj_arr = np.asarray(self._obj)
if checks.is_numba_func(apply_func):
result = combine_fns.apply_and_concat_one_nb(
ntimes, apply_func, obj_arr, *args, **kwargs)
else:
result = combine_fns.apply_and_concat_one(ntimes, apply_func,
obj_arr, *args, **kwargs)
# Build column hierarchy
if keys is not None:
new_columns = index_fns.combine_indexes(keys, self.wrapper.columns)
else:
top_columns = pd.Index(np.arange(ntimes), name='apply_idx')
new_columns = index_fns.combine_indexes(top_columns,
self.wrapper.columns)
return self.wrapper.wrap(result,
group_by=False,
**merge_dicts(dict(columns=new_columns),
wrap_kwargs))
def information_ratio(self):
"""Information ratio of a strategy.
!!! note
`factor_returns` must be set."""
checks.assert_not_none(self.factor_returns)
return self.wrapper.wrap_reduced(
nb.information_ratio_nb(self.returns.vbt.to_2d_array(),
self.factor_returns.vbt.to_2d_array()))
def down_capture(self):
"""Capture ratio for periods when the benchmark return is negative.
!!! note
`factor_returns` must be set."""
checks.assert_not_none(self.factor_returns)
return self.wrapper.wrap_reduced(
nb.down_capture_nb(self.returns.vbt.to_2d_array(),
self.factor_returns.vbt.to_2d_array(),
self.ann_factor))
def capture(self):
"""Capture ratio.
!!! note
`factor_returns` must be set."""
checks.assert_not_none(self.factor_returns)
return self.wrapper.wrap_reduced(
nb.capture_nb(self.returns.vbt.to_2d_array(),
self.factor_returns.vbt.to_2d_array(),
self.ann_factor))
def beta(self):
"""Beta.
!!! note
`factor_returns` must be set."""
checks.assert_not_none(self.factor_returns)
return self.wrapper.wrap_reduced(
nb.beta_nb(self.returns.vbt.to_2d_array(),
self.factor_returns.vbt.to_2d_array(),
risk_free=self.risk_free))
def wrap_reduced(self,
a,
index=None,
columns=None,
time_units=False,
collapse=None,
**kwargs):
"""Wrap result of reduction.
`index` can be set when reducing to an array of values (vs. one value) per column.
`columns` can be set to override object's default columns.
If `time_units` is set, calls `to_time_units`."""
checks.assert_not_none(self.ndim)
group_by = self.grouper.resolve_group_by(**kwargs)
if columns is None:
columns = self.grouper.get_columns(**kwargs)
if collapse is None:
collapse = group_by is not None and group_by is not False and self.grouped_ndim == 1
a = np.asarray(a)
if time_units:
a = self.to_time_units(a)
if a.ndim == 0:
# Scalar per Series/DataFrame
if time_units:
return pd.to_timedelta(a.item())
return a.item()
if a.ndim == 1:
if self.ndim == 1 or (self.ndim == 2 and len(columns) == 1
and collapse):
if a.shape[0] == 1:
# Scalar per Series/DataFrame with one column
if time_units:
return pd.to_timedelta(a[0])
return a[0]
# Array per Series
name = columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a, index=index, name=name)
# Scalar per column in a DataFrame
if index is None:
index = columns
return pd.Series(a, index=index)
if self.ndim == 1:
# Array per Series
name = columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a[:, 0], index=index, name=name)
# Array per column in a DataFrame
return pd.DataFrame(a, index=index, columns=columns)
def combine_with(self,
other,
*args,
combine_func=None,
to_2d=False,
broadcast_kwargs=None,
wrap_kwargs=None,
**kwargs):
"""Combine both using `combine_func` into a Series/DataFrame of the same shape.
All arguments will be broadcast using `vectorbt.base.reshape_fns.broadcast`
with `broadcast_kwargs`.
Arguments `*args` and `**kwargs` will be directly passed to `combine_func`.
If `to_2d` is True, 2-dimensional NumPy arrays will be passed, otherwise as is.
!!! note
The resulted array must have the same shape as broadcast input arrays.
## Example
```python-repl
>>> import vectorbt as vbt
>>> import pandas as pd
>>> sr = pd.Series([1, 2], index=['x', 'y'])
>>> df = pd.DataFrame([[3, 4], [5, 6]], index=['x', 'y'], columns=['a', 'b'])
>>> sr.vbt.combine_with(df, combine_func=lambda x, y: x + y)
a b
x 4 5
y 7 8
```
"""
if isinstance(other, BaseAccessor):
other = other._obj
checks.assert_not_none(combine_func)
if broadcast_kwargs is None:
broadcast_kwargs = {}
if checks.is_numba_func(combine_func):
# Numba requires writable arrays
broadcast_kwargs = merge_dicts(
dict(require_kwargs=dict(requirements='W')), broadcast_kwargs)
new_obj, new_other = reshape_fns.broadcast(self._obj, other,
**broadcast_kwargs)
# Optionally cast to 2d array
if to_2d:
new_obj_arr = reshape_fns.to_2d(new_obj, raw=True)
new_other_arr = reshape_fns.to_2d(new_other, raw=True)
else:
new_obj_arr = np.asarray(new_obj)
new_other_arr = np.asarray(new_other)
result = combine_func(new_obj_arr, new_other_arr, *args, **kwargs)
return new_obj.vbt.wrapper.wrap(result, **merge_dicts({}, wrap_kwargs))
def alpha(self):
"""Annualized alpha.
!!! note
`factor_returns` must be set."""
checks.assert_not_none(self.factor_returns)
return self.wrapper.wrap_reduced(
nb.alpha_nb(self.returns.vbt.to_2d_array(),
self.factor_returns.vbt.to_2d_array(),
reshape_fns.to_1d(self.beta, raw=True),
self.ann_factor,
risk_free=self.risk_free))
def unstack_to_df(arg, index_levels=None, column_levels=None, symmetric=False):
"""Reshape `arg` based on its multi-index into a DataFrame.
Use `index_levels` to specify what index levels will form new index, and `column_levels`
for new columns. Set `symmetric` to `True` to make DataFrame symmetric.
Example:
```python-repl
>>> import pandas as pd
>>> from vectorbt.utils.reshape_fns import unstack_to_df
>>> index = pd.MultiIndex.from_arrays(
... [[1, 1, 2, 2], [3, 4, 3, 4], ['a', 'b', 'c', 'd']],
... names=['x', 'y', 'z'])
>>> sr = pd.Series([1, 2, 3, 4], index=index)
>>> print(unstack_to_df(sr, index_levels=(0, 1), column_levels=2))
z a b c d
x y
1 3 1.0 NaN NaN NaN
1 4 NaN 2.0 NaN NaN
2 3 NaN NaN 3.0 NaN
2 4 NaN NaN NaN 4.0
```"""
# Perform checks
checks.assert_type(arg, (pd.Series, pd.DataFrame))
if checks.is_frame(arg):
if arg.shape[0] == 1:
arg = arg.iloc[0, :]
elif arg.shape[1] == 1:
arg = arg.iloc[:, 0]
checks.assert_type(arg.index, pd.MultiIndex)
sr = to_1d(arg)
if len(sr.index.levels) > 2:
checks.assert_not_none(index_levels)
checks.assert_not_none(column_levels)
else:
index_levels = 0
column_levels = 1
# Build new index and column hierarchies
new_index = index_fns.select_levels(arg.index, index_levels).unique()
new_columns = index_fns.select_levels(arg.index, column_levels).unique()
# Unstack and post-process
unstacked = unstack_to_array(sr, levels=(index_levels, column_levels))
df = pd.DataFrame(unstacked, index=new_index, columns=new_columns)
if symmetric:
return make_symmetric(df)
return df
def map_reduce_between(self,
*args,
other=None,
map_func_nb=None,
reduce_func_nb=None,
broadcast_kwargs={}):
"""See `vectorbt.signals.nb.map_reduce_between_nb`.
If `other` specified, see `vectorbt.signals.nb.map_reduce_between_two_nb`.
Arguments will be broadcasted using `vectorbt.utils.reshape_fns.broadcast`
with `broadcast_kwargs`.
Example:
Get maximum distance between signals in `signals`:
```python-repl
>>> distance_map_nb = njit(lambda col, prev_i, next_i: next_i - prev_i)
>>> max_reduce_nb = njit(lambda col, a: np.nanmax(a))
>>> print(signals.vbt.signals.map_reduce_between(
... map_func_nb=distance_map_nb, reduce_func_nb=max_reduce_nb))
a 3.0
b 3.0
c NaN
dtype: float64
```"""
checks.assert_not_none(map_func_nb)
checks.assert_not_none(reduce_func_nb)
checks.assert_numba_func(map_func_nb)
checks.assert_numba_func(reduce_func_nb)
if other is None:
# One input array
result = nb.map_reduce_between_nb(self.to_2d_array(), map_func_nb,
reduce_func_nb, *args)
if isinstance(self._obj, pd.Series):
return result[0]
return pd.Series(result, index=self.columns)
else:
# Two input arrays
obj, other = reshape_fns.broadcast(self._obj, other,
**broadcast_kwargs)
other.vbt.signals.validate()
result = nb.map_reduce_between_two_nb(self.to_2d_array(),
other.vbt.to_2d_array(),
map_func_nb, reduce_func_nb,
*args)
if isinstance(obj, pd.Series):
return result[0]
return pd.Series(result, index=obj.vbt.columns)
def apply(self,
*args,
apply_func: tp.Optional[tp.Callable] = None,
keep_pd: bool = False,
to_2d: bool = False,
wrap_kwargs: tp.KwargsLike = None,
**kwargs) -> tp.SeriesFrame:
"""Apply a function `apply_func`.
Args:
*args: Variable arguments passed to `apply_func`.
apply_func (callable): Apply function.
Can be Numba-compiled.
keep_pd (bool): Whether to keep inputs as pandas objects, otherwise convert to NumPy arrays.
to_2d (bool): Whether to reshape inputs to 2-dim arrays, otherwise keep as-is.
wrap_kwargs (dict): Keyword arguments passed to `vectorbt.base.array_wrapper.ArrayWrapper.wrap`.
**kwargs: Keyword arguments passed to `combine_func`.
!!! note
The resulted array must have the same shape as the original array.
## Example
```python-repl
>>> import vectorbt as vbt
>>> import pandas as pd
>>> sr = pd.Series([1, 2], index=['x', 'y'])
>>> sr2.vbt.apply(apply_func=lambda x: x ** 2)
i2
x2 1
y2 4
z2 9
Name: a2, dtype: int64
```
"""
checks.assert_not_none(apply_func)
# Optionally cast to 2d array
if to_2d:
obj = reshape_fns.to_2d(self.obj, raw=not keep_pd)
else:
if not keep_pd:
obj = np.asarray(self.obj)
else:
obj = self.obj
result = apply_func(obj, *args, **kwargs)
return self.wrapper.wrap(result,
group_by=False,
**merge_dicts({}, wrap_kwargs))
def map_reduce_between(self,
*args,
other=None,
map_func_nb=None,
reduce_func_nb=None,
broadcast_kwargs={}):
"""See `vectorbt.signals.nb.map_reduce_between_nb`.
If `other` specified, see `vectorbt.signals.nb.map_reduce_between_two_nb`.
Arguments will be broadcasted using `vectorbt.base.reshape_fns.broadcast`
with `broadcast_kwargs`.
Example:
Get average distance between signals in `sig`:
```python-repl
>>> distance_map_nb = njit(lambda col, from_i, to_i: to_i - from_i)
>>> mean_reduce_nb = njit(lambda col, a: np.nanmean(a))
>>> print(sig.vbt.signals.map_reduce_between(
... map_func_nb=distance_map_nb, reduce_func_nb=mean_reduce_nb))
a NaN
b 2.0
c 1.0
dtype: float64
```"""
checks.assert_not_none(map_func_nb)
checks.assert_not_none(reduce_func_nb)
checks.assert_numba_func(map_func_nb)
checks.assert_numba_func(reduce_func_nb)
if other is None:
# One input array
result = nb.map_reduce_between_nb(self.to_2d_array(), map_func_nb,
reduce_func_nb, *args)
if isinstance(self._obj, pd.Series):
return result[0]
return pd.Series(result, index=self.columns)
else:
# Two input arrays
obj, other = reshape_fns.broadcast(self._obj, other,
**broadcast_kwargs)
checks.assert_dtype(other, np.bool_)
result = nb.map_reduce_between_two_nb(self.to_2d_array(),
other.vbt.to_2d_array(),
map_func_nb, reduce_func_nb,
*args)
return self.wrap_reduced(result)
def __init__(self, records_arr, wrapper, idx_field=None):
if not isinstance(records_arr, np.ndarray):
records_arr = np.asarray(records_arr)
checks.assert_not_none(records_arr.dtype.fields)
checks.assert_value_in('col', records_arr.dtype.names)
checks.assert_type(wrapper, ArrayWrapper)
if idx_field is not None:
checks.assert_value_in(idx_field, records_arr.dtype.names)
else:
if 'idx' in records_arr.dtype.names:
idx_field = 'idx'
self.records_arr = records_arr
self.wrapper = wrapper
self.idx_field = idx_field
PandasIndexer.__init__(self, _records_indexing_func)
def apply_and_concat(self,
ntimes,
*args,
apply_func=None,
pass_2d=False,
as_columns=None,
**kwargs):
"""Apply `apply_func` `ntimes` times and concatenate the results along columns.
See `vectorbt.utils.combine_fns.apply_and_concat`.
Arguments `*args` and `**kwargs` will be directly passed to `apply_func`.
If `pass_2d` is `True`, 2-dimensional NumPy arrays will be passed, otherwise as is.
Use `as_columns` as a top-level column level.
Example:
```python-repl
>>> import pandas as pd
>>> df = pd.DataFrame([[3, 4], [5, 6]], index=['x', 'y'], columns=['a', 'b'])
>>> print(df.vbt.apply_and_concat(3, [1, 2, 3],
... apply_func=lambda i, a, b: a * b[i], as_columns=['c', 'd', 'e']))
c d e
a b a b a b
x 3 4 6 8 9 12
y 5 6 10 12 15 18
```"""
checks.assert_not_none(apply_func)
# Optionally cast to 2d array
if pass_2d:
obj_arr = reshape_fns.to_2d(np.asarray(self._obj))
else:
obj_arr = np.asarray(self._obj)
if checks.is_numba_func(apply_func):
result = combine_fns.apply_and_concat_nb(obj_arr, ntimes,
apply_func, *args,
**kwargs)
else:
result = combine_fns.apply_and_concat(obj_arr, ntimes, apply_func,
*args, **kwargs)
# Build column hierarchy
if as_columns is not None:
new_columns = index_fns.combine_indexes(as_columns, self.columns)
else:
new_columns = index_fns.tile_index(self.columns, ntimes)
return self.wrap_array(result, columns=new_columns)
def combine_with(self,
other,
*args,
combine_func=None,
pass_2d=False,
broadcast_kwargs={},
**kwargs):
"""Combine both using `combine_func` into a Series/DataFrame of the same shape.
All arguments will be broadcasted using `vectorbt.utils.reshape_fns.broadcast`
with `broadcast_kwargs`.
Arguments `*args` and `**kwargs` will be directly passed to `combine_func`.
If `pass_2d` is `True`, 2-dimensional NumPy arrays will be passed, otherwise as is.
Example:
```python-repl
>>> import pandas as pd
>>> sr = pd.Series([1, 2], index=['x', 'y'])
>>> df = pd.DataFrame([[3, 4], [5, 6]], index=['x', 'y'], columns=['a', 'b'])
>>> print(sr.vbt.combine_with(df, combine_func=lambda x, y: x + y))
a b
x 4 5
y 7 8
```"""
if isinstance(other, Base_Accessor):
other = other._obj
checks.assert_not_none(combine_func)
if checks.is_numba_func(combine_func):
# Numba requires writable arrays
broadcast_kwargs = {**dict(writeable=True), **broadcast_kwargs}
new_obj, new_other = reshape_fns.broadcast(self._obj, other,
**broadcast_kwargs)
# Optionally cast to 2d array
if pass_2d:
new_obj_arr = reshape_fns.to_2d(np.asarray(new_obj))
new_other_arr = reshape_fns.to_2d(np.asarray(new_other))
else:
new_obj_arr = np.asarray(new_obj)
new_other_arr = np.asarray(new_other)
result = combine_func(new_obj_arr, new_other_arr, *args, **kwargs)
return new_obj.vbt.wrap_array(result)
def wrap_reduced(self, a, index=None, columns=None, time_units=False):
"""Wrap result of reduction.
`index` can be set when reducing to an array of values (vs. one value) per column.
`columns` can be set to override object's default columns.
If `time_units` is set, calls `to_time_units`."""
checks.assert_not_none(self.ndim)
if columns is None:
columns = self.columns
a = np.asarray(a)
if time_units:
a = self.to_time_units(a)
if a.ndim == 0:
# Scalar value
if time_units:
return pd.to_timedelta(a.item())
return a.item()
if a.ndim == 1:
if self.ndim == 1:
if a.shape[0] == 1:
# Scalar value
if time_units:
return pd.to_timedelta(a[0])
return a[0]
# Array per series
name = columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a, index=index, name=name)
# Value per column
if index is None:
index = columns
return pd.Series(a, index=index)
if self.ndim == 1:
# Array per series
name = columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a[:, 0], index=index, name=name)
# Value per column
return pd.DataFrame(a, index=index, columns=columns)
def __init__(self, wrapper, records_arr, idx_field='auto', **kwargs):
Wrapping.__init__(self,
wrapper,
records_arr=records_arr,
idx_field=idx_field,
**kwargs)
records_arr = np.asarray(records_arr)
checks.assert_not_none(records_arr.dtype.fields)
checks.assert_in('id', records_arr.dtype.names)
checks.assert_in('col', records_arr.dtype.names)
if idx_field == 'auto':
if 'idx' in records_arr.dtype.names:
idx_field = 'idx'
elif idx_field is not None:
checks.assert_in(idx_field, records_arr.dtype.names)
self._records_arr = records_arr
self._idx_field = idx_field
self._col_mapper = ColumnMapper(wrapper, records_arr['col'])
def combine_with(self,
other,
*args,
combine_func=None,
broadcast_kwargs={},
**kwargs):
"""Broadcast with other and combine.
The returned shape is the same as broadcasted shape."""
if isinstance(other, Base_Accessor):
other = other._obj
checks.assert_not_none(combine_func)
if checks.is_numba_func(combine_func):
# Numba requires writable arrays
broadcast_kwargs = {**dict(writeable=True), **broadcast_kwargs}
new_obj, new_other = reshape_fns.broadcast(self._obj, other,
**broadcast_kwargs)
return new_obj.vbt.wrap_array(
combine_func(np.asarray(new_obj), np.asarray(new_other), *args,
**kwargs))
def __init__(self, columns, group_by=None, allow_enable=True, allow_disable=True, allow_modify=True):
Configured.__init__(
self,
columns=columns,
group_by=group_by,
allow_enable=allow_enable,
allow_disable=allow_disable,
allow_modify=allow_modify
)
checks.assert_not_none(columns)
self._columns = columns
if group_by is None or isinstance(group_by, bool):
self._group_by = None
else:
self._group_by = group_by_to_index(columns, group_by)
# Everything is allowed by default
self._allow_enable = allow_enable
self._allow_disable = allow_disable
self._allow_modify = allow_modify
def map_reduce_partitions(
self,
map_func_nb: tp.Optional[tp.SignalMapFunc] = None,
map_args: tp.Optional[tp.Args] = None,
reduce_func_nb: tp.Optional[tp.SignalReduceFunc] = None,
reduce_args: tp.Optional[tp.Args] = None,
wrap_kwargs: tp.KwargsLike = None) -> tp.MaybeSeries:
"""See `vectorbt.signals.nb.map_reduce_partitions_nb`.
## Example
Get average length of each partition in `sig`:
```python-repl
>>> distance_map_nb = njit(lambda from_i, to_i, col: to_i - from_i)
>>> mean_reduce_nb = njit(lambda col, a: np.nanmean(a))
>>> sig.vbt.signals.map_reduce_partitions(
... map_func_nb=distance_map_nb,
... reduce_func_nb=mean_reduce_nb)
a 1.0
b 1.0
c 3.0
dtype: float64
```
"""
checks.assert_not_none(map_func_nb)
checks.assert_not_none(reduce_func_nb)
checks.assert_numba_func(map_func_nb)
checks.assert_numba_func(reduce_func_nb)
if map_args is None:
map_args = ()
if reduce_args is None:
reduce_args = ()
result = nb.map_reduce_partitions_nb(self.to_2d_array(), map_func_nb,
map_args, reduce_func_nb,
reduce_args)
wrap_kwargs = merge_dicts(dict(name_or_index='map_reduce_partitions'),
wrap_kwargs)
return self.wrapper.wrap_reduced(result, **wrap_kwargs)
def wrap_reduced(self, a, index=None, name=None, time_units=False):
"""Wrap result of reduction.
Argument `index` can be passed when reducing to an array of values (vs. one value) per column.
If `time_units` is set, calls `to_time_units`."""
checks.assert_not_none(self.ndim)
a = np.asarray(a)
if time_units:
a = self.to_time_units(a)
if a.ndim == 0:
# Scalar value
if time_units:
return pd.to_timedelta(a.item())
return a.item()
elif a.ndim == 1:
if self.ndim == 1:
if a.shape[0] == 1:
# Scalar value
if time_units:
return pd.to_timedelta(a[0])
return a[0]
# Array per series
if name is None:
name = self.columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a, index=index, name=name)
# Value per column
return pd.Series(a, index=self.columns, name=name)
if self.ndim == 1:
# Array per series
if name is None:
name = self.columns[0]
if name == 0: # was a Series before
name = None
return pd.Series(a[:, 0], index=index, name=name)
# Value per column
return pd.DataFrame(a, index=index, columns=self.columns)
def __init__(self, index, columns, ndim, freq=None, column_only_select=None,
group_select=None, grouped_ndim=None, **kwargs):
config = dict(
index=index,
columns=columns,
ndim=ndim,
freq=freq,
column_only_select=column_only_select,
group_select=group_select,
grouped_ndim=grouped_ndim,
)
checks.assert_not_none(index)
checks.assert_not_none(columns)
checks.assert_not_none(ndim)
if not isinstance(index, pd.Index):
index = pd.Index(index)
if not isinstance(columns, pd.Index):
columns = pd.Index(columns)
self._index = index
self._columns = columns
self._ndim = ndim
self._freq = freq
self._column_only_select = column_only_select
self._group_select = group_select
self._grouper = ColumnGrouper(columns, **kwargs)
self._grouped_ndim = grouped_ndim
PandasIndexer.__init__(self)
Configured.__init__(self, **merge_dicts(config, self._grouper._config))
def unstack_to_df(arg, index_levels=None, column_levels=None, symmetric=False):
"""Reshape object based on multi-index into dataframe."""
# Perform checks
checks.assert_type(arg, (pd.Series, pd.DataFrame))
if checks.is_frame(arg):
if arg.shape[0] == 1:
arg = arg.iloc[0, :]
elif arg.shape[1] == 1:
arg = arg.iloc[:, 0]
checks.assert_type(arg.index, pd.MultiIndex)
sr = to_1d(arg)
if len(sr.index.levels) > 2:
checks.assert_not_none(index_levels)
checks.assert_not_none(column_levels)
else:
index_levels = 0
column_levels = 1
# Build new index and column hierarchies
new_index = np.unique(index_fns.select_levels(arg.index, index_levels))
new_columns = np.unique(index_fns.select_levels(arg.index, column_levels))
if isinstance(index_levels, (list, tuple)):
new_index = pd.MultiIndex.from_tuples(new_index, names=index_levels)
else:
new_index = pd.Index(new_index, name=index_levels)
if isinstance(column_levels, (list, tuple)):
new_columns = pd.MultiIndex.from_tuples(new_columns,
names=column_levels)
else:
new_columns = pd.Index(new_columns, name=column_levels)
# Unstack and post-process
unstacked = unstack_to_array(sr, levels=(index_levels, column_levels))
df = pd.DataFrame(unstacked, index=new_index, columns=new_columns)
if symmetric:
return make_symmetric(df)
return df