bs-python-utils 0.5.3__tar.gz → 0.8.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/PKG-INFO +72 -30
  2. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/README.md +48 -8
  3. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/__init__.py +14 -6
  4. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bivariate_quantiles.py +83 -60
  5. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_altair.py +280 -106
  6. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_logging.py +13 -7
  7. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_mem.py +6 -3
  8. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_opt.py +145 -86
  9. bs_python_utils-0.8.1/bs_python_utils/bs_plots.py +62 -0
  10. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_seaborn.py +64 -30
  11. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_sparse_gaussian.py +13 -6
  12. bs_python_utils-0.8.1/bs_python_utils/bsmplutils.py +95 -0
  13. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bsnputils.py +176 -87
  14. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bssputils.py +15 -15
  15. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bsstats.py +47 -30
  16. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bsutils.py +49 -29
  17. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/chebyshev.py +67 -54
  18. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/distance_covariances.py +1 -1
  19. bs_python_utils-0.8.1/bs_python_utils/examples/examples_altair.py +225 -0
  20. {bs_python_utils-0.5.3/bs_python_utils → bs_python_utils-0.8.1/bs_python_utils/examples}/examples_mem.py +2 -1
  21. bs_python_utils-0.8.1/bs_python_utils/examples/examples_mpl.py +13 -0
  22. bs_python_utils-0.5.3/bs_python_utils/example_opt.py → bs_python_utils-0.8.1/bs_python_utils/examples/examples_opt.py +3 -2
  23. {bs_python_utils-0.5.3/bs_python_utils → bs_python_utils-0.8.1/bs_python_utils/examples}/examples_seaborn.py +14 -10
  24. {bs_python_utils-0.5.3/bs_python_utils → bs_python_utils-0.8.1/bs_python_utils/examples}/examples_sklearn.py +2 -1
  25. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/pandas_utils.py +4 -4
  26. bs_python_utils-0.8.1/bs_python_utils/sklearn_utils.py +123 -0
  27. bs_python_utils-0.8.1/bs_python_utils/streamlit_utils.py +100 -0
  28. bs_python_utils-0.8.1/bs_python_utils.egg-info/PKG-INFO +125 -0
  29. bs_python_utils-0.8.1/bs_python_utils.egg-info/SOURCES.txt +47 -0
  30. bs_python_utils-0.8.1/bs_python_utils.egg-info/dependency_links.txt +1 -0
  31. bs_python_utils-0.8.1/bs_python_utils.egg-info/requires.txt +19 -0
  32. bs_python_utils-0.8.1/bs_python_utils.egg-info/top_level.txt +1 -0
  33. bs_python_utils-0.8.1/pyproject.toml +36 -0
  34. bs_python_utils-0.8.1/setup.cfg +4 -0
  35. bs_python_utils-0.8.1/tests/test_bivariate_quantiles.py +80 -0
  36. bs_python_utils-0.8.1/tests/test_bs_opt.py +81 -0
  37. bs_python_utils-0.8.1/tests/test_bs_sparse_gaussian.py +17 -0
  38. bs_python_utils-0.8.1/tests/test_bsnputils.py +338 -0
  39. bs_python_utils-0.8.1/tests/test_bssputils.py +29 -0
  40. bs_python_utils-0.8.1/tests/test_bsstats.py +56 -0
  41. bs_python_utils-0.8.1/tests/test_bsutils.py +169 -0
  42. bs_python_utils-0.8.1/tests/test_chebyshev.py +256 -0
  43. bs_python_utils-0.8.1/tests/test_nptri.py +51 -0
  44. bs_python_utils-0.8.1/tests/test_random_draws.py +32 -0
  45. bs_python_utils-0.8.1/tests/test_sklearn_utils.py +52 -0
  46. bs_python_utils-0.5.3/bs_python_utils/bs_plots.py +0 -3
  47. bs_python_utils-0.5.3/bs_python_utils/bsmplutils.py +0 -36
  48. bs_python_utils-0.5.3/bs_python_utils/examples_altair.py +0 -182
  49. bs_python_utils-0.5.3/bs_python_utils/sklearn_utils.py +0 -76
  50. bs_python_utils-0.5.3/pyproject.toml +0 -103
  51. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/LICENSE +0 -0
  52. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/Timer.py +0 -0
  53. {bs_python_utils-0.5.3 → bs_python_utils-0.8.1}/bs_python_utils/bs_mathstr.py +0 -0
  54. {bs_python_utils-0.5.3/bs_python_utils → bs_python_utils-0.8.1/bs_python_utils/examples}/examples_distance_covariances.py +0 -0
