wax-ot 0.1.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.

Potentially problematic release.


This version of wax-ot might be problematic. Click here for more details.

@@ -0,0 +1,13 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ .venv/
4
+ build/
5
+ dist/
6
+ *.egg-info/
7
+ .pytest_cache/
8
+ .ruff_cache/
9
+ .DS_Store
10
+
11
+ # Datasets downloaded at run time. These reach 1.8 GB (CelebA alone is 1.3 GB)
12
+ # and must never enter git history. Override the location with $WAX_CACHE.
13
+ data/cache/
@@ -0,0 +1,37 @@
1
+ cff-version: 1.2.0
2
+ message: >-
3
+ This is an independent implementation. If you use it, please cite the
4
+ original paper first, and this software second.
5
+ title: "wax-ot: an independent implementation of WaX"
6
+ abstract: >-
7
+ Attribution of Wasserstein distances to instance pairs, input features and
8
+ interpretable subspaces, by layer-wise relevance propagation on the
9
+ neuralized distance graph.
10
+ type: software
11
+ authors:
12
+ - family-names: Hassan
13
+ given-names: Omar A.
14
+ repository-code: "https://github.com/omar-A-hassan/wax-ot"
15
+ license: MIT
16
+ version: 0.1.0
17
+ date-released: 2026-08-21
18
+ references:
19
+ - type: article
20
+ title: >-
21
+ Wasserstein Distances Made Explainable: Insights Into Dataset Shifts and
22
+ Transport Phenomena
23
+ authors:
24
+ - family-names: Naumann
25
+ given-names: Philip
26
+ - family-names: Kauffmann
27
+ given-names: Jacob
28
+ - family-names: Montavon
29
+ given-names: Grégoire
30
+ journal: IEEE Transactions on Pattern Analysis and Machine Intelligence
31
+ volume: 48
32
+ issue: 6
33
+ start: 6393
34
+ end: 6406
35
+ year: 2026
36
+ doi: 10.1109/TPAMI.2026.3656947
37
+ url: "https://arxiv.org/abs/2505.06123"
wax_ot-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Omar A. Hassan
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
wax_ot-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,294 @@
1
+ Metadata-Version: 2.5
2
+ Name: wax-ot
3
+ Version: 0.1.0
4
+ Summary: WaX: explainable Wasserstein distances. Attributes a Wasserstein distance to instances, features and subspaces (Naumann, Kauffmann and Montavon, TPAMI 2026).
5
+ Project-URL: Homepage, https://github.com/omar-A-hassan/wax-ot
6
+ Project-URL: Issues, https://github.com/omar-A-hassan/wax-ot/issues
7
+ Project-URL: Paper, https://arxiv.org/abs/2505.06123
8
+ Author-email: "Omar A. Hassan" <quickyomar@gmail.com>
9
+ License-Expression: MIT
10
+ License-File: LICENSE
11
+ Keywords: dataset-shift,explainable-ai,layer-wise-relevance-propagation,optimal-transport,wasserstein
12
+ Classifier: Development Status :: 4 - Beta
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
16
+ Requires-Python: >=3.10
17
+ Requires-Dist: numpy>=2.2.6
18
+ Requires-Dist: pot>=0.9.7.post1
19
+ Requires-Dist: scipy>=1.15.3
20
+ Provides-Extra: repro
21
+ Requires-Dist: matplotlib>=3.10.9; extra == 'repro'
22
+ Requires-Dist: open-clip-torch>=3.3.0; extra == 'repro'
23
+ Requires-Dist: scikit-learn>=1.7.2; extra == 'repro'
24
+ Requires-Dist: torch>=2.13.0; extra == 'repro'
25
+ Provides-Extra: torch
26
+ Requires-Dist: torch>=2.13.0; extra == 'torch'
27
+ Provides-Extra: viz
28
+ Requires-Dist: matplotlib>=3.10.9; extra == 'viz'
29
+ Description-Content-Type: text/markdown
30
+
31
+ # wax-ot
32
+
33
+ WaX: an explanation method for Wasserstein distances.
34
+
35
+ This package is an independent implementation of the method in this article:
36
+
37
+ > **Wasserstein Distances Made Explainable: Insights Into Dataset Shifts and
38
+ > Transport Phenomena.** P. Naumann, J. Kauffmann, G. Montavon.
39
+ > IEEE Transactions on Pattern Analysis and Machine Intelligence,
40
+ > volume 48, number 6, pages 6393 to 6406, 2026.
41
+ > doi:[10.1109/TPAMI.2026.3656947](https://doi.org/10.1109/TPAMI.2026.3656947).
42
+ > Preprint with the supplementary notes:
43
+ > [arXiv:2505.06123](https://arxiv.org/abs/2505.06123).
44
+
45
+ The authors of the article did not write this package. They are not related to
46
+ it. On the date of this release, no official implementation is public.
47
+
48
+ Section [Verification](#verification) gives the tests and the results.
49
+
50
+ ## Description
51
+
52
+ A Wasserstein distance gives the quantity of the difference between two
53
+ datasets. It does not give the cause of the difference.
54
+
55
+ WaX gives the cause. The method writes the distance as a network of two layers
56
+ and keeps the coupling constant. The method then does a backward pass with
57
+ layer-wise relevance propagation rules. The result is a division of the
58
+ distance into relevance values for:
59
+
60
+ - instance pairs, `R_kl`: the source points and the target points that cause
61
+ the distance;
62
+ - input features, `R_i`: the measured quantities that cause the distance;
63
+ - subspaces (U-WaX): the directions in the input space that cause the distance.
64
+
65
+ The division keeps the total. Thus `sum(R_i)` and `sum(R_kl)` are both equal to
66
+ the distance `W_p`. The test suite examines this property.
67
+
68
+ ## Installation
69
+
70
+ ```
71
+ pip install wax-ot
72
+ ```
73
+
74
+ The core package needs numpy, scipy and POT only. Four optional extras are
75
+ available:
76
+
77
+ | Extra | Adds | Necessary for |
78
+ |---|---|---|
79
+ | `wax-ot[viz]` | matplotlib | the `wax.plotting` module |
80
+ | `wax-ot[torch]` | torch | the comparison with the reference implementation |
81
+ | `wax-ot[repro]` | the two extras above, scikit-learn and open-clip | the experiment scripts |
82
+
83
+ ## Use
84
+
85
+ ```python
86
+ import numpy as np
87
+ from wax import attribute, exact
88
+
89
+ X = np.random.randn(300, 8)
90
+ Y = np.random.randn(300, 8) + np.array([0.6, 0, 0, 0.2, 0, 0, -0.4, 0])
91
+
92
+ coupling = exact(X, Y, p=2, q=2) # also: sinkhorn(...) or uniform(...)
93
+ a = attribute(X, Y, coupling, p=2, q=2)
94
+
95
+ a.W # the Wasserstein distance
96
+ a.R_i # relevance of each feature, shape (8,)
97
+ a.R_kl # relevance of each instance pair, shape (300, 300)
98
+ a.conserved # True if sum(R_i) and sum(R_kl) are equal to W
99
+ ```
100
+
101
+ ### How to read the relevance values
102
+
103
+ The relevance values use the same units as the distance `W`. They are not
104
+ percentages. If `W` is 1.0, and one feature has an `R_i` of 0.4, then that
105
+ feature causes 40 percent of the distance. The sum of the values is always
106
+ equal to `W`.
107
+
108
+ The value `R_kl[k, l]` is the part of the distance that comes from the movement
109
+ of source point `k` to target point `l`. Large values show unusual pairs.
110
+
111
+ Make the scale of all features equal before you calculate the relevance values.
112
+ If you do not do this, the feature with the largest values controls the result.
113
+ The function `wax.datasets.standardize` does this operation.
114
+
115
+ ### Subspaces
116
+
117
+ ```python
118
+ from wax import uwax_search, uwax_attribute
119
+
120
+ U, history = uwax_search(X, Y, coupling, r=4, dims=[1, 1, 1], seed=0)
121
+ s = uwax_attribute(X, Y, coupling, U, r=4)
122
+
123
+ s.R_c # relevance of each subspace
124
+ s.captured # the part of W that the subspaces explain
125
+ ```
126
+
127
+ ## Modules
128
+
129
+ | Module | Part of the article |
130
+ |---|---|
131
+ | `wax.coupling` | the couplings: exact transport (eq. 1), Sinkhorn, uniform |
132
+ | `wax.forward` | the network `z_kl` to `W_p` (eq. 2) |
133
+ | `wax.wax` | the relevance rules (eq. 3) and the reference implementation |
134
+ | `wax.uwax` | the subspace search and its relevance values (eqs. 9 to 12) |
135
+ | `wax.metrics` | Symmetric Relevance Gain (eq. 5), cosine similarity (eq. 6) |
136
+ | `wax.baselines` | MeanShift, Occlusion, Coupling, Uniform, Logistic |
137
+ | `wax.datasets` | the UCI loaders and the preprocessing of Supplementary Note E |
138
+
139
+ ## Verification
140
+
141
+ The IEEE article does not contain the supplementary notes. The arXiv preprint
142
+ contains all of them. Supplementary Note D contains the implementation of WaX
143
+ by the authors. This package follows Notes D, E, F and G.
144
+
145
+ ### Comparison with the implementation by the authors
146
+
147
+ The function `wax.attribute` is an independent closed form of the relevance
148
+ rules. The function `wax.torch_gradient_attribution` is a copy of the code in
149
+ Note D, Figure S3. That code uses the detach operation two times. Therefore
150
+ automatic differentiation gives the same result as the rules for all values of
151
+ `alpha` and `beta`.
152
+
153
+ The two functions agree to 2.4e-15. The test uses three couplings and eight
154
+ sets of the parameters `(p, q, alpha, beta)`. Four of the eight sets have an
155
+ `alpha` that is different from `p`, or a `beta` that is different from `q`. To
156
+ run this test:
157
+
158
+ ```
159
+ pytest tests/test_wax.py -k matches_authors_reference_implementation
160
+ ```
161
+
162
+ The test suite has 64 tests. This comparison is 24 of them. The other 40 tests
163
+ examine the conservation property and the closed form of Supplementary Note J.
164
+
165
+ ### Results for Table III
166
+
167
+ The table gives the mean absolute error against the published table. Each value
168
+ is a mean over the six delays. File
169
+ [results/02_table3.md](results/02_table3.md) gives all the numbers.
170
+
171
+ | Dataset | Uniform | MeanShift | Logistic L2 | Logistic GI | WaX | WaX reg | Mean |
172
+ |---|---|---|---|---|---|---|---|
173
+ | air-quality | 0.032 | 0.096 | 0.016 | 0.052 | 0.062 | 0.030 | 0.048 |
174
+ | electricity | 0.160 | 0.058 | 0.227 | 0.152 | 0.041 | 0.081 | 0.120 |
175
+ | appliances | 0.020 | 0.061 | 0.115 | 0.139 | 0.035 | 0.007 | 0.063 |
176
+ | **all** | 0.071 | 0.072 | 0.119 | 0.114 | **0.046** | **0.039** | **0.077** |
177
+
178
+ The two WaX columns have the smallest error. The datasets also have the
179
+ dimensions of the published table. The feature counts 9, 7 and 25 are correct.
180
+ The subset sizes for air-quality are within five days of the published values
181
+ in each row.
182
+
183
+ ### Results for Table I
184
+
185
+ The scripts use three of the six datasets in Table I. They do not use Crime,
186
+ Mice or Robot.
187
+
188
+ The preprocessing of Supplementary Note E gives almost the same dataset
189
+ dimensions as the article:
190
+
191
+ | Dataset | This package (N/M/d) | Article (N/M/d) |
192
+ |---|---|---|
193
+ | wisconsin | 332 / 163 / 30 | 332 / 163 / 30 |
194
+ | musk1 | 133 / 167 / 166 | 133 / 167 / 166 |
195
+ | wine-quality | 3783 / 1086 / 12 | 3781 / 1086 / 12 |
196
+
197
+ The next table gives the mean absolute error against Table I. File
198
+ [results/01_table1.md](results/01_table1.md) gives all the numbers.
199
+
200
+ | Dataset | MeanShift | Occlusion | Coupling | WaX | Mean |
201
+ |---|---|---|---|---|---|
202
+ | wine-quality | 0.004 | 0.005 | 0.012 | 0.007 | **0.007** |
203
+ | wisconsin | 0.075 | 0.084 | 0.080 | 0.093 | 0.083 |
204
+ | musk1 | 0.325 | 0.198 | 0.222 | 0.174 | 0.230 |
205
+ | **all** | 0.135 | 0.095 | 0.105 | 0.091 | **0.107** |
206
+
207
+ WaX has the best value, or an equal value, in 13 of these 15 rows. In the
208
+ article, WaX has the best value, or an equal value, in 14 of the same 15 rows.
209
+
210
+ The error increases when the number of samples decreases. Supplementary Note F
211
+ gives the same relation. Wine has 3783 source rows and an error of 0.007. Musk1
212
+ has 133 source rows and 166 features, and an error of 0.230.
213
+
214
+ ### Known differences from the article
215
+
216
+ - **Number of bootstrap trials.** Note F uses 100 trials. On a laptop with 8 GB
217
+ of memory, 100 trials of the Wine dataset need approximately 18 hours. Thus
218
+ the default value of `--reps` is 3. This is the largest cause of error in
219
+ Table I. Use `--reps 100` for the protocol of the article.
220
+ - **Selection of the electricity clients.** The scripts use the first seven
221
+ client columns with data for the full period. The notes do not say which
222
+ seven of the 370 clients the article uses. This dataset has the largest error
223
+ in Table III.
224
+ - **Selection of the air-quality channels.** The scripts keep the nine channels
225
+ with a minimum of 90 percent of the values present. This gives the nine
226
+ channels of the article. The notes do not give the threshold, and other
227
+ thresholds also give nine channels.
228
+ - **Tests that are not done.** The PLISM rows of Table III need three
229
+ histopathology models. The confidence intervals of Table S2 are not
230
+ calculated.
231
+ - **Results that are not compared.** Scripts 03, 04 and 05 make figures. The
232
+ figures agree with the article, but no script compares them with a published
233
+ number. Script 05 uses a maximum of 500 images from each face dataset,
234
+ because of the memory limit.
235
+
236
+ A laptop with an Apple M3 processor and 8 GB of memory made all the results
237
+ above. The scripts used the processor only.
238
+
239
+ ## How to run the experiments
240
+
241
+ ```
242
+ pip install "wax-ot[repro]"
243
+ python scripts/01_faithfulness_srg.py # Table I. Writes results/01_table1.*
244
+ python scripts/02_characterize_shift.py # Table III. Writes results/02_table3.*
245
+ python scripts/03_uwax_abalone.py # Section VI. U-WaX subspaces
246
+ python scripts/04_domain_alignment.py # Section V. Domain alignment
247
+ python scripts/05_celeba_lfw.py # Section VIII. CelebA and LFW
248
+ ```
249
+
250
+ Scripts 02 and 05 download datasets at the first use. The electricity dataset
251
+ is approximately 260 MB. The CelebA dataset is approximately 1.3 GB. The
252
+ scripts keep the files in `~/.cache/wax`. To use a different directory, set the
253
+ environment variable `WAX_CACHE`.
254
+
255
+ ## Development
256
+
257
+ ```
258
+ uv sync --all-extras
259
+ uv run pytest
260
+ uv run ruff check
261
+ ```
262
+
263
+ ## Data and licenses
264
+
265
+ The code has the MIT license. The method and the text of the article belong to
266
+ the authors. This repository does not contain the article. It gives a link to
267
+ the article.
268
+
269
+ Two data files are in the package. They keep their original licenses: the UCI
270
+ abalone data (Nash et al., 1995) and the google-10000-english word list. The
271
+ scripts download the other datasets at run time. This repository does not
272
+ contain them.
273
+
274
+ ## Citation
275
+
276
+ Cite the article first. This package is an implementation of the method in the
277
+ article. It is not new research.
278
+
279
+ ```bibtex
280
+ @article{naumann2026wasserstein,
281
+ author = {Naumann, Philip and Kauffmann, Jacob and Montavon, Gr{\'e}goire},
282
+ title = {Wasserstein Distances Made Explainable: Insights Into Dataset
283
+ Shifts and Transport Phenomena},
284
+ journal = {IEEE Transactions on Pattern Analysis and Machine Intelligence},
285
+ volume = {48},
286
+ number = {6},
287
+ pages = {6393--6406},
288
+ year = {2026},
289
+ doi = {10.1109/TPAMI.2026.3656947}
290
+ }
291
+ ```
292
+
293
+ To cite this package in addition to the article, use the metadata in
294
+ [CITATION.cff](CITATION.cff).
wax_ot-0.1.0/README.md ADDED
@@ -0,0 +1,264 @@
1
+ # wax-ot
2
+
3
+ WaX: an explanation method for Wasserstein distances.
4
+
5
+ This package is an independent implementation of the method in this article:
6
+
7
+ > **Wasserstein Distances Made Explainable: Insights Into Dataset Shifts and
8
+ > Transport Phenomena.** P. Naumann, J. Kauffmann, G. Montavon.
9
+ > IEEE Transactions on Pattern Analysis and Machine Intelligence,
10
+ > volume 48, number 6, pages 6393 to 6406, 2026.
11
+ > doi:[10.1109/TPAMI.2026.3656947](https://doi.org/10.1109/TPAMI.2026.3656947).
12
+ > Preprint with the supplementary notes:
13
+ > [arXiv:2505.06123](https://arxiv.org/abs/2505.06123).
14
+
15
+ The authors of the article did not write this package. They are not related to
16
+ it. On the date of this release, no official implementation is public.
17
+
18
+ Section [Verification](#verification) gives the tests and the results.
19
+
20
+ ## Description
21
+
22
+ A Wasserstein distance gives the quantity of the difference between two
23
+ datasets. It does not give the cause of the difference.
24
+
25
+ WaX gives the cause. The method writes the distance as a network of two layers
26
+ and keeps the coupling constant. The method then does a backward pass with
27
+ layer-wise relevance propagation rules. The result is a division of the
28
+ distance into relevance values for:
29
+
30
+ - instance pairs, `R_kl`: the source points and the target points that cause
31
+ the distance;
32
+ - input features, `R_i`: the measured quantities that cause the distance;
33
+ - subspaces (U-WaX): the directions in the input space that cause the distance.
34
+
35
+ The division keeps the total. Thus `sum(R_i)` and `sum(R_kl)` are both equal to
36
+ the distance `W_p`. The test suite examines this property.
37
+
38
+ ## Installation
39
+
40
+ ```
41
+ pip install wax-ot
42
+ ```
43
+
44
+ The core package needs numpy, scipy and POT only. Four optional extras are
45
+ available:
46
+
47
+ | Extra | Adds | Necessary for |
48
+ |---|---|---|
49
+ | `wax-ot[viz]` | matplotlib | the `wax.plotting` module |
50
+ | `wax-ot[torch]` | torch | the comparison with the reference implementation |
51
+ | `wax-ot[repro]` | the two extras above, scikit-learn and open-clip | the experiment scripts |
52
+
53
+ ## Use
54
+
55
+ ```python
56
+ import numpy as np
57
+ from wax import attribute, exact
58
+
59
+ X = np.random.randn(300, 8)
60
+ Y = np.random.randn(300, 8) + np.array([0.6, 0, 0, 0.2, 0, 0, -0.4, 0])
61
+
62
+ coupling = exact(X, Y, p=2, q=2) # also: sinkhorn(...) or uniform(...)
63
+ a = attribute(X, Y, coupling, p=2, q=2)
64
+
65
+ a.W # the Wasserstein distance
66
+ a.R_i # relevance of each feature, shape (8,)
67
+ a.R_kl # relevance of each instance pair, shape (300, 300)
68
+ a.conserved # True if sum(R_i) and sum(R_kl) are equal to W
69
+ ```
70
+
71
+ ### How to read the relevance values
72
+
73
+ The relevance values use the same units as the distance `W`. They are not
74
+ percentages. If `W` is 1.0, and one feature has an `R_i` of 0.4, then that
75
+ feature causes 40 percent of the distance. The sum of the values is always
76
+ equal to `W`.
77
+
78
+ The value `R_kl[k, l]` is the part of the distance that comes from the movement
79
+ of source point `k` to target point `l`. Large values show unusual pairs.
80
+
81
+ Make the scale of all features equal before you calculate the relevance values.
82
+ If you do not do this, the feature with the largest values controls the result.
83
+ The function `wax.datasets.standardize` does this operation.
84
+
85
+ ### Subspaces
86
+
87
+ ```python
88
+ from wax import uwax_search, uwax_attribute
89
+
90
+ U, history = uwax_search(X, Y, coupling, r=4, dims=[1, 1, 1], seed=0)
91
+ s = uwax_attribute(X, Y, coupling, U, r=4)
92
+
93
+ s.R_c # relevance of each subspace
94
+ s.captured # the part of W that the subspaces explain
95
+ ```
96
+
97
+ ## Modules
98
+
99
+ | Module | Part of the article |
100
+ |---|---|
101
+ | `wax.coupling` | the couplings: exact transport (eq. 1), Sinkhorn, uniform |
102
+ | `wax.forward` | the network `z_kl` to `W_p` (eq. 2) |
103
+ | `wax.wax` | the relevance rules (eq. 3) and the reference implementation |
104
+ | `wax.uwax` | the subspace search and its relevance values (eqs. 9 to 12) |
105
+ | `wax.metrics` | Symmetric Relevance Gain (eq. 5), cosine similarity (eq. 6) |
106
+ | `wax.baselines` | MeanShift, Occlusion, Coupling, Uniform, Logistic |
107
+ | `wax.datasets` | the UCI loaders and the preprocessing of Supplementary Note E |
108
+
109
+ ## Verification
110
+
111
+ The IEEE article does not contain the supplementary notes. The arXiv preprint
112
+ contains all of them. Supplementary Note D contains the implementation of WaX
113
+ by the authors. This package follows Notes D, E, F and G.
114
+
115
+ ### Comparison with the implementation by the authors
116
+
117
+ The function `wax.attribute` is an independent closed form of the relevance
118
+ rules. The function `wax.torch_gradient_attribution` is a copy of the code in
119
+ Note D, Figure S3. That code uses the detach operation two times. Therefore
120
+ automatic differentiation gives the same result as the rules for all values of
121
+ `alpha` and `beta`.
122
+
123
+ The two functions agree to 2.4e-15. The test uses three couplings and eight
124
+ sets of the parameters `(p, q, alpha, beta)`. Four of the eight sets have an
125
+ `alpha` that is different from `p`, or a `beta` that is different from `q`. To
126
+ run this test:
127
+
128
+ ```
129
+ pytest tests/test_wax.py -k matches_authors_reference_implementation
130
+ ```
131
+
132
+ The test suite has 64 tests. This comparison is 24 of them. The other 40 tests
133
+ examine the conservation property and the closed form of Supplementary Note J.
134
+
135
+ ### Results for Table III
136
+
137
+ The table gives the mean absolute error against the published table. Each value
138
+ is a mean over the six delays. File
139
+ [results/02_table3.md](results/02_table3.md) gives all the numbers.
140
+
141
+ | Dataset | Uniform | MeanShift | Logistic L2 | Logistic GI | WaX | WaX reg | Mean |
142
+ |---|---|---|---|---|---|---|---|
143
+ | air-quality | 0.032 | 0.096 | 0.016 | 0.052 | 0.062 | 0.030 | 0.048 |
144
+ | electricity | 0.160 | 0.058 | 0.227 | 0.152 | 0.041 | 0.081 | 0.120 |
145
+ | appliances | 0.020 | 0.061 | 0.115 | 0.139 | 0.035 | 0.007 | 0.063 |
146
+ | **all** | 0.071 | 0.072 | 0.119 | 0.114 | **0.046** | **0.039** | **0.077** |
147
+
148
+ The two WaX columns have the smallest error. The datasets also have the
149
+ dimensions of the published table. The feature counts 9, 7 and 25 are correct.
150
+ The subset sizes for air-quality are within five days of the published values
151
+ in each row.
152
+
153
+ ### Results for Table I
154
+
155
+ The scripts use three of the six datasets in Table I. They do not use Crime,
156
+ Mice or Robot.
157
+
158
+ The preprocessing of Supplementary Note E gives almost the same dataset
159
+ dimensions as the article:
160
+
161
+ | Dataset | This package (N/M/d) | Article (N/M/d) |
162
+ |---|---|---|
163
+ | wisconsin | 332 / 163 / 30 | 332 / 163 / 30 |
164
+ | musk1 | 133 / 167 / 166 | 133 / 167 / 166 |
165
+ | wine-quality | 3783 / 1086 / 12 | 3781 / 1086 / 12 |
166
+
167
+ The next table gives the mean absolute error against Table I. File
168
+ [results/01_table1.md](results/01_table1.md) gives all the numbers.
169
+
170
+ | Dataset | MeanShift | Occlusion | Coupling | WaX | Mean |
171
+ |---|---|---|---|---|---|
172
+ | wine-quality | 0.004 | 0.005 | 0.012 | 0.007 | **0.007** |
173
+ | wisconsin | 0.075 | 0.084 | 0.080 | 0.093 | 0.083 |
174
+ | musk1 | 0.325 | 0.198 | 0.222 | 0.174 | 0.230 |
175
+ | **all** | 0.135 | 0.095 | 0.105 | 0.091 | **0.107** |
176
+
177
+ WaX has the best value, or an equal value, in 13 of these 15 rows. In the
178
+ article, WaX has the best value, or an equal value, in 14 of the same 15 rows.
179
+
180
+ The error increases when the number of samples decreases. Supplementary Note F
181
+ gives the same relation. Wine has 3783 source rows and an error of 0.007. Musk1
182
+ has 133 source rows and 166 features, and an error of 0.230.
183
+
184
+ ### Known differences from the article
185
+
186
+ - **Number of bootstrap trials.** Note F uses 100 trials. On a laptop with 8 GB
187
+ of memory, 100 trials of the Wine dataset need approximately 18 hours. Thus
188
+ the default value of `--reps` is 3. This is the largest cause of error in
189
+ Table I. Use `--reps 100` for the protocol of the article.
190
+ - **Selection of the electricity clients.** The scripts use the first seven
191
+ client columns with data for the full period. The notes do not say which
192
+ seven of the 370 clients the article uses. This dataset has the largest error
193
+ in Table III.
194
+ - **Selection of the air-quality channels.** The scripts keep the nine channels
195
+ with a minimum of 90 percent of the values present. This gives the nine
196
+ channels of the article. The notes do not give the threshold, and other
197
+ thresholds also give nine channels.
198
+ - **Tests that are not done.** The PLISM rows of Table III need three
199
+ histopathology models. The confidence intervals of Table S2 are not
200
+ calculated.
201
+ - **Results that are not compared.** Scripts 03, 04 and 05 make figures. The
202
+ figures agree with the article, but no script compares them with a published
203
+ number. Script 05 uses a maximum of 500 images from each face dataset,
204
+ because of the memory limit.
205
+
206
+ A laptop with an Apple M3 processor and 8 GB of memory made all the results
207
+ above. The scripts used the processor only.
208
+
209
+ ## How to run the experiments
210
+
211
+ ```
212
+ pip install "wax-ot[repro]"
213
+ python scripts/01_faithfulness_srg.py # Table I. Writes results/01_table1.*
214
+ python scripts/02_characterize_shift.py # Table III. Writes results/02_table3.*
215
+ python scripts/03_uwax_abalone.py # Section VI. U-WaX subspaces
216
+ python scripts/04_domain_alignment.py # Section V. Domain alignment
217
+ python scripts/05_celeba_lfw.py # Section VIII. CelebA and LFW
218
+ ```
219
+
220
+ Scripts 02 and 05 download datasets at the first use. The electricity dataset
221
+ is approximately 260 MB. The CelebA dataset is approximately 1.3 GB. The
222
+ scripts keep the files in `~/.cache/wax`. To use a different directory, set the
223
+ environment variable `WAX_CACHE`.
224
+
225
+ ## Development
226
+
227
+ ```
228
+ uv sync --all-extras
229
+ uv run pytest
230
+ uv run ruff check
231
+ ```
232
+
233
+ ## Data and licenses
234
+
235
+ The code has the MIT license. The method and the text of the article belong to
236
+ the authors. This repository does not contain the article. It gives a link to
237
+ the article.
238
+
239
+ Two data files are in the package. They keep their original licenses: the UCI
240
+ abalone data (Nash et al., 1995) and the google-10000-english word list. The
241
+ scripts download the other datasets at run time. This repository does not
242
+ contain them.
243
+
244
+ ## Citation
245
+
246
+ Cite the article first. This package is an implementation of the method in the
247
+ article. It is not new research.
248
+
249
+ ```bibtex
250
+ @article{naumann2026wasserstein,
251
+ author = {Naumann, Philip and Kauffmann, Jacob and Montavon, Gr{\'e}goire},
252
+ title = {Wasserstein Distances Made Explainable: Insights Into Dataset
253
+ Shifts and Transport Phenomena},
254
+ journal = {IEEE Transactions on Pattern Analysis and Machine Intelligence},
255
+ volume = {48},
256
+ number = {6},
257
+ pages = {6393--6406},
258
+ year = {2026},
259
+ doi = {10.1109/TPAMI.2026.3656947}
260
+ }
261
+ ```
262
+
263
+ To cite this package in addition to the article, use the metadata in
264
+ [CITATION.cff](CITATION.cff).