function_base.py 文件源码

def _median(a, axis=None, out=None, overwrite_input=False):
    # can't be reasonably be implemented in terms of percentile as we have to
    # call mean to not break astropy
    a = np.asanyarray(a)

    # Set the partition indexes
    if axis is None:
        sz = a.size
    else:
        sz = a.shape[axis]
    if sz % 2 == 0:
        szh = sz // 2
        kth = [szh - 1, szh]
    else:
        kth = [(sz - 1) // 2]
    # Check if the array contains any nan's
    if np.issubdtype(a.dtype, np.inexact):
        kth.append(-1)

    if overwrite_input:
        if axis is None:
            part = a.ravel()
            part.partition(kth)
        else:
            a.partition(kth, axis=axis)
            part = a
    else:
        part = partition(a, kth, axis=axis)

    if part.shape == ():
        # make 0-D arrays work
        return part.item()
    if axis is None:
        axis = 0

    indexer = [slice(None)] * part.ndim
    index = part.shape[axis] // 2
    if part.shape[axis] % 2 == 1:
        # index with slice to allow mean (below) to work
        indexer[axis] = slice(index, index+1)
    else:
        indexer[axis] = slice(index-1, index+1)

    # Check if the array contains any nan's
    if np.issubdtype(a.dtype, np.inexact) and sz > 0:
        # warn and return nans like mean would
        rout = mean(part[indexer], axis=axis, out=out)
        return np.lib.utils._median_nancheck(part, rout, axis, out)
    else:
        # if there are no nans
        # Use mean in odd and even case to coerce data type
        # and check, use out array.
        return mean(part[indexer], axis=axis, out=out)