bs-python-utils 0.5.2__tar.gz → 0.8.0__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 (55) hide show
  1. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/PKG-INFO +65 -28
  2. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/README.md +45 -6
  3. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/Timer.py +16 -11
  4. bs_python_utils-0.8.0/bs_python_utils/__init__.py +20 -0
  5. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bivariate_quantiles.py +98 -87
  6. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_altair.py +288 -122
  7. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_logging.py +15 -23
  8. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_mathstr.py +26 -35
  9. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_mem.py +28 -22
  10. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_opt.py +155 -102
  11. bs_python_utils-0.8.0/bs_python_utils/bs_plots.py +62 -0
  12. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_seaborn.py +74 -40
  13. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bs_sparse_gaussian.py +8 -5
  14. bs_python_utils-0.8.0/bs_python_utils/bsmplutils.py +94 -0
  15. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bsnputils.py +234 -115
  16. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bssputils.py +19 -16
  17. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bsstats.py +59 -42
  18. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/bsutils.py +78 -46
  19. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/chebyshev.py +88 -54
  20. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/distance_covariances.py +7 -2
  21. bs_python_utils-0.8.0/bs_python_utils/examples/examples_altair.py +225 -0
  22. {bs_python_utils-0.5.2/bs_python_utils → bs_python_utils-0.8.0/bs_python_utils/examples}/examples_mem.py +2 -1
  23. bs_python_utils-0.8.0/bs_python_utils/examples/examples_mpl.py +13 -0
  24. bs_python_utils-0.5.2/bs_python_utils/example_opt.py → bs_python_utils-0.8.0/bs_python_utils/examples/examples_opt.py +3 -2
  25. {bs_python_utils-0.5.2/bs_python_utils → bs_python_utils-0.8.0/bs_python_utils/examples}/examples_seaborn.py +14 -10
  26. {bs_python_utils-0.5.2/bs_python_utils → bs_python_utils-0.8.0/bs_python_utils/examples}/examples_sklearn.py +2 -1
  27. {bs_python_utils-0.5.2 → bs_python_utils-0.8.0}/bs_python_utils/pandas_utils.py +13 -7
  28. bs_python_utils-0.8.0/bs_python_utils/sklearn_utils.py +123 -0
  29. bs_python_utils-0.8.0/bs_python_utils/streamlit_utils.py +100 -0
  30. bs_python_utils-0.8.0/bs_python_utils.egg-info/PKG-INFO +118 -0
  31. bs_python_utils-0.8.0/bs_python_utils.egg-info/SOURCES.txt +46 -0
  32. bs_python_utils-0.8.0/bs_python_utils.egg-info/dependency_links.txt +1 -0
  33. bs_python_utils-0.8.0/bs_python_utils.egg-info/requires.txt +17 -0
  34. bs_python_utils-0.8.0/bs_python_utils.egg-info/top_level.txt +1 -0
  35. bs_python_utils-0.8.0/pyproject.toml +31 -0
  36. bs_python_utils-0.8.0/setup.cfg +4 -0
  37. bs_python_utils-0.8.0/tests/test_bivariate_quantiles.py +80 -0
  38. bs_python_utils-0.8.0/tests/test_bs_opt.py +81 -0
  39. bs_python_utils-0.8.0/tests/test_bs_sparse_gaussian.py +17 -0
  40. bs_python_utils-0.8.0/tests/test_bsnputils.py +338 -0
  41. bs_python_utils-0.8.0/tests/test_bssputils.py +29 -0
  42. bs_python_utils-0.8.0/tests/test_bsstats.py +56 -0
  43. bs_python_utils-0.8.0/tests/test_bsutils.py +169 -0
  44. bs_python_utils-0.8.0/tests/test_chebyshev.py +256 -0
  45. bs_python_utils-0.8.0/tests/test_nptri.py +51 -0
  46. bs_python_utils-0.8.0/tests/test_random_draws.py +32 -0
  47. bs_python_utils-0.8.0/tests/test_sklearn_utils.py +52 -0
  48. bs_python_utils-0.5.2/LICENSE +0 -21
  49. bs_python_utils-0.5.2/bs_python_utils/__init__.py +0 -0
  50. bs_python_utils-0.5.2/bs_python_utils/bs_plots.py +0 -2
  51. bs_python_utils-0.5.2/bs_python_utils/bsmplutils.py +0 -34
  52. bs_python_utils-0.5.2/bs_python_utils/examples_altair.py +0 -182
  53. bs_python_utils-0.5.2/bs_python_utils/sklearn_utils.py +0 -74
  54. bs_python_utils-0.5.2/pyproject.toml +0 -103
  55. {bs_python_utils-0.5.2/bs_python_utils → bs_python_utils-0.8.0/bs_python_utils/examples}/examples_distance_covariances.py +0 -0
