bs-python-utils 0.8.5__tar.gz → 0.9.2__tar.gz

{bs_python_utils-0.8.5 → bs_python_utils-0.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bs-python-utils
-Version: 0.8.5
+Version: 0.9.2
 Summary: Utilities programs for my Python code
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown
@@ -25,6 +25,7 @@ Requires-Dist: statsmodels>=0.14.5
 Requires-Dist: streamlit>=1.49.1
 Requires-Dist: vega-datasets>=0.9.0
 Requires-Dist: colorama>=0.4.6
+Requires-Dist: pyarrow>=21.0.0
 Dynamic: license-file
 
 ## bs-python-utils
@@ -42,6 +43,9 @@ Dynamic: license-file
 
 ### Release notes
 
+#### 0.9 (March 18, 2026)
+Improved bivariate quantiles.
+
 #### 0.8.5 (January 25, 2026)
 Added colored logger.
 

{bs_python_utils-0.8.5 → bs_python_utils-0.9.2}/README.md

@@ -13,6 +13,9 @@
 
 ### Release notes
 
+#### 0.9 (March 18, 2026)
+Improved bivariate quantiles.
+
 #### 0.8.5 (January 25, 2026)
 Added colored logger.
 

bs_python_utils-0.9.2/bs_python_utils/bivariate_quantiles.py

@@ -0,0 +1,238 @@
+"""This takes in observations of a bivariate random variable `y`
+and computes vector quantiles and vector ranks à la
+[Chernozhukov-Galichon-Hallin-Henry (*Ann. Stats.* 2017)](
+https://projecteuclid.org/journals/annals-of-statistics/volume-45/
+issue-1/MongeKantorovich-depth-quantiles-ranks-and-signs/10.1214/16-AOS1450.full).
+
+
+Note:
+    if the math looks strange in the documentation, just reload the page.
+
+The sequence of steps is as follows:
+
+1.  optimize the weights: `v = solve_for_v(y, n_nodes)` given  `n_nodes` Chebyshev nodes for numerical integration
+2. to obtain the $(u_1,u_2)$ quantiles for $(u_1, u_2)\\in [0,1]$, run
+`qtiles_y = bivariate_quantiles_v(y, v, u1, u2)`
+3. to compute the vector ranks for all points in the sample (the barycenters
+  of the cells in the power diagram):
+`ranks_y = bivariate_ranks_v(y, v, n_nodes)`
+
+Steps 1 and 2 can be combined: `qtiles_y = bivariate_quantiles(y, v, u1, u2, n_nodes)`
+
+Steps 1 and 3 can be combined: `ranks_y = bivariate_ranks(y, n_nodes)`
+"""
+
+from typing import cast
+
+import numpy as np
+
+from bs_python_utils.bs_opt import minimize_free, print_optimization_results
+from bs_python_utils.bsnputils import TwoArrays, npmaxabs
+from bs_python_utils.bsutils import bs_error_abort
+from bs_python_utils.chebyshev import Interval, cheb_get_nodes_1d
+
+
+def _compute_ad(y: np.ndarray) -> TwoArrays:
+    """Build the `A` and `dy2` matrices used in the dual optimisation."""
+    y1 = y[:, 0]
+    dy1 = np.subtract.outer(y1, y1)
+    y2 = y[:, 1]
+    dy2 = np.subtract.outer(y2, y2)
+    np.fill_diagonal(dy2, 1.0)
+    dy2 = dy2.T
+    a_mat = np.divide(dy1, dy2)
+    return a_mat, dy2
+
+
+def _compute_m_M(
+    v: np.ndarray, a_mat: np.ndarray, dy2: np.ndarray, tau1_nodes: np.ndarray
+) -> TwoArrays:
+    """Build the `m` and `M` matrices used in the dual optimisation."""
+    dv = np.subtract.outer(v, v)
+    b_mat = dv / dy2
+    np.fill_diagonal(dy2, 0.0)
+    EPS = 1e-12
+    maskp = dy2 < EPS
+    maskm = dy2 > -EPS
+    n, n_nodes = v.size, tau1_nodes.size
+    m_low = np.empty((n, n_nodes))
+    m_high = np.empty((n, n_nodes))
+    for i, tau1 in enumerate(tau1_nodes):
+        f_mat = tau1 * a_mat - b_mat
+        f_matp = f_mat.copy()
+        f_matm = f_mat.copy()
+        f_matp[maskp] = 1
+        f_matm[maskm] = 0
+        m_low[:, i] = np.max(f_matm, axis=1)
+        m_high[:, i] = np.min(f_matp, axis=1)
+    return np.clip(m_low, 0.0, 1.0), np.clip(m_high, 0.0, 1.0)
+
+
+def bivariate_quantiles_v(y: np.ndarray, tau: np.ndarray, v: np.ndarray) -> np.ndarray:
+    """Evaluate vector quantiles for a given set of dual weights.
+
+    Args:
+        y: Observations with shape ``(n, 2)``.
+        tau: Evaluation points in ``[0, 1]^2`` (shape ``(m, 2)``).
+        v: Dual weights solving the optimal transport problem (length ``n``).
+
+    Returns:
+        Array of quantile locations with shape ``(m, 2)``.
+    """
+    if tau.shape[1] != 2:
+        bs_error_abort("tau must have two columns")
+    q = y[np.argmax(tau @ y.T - v, axis=1), :]
+    return cast(np.ndarray, q)
+
+
+def _objgrad(
+    v: np.ndarray, args: list, gr: bool = False
+) -> float | tuple[float, np.ndarray, np.ndarray]:
+    """computes the expectation of $\\psi(U, v)$ and perhaps its gradient wrt `v` or the bivariate ranks
+
+    Args:
+        v: an `(n-1)`-vector
+        args: a list of other arguments `[y, a_mat, dy2, tau1_nodes, tau1_weights, verbose]`
+        gr: if `False`, we only return the value of the objective function
+            if `True`, we also return the gradient and the bivariate ranks
+
+    Returns:
+        the value of the expectation and perhaps its gradient and the bivariate ranks
+    """
+    y = args[0]
+    y1 = y[:, 0]
+    y2 = y[:, 1]
+    n = y.shape[0]
+    a_mat, dy2 = args[1], args[2]
+    tau1_nodes = args[3]
+    tau1_weights = args[4]
+    vs1 = np.append(v, -np.sum(v))
+    m, M = _compute_m_M(vs1, a_mat, dy2, tau1_nodes)
+    # print(f"m is {m}")
+    # print(f"M  is {M}")
+    # import sys
+
+    # sys.exit(1)
+
+    EPS = 1e-12
+    obj_val = 0.0
+    probs = np.zeros(n)
+    bivrank = np.zeros((n, 2))
+    for k in range(n):
+        Mk = M[k, :]
+        mk = m[k, :]
+        pos_diffs = np.maximum(Mk - mk, 0.0)
+        # print(f"pos_diffs for k={k} are {pos_diffs}")
+        pos_diffs_sq = np.maximum(Mk * Mk - mk * mk, 0.0)
+        probs[k] = pos_diffs @ tau1_weights
+        # print(f"probs[{k}] = {probs[k]}")
+        factor1 = (tau1_nodes * pos_diffs) @ tau1_weights
+        factor2 = (pos_diffs_sq @ tau1_weights) / 2.0
+        obj_val += y1[k] * factor1 + y2[k] * factor2 - vs1[k] * probs[k]
+        if probs[k] > EPS:
+            bivrank[k, 0] = factor1 / probs[k]
+            bivrank[k, 1] = factor2 / probs[k]
+
+    # print(f"{np.min(probs)=}")
+
+    if gr:
+        grad_val = probs[-1] - probs[:-1]
+        return obj_val, grad_val, bivrank
+    else:
+        return obj_val
+
+
+def _obj(v: np.ndarray, args: list):
+    return _objgrad(v, args)
+
+
+def _grad(v: np.ndarray, args: list):
+    res_objg = cast(tuple[float, np.ndarray], _objgrad(v, args, gr=True))
+    grad_val = res_objg[1]
+    verbose = args[-1]
+    if verbose:
+        print(f"The error on the gradient is {npmaxabs(grad_val)}")
+    return grad_val
+
+
+def _solve_for_v(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> TwoArrays:
+    """Solve the dual optimisation to obtain the optimal weights ``v`` and the bivariate ranks
+
+    Args:
+        y: Observations with shape ``(n, 2)``.
+        n_nodes: Number of Chebyshev nodes for the quadrature.
+        verbose: Print optimisation diagnostics when ``True``.
+
+    Returns:
+        Array of length ``n`` containing the optimal weights (including the
+            residual term).
+        Array of shape ``(n, 2)`` containing the bivariate ranks (the barycenters of the cells in the power diagram).
+    """
+    d = y.shape[1]
+
+    if d != 2:
+        bs_error_abort(f"only works for 2-dimensional y, not for {d}")
+
+    v0 = np.mean(y[:-1, :], 1)
+
+    interval01 = Interval(0.0, 1.0)
+    tau1_nodes, tau1_weights = cheb_get_nodes_1d(interval01, n_nodes)
+
+    a_mat, dy2 = _compute_ad(y)
+
+    argsog = [y, a_mat, dy2, tau1_nodes, tau1_weights, verbose]
+
+    res = minimize_free(_obj, _grad, v0, args=argsog)
+    if verbose:
+        print_optimization_results(res, "Minimizing over v")
+
+    if not res.success:
+        bs_error_abort("Problem! the optimization failed.")
+    vstar = res.x
+    if verbose:
+        print(f"The final gradient over v is close to 0: error {npmaxabs(res.jac)}")
+    _, _, bivranks = cast(tuple, _objgrad(vstar, argsog, gr=True))
+    vstar = np.append(vstar, -np.sum(vstar))
+    return cast(np.ndarray, vstar), cast(np.ndarray, bivranks)
+
+
+def bivariate_ranks(
+    y: np.ndarray,
+    n_nodes: int = 32,
+    verbose: bool = False,
+) -> np.ndarray:
+    """Compute the barycentric ranks of each observation.
+
+    Args:
+        y: Observations with shape ``(n, 2)``.
+        n_nodes: Number of Chebyshev nodes used in the quadrature.
+        verbose: Print diagnostics when ``True``.
+
+    Returns:
+        Array of average ranks (shape ``(n, 2)``) with ``nan`` for zero-mass cells.
+    """
+    d = y.shape[1]
+
+    if d != 2:
+        bs_error_abort(f"only works for 2-dimensional y, not for {d}")
+
+    _, bivranks = _solve_for_v(y, n_nodes, verbose)
+    return cast(np.ndarray, bivranks)
+
+
+def bivariate_quantiles(
+    y: np.ndarray, tau: np.ndarray, n_nodes: int = 32, verbose: bool = False
+) -> np.ndarray:
+    """Solve for the dual weights then evaluate bivariate quantiles.
+
+    Args:
+        y: Observations, shape ``(n, 2)``.
+        tau: Query points in ``[0, 1]^2`` (shape ``(m, 2)``).
+        n_nodes: Number of Chebyshev nodes for the quadrature.
+        verbose: Print optimisation diagnostics when ``True``.
+
+    Returns:
+        Bivariate quantiles at ``u``.
+    """
+    v, _ = _solve_for_v(y, n_nodes, verbose)
+    return bivariate_quantiles_v(y, tau, v)

{bs_python_utils-0.8.5 → bs_python_utils-0.9.2}/bs_python_utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bs-python-utils
-Version: 0.8.5
+Version: 0.9.2
 Summary: Utilities programs for my Python code
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown
@@ -25,6 +25,7 @@ Requires-Dist: statsmodels>=0.14.5
 Requires-Dist: streamlit>=1.49.1
 Requires-Dist: vega-datasets>=0.9.0
 Requires-Dist: colorama>=0.4.6
+Requires-Dist: pyarrow>=21.0.0
 Dynamic: license-file
 
 ## bs-python-utils
@@ -42,6 +43,9 @@ Dynamic: license-file
 
 ### Release notes
 
+#### 0.9 (March 18, 2026)
+Improved bivariate quantiles.
+
 #### 0.8.5 (January 25, 2026)
 Added colored logger.
 

{bs_python_utils-0.8.5 → bs_python_utils-0.9.2}/bs_python_utils.egg-info/requires.txt

@@ -18,3 +18,4 @@ statsmodels>=0.14.5
 streamlit>=1.49.1
 vega-datasets>=0.9.0
 colorama>=0.4.6
+pyarrow>=21.0.0

{bs_python_utils-0.8.5 → bs_python_utils-0.9.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bs-python-utils"
-version = "0.8.5"
+version = "0.9.2"
 description = "Utilities programs for my Python code"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -25,6 +25,7 @@ dependencies = [
     "streamlit>=1.49.1",
     "vega-datasets>=0.9.0",
     "colorama>=0.4.6",
+    "pyarrow>=21.0.0",
 ]
 
 [tool.setuptools]

bs_python_utils-0.9.2/tests/test_bivariate_quantiles.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from bs_python_utils.bivariate_quantiles import (
+    _compute_ad,
+    _solve_for_v,
+)
+from bs_python_utils.bsnputils import npmaxabs
+
+
+def test_bivariate_ranks():
+    rng = np.random.default_rng(123)
+    n = 500
+    y = rng.normal(size=(n, 2))
+    nodes = 32
+
+    _, bivranks = _solve_for_v(y, n_nodes=nodes, verbose=False)
+
+    bivranks_th = np.empty_like(bivranks)
+    for k in range(n):
+        bivranks_th[k, 0] = np.mean(y[:, 0] < y[k, 0]) + 0.5 * np.mean(
+            y[:, 0] == y[k, 0]
+        )
+        bivranks_th[k, 1] = np.mean(y[:, 1] < y[k, 1]) + 0.5 * np.mean(
+            y[:, 1] == y[k, 1]
+        )
+
+    assert npmaxabs(bivranks - bivranks_th) < 0.01
+
+
+def test_compute_ad():
+    y = np.array([[0.0, 1.0], [1.0, 1.3], [-1.0, 3.0]])
+    a_mat, dy2 = _compute_ad(y)
+
+    a_mat_th = np.array(
+        [[0.0, -1.0 / 0.3, 0.5], [-1.0 / 0.3, 0.0, 2.0 / 1.7], [0.5, 2.0 / 1.7, 0.0]]
+    )
+    dy2_th = np.array([[1.0, 0.3, 2.0], [-0.3, 1.0, 1.7], [-2.0, -1.7, 1.0]])
+
+    assert np.allclose(a_mat, a_mat_th)
+    assert np.allclose(dy2, dy2_th)

bs_python_utils-0.8.5/bs_python_utils/bivariate_quantiles.py

@@ -1,306 +0,0 @@
-"""This takes in observations of a bivariate random variable `y`
-and computes vector quantiles and vector ranks à la
-[Chernozhukov-Galichon-Hallin-Henry (*Ann. Stats.* 2017)](
-https://projecteuclid.org/journals/annals-of-statistics/volume-45/
-issue-1/MongeKantorovich-depth-quantiles-ranks-and-signs/10.1214/16-AOS1450.full).
-
-
-Note:
-    if the math looks strange in the documentation, just reload the page.
-
-The sequence of steps is as follows:
-
-* choose a  number of Chebyshev nodes for numerical integration and optimize
-  the weights: `v = solve_for_v(y, n_nodes)`
-* to obtain the $(u_1,u_2)$ quantiles for $(u_1, u_2)\\in [0,1]$, run
-`qtiles_y = bivariate_quantiles_v(y, v, u1, u2)`
-* to compute the vector ranks for all points in the sample (the barycenters
-  of the cells in the power diagram):
-`ranks_y = bivariate_ranks_v(y, v, n_nodes)`
-
-Steps 1 and 2 can be combined: `qtiles_y = bivariate_quantiles(y, v, u1, u2, n_nodes)`
-
-Steps 1 and 3 can be combined: `ranks_y = bivariate_ranks(y, n_nodes)`
-"""
-
-from typing import cast
-
-import numpy as np
-
-from bs_python_utils.bs_opt import minimize_free, print_optimization_results
-from bs_python_utils.bsnputils import TwoArrays, npmaxabs
-from bs_python_utils.bsutils import bs_error_abort
-from bs_python_utils.chebyshev import Interval, cheb_get_nodes_1d
-
-
-def _compute_ab(y_sorted: np.ndarray, v_sorted: np.ndarray) -> TwoArrays:
-    """Build the `A` and `B` matrices used in the dual optimisation."""
-    y1 = y_sorted[:, 0]
-    dy1 = np.subtract.outer(y1, y1)
-    y2 = y_sorted[:, 1]
-    dy2 = np.subtract.outer(y2, y2)
-    np.fill_diagonal(dy2, 1.0)
-    dv = np.subtract.outer(v_sorted, v_sorted)
-    with np.errstate(divide="ignore", invalid="ignore"):
-        a_mat = np.divide(dy1.T, dy2, where=np.abs(dy2) > 1e-12)
-        b_mat = np.divide(dv.T, dy2, where=np.abs(dy2) > 1e-12)
-    a_mat = np.nan_to_num(a_mat, nan=0.0, posinf=0.0, neginf=0.0)
-    b_mat = np.nan_to_num(b_mat, nan=0.0, posinf=0.0, neginf=0.0)
-    return a_mat, b_mat
-
-
-def _compute_u2_bounds(
-    k: int, u1: np.ndarray, a_mat: np.ndarray, b_mat: np.ndarray
-) -> TwoArrays:
-    """Return the admissible interval of ``u2`` that selects index ``k``."""
-    n = a_mat.shape[0]
-    m = u1.size
-    if k == 0:
-        left_bound = np.zeros(m)
-        a_right = a_mat[0, 1:]
-        b_right = b_mat[0, 1:]
-        if a_right.size:
-            right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
-        else:
-            right_bound = np.ones(m)
-    elif 1 <= k < n - 1:
-        a_left = a_mat[k, :k]
-        b_left = b_mat[k, :k]
-        if a_left.size:
-            left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
-        else:
-            left_bound = np.zeros(m)
-        a_right = a_mat[k, (k + 1) :]
-        b_right = b_mat[k, (k + 1) :]
-        if a_right.size:
-            right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
-        else:
-            right_bound = np.ones(m)
-    elif k == n - 1:
-        a_left = a_mat[-1, :-1]
-        b_left = b_mat[-1, :-1]
-        if a_left.size:
-            left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
-        else:
-            left_bound = np.zeros(m)
-        right_bound = np.ones(m)
-    else:
-        bs_error_abort(f"{k=} is not compatible with {n=}")
-    left_bound = np.clip(left_bound, 0.0, 1.0)
-    right_bound = np.clip(right_bound, 0.0, 1.0)
-
-    return left_bound, right_bound
-
-
-def bivariate_quantiles_v(y: np.ndarray, u: np.ndarray, v: np.ndarray) -> np.ndarray:
-    """Evaluate vector quantiles for a given set of dual weights.
-
-    Args:
-        y: Observations with shape ``(n, 2)``.
-        u: Evaluation points in ``[0, 1]^2`` (shape ``(m, 2)``).
-        v: Dual weights solving the optimal transport problem (length ``n``).
-
-    Returns:
-        Array of quantile locations with shape ``(m, 2)``.
-    """
-    u = np.atleast_2d(u)
-    if u.shape[1] != 2:
-        bs_error_abort("u must have two columns")
-    m = u.shape[0]
-    q = np.empty((m, 2))
-    block = max(1, min(m, 5_000))
-    for start in range(0, m, block):
-        stop = min(start + block, m)
-        chunk = u[start:stop]
-        net_val = chunk @ y.T - v
-        k_max = np.argmax(net_val, axis=1)
-        q[start:stop] = y[k_max]
-    return cast(np.ndarray, q)
-
-
-def bivariate_ranks_v(
-    y: np.ndarray, v: np.ndarray, n_nodes: int = 32, presorted: bool = False
-) -> np.ndarray:
-    """Compute the barycentric ranks of each observation given optimal weights.
-
-    Args:
-        y: Observations with shape ``(n, 2)``.
-        v: Dual weights returned by ``solve_for_v_``.
-        n_nodes: Number of Chebyshev nodes used in the quadrature.
-        presorted: Set to ``True`` when ``y``/``v`` are pre-sorted by the
-            second coordinate.
-
-    Returns:
-        Array of average ranks (shape ``(n, 2)``) with ``nan`` for zero-mass cells.
-    """
-    n, d = y.shape
-
-    if d != 2:
-        bs_error_abort(f"only works for 2-dimensional y, not for {d}")
-
-    interval01 = Interval(0.0, 1.0)
-    u1_nodes, u1_weights = cheb_get_nodes_1d(interval01, n_nodes)
-
-    if presorted:
-        sort_order = np.arange(n)
-        y_sorted = y
-        v_sorted = v
-    else:
-        sort_order = np.argsort(y[:, 1])
-        y_sorted = y[sort_order, :]
-        v_sorted = v[sort_order]
-
-    a_mat, b_mat = _compute_ab(y_sorted, v_sorted)
-
-    average_ranks = np.zeros((n, 2))
-
-    for k in range(n):
-        left_bounds, right_bounds = _compute_u2_bounds(k, u1_nodes, a_mat, b_mat)
-        pos_diffs = np.maximum(right_bounds - left_bounds, 0.0)
-        pos_diffs_sq = np.maximum(
-            right_bounds * right_bounds - left_bounds * left_bounds, 0.0
-        )
-        prob_k = pos_diffs @ u1_weights
-        if prob_k <= 1e-12:
-            average_ranks[sort_order[k], :] = np.array([np.nan, np.nan])
-            continue
-        average_ranks[sort_order[k], 0] = ((u1_nodes * pos_diffs) @ u1_weights) / prob_k
-        average_ranks[sort_order[k], 1] = ((pos_diffs_sq @ u1_weights) / 2.0) / prob_k
-
-    return average_ranks
-
-
-def _objgrad(
-    v_sorted: np.ndarray, args: list, gr: bool = False
-) -> float | tuple[float, np.ndarray]:
-    """computes the expectation of $\\psi(U, v)$ and perhaps its gradient wrt `v`
-
-    Args:
-        v_sorted: an `n`-vector of weights, sorted by increasing `y[:, 1]`
-        args: a list of other arguments `[y_sorted, u1_nodes, u1_weights, verbose]`
-        gr: if `True`, we also evaluate the gradient
-
-    Returns:
-        the value of the expectation and perhaps its gradient
-    """
-    y_sorted = args[0]
-    n = y_sorted.shape[0]
-    u1_nodes = args[1]
-    u1_weights = args[2]
-    vs1 = np.append(v_sorted, -np.sum(v_sorted))
-    a_mat, b_mat = _compute_ab(y_sorted, vs1)
-
-    obj_val = 0.0
-    probs = np.zeros(n)
-    for k in range(n):
-        left_bounds, right_bounds = _compute_u2_bounds(k, u1_nodes, a_mat, b_mat)
-        pos_diffs = np.maximum(right_bounds - left_bounds, 0.0)
-        pos_diffs_sq = np.maximum(
-            right_bounds * right_bounds - left_bounds * left_bounds, 0.0
-        )
-        obj_val += (
-            y_sorted[k, 0] * ((u1_nodes * pos_diffs) @ u1_weights)
-            + y_sorted[k, 1] * (pos_diffs_sq @ u1_weights) / 2.0
-        )
-        probs[k] = pos_diffs @ u1_weights
-        obj_val -= vs1[k] * probs[k]
-
-    if gr:
-        grad_val = probs[-1] - probs[:-1]
-        return obj_val, grad_val
-    else:
-        return cast(float, obj_val)
-
-
-def _obj(v_sorted: np.ndarray, args: list):
-    return _objgrad(v_sorted, args)
-
-
-def _grad(v_sorted: np.ndarray, args: list):
-    res_objg = cast(tuple[float, np.ndarray], _objgrad(v_sorted, args, gr=True))
-    grad_val = res_objg[1]
-    verbose = args[3]
-    if verbose:
-        print(f"The error on the gradient is {npmaxabs(grad_val)}")
-    return grad_val
-
-
-def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.ndarray:
-    """Solve the dual optimisation to obtain the optimal weights ``v``.
-
-    Args:
-        y: Observations with shape ``(n, 2)``.
-        n_nodes: Number of Chebyshev nodes for the quadrature.
-        verbose: Print optimisation diagnostics when ``True``.
-
-    Returns:
-        Array of length ``n`` containing the optimal weights (including the
-            residual term).
-    """
-    n, d = y.shape
-
-    if d != 2:
-        bs_error_abort(f"only works for 2-dimensional y, not for {d}")
-
-    # sort by increasing y[:, 1]
-    sort_order = np.argsort(y[:, 1])
-    y_sorted = y[sort_order, :]
-
-    v0 = np.mean(y_sorted[:-1, :], 1)
-
-    interval01 = Interval(0.0, 1.0)
-    u1_nodes, u1_weights = cheb_get_nodes_1d(interval01, n_nodes)
-
-    argsog = [y_sorted, u1_nodes, u1_weights, verbose]
-
-    res = minimize_free(_obj, _grad, v0, args=argsog)
-    if verbose:
-        print_optimization_results(res, "Minimizing over v")
-
-    if not res.success:
-        bs_error_abort("Problem! the optimization failed.")
-    vstar = res.x
-    if verbose:
-        print(f"The final gradient over v is close to 0: error {npmaxabs(res.jac)}")
-    vstar1_sorted = np.append(vstar, -np.sum(vstar))
-
-    # revert to original order
-    vstar1 = np.zeros_like(vstar1_sorted)
-    vstar1[sort_order] = vstar1_sorted
-
-    return vstar1
-
-
-def bivariate_quantiles(
-    y: np.ndarray, u: np.ndarray, n_nodes: int = 32, verbose: bool = False
-) -> np.ndarray:
-    """Solve for the dual weights then evaluate bivariate quantiles.
-
-    Args:
-        y: Observations, shape ``(n, 2)``.
-        u: Query points in ``[0, 1]^2`` (shape ``(m, 2)``).
-        n_nodes: Number of Chebyshev nodes for the quadrature.
-        verbose: Print optimisation diagnostics when ``True``.
-
-    Returns:
-        Bivariate quantiles at ``u``.
-    """
-    v = solve_for_v_(y, n_nodes, verbose)
-    return bivariate_quantiles_v(y, u, v)
-
-
-def bivariate_ranks(
-    y: np.ndarray, n_nodes: int = 32, verbose: bool = False
-) -> np.ndarray:
-    """Compute ranks by first solving for the optimal weights ``v``.
-
-    Args:
-        y: Observations, shape ``(n, 2)``.
-        n_nodes: Number of Chebyshev nodes for the quadrature.
-        verbose: Print optimisation diagnostics when ``True``.
-
-    Returns:
-        Average ranks with shape ``(n, 2)``.
-    """
-    v = solve_for_v_(y, n_nodes, verbose)
-    return bivariate_ranks_v(y, v, n_nodes)

bs_python_utils-0.8.5/tests/test_bivariate_quantiles.py

@@ -1,80 +0,0 @@
-import numpy as np
-import pytest
-
-from bs_python_utils.bivariate_quantiles import (
-    _compute_ab,
-    _compute_u2_bounds,
-    bivariate_quantiles,
-    bivariate_quantiles_v,
-    bivariate_ranks,
-    bivariate_ranks_v,
-    solve_for_v_,
-)
-from bs_python_utils.bsnputils import bsgrid, ecdf, npmaxabs
-
-
-def test_bivariate_quantiles_recompute_matches():
-    rng = np.random.default_rng(123)
-    y = rng.normal(size=(64, 2))
-    nodes = 16
-
-    weights = solve_for_v_(y, n_nodes=nodes, verbose=False)
-
-    q = np.linspace(0.0, 1.0, 5)
-    u_points = bsgrid(q, q)
-
-    direct = bivariate_quantiles_v(y, u_points, weights)
-    via_solver = bivariate_quantiles(y, u_points, n_nodes=nodes, verbose=False)
-
-    ranks_direct = bivariate_ranks_v(y, weights, n_nodes=nodes)
-    ranks_via_solver = bivariate_ranks(y, n_nodes=nodes, verbose=False)
-
-    assert np.allclose(direct, via_solver)
-    assert np.allclose(ranks_direct, ranks_via_solver)
-    assert (
-        npmaxabs(ranks_direct - np.column_stack((ecdf(y[:, 0]), ecdf(y[:, 1])))) < 0.2
-    )
-
-
-def test_chunked_quantiles_matches_numpy_argmax():
-    rng = np.random.default_rng(7)
-    y = rng.normal(size=(10, 2))
-    v = rng.normal(size=10)
-    u = rng.uniform(0.0, 1.0, size=(6000, 2))
-
-    chunked = bivariate_quantiles_v(y, u, v)
-    expected = y[np.argmax(u @ y.T - v, axis=1)]
-
-    assert np.allclose(chunked, expected)
-
-
-def test_compute_ab_handles_ties():
-    y = np.array([[0.0, 1.0], [1.0, 1.0]])
-    v = np.array([0.2, -0.2])
-    a_mat, b_mat = _compute_ab(y, v)
-
-    assert np.isfinite(a_mat).all()
-    assert np.isfinite(b_mat).all()
-
-
-@pytest.mark.parametrize("k", [0, 1])
-def test_compute_u2_bounds_degenerate(k):
-    y = np.array([[0.0, 0.0], [1.0, 1.0]])
-    v = np.array([0.1, -0.1])
-    a_mat, b_mat = _compute_ab(y, v)
-    nodes = np.linspace(0.0, 1.0, 5)
-
-    left, right = _compute_u2_bounds(k, nodes, a_mat, b_mat)
-
-    assert np.isfinite(left).all()
-    assert np.isfinite(right).all()
-    assert (right >= left).all()
-
-
-def test_bivariate_ranks_zero_mass_returns_nan():
-    y = np.array([[0.0, 0.0], [1.0, 1.0]])
-    v = np.array([0.1, -0.1])
-    ranks = bivariate_ranks_v(y, v, n_nodes=8, presorted=True)
-
-    assert np.isnan(ranks[0]).all()
-    assert np.isfinite(ranks[1]).all()