Generic function for cumulative reduction Parameters ---------- func: callable Cumulative function like np.cumsum or np.cumprod binop: callable Associated binary operator like ``np.cumsum->add`` or ``np.cumprod->mul`` ident: Number Associated identity lik
(
func,
binop,
ident,
x,
axis=None,
dtype=None,
out=None,
method="sequential",
preop=None,
)
| 1153 | |
| 1154 | |
| 1155 | def cumreduction( |
| 1156 | func, |
| 1157 | binop, |
| 1158 | ident, |
| 1159 | x, |
| 1160 | axis=None, |
| 1161 | dtype=None, |
| 1162 | out=None, |
| 1163 | method="sequential", |
| 1164 | preop=None, |
| 1165 | ): |
| 1166 | """Generic function for cumulative reduction |
| 1167 | |
| 1168 | Parameters |
| 1169 | ---------- |
| 1170 | func: callable |
| 1171 | Cumulative function like np.cumsum or np.cumprod |
| 1172 | binop: callable |
| 1173 | Associated binary operator like ``np.cumsum->add`` or ``np.cumprod->mul`` |
| 1174 | ident: Number |
| 1175 | Associated identity like ``np.cumsum->0`` or ``np.cumprod->1`` |
| 1176 | x: dask Array |
| 1177 | axis: int |
| 1178 | dtype: dtype |
| 1179 | method : {'sequential', 'blelloch'}, optional |
| 1180 | Choose which method to use to perform the cumsum. Default is 'sequential'. |
| 1181 | |
| 1182 | * 'sequential' performs the scan of each prior block before the current block. |
| 1183 | * 'blelloch' is a work-efficient parallel scan. It exposes parallelism by first |
| 1184 | calling ``preop`` on each block and combines the values via a binary tree. |
| 1185 | This method may be faster or more memory efficient depending on workload, |
| 1186 | scheduler, and hardware. More benchmarking is necessary. |
| 1187 | preop: callable, optional |
| 1188 | Function used by 'blelloch' method, |
| 1189 | like ``np.cumsum->np.sum`` or ``np.cumprod->np.prod`` |
| 1190 | |
| 1191 | Returns |
| 1192 | ------- |
| 1193 | dask array |
| 1194 | |
| 1195 | See also |
| 1196 | -------- |
| 1197 | cumsum |
| 1198 | cumprod |
| 1199 | """ |
| 1200 | if method == "blelloch": |
| 1201 | if preop is None: |
| 1202 | raise TypeError( |
| 1203 | 'cumreduction with "blelloch" method required `preop=` argument' |
| 1204 | ) |
| 1205 | return prefixscan_blelloch(func, preop, binop, x, axis, dtype, out=out) |
| 1206 | elif method != "sequential": |
| 1207 | raise ValueError( |
| 1208 | f'Invalid method for cumreduction. Expected "sequential" or "blelloch". Got: {method!r}' |
| 1209 | ) |
| 1210 | |
| 1211 | if axis is None: |
| 1212 | if x.ndim > 1: |