@@ -1,28 +1,30 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: bs-python-utils
3
- Version: 0.5.3
4
- Summary: my Python utilities
5
- Home-page: https://github.com/bsalanie/bs-python-utils
6
- Author: Bernard Salanie
7
- Author-email: bsalanie@columbia.edu
8
- Requires-Python: >=3.8,<4.0
9
- Classifier: Programming Language :: Python :: 3
10
- Classifier: Programming Language :: Python :: 3.8
11
- Classifier: Programming Language :: Python :: 3.9
12
- Classifier: Programming Language :: Python :: 3.10
13
- Classifier: Programming Language :: Python :: 3.11
14
- Requires-Dist: altair (>=4.2.2,<5.0.0)
15
- Requires-Dist: altair-saver (>=0.5.0,<0.6.0)
16
- Requires-Dist: emcee (>=3.1.4,<4.0.0)
17
- Requires-Dist: griffe (>=0.32.3,<0.33.0)
18
- Requires-Dist: matplotlib (>=3.7.1,<4.0.0)
19
- Requires-Dist: scikit-learn (>=1.2.2,<2.0.0)
20
- Requires-Dist: seaborn (>=0.12.2,<0.13.0)
21
- Requires-Dist: statsmodels (>=0.13.5,<0.14.0)
22
- Requires-Dist: vega-datasets (>=0.9.0,<0.10.0)
23
- Project-URL: Documentation, https://bsalanie.github.io/bs-python-utils/
24
- Project-URL: Repository, https://github.com/bsalanie/bs-python-utils
3
+ Version: 0.8.1
4
+ Summary: Utilities programs for my Python code
5
+ Requires-Python: >=3.12
25
6
  Description-Content-Type: text/markdown
7
+ License-File: LICENSE
8
+ Requires-Dist: altair>=5.5.0
9
+ Requires-Dist: altair-saver>=0.5.0
10
+ Requires-Dist: emcee>=3.1.6
11
+ Requires-Dist: matplotlib>=3.10.6
12
+ Requires-Dist: mkdocs>=1.6.1
13
+ Requires-Dist: mkdocs-material>=9.6.20
14
+ Requires-Dist: mkdocstrings-python>=1.18.2
15
+ Requires-Dist: numpy>=2.3.3
16
+ Requires-Dist: pandas>=2.3.2
17
+ Requires-Dist: pre-commit>=4.3.0
18
+ Requires-Dist: pytest>=8.4.2
19
+ Requires-Dist: mypy
20
+ Requires-Dist: ruff>=0.14.0
21
+ Requires-Dist: scikit-learn>=1.7.2
22
+ Requires-Dist: scipy>=1.16.2
23
+ Requires-Dist: seaborn>=0.13.2
24
+ Requires-Dist: statsmodels>=0.14.5
25
+ Requires-Dist: streamlit>=1.49.1
26
+ Requires-Dist: vega-datasets>=0.9.0
27
+ Dynamic: license-file
26
28
 
27
29
  # bs-python-utils
28
30
 
@@ -32,35 +34,75 @@ Description-Content-Type: text/markdown
32
34
 