@@ -1,28 +1,26 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: bs-python-utils
3
- Version: 0.5.2
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.0
4
+ Summary: Add your description here
5
+ Requires-Python: >=3.12
25
6
  Description-Content-Type: text/markdown
7
+ Requires-Dist: altair>=5.5.0
8
+ Requires-Dist: altair-saver>=0.5.0
9
+ Requires-Dist: emcee>=3.1.6
10
+ Requires-Dist: matplotlib>=3.10.6
11
+ Requires-Dist: mkdocs>=1.6.1
12
+ Requires-Dist: mkdocs-material>=9.6.20
13
+ Requires-Dist: mkdocstrings-python>=1.18.2
14
+ Requires-Dist: numpy>=2.3.3
15
+ Requires-Dist: pandas>=2.3.2
16
+ Requires-Dist: pre-commit>=4.3.0
17
+ Requires-Dist: pytest>=8.4.2
18
+ Requires-Dist: scikit-learn>=1.7.2
19
+ Requires-Dist: scipy>=1.16.2
20
+ Requires-Dist: seaborn>=0.13.2
21
+ Requires-Dist: statsmodels>=0.14.5
22
+ Requires-Dist: streamlit>=1.49.1
23
+ Requires-Dist: vega-datasets>=0.9.0
26
24
 
27
25
  # bs-python-utils
28
26
 
@@ -32,33 +30,72 @@ Description-Content-Type: text/markdown
32
30
 
