tests: fix gauge invariance test

tests/test_basic.py

@@ -102,7 +102,7 @@ def test_gauge_invariance():
         nx, ny, 2, 2
     )
     s.omega = 0.987
-    s.eta = 0.789 * 0
+    s.eta = 0.789
     s.D = 1.0
     s.Lx = 10
     s.Ly = 10
@@ -143,20 +143,32 @@ def test_gauge_invariance():
     # Should be invariant to numerical accuracy
     assert_allclose(S1, S0, atol=1e-9)
 
-    # Current should be conserved to numerical accuracy
+    # Current should obey a continuity law to numerical accuracy
     # due to the gauge symmetry
     from usadelndsoc.solver import Result, tr
 
-    res = Result(s, s.omega, Phi2)
+    # Compute current divergence
+    res = Result(s, s.omega, Phi)
 
     t3 = np.diag([1, 1, -1, -1])
     J = tr(res.J[i, j] @ t3)
 
-    print(np.where(res.J))
-
     div_J_x = J[0] - J[2]
     div_J_y = J[1] - J[3]
     div_J = div_J_x + div_J_y
 
-    print(J.tolist())
-    assert_allclose(div_J, 0, atol=1e-9)
+    # Compute source term (numdiff)
+    Sa = s._core.eval(Phi)
+
+    Phi2 = Phi.copy()
+    ds = 1e-7
+    dW = expm(ds * t3)
+    idW = np.linalg.inv(dW)
+    Phi2[i, j, 0] = dW[:2, :2] @ Phi[i, j, 0] @ idW[2:, 2:]
+    Phi2[i, j, 1] = dW[2:, 2:] @ Phi[i, j, 1] @ idW[:2, :2]
+
+    Sb = s._core.eval(Phi2)
+    dS_ds = (Sb - Sa) / ds / 16
+
+    # They must match
+    assert_allclose(div_J, dS_ds, rtol=1e-6)