33
35
  [![Build status](https://img.shields.io/github/actions/workflow/status/bsalanie/bs-python-utils/main.yml?branch=main)](https://github.com/bsalanie/bs-python-utils/actions/workflows/main.yml?query=branch%3Amain) <!-- [![codecov](https://codecov.io/gh/bsalanie/bs-python-utils/branch/main/graph/badge.svg)](https://codecov.io/gh/bsalanie/bs-python-utils) --> <!-- [![Commit activity](https://img.shields.io/github/commit-activity/m/bsalanie/bs-python-utils)](https://img.shields.io/github/commit-activity/m/bsalanie/bs-python-utils) --> [![License](https://img.shields.io/github/license/bsalanie/bs-python-utils)](https://img.shields.io/github/license/bsalanie/bs-python-utils)
34
36
 
35
- My Python utilities.
37
+ **My Python utilities**.
36
38
 
37
- - **Github repository**: <https://github.com/bsalanie/bs-python-utils/>
38
- - **Documentation** <https://bsalanie.github.io/bs-python-utils/>
39
+ - **Github repository**: <https://github.com/bsalanie/bs-python-utils/>
40
+ - **Documentation** <https://bsalanie.github.io/bs-python-utils/>
39
41
 
40
42
  ### Release notes
41
- #### 0.5.3 (August 19, 2023)
42
- Improved the documentation.
43
+
44
+ #### 0.8.1 (October 18, 2025)
45
+ Deleted extra examples scripts.
46
+
47
+ #### 0.8 (October 14, 2025)
48
+ Switched to `uv` for project management. Used Codex to clean up the code and generate more tests. Added the `examples` subdirectory.
49
+
50
+ #### 0.7 (June 18, 2024)
51
+ Added `make_lowertri` in `bsnputils`, and some utility code in `bs_plots`.
52
+
53
+ #### 0.6.4 (November 26, 2023)
54
+ Moved to Altair 5.0; made timing optional in `bs_opt`.
55
+
56
+ #### 0.6.3 (November 26, 2023)
57
+ Updated statsmodels dependency.
58
+
59
+ #### 0.6.1 and 0.6.2 (November 18, 2023)
60
+ Added a Matplotlib plot for discrete choice models, a boxplot for Altair, and a density plot for Seaborn.
61
+
62
+ #### 0.6 (November 9, 2023)
63
+
64
+ Added matrix heatmap plot in Altair and Streamlit utility code.
65
+
66
+ #### 0.5.5 (October 23, 2023)
67
+
68
+ Added `fstring***` in `bsutils`.
69
+
43
70
  #### 0.5.2 (August 16, 2023)
71
+
44
72
  Added two interpolation routines from values at the Chebyshev nodes in `chebyshev` module.
73
+
45
74
  #### 0.5.1 (August 14, 2023)
46
- Added two functions to draw random samples in `bsstats` module.
75
+
76
+ Added two functions to draw random samples in `bsstats` module.
77
+
47
78
  #### 0.5 (August 7, 2023)
79
+
48
80
  Added 1-dimensional root finding in `chebyshev` module.
81
+
49
82
  #### 0.4.2 (August 2, 2023)
83
+
50
84
  Updated the docs.
85
+
51
86
  #### 0.4.1 (July 23, 2023)
87
+
52
88
  Only print if verbose in bivariate quantiles.
89
+
53
90
  #### 0.4 (July 22, 2023)
91
+
54
92
  Added bivariate quantiles and ranks à la optimal transportation.
93
+
55
94
  #### 0.3 (July 21, 2023)
95
+
56
96
  Added in Numpy utils a function to set upper and lower triangle to a scalar;
57
97
  and `minimize_free` in `bs_opt`.
58
98
 
59
99
  #### 0.2 (July 17, 2023)
100
+
60
101
  Fixed printing in accelerated gradient descent.
61
102
 
62
103
  #### 0.1 (June 20, 2023)
63
- Added `grid_function` and the `chebyshev` module.
104
+
105
+ Added `grid_function` and the `chebyshev` module.
64
106
 
65
107
  #### 0.0.6 (May 9, 2023)
66
108
 
@@ -6,35 +6,75 @@
6
6
 
7
7
  [![Build status](https://img.shields.io/github/actions/workflow/status/bsalanie/bs-python-utils/main.yml?branch=main)](https://github.com/bsalanie/bs-python-utils/actions/workflows/main.yml?query=branch%3Amain) <!-- [![codecov](https://codecov.io/gh/bsalanie/bs-python-utils/branch/main/graph/badge.svg)](https://codecov.io/gh/bsalanie/bs-python-utils) --> <!-- [![Commit activity](https://img.shields.io/github/commit-activity/m/bsalanie/bs-python-utils)](https://img.shields.io/github/commit-activity/m/bsalanie/bs-python-utils) --> [![License](https://img.shields.io/github/license/bsalanie/bs-python-utils)](https://img.shields.io/github/license/bsalanie/bs-python-utils)
8
8
 
9
- My Python utilities.
9
+ **My Python utilities**.
10
10
 
11
- - **Github repository**: <https://github.com/bsalanie/bs-python-utils/>
12
- - **Documentation** <https://bsalanie.github.io/bs-python-utils/>
11
+ - **Github repository**: <https://github.com/bsalanie/bs-python-utils/>
12
+ - **Documentation** <https://bsalanie.github.io/bs-python-utils/>
13
13
 
14
14
  ### Release notes
15
- #### 0.5.3 (August 19, 2023)
16
- Improved the documentation.
15
+
16
+ #### 0.8.1 (October 18, 2025)
17
+ Deleted extra examples scripts.
18
+
19
+ #### 0.8 (October 14, 2025)
20
+ Switched to `uv` for project management. Used Codex to clean up the code and generate more tests. Added the `examples` subdirectory.
21
+
22
+ #### 0.7 (June 18, 2024)
23
+ Added `make_lowertri` in `bsnputils`, and some utility code in `bs_plots`.
24
+
25
+ #### 0.6.4 (November 26, 2023)
26
+ Moved to Altair 5.0; made timing optional in `bs_opt`.
27
+
28
+ #### 0.6.3 (November 26, 2023)
29
+ Updated statsmodels dependency.
30
+
31
+ #### 0.6.1 and 0.6.2 (November 18, 2023)
32
+ Added a Matplotlib plot for discrete choice models, a boxplot for Altair, and a density plot for Seaborn.
33
+
34
+ #### 0.6 (November 9, 2023)
35
+
36
+ Added matrix heatmap plot in Altair and Streamlit utility code.
37
+
38
+ #### 0.5.5 (October 23, 2023)
39
+
40
+ Added `fstring***` in `bsutils`.
41
+
17
42
  #### 0.5.2 (August 16, 2023)
43
+
18
44
  Added two interpolation routines from values at the Chebyshev nodes in `chebyshev` module.
45
+
19
46
  #### 0.5.1 (August 14, 2023)
20
- Added two functions to draw random samples in `bsstats` module.
47
+
48
+ Added two functions to draw random samples in `bsstats` module.
49
+
21
50
  #### 0.5 (August 7, 2023)
51
+
22
52
  Added 1-dimensional root finding in `chebyshev` module.
53
+
23
54
  #### 0.4.2 (August 2, 2023)
55
+
24
56
  Updated the docs.
57
+
25
58
  #### 0.4.1 (July 23, 2023)
59
+
26
60
  Only print if verbose in bivariate quantiles.
61
+
27
62
  #### 0.4 (July 22, 2023)
63
+
28
64
  Added bivariate quantiles and ranks à la optimal transportation.
65
+
29
66
  #### 0.3 (July 21, 2023)
67
+
30
68
  Added in Numpy utils a function to set upper and lower triangle to a scalar;
31
69
  and `minimize_free` in `bs_opt`.
32
70
 
33
71
  #### 0.2 (July 17, 2023)
72
+
34
73
  Fixed printing in accelerated gradient descent.
35
74
 
36
75
  #### 0.1 (June 20, 2023)
37
- Added `grid_function` and the `chebyshev` module.
76
+
77
+ Added `grid_function` and the `chebyshev` module.
38
78
 
39
79
  #### 0.0.6 (May 9, 2023)
40
80
 
@@ -54,4 +94,4 @@ Satisfied mypy.
54
94
 
55
95
  #### 0.0.2 (April 24, 2023)
56
96
 
57
- Fixed main PyPI page.
97
+ Fixed main PyPI page.
@@ -1,6 +1,8 @@
1
- """This package contains a number of functions that I have found useful in my programming.
1
+ """This package contains a number of functions that I have found useful in my
2
+ programming.
2
3
 
3
- * `bsutils` has (inter alia) some I/O functions, error reporting, and $C^2$ extensions of log and exp
4
+ * `bsutils` has (inter alia) some I/O functions, error reporting, and $C^2$
5
+ extensions of log and exp
4
6
  * `bs_logging` has customized logging
5
7
  * `bs_mathstr` has Unicode for math strings
6
8
  * `bs_mem` reports memory usage
@@ -9,12 +11,18 @@
9
11
  * `bssputils` has Scipy functions
10
12
  * `sklearn_utils` has Sklearn functions
11
13
  * `pandas_utils` has Pandas functions
12
- * `bsstats` has TSLS code; nonparametric and flexible estimation; and code to draw random samples
14
+ * `bsstats` has TSLS code; nonparametric and flexible estimation; and code to
15
+ draw random samples
13
16
  * `bs_opt` interfaces with `scipy.optimize`
14
- * `bs_sparse_gaussian` uses sparse integration to evaluate $E f(X)$ for $X \\simeq N(0,1)$
17
+ * `bs_sparse_gaussian` uses sparse integration to evaluate $E f(X)$ for
18
+ $X \\simeq N(0,1)$
15
19
  * `chebyshev` has Chebyshev interpolation and integration in dimensions 1 and 2
16
20
  * `distance_covariances` has measures of nonlinear dependence between random variables
17
- * `bivariate_quantiles` computes quantiles and ranks for 2-dimensional random variables à 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).
21
+ * `bivariate_quantiles` computes quantiles and ranks for 2-dimensional random
22
+ variables à la [Chernozhukov-Galichon-Hallin-Henry (*Ann. Stats.* 2017)](
23
+ https://projecteuclid.org/journals/annals-of-statistics/volume-45/issue-1/
24
+ MongeKantorovich-depth-quantiles-ranks-and-signs/10.1214/16-AOS1450.full).
18
25
 
19
- * `bs_plots` gathers plotting routines in Matplotlib, Seaborn and Altair from `bsmplutils`, `bs_seaborn`, and `bs_altair`.
26
+ * `bs_plots` gathers plotting routines in Matplotlib, Seaborn and Altair from
27
+ `bsmplutils`, `bs_seaborn`, and `bs_altair`.
20
28
  """
@@ -1,5 +1,8 @@
1
1
  """This takes in observations of a bivariate random variable `y`
2
- 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).
2
+ and computes vector quantiles and vector ranks à la
3
+ [Chernozhukov-Galichon-Hallin-Henry (*Ann. Stats.* 2017)](
4
+ https://projecteuclid.org/journals/annals-of-statistics/volume-45/
5
+ issue-1/MongeKantorovich-depth-quantiles-ranks-and-signs/10.1214/16-AOS1450.full).
3
6
 
4
7
 
5
8
  Note:
@@ -7,16 +10,19 @@ Note:
7
10
 
8
11
  The sequence of steps is as follows:
9
12
 
10
- * choose a number of Chebyshev nodes for numerical integration and optimize the weights: `v = solve_for_v(y, n_nodes)`
13
+ * choose a number of Chebyshev nodes for numerical integration and optimize
14
+ the weights: `v = solve_for_v(y, n_nodes)`
11
15
  * to obtain the $(u_1,u_2)$ quantiles for $(u_1, u_2)\\in [0,1]$, run
12
16
  `qtiles_y = bivariate_quantiles_v(y, v, u1, u2)`
13
- * to compute the vector ranks for all points in the sample (the barycenters of the cells in the power diagram):
17
+ * to compute the vector ranks for all points in the sample (the barycenters
18
+ of the cells in the power diagram):
14
19
  `ranks_y = bivariate_ranks_v(y, v, n_nodes)`
15
20
 
16
21
  Steps 1 and 2 can be combined: `qtiles_y = bivariate_quantiles(y, v, u1, u2, n_nodes)`
17
22
 
18
23
  Steps 1 and 3 can be combined: `ranks_y = bivariate_ranks(y, n_nodes)`
19
24
  """
25
+
20
26
  from typing import cast
21
27
 
22
28
  import numpy as np
@@ -28,62 +34,55 @@ from bs_python_utils.chebyshev import Interval, cheb_get_nodes_1d
28
34
 
29
35
 
30
36
  def _compute_ab(y_sorted: np.ndarray, v_sorted: np.ndarray) -> TwoArrays:
31
- """evaluates the A and B matrices
32
-
33
- Args:
34
- y_sorted: the observations, an `(n,2)` matrix sorted by increasing `y[:, 1]`
35
- v_sorted: the weights, an `n`-vector in the same sort order
36
-
37
- Returns:
38
- two `(n, n)` matrices
39
-
40
- Warning:
41
- the user must make sure that `y` and `v` have been sorted beforehand
42
- """
43
- n = v_sorted.size
37
+ """Build the `A` and `B` matrices used in the dual optimisation."""
44
38
  y1 = y_sorted[:, 0]
45
39
  dy1 = np.subtract.outer(y1, y1)
46
40
  y2 = y_sorted[:, 1]
47
41
  dy2 = np.subtract.outer(y2, y2)
48
- np.fill_diagonal(dy2, np.ones(n)) # to avoid division by 0
42
+ np.fill_diagonal(dy2, 1.0)
49
43
  dv = np.subtract.outer(v_sorted, v_sorted)
50
- a_mat = dy1.T / dy2
51
- b_mat = dv.T / dy2
44
+ with np.errstate(divide="ignore", invalid="ignore"):
45
+ a_mat = np.divide(dy1.T, dy2, where=np.abs(dy2) > 1e-12)
46
+ b_mat = np.divide(dv.T, dy2, where=np.abs(dy2) > 1e-12)
47
+ a_mat = np.nan_to_num(a_mat, nan=0.0, posinf=0.0, neginf=0.0)
48
+ b_mat = np.nan_to_num(b_mat, nan=0.0, posinf=0.0, neginf=0.0)
52
49
  return a_mat, b_mat
53
50
 
54
51
 
55
52
  def _compute_u2_bounds(
56
53
  k: int, u1: np.ndarray, a_mat: np.ndarray, b_mat: np.ndarray
57
54
  ) -> TwoArrays:
58
- """for given u1, calculates the bounds on u2 that make k the chosen observation
59
-
60
- Args:
61
- k: an integer between 0 and (n-1)
62
- u1: a vector of size `m`
63
- a_mat: the `A` matrix of size `(n,n)`
64
- b_mat: the `B` matrix of size `(n,n)`
65
-
66
- Returns:
67
- two vectors of size `m`, the left and right bounds in [0,1]
68
- """
55
+ """Return the admissible interval of ``u2`` that selects index ``k``."""
69
56
  n = a_mat.shape[0]
70
57
  m = u1.size
71
58
  if k == 0:
72
59
  left_bound = np.zeros(m)
73
60
  a_right = a_mat[0, 1:]
74
61
  b_right = b_mat[0, 1:]
75
- right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
62
+ if a_right.size:
63
+ right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
64
+ else:
65
+ right_bound = np.ones(m)
76
66
  elif 1 <= k < n - 1:
77
67
  a_left = a_mat[k, :k]
78
68
  b_left = b_mat[k, :k]
79
- left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
69
+ if a_left.size:
70
+ left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
71
+ else:
72
+ left_bound = np.zeros(m)
80
73
  a_right = a_mat[k, (k + 1) :]
81
74
  b_right = b_mat[k, (k + 1) :]
82
- right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
75
+ if a_right.size:
76
+ right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
77
+ else:
78
+ right_bound = np.ones(m)
83
79
  elif k == n - 1:
84
80
  a_left = a_mat[-1, :-1]
85
81
  b_left = b_mat[-1, :-1]
86
- left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
82
+ if a_left.size:
83
+ left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
84
+ else:
85
+ left_bound = np.zeros(m)
87
86
  right_bound = np.ones(m)
88
87
  else:
89
88
  bs_error_abort(f"{k=} is not compatible with {n=}")
@@ -94,34 +93,45 @@ def _compute_u2_bounds(
94
93
 
95
94
 
96
95
  def bivariate_quantiles_v(y: np.ndarray, u: np.ndarray, v: np.ndarray) -> np.ndarray:
97
- """computes the vector quantiles of `y` at values `u`, given the converged `v`
96
+ """Evaluate vector quantiles for a given set of dual weights.
98
97
 
99
98
  Args:
100
- y: the observations, an `(n,2)` matrix
101
- u: the values where we want the quantiles, an `(m,2)` matrix in $[0,1]$
102
- v: the converged values of the weights, an `n`-vector
99
+ y: Observations with shape ``(n, 2)``.
100
+ u: Evaluation points in ``[0, 1]^2`` (shape ``(m, 2)``).
101
+ v: Dual weights solving the optimal transport problem (length ``n``).
103
102
 
104
103
  Returns:
105
- an `(m,2)` matrix with the quantiles of `y` at the values `u`
104
+ Array of quantile locations with shape ``(m, 2)``.
106
105
  """
107
- net_val = u @ y.T - v
108
- k_max = np.argmax(net_val, 1)
109
- return cast(np.ndarray, y[k_max])
106
+ u = np.atleast_2d(u)
107
+ if u.shape[1] != 2:
108
+ bs_error_abort("u must have two columns")
109
+ m = u.shape[0]
110
+ q = np.empty((m, 2))
111
+ block = max(1, min(m, 5_000))
112
+ for start in range(0, m, block):
113
+ stop = min(start + block, m)
114
+ chunk = u[start:stop]
115
+ net_val = chunk @ y.T - v
116
+ k_max = np.argmax(net_val, axis=1)
117
+ q[start:stop] = y[k_max]
118
+ return cast(np.ndarray, q)
110
119
 
111
120
 
112
121
  def bivariate_ranks_v(
113
122
  y: np.ndarray, v: np.ndarray, n_nodes: int = 32, presorted: bool = False
114
123
  ) -> np.ndarray:
115
- """computes the vector ranks of `y`, given the converged `v`
124
+ """Compute the barycentric ranks of each observation given optimal weights.
116
125
 
117
126
  Args:
118
- y: the observations, an `(n,2)` matrix
119
- v: the converged values of the weights, an `n`-vector
120
- n_nodes: the number of nodes for Chebyshev integration
121
- presorted: if `True`, then `y` and `v` are sorted by increasing `y[:, 1]`.
127
+ y: Observations with shape ``(n, 2)``.
128
+ v: Dual weights returned by ``solve_for_v_``.
129
+ n_nodes: Number of Chebyshev nodes used in the quadrature.
130
+ presorted: Set to ``True`` when ``y``/``v`` are pre-sorted by the
131
+ second coordinate.
122
132
 
123
133
  Returns:
124
- an `(n,2)` matrix with the average ranks of `y`
134
+ Array of average ranks (shape ``(n, 2)``) with ``nan`` for zero-mass cells.
125
135
  """
126
136
  n, d = y.shape
127
137
 
@@ -151,6 +161,9 @@ def bivariate_ranks_v(
151
161
  right_bounds * right_bounds - left_bounds * left_bounds, 0.0
152
162
  )
153
163
  prob_k = pos_diffs @ u1_weights
164
+ if prob_k <= 1e-12:
165
+ average_ranks[sort_order[k], :] = np.array([np.nan, np.nan])
166
+ continue
154
167
  average_ranks[sort_order[k], 0] = ((u1_nodes * pos_diffs) @ u1_weights) / prob_k
155
168
  average_ranks[sort_order[k], 1] = ((pos_diffs_sq @ u1_weights) / 2.0) / prob_k
156
169
 
@@ -213,6 +226,17 @@ def _grad(v_sorted: np.ndarray, args: list):
213
226
 
214
227
 
215
228
  def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.ndarray:
229
+ """Solve the dual optimisation to obtain the optimal weights ``v``.
230
+
231
+ Args:
232
+ y: Observations with shape ``(n, 2)``.
233
+ n_nodes: Number of Chebyshev nodes for the quadrature.
234
+ verbose: Print optimisation diagnostics when ``True``.
235
+
236
+ Returns:
237
+ Array of length ``n`` containing the optimal weights (including the
238
+ residual term).
239
+ """
216
240
  n, d = y.shape
217
241
 
218
242
  if d != 2:
@@ -250,17 +274,16 @@ def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.
250
274
  def bivariate_quantiles(
251
275
  y: np.ndarray, u: np.ndarray, n_nodes: int = 32, verbose: bool = False
252
276
  ) -> np.ndarray:
253
- """computes the bivariate quantiles of `y` at the quantiles `u`
277
+ """Solve for the dual weights then evaluate bivariate quantiles.
254
278
 
255
279
  Args:
256
- y: the observations, an `(n, 2)` matrix
257
- u: the quantiles at which to compute the bivariate quantiles,
258
- an `(m, 2)` matrix
259
- n_nodes: the number of nodes to use for the quadrature
260
- verbose: if `True`, print some information
280
+ y: Observations, shape ``(n, 2)``.
281
+ u: Query points in ``[0, 1]^2`` (shape ``(m, 2)``).
282
+ n_nodes: Number of Chebyshev nodes for the quadrature.
283
+ verbose: Print optimisation diagnostics when ``True``.
261
284
 
262
285
  Returns:
263
- an `(m, 2)` matrix of bivariate quantiles
286
+ Bivariate quantiles at ``u``.
264
287
  """
265
288
  v = solve_for_v_(y, n_nodes, verbose)
266
289
  return bivariate_quantiles_v(y, u, v)
@@ -269,15 +292,15 @@ def bivariate_quantiles(
269
292
  def bivariate_ranks(
270
293
  y: np.ndarray, n_nodes: int = 32, verbose: bool = False
271
294
  ) -> np.ndarray:
272
- """computes the bivariate ranks of `y`
295
+ """Compute ranks by first solving for the optimal weights ``v``.
273
296
 
274
297
  Args:
275
- y: the observations, an `(n, 2)` matrix
276
- n_nodes: the number of nodes to use for the quadrature
277
- verbose: if `True`, print some information
298
+ y: Observations, shape ``(n, 2)``.
299
+ n_nodes: Number of Chebyshev nodes for the quadrature.
300
+ verbose: Print optimisation diagnostics when ``True``.
278
301
 
279
302
  Returns:
280
- the `(n, 2)` matrix of bivariate average ranks
303
+ Average ranks with shape ``(n, 2)``.
281
304
  """
282
305
  v = solve_for_v_(y, n_nodes, verbose)
283
306
  return bivariate_ranks_v(y, v, n_nodes)