33
31
  [![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
32
 
35
- My Python utilities.
33
+ **My Python utilities**.
36
34
 
37
- - **Github repository**: <https://github.com/bsalanie/bs-python-utils/>
38
- - **Documentation** <https://bsalanie.github.io/bs-python-utils/>
35
+ - **Github repository**: <https://github.com/bsalanie/bs-python-utils/>
36
+ - **Documentation** <https://bsalanie.github.io/bs-python-utils/>
39
37
 
40
38
  ### Release notes
39
+
40
+ #### 0.8 (October 14, 2025)
41
+ Switched to `uv` for project management. Used Codex to clean up the code and generate more tests. Added the `examples` subdirectory.
42
+
43
+ #### 0.7 (June 18, 2024)
44
+ Added `make_lowertri` in `bsnputils`, and some utility code in `bs_plots`.
45
+
46
+ #### 0.6.4 (November 26, 2023)
47
+ Moved to Altair 5.0; made timing optional in `bs_opt`.
48
+
49
+ #### 0.6.3 (November 26, 2023)
50
+ Updated statsmodels dependency.
51
+
52
+ #### 0.6.1 and 0.6.2 (November 18, 2023)
53
+ Added a Matplotlib plot for discrete choice models, a boxplot for Altair, and a density plot for Seaborn.
54
+
55
+ #### 0.6 (November 9, 2023)
56
+
57
+ Added matrix heatmap plot in Altair and Streamlit utility code.
58
+
59
+ #### 0.5.5 (October 23, 2023)
60
+
61
+ Added `fstring***` in `bsutils`.
62
+
41
63
  #### 0.5.2 (August 16, 2023)
64
+
42
65
  Added two interpolation routines from values at the Chebyshev nodes in `chebyshev` module.
66
+
43
67
  #### 0.5.1 (August 14, 2023)
44
- Added two functions to draw random samples in `bsstats` module.
68
+
69
+ Added two functions to draw random samples in `bsstats` module.
70
+
45
71
  #### 0.5 (August 7, 2023)
72
+
46
73
  Added 1-dimensional root finding in `chebyshev` module.
74
+
47
75
  #### 0.4.2 (August 2, 2023)
76
+
48
77
  Updated the docs.
78
+
49
79
  #### 0.4.1 (July 23, 2023)
80
+
50
81
  Only print if verbose in bivariate quantiles.
82
+
51
83
  #### 0.4 (July 22, 2023)
84
+
52
85
  Added bivariate quantiles and ranks à la optimal transportation.
86
+
53
87
  #### 0.3 (July 21, 2023)
88
+
54
89
  Added in Numpy utils a function to set upper and lower triangle to a scalar;
55
90
  and `minimize_free` in `bs_opt`.
56
91
 
57
92
  #### 0.2 (July 17, 2023)
93
+
58
94
  Fixed printing in accelerated gradient descent.
59
95
 
60
96
  #### 0.1 (June 20, 2023)
61
- Added `grid_function` and the `chebyshev` module.
97
+
98
+ Added `grid_function` and the `chebyshev` module.
62
99
 
63
100
  #### 0.0.6 (May 9, 2023)
64
101
 
@@ -6,33 +6,72 @@
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
+
16
+ #### 0.8 (October 14, 2025)
17
+ Switched to `uv` for project management. Used Codex to clean up the code and generate more tests. Added the `examples` subdirectory.
18
+
19
+ #### 0.7 (June 18, 2024)
20
+ Added `make_lowertri` in `bsnputils`, and some utility code in `bs_plots`.
21
+
22
+ #### 0.6.4 (November 26, 2023)
23
+ Moved to Altair 5.0; made timing optional in `bs_opt`.
24
+
25
+ #### 0.6.3 (November 26, 2023)
26
+ Updated statsmodels dependency.
27
+
28
+ #### 0.6.1 and 0.6.2 (November 18, 2023)
29
+ Added a Matplotlib plot for discrete choice models, a boxplot for Altair, and a density plot for Seaborn.
30
+
31
+ #### 0.6 (November 9, 2023)
32
+
33
+ Added matrix heatmap plot in Altair and Streamlit utility code.
34
+
35
+ #### 0.5.5 (October 23, 2023)
36
+
37
+ Added `fstring***` in `bsutils`.
38
+
15
39
  #### 0.5.2 (August 16, 2023)
40
+
16
41
  Added two interpolation routines from values at the Chebyshev nodes in `chebyshev` module.
42
+
17
43
  #### 0.5.1 (August 14, 2023)
18
- Added two functions to draw random samples in `bsstats` module.
44
+
45
+ Added two functions to draw random samples in `bsstats` module.
46
+
19
47
  #### 0.5 (August 7, 2023)
48
+
20
49
  Added 1-dimensional root finding in `chebyshev` module.
50
+
21
51
  #### 0.4.2 (August 2, 2023)
52
+
22
53
  Updated the docs.
54
+
23
55
  #### 0.4.1 (July 23, 2023)
56
+
24
57
  Only print if verbose in bivariate quantiles.
58
+
25
59
  #### 0.4 (July 22, 2023)
60
+
26
61
  Added bivariate quantiles and ranks à la optimal transportation.
62
+
27
63
  #### 0.3 (July 21, 2023)
64
+
28
65
  Added in Numpy utils a function to set upper and lower triangle to a scalar;
29
66
  and `minimize_free` in `bs_opt`.
30
67
 
31
68
  #### 0.2 (July 17, 2023)
69
+
32
70
  Fixed printing in accelerated gradient descent.
33
71
 
34
72
  #### 0.1 (June 20, 2023)
35
- Added `grid_function` and the `chebyshev` module.
73
+
74
+ Added `grid_function` and the `chebyshev` module.
36
75
 
37
76
  #### 0.0.6 (May 9, 2023)
38
77
 
@@ -52,4 +91,4 @@ Satisfied mypy.
52
91
 
53
92
  #### 0.0.2 (April 24, 2023)
54
93
 
55
- Fixed main PyPI page.
94
+ Fixed main PyPI page.
@@ -1,5 +1,8 @@
1
1
  """
2
- utilities to time code
2
+ Utilities to time code:
3
+
4
+ * a `Timer` class that can be used as a context manager
5
+ * a `timeit` decorator for functions.
3
6
  """
4
7
 
5
8
  import time
@@ -28,20 +31,22 @@ class Timer:
28
31
  A timer that can be started, stopped, and reset as needed by the user.
29
32
  It keeps track of the total elapsed time in the `elapsed` attribute::
30
33
 
31
- with Timer() as t:
32
- ....
33
- print(f"... took {t.elapsed} seconds")
34
+ Examples:
35
+ >>> with Timer() as t:
36
+ >>> ....
37
+ >>> print(f"... took {t.elapsed} seconds")
34
38
 
35
39
  use `Timer(time.process_time)` to get only CPU time.
36
40
 
37
- can also do::
41
+ can also do:
38
42
 
39
- t = Timer()
40
- t.start()
41
- t.stop()
42
- t.start() # will add to the same counter
43
- t.stop()
44
- print(f"{t.elapsed} seconds total")
43
+ Examples:
44
+ >>> t = Timer()
45
+ >>> t.start()
46
+ >>> t.stop()
47
+ >>> t.start() # will add to the same counter
48
+ >>> t.stop()
49
+ >>> print(f"{t.elapsed} seconds total")
45
50
  """
46
51
 
47
52
  def __init__(self, func: Callable = time.perf_counter) -> None:
@@ -0,0 +1,20 @@
1
+ """This package contains a number of functions that I have found useful in my programming.
2
+
3
+ * `bsutils` has (inter alia) some I/O functions, error reporting, and $C^2$ extensions of log and exp
4
+ * `bs_logging` has customized logging
5
+ * `bs_mathstr` has Unicode for math strings
6
+ * `bs_mem` reports memory usage
7
+ * `Timer` has a `Timer` class to time code execution.
8
+ * `bsnputils` has Numpy functions
9
+ * `bssputils` has Scipy functions
10
+ * `sklearn_utils` has Sklearn functions
11
+ * `pandas_utils` has Pandas functions
12
+ * `bsstats` has TSLS code; nonparametric and flexible estimation; and code to draw random samples
13
+ * `bs_opt` interfaces with `scipy.optimize`
14
+ * `bs_sparse_gaussian` uses sparse integration to evaluate $E f(X)$ for $X \\simeq N(0,1)$
15
+ * `chebyshev` has Chebyshev interpolation and integration in dimensions 1 and 2
16
+ * `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).
18
+
19
+ * `bs_plots` gathers plotting routines in Matplotlib, Seaborn and Altair from `bsmplutils`, `bs_seaborn`, and `bs_altair`.
20
+ """
@@ -1,94 +1,83 @@
1
1
  """This takes in observations of a bivariate random variable `y`
2
- and computes vector quantiles and vector ranks.
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).
3
3
 
4
- The sequence of steps is as follows:
5
4
 
6
- 1. choose a number of Chebyshev nodes for numerical integration and
7
- optimize the weights: `v = solve_for_v(y, n_nodes)`
5
+ Note:
6
+ if the math looks strange in the documentation, just reload the page.
8
7
 
9
- 2a. to obtain the $(u_1,u_2)$ quantiles for $u_1, u_2\\in [0,1]$, run
10
- `qtiles_y = bivariate_quantiles_v(y, v, u1, u2)`
8
+ The sequence of steps is as follows:
11
9
 
12
- 2b. to compute the vector ranks for all points in the sample (the barycenters of the cells in the power diagram):
10
+ * choose a number of Chebyshev nodes for numerical integration and optimize the weights: `v = solve_for_v(y, n_nodes)`
11
+ * to obtain the $(u_1,u_2)$ quantiles for $(u_1, u_2)\\in [0,1]$, run
12
+ `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):
13
14
  `ranks_y = bivariate_ranks_v(y, v, n_nodes)`
14
15
 
15
- Steps 1 and 2a can be combined: `qtiles_y = bivariate_quantiles(y, v, u1, u2, n_nodes)`
16
+ Steps 1 and 2 can be combined: `qtiles_y = bivariate_quantiles(y, v, u1, u2, n_nodes)`
16
17
 
17
- Steps 1 and 2b can be combined: `ranks_y = bivariate_ranks(y, n_nodes)`
18
+ Steps 1 and 3 can be combined: `ranks_y = bivariate_ranks(y, n_nodes)`
18
19
  """
20
+
19
21
  from typing import cast
20
22
 
21
23
  import numpy as np
22
24
 
23
- from bs_python_utils.bs_opt import (
24
- minimize_free,
25
- print_optimization_results,
26
- )
27
- from bs_python_utils.bsnputils import (
28
- TwoArrays,
29
- npmaxabs,
30
- )
25
+ from bs_python_utils.bs_opt import minimize_free, print_optimization_results
26
+ from bs_python_utils.bsnputils import TwoArrays, npmaxabs
31
27
  from bs_python_utils.bsutils import bs_error_abort
32
28
  from bs_python_utils.chebyshev import Interval, cheb_get_nodes_1d
33
29
 
34
30
 
35
- def compute_ab_(y_sorted: np.ndarray, v_sorted: np.ndarray) -> TwoArrays:
36
- """evaluates the A and B matrices
37
-
38
- Args:
39
- y_sorted: the observations, an `(n,2)` matrix sorted by increasing `y[:, 1]`
40
- v_sorted: the weights, an `n`-vector in the same sort order
41
-
42
- Returns:
43
- two `(n, n)` matrices
44
-
45
- Warning:
46
- the user must make sure that `y` and `v` have been sorted beforehand
47
- """
48
- n = v_sorted.size
31
+ def _compute_ab(y_sorted: np.ndarray, v_sorted: np.ndarray) -> TwoArrays:
32
+ """Build the `A` and `B` matrices used in the dual optimisation."""
49
33
  y1 = y_sorted[:, 0]
50
34
  dy1 = np.subtract.outer(y1, y1)
51
35
  y2 = y_sorted[:, 1]
52
36
  dy2 = np.subtract.outer(y2, y2)
53
- np.fill_diagonal(dy2, np.ones(n)) # to avoid division by 0
37
+ np.fill_diagonal(dy2, 1.0)
54
38
  dv = np.subtract.outer(v_sorted, v_sorted)
55
- a_mat = dy1.T / dy2
56
- b_mat = dv.T / dy2
39
+ with np.errstate(divide="ignore", invalid="ignore"):
40
+ a_mat = np.divide(dy1.T, dy2, where=np.abs(dy2) > 1e-12)
41
+ b_mat = np.divide(dv.T, dy2, where=np.abs(dy2) > 1e-12)
42
+ a_mat = np.nan_to_num(a_mat, nan=0.0, posinf=0.0, neginf=0.0)
43
+ b_mat = np.nan_to_num(b_mat, nan=0.0, posinf=0.0, neginf=0.0)
57
44
  return a_mat, b_mat
58
45
 
59
46
 
60
- def compute_u2_bounds_(
47
+ def _compute_u2_bounds(
61
48
  k: int, u1: np.ndarray, a_mat: np.ndarray, b_mat: np.ndarray
62
49
  ) -> TwoArrays:
63
- """for given u1, calculates the bounds on u2 that make k the chosen observation
64
-
65
- Args:
66
- k: an integer between 0 and (n-1)
67
- u1: a vector of size `m`
68
- a_mat: the `A` matrix of size `(n,n)`
69
- b_mat: the `B` matrix of size `(n,n)`
70
-
71
- Returns:
72
- two vectors of size `m`, the left and right bounds in [0,1]
73
- """
50
+ """Return the admissible interval of ``u2`` that selects index ``k``."""
74
51
  n = a_mat.shape[0]
75
52
  m = u1.size
76
53
  if k == 0:
77
54
  left_bound = np.zeros(m)
78
55
  a_right = a_mat[0, 1:]
79
56
  b_right = b_mat[0, 1:]
80
- right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
57
+ if a_right.size:
58
+ right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
59
+ else:
60
+ right_bound = np.ones(m)
81
61
  elif 1 <= k < n - 1:
82
62
  a_left = a_mat[k, :k]
83
63
  b_left = b_mat[k, :k]
84
- left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
64
+ if a_left.size:
65
+ left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
66
+ else:
67
+ left_bound = np.zeros(m)
85
68
  a_right = a_mat[k, (k + 1) :]
86
69
  b_right = b_mat[k, (k + 1) :]
87
- right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
70
+ if a_right.size:
71
+ right_bound = np.min(np.outer(u1, a_right) - b_right, 1)
72
+ else:
73
+ right_bound = np.ones(m)
88
74
  elif k == n - 1:
89
75
  a_left = a_mat[-1, :-1]
90
76
  b_left = b_mat[-1, :-1]
91
- left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
77
+ if a_left.size:
78
+ left_bound = np.max(np.outer(u1, a_left) - b_left, 1)
79
+ else:
80
+ left_bound = np.zeros(m)
92
81
  right_bound = np.ones(m)
93
82
  else:
94
83
  bs_error_abort(f"{k=} is not compatible with {n=}")
@@ -99,34 +88,44 @@ def compute_u2_bounds_(
99
88
 
100
89
 
101
90
  def bivariate_quantiles_v(y: np.ndarray, u: np.ndarray, v: np.ndarray) -> np.ndarray:
102
- """computes the vector quantiles of `y` at values `u`, given the converged `v`
91
+ """Evaluate vector quantiles for a given set of dual weights.
103
92
 
104
93
  Args:
105
- y: the observations, an `(n,2)` matrix
106
- u: the values where we want the quantiles, an `(m,2)` matrix in $[0,1]$
107
- v: the converged values of the weights, an `n`-vector
94
+ y: Observations with shape ``(n, 2)``.
95
+ u: Evaluation points in ``[0, 1]^2`` (shape ``(m, 2)``).
96
+ v: Dual weights solving the optimal transport problem (length ``n``).
108
97
 
109
98
  Returns:
110
- an `(m,2)` matrix with the quantiles of `y` at the values `u`
99
+ Array of quantile locations with shape ``(m, 2)``.
111
100
  """
112
- net_val = u @ y.T - v
113
- k_max = np.argmax(net_val, 1)
114
- return cast(np.ndarray, y[k_max])
101
+ u = np.atleast_2d(u)
102
+ if u.shape[1] != 2:
103
+ bs_error_abort("u must have two columns")
104
+ m = u.shape[0]
105
+ q = np.empty((m, 2))
106
+ block = max(1, min(m, 5_000))
107
+ for start in range(0, m, block):
108
+ stop = min(start + block, m)
109
+ chunk = u[start:stop]
110
+ net_val = chunk @ y.T - v
111
+ k_max = np.argmax(net_val, axis=1)
112
+ q[start:stop] = y[k_max]
113
+ return cast(np.ndarray, q)
115
114
 
116
115
 
117
116
  def bivariate_ranks_v(
118
117
  y: np.ndarray, v: np.ndarray, n_nodes: int = 32, presorted: bool = False
119
118
  ) -> np.ndarray:
120
- """computes the vector ranks of `y`, given the converged `v`
119
+ """Compute the barycentric ranks of each observation given optimal weights.
121
120
 
122
121
  Args:
123
- y: the observations, an `(n,2)` matrix
124
- v: the converged values of the weights, an `n`-vector
125
- n_nodes: the number of nodes for Chebyshev integration
126
- presorted: if `True`, then `y` and `v` are sorted by increasing `y[:, 1]`.
122
+ y: Observations with shape ``(n, 2)``.
123
+ v: Dual weights returned by ``solve_for_v_``.
124
+ n_nodes: Number of Chebyshev nodes used in the quadrature.
125
+ presorted: Set to ``True`` when ``y``/``v`` are pre-sorted by the second coordinate.
127
126
 
128
127
  Returns:
129
- an `(n,2)` matrix with the average ranks of `y`
128
+ Array of average ranks (shape ``(n, 2)``) with ``nan`` for zero-mass cells.
130
129
  """
131
130
  n, d = y.shape
132
131
 
@@ -145,24 +144,27 @@ def bivariate_ranks_v(
145
144
  y_sorted = y[sort_order, :]
146
145
  v_sorted = v[sort_order]
147
146
 
148
- a_mat, b_mat = compute_ab_(y_sorted, v_sorted)
147
+ a_mat, b_mat = _compute_ab(y_sorted, v_sorted)
149
148
 
150
149
  average_ranks = np.zeros((n, 2))
151
150
 
152
151
  for k in range(n):
153
- left_bounds, right_bounds = compute_u2_bounds_(k, u1_nodes, a_mat, b_mat)
152
+ left_bounds, right_bounds = _compute_u2_bounds(k, u1_nodes, a_mat, b_mat)
154
153
  pos_diffs = np.maximum(right_bounds - left_bounds, 0.0)
155
154
  pos_diffs_sq = np.maximum(
156
155
  right_bounds * right_bounds - left_bounds * left_bounds, 0.0
157
156
  )
158
157
  prob_k = pos_diffs @ u1_weights
158
+ if prob_k <= 1e-12:
159
+ average_ranks[sort_order[k], :] = np.array([np.nan, np.nan])
160
+ continue
159
161
  average_ranks[sort_order[k], 0] = ((u1_nodes * pos_diffs) @ u1_weights) / prob_k
160
162
  average_ranks[sort_order[k], 1] = ((pos_diffs_sq @ u1_weights) / 2.0) / prob_k
161
163
 
162
164
  return average_ranks
163
165
 
164
166
 
165
- def objgrad_(
167
+ def _objgrad(
166
168
  v_sorted: np.ndarray, args: list, gr: bool = False
167
169
  ) -> float | tuple[float, np.ndarray]:
168
170
  """computes the expectation of $\\psi(U, v)$ and perhaps its gradient wrt `v`
@@ -180,12 +182,12 @@ def objgrad_(
180
182
  u1_nodes = args[1]
181
183
  u1_weights = args[2]
182
184
  vs1 = np.append(v_sorted, -np.sum(v_sorted))
183
- a_mat, b_mat = compute_ab_(y_sorted, vs1)
185
+ a_mat, b_mat = _compute_ab(y_sorted, vs1)
184
186
 
185
187
  obj_val = 0.0
186
188
  probs = np.zeros(n)
187
189
  for k in range(n):
188
- left_bounds, right_bounds = compute_u2_bounds_(k, u1_nodes, a_mat, b_mat)
190
+ left_bounds, right_bounds = _compute_u2_bounds(k, u1_nodes, a_mat, b_mat)
189
191
  pos_diffs = np.maximum(right_bounds - left_bounds, 0.0)
190
192
  pos_diffs_sq = np.maximum(
191
193
  right_bounds * right_bounds - left_bounds * left_bounds, 0.0
@@ -204,20 +206,30 @@ def objgrad_(
204
206
  return cast(float, obj_val)
205
207
 
206
208
 
207
- def obj_(v_sorted: np.ndarray, args: list):
208
- return objgrad_(v_sorted, args)
209
+ def _obj(v_sorted: np.ndarray, args: list):
210
+ return _objgrad(v_sorted, args)
209
211
 
210
212
 
211
- def grad_(v_sorted: np.ndarray, args: list):
212
- res_objg = cast(tuple[float, np.ndarray], objgrad_(v_sorted, args, gr=True))
213
+ def _grad(v_sorted: np.ndarray, args: list):
214
+ res_objg = cast(tuple[float, np.ndarray], _objgrad(v_sorted, args, gr=True))
213
215
  grad_val = res_objg[1]
214
216
  verbose = args[3]
215
217
  if verbose:
216
- print(f"grad_err is {npmaxabs(grad_val)}")
218
+ print(f"The error on the gradient is {npmaxabs(grad_val)}")
217
219
  return grad_val
218
220
 
219
221
 
220
222
  def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.ndarray:
223
+ """Solve the dual optimisation to obtain the optimal weights ``v``.
224
+
225
+ Args:
226
+ y: Observations with shape ``(n, 2)``.
227
+ n_nodes: Number of Chebyshev nodes for the quadrature.
228
+ verbose: Print optimisation diagnostics when ``True``.
229
+
230
+ Returns:
231
+ Array of length ``n`` containing the optimal weights (including the residual term).
232
+ """
221
233
  n, d = y.shape
222
234
 
223
235
  if d != 2:
@@ -234,7 +246,7 @@ def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.
234
246
 
235
247
  argsog = [y_sorted, u1_nodes, u1_weights, verbose]
236
248
 
237
- res = minimize_free(obj_, grad_, v0, args=argsog)
249
+ res = minimize_free(_obj, _grad, v0, args=argsog)
238
250
  if verbose:
239
251
  print_optimization_results(res, "Minimizing over v")
240
252
 
@@ -255,17 +267,16 @@ def solve_for_v_(y: np.ndarray, n_nodes: int = 32, verbose: bool = False) -> np.
255
267
  def bivariate_quantiles(
256
268
  y: np.ndarray, u: np.ndarray, n_nodes: int = 32, verbose: bool = False
257
269
  ) -> np.ndarray:
258
- """computes the bivariate quantiles of `y` at the quantiles `u`
270
+ """Solve for the dual weights then evaluate bivariate quantiles.
259
271
 
260
272
  Args:
261
- y: the observations, an `(n, 2)` matrix
262
- u: the quantiles at which to compute the bivariate quantiles,
263
- an `(m, 2)` matrix
264
- n_nodes: the number of nodes to use for the quadrature
265
- verbose: if `True`, print some information
273
+ y: Observations, shape ``(n, 2)``.
274
+ u: Query points in ``[0, 1]^2`` (shape ``(m, 2)``).
275
+ n_nodes: Number of Chebyshev nodes for the quadrature.
276
+ verbose: Print optimisation diagnostics when ``True``.
266
277
 
267
278
  Returns:
268
- an `(m, 2)` matrix of bivariate quantiles
279
+ Bivariate quantiles at ``u``.
269
280
  """
270
281
  v = solve_for_v_(y, n_nodes, verbose)
271
282
  return bivariate_quantiles_v(y, u, v)
@@ -274,15 +285,15 @@ def bivariate_quantiles(
274
285
  def bivariate_ranks(
275
286
  y: np.ndarray, n_nodes: int = 32, verbose: bool = False
276
287
  ) -> np.ndarray:
277
- """computes the bivariate ranks of `y`
288
+ """Compute ranks by first solving for the optimal weights ``v``.
278
289
 
279
290
  Args:
280
- y: the observations, an `(n, 2)` matrix
281
- n_nodes: the number of nodes to use for the quadrature
282
- verbose: if `True`, print some information
291
+ y: Observations, shape ``(n, 2)``.
292
+ n_nodes: Number of Chebyshev nodes for the quadrature.
293
+ verbose: Print optimisation diagnostics when ``True``.
283
294
 
284
295
  Returns:
285
- the `(n, 2)` matrix of bivariate average ranks
296
+ Average ranks with shape ``(n, 2)``.
286
297
  """
287
298
  v = solve_for_v_(y, n_nodes, verbose)
288
299
  return bivariate_ranks_v(y, v, n_nodes)