Understand wquantiles module in Python

wquantiles(github)

Weighted quantiles with Python, including weighted median.

The main methods are quantile and median. The input of quantile is a numpy array (data), a numpy array of weights of one dimension and the value of the quantile (between 0 and 1) to compute. The weighting is applied along the last axis.

The method median is an alias to quantile(data, weights, 0.5).

Code

""" Library to compute weighted quantiles, including the weighted median, of numpy arrays. """ from __future__ import print_function import numpy as np __version__ = "0.4" def quantile_1D(data, weights, quantile): """ Compute the weighted quantile of a 1D numpy array. Parameters ---------- data : ndarray Input array (one dimension). weights : ndarray Array with the weights of the same size of `data`. quantile : float Quantile to compute. It must have a value between 0 and 1. Returns ------- quantile_1D : float The output value. """ # Check the data if not isinstance(data, np.matrix): data = np.asarray(data) if not isinstance(weights, np.matrix): weights = np.asarray(weights) nd = data.ndim if nd != 1: raise TypeError("data must be a one dimensional array") ndw = weights.ndim if ndw != 1: raise TypeError("weights must be a one dimensional array") if data.shape != weights.shape: raise TypeError("the length of data and weights must be the same") if ((quantile > 1.) or (quantile < 0.)): raise ValueError("quantile must have a value between 0. and 1.") # Sort the data ind_sorted = np.argsort(data) sorted_data = data[ind_sorted] sorted_weights = weights[ind_sorted] # Compute the auxiliary arrays Sn = np.cumsum(sorted_weights) # TODO: Check that the weights do not sum zero #assert Sn != 0, "The sum of the weights must not be zero" Pn = (Sn-0.5*sorted_weights)/np.sum(sorted_weights) # Get the value of the weighted median return np.interp(quantile, Pn, sorted_data) def quantile(data, weights, quantile): """ Weighted quantile of an array with respect to the last axis. Parameters ---------- data : ndarray Input array. weights : ndarray Array with the weights. It must have the same size of the last axis of `data`. quantile : float Quantile to compute. It must have a value between 0 and 1. Returns ------- quantile : float The output value. """ # TODO: Allow to specify the axis nd = data.ndim if nd == 0: TypeError("data must have at least one dimension") elif nd == 1: return quantile_1D(data, weights, quantile) elif nd > 1: n = data.shape imr = data.reshape((np.prod(n[:-1]), n[-1])) result = np.apply_along_axis(quantile_1D, -1, imr, weights, quantile) return result.reshape(n[:-1]) def median(data, weights): """ Weighted median of an array with respect to the last axis. Alias for `quantile(data, weights, 0.5)`. """ return quantile(data, weights, 0.5)

Functions

np.argsort(data)

np.cumsum()