| 346 | return h.clip(1, n) |
| 347 | |
| 348 | def _weighted_quantile_1d( |
| 349 | data: np.ndarray, |
| 350 | weights: np.ndarray, |
| 351 | q: np.ndarray, |
| 352 | skipna: bool, |
| 353 | method: QUANTILE_METHODS = "linear", |
| 354 | ) -> np.ndarray: |
| 355 | # This algorithm has been adapted from: |
| 356 | # https://aakinshin.net/posts/weighted-quantiles/#reference-implementation |
| 357 | is_nan = np.isnan(data) |
| 358 | if skipna: |
| 359 | # Remove nans from data and weights |
| 360 | not_nan = ~is_nan |
| 361 | data = data[not_nan] |
| 362 | weights = weights[not_nan] |
| 363 | elif is_nan.any(): |
| 364 | # Return nan if data contains any nan |
| 365 | return np.full(q.size, np.nan) |
| 366 | |
| 367 | # Filter out data (and weights) associated with zero weights, which also flattens them |
| 368 | nonzero_weights = weights != 0 |
| 369 | data = data[nonzero_weights] |
| 370 | weights = weights[nonzero_weights] |
| 371 | n = data.size |
| 372 | |
| 373 | if n == 0: |
| 374 | # Possibly empty after nan or zero weight filtering above |
| 375 | return np.full(q.size, np.nan) |
| 376 | |
| 377 | # Kish's effective sample size |
| 378 | nw = weights.sum() ** 2 / (weights**2).sum() |
| 379 | |
| 380 | # Sort data and weights |
| 381 | sorter = np.argsort(data) |
| 382 | data = data[sorter] |
| 383 | weights = weights[sorter] |
| 384 | |
| 385 | # Normalize and sum the weights |
| 386 | weights = weights / weights.sum() |
| 387 | weights_cum = np.append(0, weights.cumsum()) |
| 388 | |
| 389 | # Vectorize the computation by transposing q with respect to weights |
| 390 | q = np.atleast_2d(q).T |
| 391 | |
| 392 | # Get the interpolation parameter for each q |
| 393 | h = _get_h(nw, q, method) |
| 394 | |
| 395 | # Find the samples contributing to the quantile computation (at *positions* between (h-1)/nw and h/nw) |
| 396 | u = np.maximum((h - 1) / nw, np.minimum(h / nw, weights_cum)) |
| 397 | |
| 398 | # Compute their relative weight |
| 399 | v = u * nw - h + 1 |
| 400 | w = np.diff(v) |
| 401 | |
| 402 | # Apply the weights |
| 403 | return (data * w).sum(axis=1) |
| 404 | |
| 405 | if skipna is None and da.dtype.kind in "cfO": |