bs-python-utils 0.8.3__py3-none-any.whl → 0.9__py3-none-any.whl

bs_python_utils/bivariate_quantiles.py

@@ -10,11 +10,10 @@ Note:
 
 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
+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)`
-* to compute the vector ranks for all points in the sample (the barycenters
+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)`
 
@@ -33,200 +32,131 @@ 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]
+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_sorted[:, 1]
+    y2 = y[:, 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
+    dy2 = dy2.T
+    a_mat = np.divide(dy1, dy2)
+    return a_mat, dy2
 
 
-def _compute_u2_bounds(
-    k: int, u1: np.ndarray, a_mat: np.ndarray, b_mat: np.ndarray
+def _compute_m_M(
+    v: np.ndarray, a_mat: np.ndarray, dy2: np.ndarray, tau1_nodes: 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:
+    """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)``.
-        u: Evaluation points in ``[0, 1]^2`` (shape ``(m, 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)``.
     """
-    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]
+    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 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`
+    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_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
+        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
+        the value of the expectation and perhaps its gradient and the bivariate ranks
     """
-    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)
-
+    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):
-        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]
+        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
+        return obj_val, grad_val, bivrank
     else:
-        return cast(float, obj_val)
+        return obj_val
 
 
-def _obj(v_sorted: np.ndarray, args: list):
-    return _objgrad(v_sorted, args)
+def _obj(v: np.ndarray, args: list):
+    return _objgrad(v, args)
 
 
-def _grad(v_sorted: np.ndarray, args: list):
-    res_objg = cast(tuple[float, np.ndarray], _objgrad(v_sorted, args, gr=True))
+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[3]
+    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) -> np.ndarray:
-    """Solve the dual optimisation to obtain the optimal weights ``v``.
+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)``.
@@ -236,22 +166,21 @@ def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.
     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).
     """
-    n, d = y.shape
+    d = y.shape[1]
 
     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)
+    v0 = np.mean(y[:-1, :], 1)
 
     interval01 = Interval(0.0, 1.0)
-    u1_nodes, u1_weights = cheb_get_nodes_1d(interval01, n_nodes)
+    tau1_nodes, tau1_weights = cheb_get_nodes_1d(interval01, n_nodes)
 
-    argsog = [y_sorted, u1_nodes, u1_weights, verbose]
+    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:
@@ -262,45 +191,48 @@ def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.
     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
+    _, _, 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_quantiles(
-    y: np.ndarray, u: np.ndarray, n_nodes: int = 32, verbose: bool = False
+def bivariate_ranks(
+    y: np.ndarray,
+    n_nodes: int = 32,
+    verbose: bool = False,
 ) -> np.ndarray:
-    """Solve for the dual weights then evaluate bivariate quantiles.
+    """Compute the barycentric ranks of each observation.
 
     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``.
+        y: Observations with shape ``(n, 2)``.
+        n_nodes: Number of Chebyshev nodes used in the quadrature.
+        verbose: Print diagnostics when ``True``.
 
     Returns:
-        Bivariate quantiles at ``u``.
+        Array of average ranks (shape ``(n, 2)``) with ``nan`` for zero-mass cells.
     """
-    v = solve_for_v_(y, n_nodes, verbose)
-    return bivariate_quantiles_v(y, u, v)
+    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_ranks(
-    y: np.ndarray, n_nodes: int = 32, verbose: bool = False
+
+def bivariate_quantiles(
+    y: np.ndarray, tau: np.ndarray, n_nodes: int = 32, verbose: bool = False
 ) -> np.ndarray:
-    """Compute ranks by first solving for the optimal weights ``v``.
+    """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:
-        Average ranks with shape ``(n, 2)``.
+        Bivariate quantiles at ``u``.
     """
-    v = solve_for_v_(y, n_nodes, verbose)
-    return bivariate_ranks_v(y, v, n_nodes)
+    v, _ = _solve_for_v(y, n_nodes, verbose)
+    return bivariate_quantiles_v(y, tau, v)

bs_python_utils/bs_logging.py

@@ -46,7 +46,7 @@ class ColorFormatter(logging.Formatter):
 
 def get_logger(
     name: str,
-    level=logging.INFO,
+    level: int = logging.INFO,
     log_to_file: bool = False,
 ) -> logging.Logger:
     """Create and return a logger with colored output.

{bs_python_utils-0.8.3.dist-info → bs_python_utils-0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bs-python-utils
-Version: 0.8.3
+Version: 0.9
 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.8.5 (January 25, 2026)
+Added colored logger.
+
 #### 0.8.2 (October 26, 2025)
 Included grids for sparse Gaussian integration.
 

{bs_python_utils-0.8.3.dist-info → bs_python_utils-0.9.dist-info}/RECORD

@@ -1,8 +1,8 @@
 bs_python_utils/Timer.py,sha256=Rwj0Gec6VUq9Y6Rmykc5wL3Xd9dIj7Un_QgbJc2cTxo,1966
 bs_python_utils/__init__.py,sha256=XJ62LhKlvBaZH5ItthEpR7W6Sb9FEKi1d-6TZhP8DLc,1353
-bs_python_utils/bivariate_quantiles.py,sha256=Z-KGBUE6TlgayemskZb9C91ISSDwf3_PPoGMlrBr9mo,10231
+bs_python_utils/bivariate_quantiles.py,sha256=NaoNvt0QHEssFvju9YTD6dcHycd8NyWRG3cJADuk_W0,7859
 bs_python_utils/bs_altair.py,sha256=4-MyCBrpuem48BhZV0dBhPZBgrxyl9PivdY3gaSUdN8,35127
-bs_python_utils/bs_logging.py,sha256=qRfb1B6LLLsgMTl4hMsUnfpE3H_-CKlHXawtEhZUT_Q,5043
+bs_python_utils/bs_logging.py,sha256=CGi_pLeBSVd7FXsoafCI9xDu9eXLMnJebgY8Ku3-qP8,5050
 bs_python_utils/bs_mathstr.py,sha256=zW2ECxDD-29H81s_PNE1UKbIXJZtuhReJwcYim_4YTg,3595
 bs_python_utils/bs_mem.py,sha256=WaTikTlMBDBAakS6ZBh9EZuLYtNO_1XqcFXL33ws4vY,4668
 bs_python_utils/bs_opt.py,sha256=iq1fzPFm20kylg21uWXVDcfQWY7c7OpE9OYNBBaxbYU,17797
@@ -42,8 +42,8 @@ bs_python_utils/examples/examples_opt.py,sha256=pcWuTzOtTDWVJtKJJSUmaRIaPwDnAY3G
 bs_python_utils/examples/examples_seaborn.py,sha256=pwGx0pDVNGEtLMKgTqOEzoXkOsXgRND8-cJIxJQEY5o,783
 bs_python_utils/examples/examples_sklearn.py,sha256=TAnPRdRIiokrWBUUqVKkji5KXGCXyy_AjTIAEcmhtIY,792
 bs_python_utils/examples/examples_sparse_gaussian.py,sha256=xz3U2z97h-sxoxFmMDbtETRynVcGDUVahr9HCK42BtU,911
-bs_python_utils-0.8.3.dist-info/licenses/LICENSE,sha256=J03S1L_wQN4qquuluScbRUE09on78nCnNRpaUpOvHcw,1073
-bs_python_utils-0.8.3.dist-info/METADATA,sha256=ronR1HpDM8IF8ELj9QZqqoMDoSGzWj443RJF6RC35dA,3933
-bs_python_utils-0.8.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-bs_python_utils-0.8.3.dist-info/top_level.txt,sha256=4sbCpeuqAsdWmoAmZju1vf1ZFMJrSx310eS5SCyrAeQ,16
-bs_python_utils-0.8.3.dist-info/RECORD,,
+bs_python_utils-0.9.dist-info/licenses/LICENSE,sha256=J03S1L_wQN4qquuluScbRUE09on78nCnNRpaUpOvHcw,1073
+bs_python_utils-0.9.dist-info/METADATA,sha256=EuB4Uk_EJb6_RmwvLN6G7i9NRh3vHCROA9lbC-jsGCg,4015
+bs_python_utils-0.9.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+bs_python_utils-0.9.dist-info/top_level.txt,sha256=4sbCpeuqAsdWmoAmZju1vf1ZFMJrSx310eS5SCyrAeQ,16
+bs_python_utils-0.9.dist-info/RECORD,,

{bs_python_utils-0.8.3.dist-info → bs_python_utils-0.9.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.2)
+Generator: setuptools (82.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any