######################################################## Copyright (c) 2015, ArrayFire# All rights reserved.## This file is distributed under 3-clause BSD license.# The complete license agreement can be obtained at:# http://arrayfire.com/licenses/BSD-3-Clause########################################################"""Statistical algorithms (mean, var, stdev, etc)."""from.libraryimport*from.arrayimport*
[docs]defmean(a,weights=None,dim=None):"""    Calculate mean along a given dimension.    Parameters    ----------    a: af.Array        The input array.    weights: optional: af.Array. default: None.        Array to calculate the weighted mean. Must match size of the        input array.    dim: optional: int. default: None.        The dimension for which to obtain the mean from input data.    Returns    -------    output: af.Array        Array containing the mean of the input array along a given        dimension.    """ifdimisnotNone:out=Array()ifweightsisNone:safe_call(backend.get().af_mean(c_pointer(out.arr),a.arr,c_int_t(dim)))else:safe_call(backend.get().af_mean_weighted(c_pointer(out.arr),a.arr,weights.arr,c_int_t(dim)))returnoutelse:real=c_double_t(0)imag=c_double_t(0)ifweightsisNone:safe_call(backend.get().af_mean_all(c_pointer(real),c_pointer(imag),a.arr))else:safe_call(backend.get().af_mean_all_weighted(c_pointer(real),c_pointer(imag),a.arr,weights.arr))real=real.valueimag=imag.valuereturnrealifimag==0elsereal+imag*1j
[docs]defvar(a,isbiased=False,weights=None,dim=None):"""    Calculate variance along a given dimension.    Parameters    ----------    a: af.Array        The input array.    isbiased: optional: Boolean. default: False.        Boolean denoting population variance (false) or sample        variance (true).    weights: optional: af.Array. default: None.        Array to calculate for the weighted mean. Must match size of        the input array.    dim: optional: int. default: None.        The dimension for which to obtain the variance from input data.    Returns    -------    output: af.Array        Array containing the variance of the input array along a given        dimension.    """ifdimisnotNone:out=Array()ifweightsisNone:safe_call(backend.get().af_var(c_pointer(out.arr),a.arr,isbiased,c_int_t(dim)))else:safe_call(backend.get().af_var_weighted(c_pointer(out.arr),a.arr,weights.arr,c_int_t(dim)))returnoutelse:real=c_double_t(0)imag=c_double_t(0)ifweightsisNone:safe_call(backend.get().af_var_all(c_pointer(real),c_pointer(imag),a.arr,isbiased))else:safe_call(backend.get().af_var_all_weighted(c_pointer(real),c_pointer(imag),a.arr,weights.arr))real=real.valueimag=imag.valuereturnrealifimag==0elsereal+imag*1j
[docs]defmeanvar(a,weights=None,bias=VARIANCE.DEFAULT,dim=-1):"""    Calculate mean and variance along a given dimension.    Parameters    ----------    a: af.Array        The input array.    weights: optional: af.Array. default: None.        Array to calculate for the weighted mean. Must match size of        the input array.    bias: optional: af.VARIANCE. default: DEFAULT.        population variance(VARIANCE.POPULATION) or        sample variance(VARIANCE.SAMPLE).    dim: optional: int. default: -1.        The dimension for which to obtain the variance from input data.    Returns    -------    mean: af.Array        Array containing the mean of the input array along a given        dimension.    variance: af.Array        Array containing the variance of the input array along a given        dimension.    """mean_out=Array()var_out=Array()ifweightsisNone:weights=Array()safe_call(backend.get().af_meanvar(c_pointer(mean_out.arr),c_pointer(var_out.arr),a.arr,weights.arr,bias.value,c_int_t(dim)))returnmean_out,var_out
[docs]defstdev(a,dim=None):"""    Calculate standard deviation along a given dimension.    Parameters    ----------    a: af.Array        The input array.    dim: optional: int. default: None.        The dimension for which to obtain the standard deviation from        input data.    Returns    -------    output: af.Array        Array containing the standard deviation of the input array        along a given dimension.    """ifdimisnotNone:out=Array()safe_call(backend.get().af_stdev(c_pointer(out.arr),a.arr,c_int_t(dim)))returnoutelse:real=c_double_t(0)imag=c_double_t(0)safe_call(backend.get().af_stdev_all(c_pointer(real),c_pointer(imag),a.arr))real=real.valueimag=imag.valuereturnrealifimag==0elsereal+imag*1j
[docs]defcov(a,isbiased=False,dim=None):"""    Calculate covariance along a given dimension.    Parameters    ----------    a: af.Array        The input array.    isbiased: optional: Boolean. default: False.        Boolean denoting whether biased estimate should be taken.    dim: optional: int. default: None.        The dimension for which to obtain the covariance from input data.    Returns    -------    output: af.Array        Array containing the covariance of the input array along a        given dimension.    """ifdimisnotNone:out=Array()safe_call(backend.get().af_cov(c_pointer(out.arr),a.arr,isbiased,c_int_t(dim)))returnoutelse:real=c_double_t(0)imag=c_double_t(0)safe_call(backend.get().af_cov_all(c_pointer(real),c_pointer(imag),a.arr,isbiased))real=real.valueimag=imag.valuereturnrealifimag==0elsereal+imag*1j
[docs]defmedian(a,dim=None):"""    Calculate median along a given dimension.    Parameters    ----------    a: af.Array        The input array.    dim: optional: int. default: None.        The dimension for which to obtain the median from input data.    Returns    -------    output: af.Array        Array containing the median of the input array along a        given dimension.    """ifdimisnotNone:out=Array()safe_call(backend.get().af_median(c_pointer(out.arr),a.arr,c_int_t(dim)))returnoutelse:real=c_double_t(0)imag=c_double_t(0)safe_call(backend.get().af_median_all(c_pointer(real),c_pointer(imag),a.arr))real=real.valueimag=imag.valuereturnrealifimag==0elsereal+imag*1j
[docs]defcorrcoef(x,y):"""    Calculate the correlation coefficient of the input arrays.    Parameters    ----------    x: af.Array        The first input array.    y: af.Array        The second input array.    Returns    -------    output: af.Array        Array containing the correlation coefficient of the input arrays.    """real=c_double_t(0)imag=c_double_t(0)safe_call(backend.get().af_corrcoef(c_pointer(real),c_pointer(imag),x.arr,y.arr))real=real.valueimag=imag.valuereturnrealifimag==0elsereal+imag*1j
[docs]deftopk(data,k,dim=0,order=TOPK.DEFAULT):"""    Return top k elements along a single dimension.    Parameters    ----------    data: af.Array          Input array to return k elements from.    k: scalar. default: 0       The number of elements to return from input array.    dim: optional: scalar. default: 0         The dimension along which the top k elements are         extracted. Note: at the moment, topk() only supports the         extraction of values along the first dimension.    order: optional: af.TOPK. default: af.TOPK.DEFAULT           The ordering of k extracted elements. Defaults to top k max values.    Returns    -------    values: af.Array            Top k elements from input array.    indices: af.Array             Corresponding index array to top k elements.    """values=Array()indices=Array()safe_call(backend.get().af_topk(c_pointer(values.arr),c_pointer(indices.arr),data.arr,k,c_int_t(dim),order.value))returnvalues,indices