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)