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.
- wax_ot-0.1.0/.gitignore +13 -0
- wax_ot-0.1.0/CITATION.cff +37 -0
- wax_ot-0.1.0/LICENSE +21 -0
- wax_ot-0.1.0/PKG-INFO +294 -0
- wax_ot-0.1.0/README.md +264 -0
- wax_ot-0.1.0/pyproject.toml +71 -0
- wax_ot-0.1.0/scripts/01_faithfulness_srg.py +199 -0
- wax_ot-0.1.0/scripts/02_characterize_shift.py +242 -0
- wax_ot-0.1.0/scripts/03_uwax_abalone.py +91 -0
- wax_ot-0.1.0/scripts/04_domain_alignment.py +99 -0
- wax_ot-0.1.0/scripts/05_celeba_lfw.py +301 -0
- wax_ot-0.1.0/src/wax/__init__.py +75 -0
- wax_ot-0.1.0/src/wax/baselines.py +102 -0
- wax_ot-0.1.0/src/wax/coupling.py +103 -0
- wax_ot-0.1.0/src/wax/data/abalone.data +4177 -0
- wax_ot-0.1.0/src/wax/data/english_words.txt +10000 -0
- wax_ot-0.1.0/src/wax/datasets.py +558 -0
- wax_ot-0.1.0/src/wax/forward.py +50 -0
- wax_ot-0.1.0/src/wax/metrics.py +82 -0
- wax_ot-0.1.0/src/wax/plotting.py +138 -0
- wax_ot-0.1.0/src/wax/uwax.py +301 -0
- wax_ot-0.1.0/src/wax/wax.py +194 -0
- wax_ot-0.1.0/src/wax/words.py +31 -0
- wax_ot-0.1.0/tests/test_wax.py +278 -0
wax_ot-0.1.0/.gitignore
ADDED
|
@@ -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).
|