Source code for mars.tensor.statistics.average

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

from ..datasource import tensor as astensor
from ..base.broadcast_to import broadcast_to
from ..base.swapaxes import swapaxes

[docs]def average(a, axis=None, weights=None, returned=False): """ Compute the weighted average along the specified axis. Parameters ---------- a : array_like Tensor containing data to be averaged. If `a` is not a tensor, a conversion is attempted. axis : None or int or tuple of ints, optional Axis or axes along which to average `a`. The default, axis=None, will average over all of the elements of the input tensor. If axis is negative it counts from the last to the first axis. If axis is a tuple of ints, averaging is performed on all of the axes specified in the tuple instead of a single axis or all the axes as before. weights : array_like, optional A tensor of weights associated with the values in `a`. Each value in `a` contributes to the average according to its associated weight. The weights tensor can either be 1-D (in which case its length must be the size of `a` along the given axis) or of the same shape as `a`. If `weights=None`, then all data in `a` are assumed to have a weight equal to one. returned : bool, optional Default is `False`. If `True`, the tuple (`average`, `sum_of_weights`) is returned, otherwise only the average is returned. If `weights=None`, `sum_of_weights` is equivalent to the number of elements over which the average is taken. Returns ------- average, [sum_of_weights] : tensor_type or double Return the average along the specified axis. When returned is `True`, return a tuple with the average as the first element and the sum of the weights as the second element. The return type is `Float` if `a` is of integer type, otherwise it is of the same type as `a`. `sum_of_weights` is of the same type as `average`. Raises ------ ZeroDivisionError When all weights along axis are zero. See `` for a version robust to this type of error. TypeError When the length of 1D `weights` is not the same as the shape of `a` along axis. See Also -------- mean Examples -------- >>> import mars.tensor as mt >>> data = list(range(1,5)) >>> data [1, 2, 3, 4] >>> mt.average(data).execute() 2.5 >>> mt.average(range(1,11), weights=range(10,0,-1)).execute() 4.0 >>> data = mt.arange(6).reshape((3,2)) >>> data.execute() array([[0, 1], [2, 3], [4, 5]]) >>> mt.average(data, axis=1, weights=[1./4, 3./4]).execute() array([ 0.75, 2.75, 4.75]) >>> mt.average(data, weights=[1./4, 3./4]).execute() Traceback (most recent call last): ... TypeError: Axis must be specified when shapes of a and weights differ. """ from ..arithmetic import truediv, multiply a = astensor(a) if weights is None: avg = a.mean(axis) scl = avg.dtype.type(a.size / avg.size) else: wgt = astensor(weights) if issubclass(a.dtype.type, (np.integer, np.bool_)): result_dtype = np.result_type(a.dtype, wgt.dtype, "f8") else: result_dtype = np.result_type(a.dtype, wgt.dtype) # sanity checks if a.shape != wgt.shape: if axis is None: raise TypeError( "Axis must be specified when shapes of a and weights differ." ) if wgt.ndim != 1: raise TypeError( "1D weights expected when shapes of a and weights differ." ) if wgt.shape[0] != a.shape[axis]: raise ValueError( "Length of weights not compatible with specified axis." ) # setup wgt to broadcast along axis wgt = broadcast_to(wgt, (a.ndim - 1) * (1,) + wgt.shape) wgt = swapaxes(wgt, -1, axis) scl = wgt.sum(axis=axis, dtype=result_dtype) with np.errstate(divide="raise"): avg = truediv(multiply(a, wgt, dtype=result_dtype).sum(axis), scl) if returned: if scl.shape != avg.shape: scl = broadcast_to(scl, avg.shape) return avg, scl else: return avg