diff --git a/src/action.hpp b/src/action.hpp
index 2531af08a0b804a0f05b26615569596cc48d7502..691f3511dd35ed74273f0796374adcf227b36dbd 100644
--- a/src/action.hpp
+++ b/src/action.hpp
@@ -80,8 +80,8 @@ inline Scalar S_2(const array::ArrayView<const Mask, Shape0>& mask,
 
     const double dd[2] = {Lx / nx, Ly / ny};
 
-    const int N_i[2] = {1, 0};
-    const int N_j[2] = {0, 1};
+    const int N_i[4] = {1, 0, -1, 0};
+    const int N_j[4] = {0, 1, 0, -1};
 
     /*
      * Directions cycle (left->right, down->up, right->left, up->down)
@@ -102,10 +102,13 @@ inline Scalar S_2(const array::ArrayView<const Mask, Shape0>& mask,
             if (mask(i,j) == Mask::VACUUM)
                 continue;
 
+            if (i != 3 || j != 4)
+                continue;
+
             auto Q1 = Qmat.part(i,j).matrix();
 
             /* Sum over neighbors */
-            for (size_t k = 0; k < 2; ++k) {
+            for (size_t k = 1; k < 1+1; ++k) {
                 int i2 = i + N_i[k];
                 int j2 = j + N_j[k];
                 
@@ -312,7 +315,7 @@ inline Scalar S(const array::ArrayView<const Mask, Shape0>& mask,
         }
     }
 
-    S = S + S_2<Scalar>(mask, Qmat, U, Lx, Ly, D) + S_Omega<Scalar>(mask, Qmat, omega, Omega, Lx, Ly) + S_H<Scalar>(mask, Qmat, U, Lx, Ly, eta);
+    S = S + S_2<Scalar>(mask, Qmat, U, Lx, Ly, D);// + S_Omega<Scalar>(mask, Qmat, omega, Omega, Lx, Ly) + S_H<Scalar>(mask, Qmat, U, Lx, Ly, eta);
     return S;
 }
 
diff --git a/tests/test_basic.py b/tests/test_basic.py
index faaf5669a7bf41b6bb420efa5eac42f8ecf559ad..37c18555db6cdac6d73fe6370989ecbae7eaacaf 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -102,7 +102,7 @@ def test_gauge_invariance():
         nx, ny, 2, 2
     )
     s.omega = 0.987
-    s.eta = 0.789
+    s.eta = 0.789 * 0
     s.D = 1.0
     s.Lx = 10
     s.Ly = 10
@@ -142,3 +142,21 @@ def test_gauge_invariance():
 
     # Should be invariant to numerical accuracy
     assert_allclose(S1, S0, atol=1e-9)
+
+    # Current should be conserved to numerical accuracy
+    # due to the gauge symmetry
+    from usadelndsoc.solver import Result, tr
+
+    res = Result(s, s.omega, Phi2)
+
+    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)
diff --git a/tests/test_solver.py b/tests/test_solver.py
index 4d9ed657bd3533d90c940b92ba818062418ca088..ac8f535568d94523a11b561137bc6c3b0d5e5438 100644
--- a/tests/test_solver.py
+++ b/tests/test_solver.py
@@ -326,16 +326,14 @@ def test_soc_analytic(T, n, alpha_soc):
         res0 = sol0.solve(omega=ww)
         sol.omega = ww
         J = sol._core.grad_A(res0.Phi.copy()).transpose(0, 1, 2, 4, 3) * (1 / 16)
-        Jtot += pi * aa * J
+        Jtot += -1j * pi * aa * J
 
         gp = (ww - 1j * h) / np.sqrt((ww - 1j * h) ** 2 + 1)
         gm = (ww + 1j * h) / np.sqrt((ww + 1j * h) ** 2 + 1)
         fp = 1 / np.sqrt((ww - 1j * h) ** 2 + 1)
         fm = 1 / np.sqrt((ww + 1j * h) ** 2 + 1)
-        # XXX: check the 32*eta
-        Jy_an = (
-            (gp - gm) * (1 + gp * gm + fp * fm) * (-1j * alpha_soc**3) * (32 * eta)
-        )
+        # XXX: check the 8*eta
+        Jy_an = (gp - gm) * (1 + gp * gm + fp * fm) * (-1j * alpha_soc**3) * (8 * eta)
         Jy_an *= dd[1]  # multiply current density by the width of a single cell
         Jantot += pi * aa * Jy_an