| 136 | |
| 137 | |
| 138 | def test_emd2_gradients(): |
| 139 | n_samples = 100 |
| 140 | n_features = 2 |
| 141 | rng = np.random.RandomState(0) |
| 142 | |
| 143 | x = rng.randn(n_samples, n_features) |
| 144 | y = rng.randn(n_samples, n_features) |
| 145 | a = ot.utils.unif(n_samples) |
| 146 | |
| 147 | M = ot.dist(x, y) |
| 148 | |
| 149 | if torch: |
| 150 | a1 = torch.tensor(a, requires_grad=True) |
| 151 | b1 = torch.tensor(a, requires_grad=True) |
| 152 | M1 = torch.tensor(M, requires_grad=True) |
| 153 | |
| 154 | val, log = ot.emd2(a1, b1, M1, log=True) |
| 155 | |
| 156 | val.backward() |
| 157 | |
| 158 | assert a1.shape == a1.grad.shape |
| 159 | assert b1.shape == b1.grad.shape |
| 160 | assert M1.shape == M1.grad.shape |
| 161 | |
| 162 | assert np.allclose( |
| 163 | a1.grad.cpu().detach().numpy(), |
| 164 | log["u"].cpu().detach().numpy() - log["u"].cpu().detach().numpy().mean(), |
| 165 | ) |
| 166 | |
| 167 | assert np.allclose( |
| 168 | b1.grad.cpu().detach().numpy(), |
| 169 | log["v"].cpu().detach().numpy() - log["v"].cpu().detach().numpy().mean(), |
| 170 | ) |
| 171 | |
| 172 | # Testing for bug #309, checking for scaling of gradient |
| 173 | a2 = torch.tensor(a, requires_grad=True) |
| 174 | b2 = torch.tensor(a, requires_grad=True) |
| 175 | M2 = torch.tensor(M, requires_grad=True) |
| 176 | |
| 177 | val = 10.0 * ot.emd2(a2, b2, M2) |
| 178 | |
| 179 | val.backward() |
| 180 | |
| 181 | assert np.allclose( |
| 182 | 10.0 * a1.grad.cpu().detach().numpy(), a2.grad.cpu().detach().numpy() |
| 183 | ) |
| 184 | assert np.allclose( |
| 185 | 10.0 * b1.grad.cpu().detach().numpy(), b2.grad.cpu().detach().numpy() |
| 186 | ) |
| 187 | assert np.allclose( |
| 188 | 10.0 * M1.grad.cpu().detach().numpy(), M2.grad.cpu().detach().numpy() |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | def test_emd_emd2(): |