def average(
a: BlockArray,
axis: Union[None, int, List[int]] = None,
weights: Optional[BlockArray] = None,
returned: bool = False,
):
if axis and not isinstance(axis, int):
raise NotImplementedError("Only single 'axis' is currently supported.")
if weights is None:
avg = mean(a, axis)
if not returned:
return avg
weights_sum = BlockArray.from_scalar(a.size / avg.size, a.cm)
return avg, weights_sum
if a.shape != weights.shape:
raise NotImplementedError(
"1D weights broadcasting is currently not supported; "
"dimensions of 'a' and 'weights' must match.")
weights_sum = sum(weights, axis=axis)
if not all(weights_sum):
raise ZeroDivisionError("Weights along one or more axes sum to zero.")
avg = divide(sum(multiply(a, weights), axis=axis), weights_sum)
if not returned:
return avg
if avg.shape != weights_sum.shape:
weights_sum = weights_sum.broadcast_to(avg.shape)
return avg, weights_sum
def average(
a: BlockArray,
axis: Union[None, int] = None,
weights: Optional[BlockArray] = None,
returned: bool = False,
) -> Union[BlockArray, Tuple[BlockArray, BlockArray]]:
"""Compute the weighted average along the specified axis.
Args:
a: BlockArray to be averaged.
axis: Axis along which to average `a`.
weights: BlockArray of weights associated with `a`.
returned: Whether to return the sum of the weights.
Returns:
The average along the specified axis. If `returned` is True, return a tuple with the
average as the first element and the sum of the weights as the second element.
"""
if axis and not isinstance(axis, int):
raise NotImplementedError("Only single 'axis' is currently supported.")
if weights is None:
avg = mean(a, axis=axis)
if not returned:
return avg
weights_sum = BlockArray.from_scalar(a.size / avg.size, a.cm)
return avg, weights_sum
if a.shape != weights.shape:
raise NotImplementedError(
"1D weights broadcasting is currently not supported; "
"dimensions of 'a' and 'weights' must match."
)
weights_sum = sum(weights, axis=axis)
if not all(weights_sum):
raise ZeroDivisionError("Weights along one or more axes sum to zero.")
avg = divide(sum(multiply(a, weights), axis=axis), weights_sum)
if not returned:
return avg
if avg.shape != weights_sum.shape:
weights_sum = weights_sum.broadcast_to(avg.shape)
return avg, weights_sum
def scalar(self, value):
return BlockArray.from_scalar(value, self.cm)
