solver: fix self-consistent iteration to do singlet correctly

usadelndsoc/plotting.py

@@ -16,7 +16,7 @@ __all__ = [
 ]
 
 
-def plot_Delta_J_2d(x, y, Delta, J, ax=None):
+def plot_Delta_J_2d(x, y, Delta, xc, yc, Ix, Iy, ax=None):
     """
     Color plot of order parameter and supercurrents.
     """
@@ -24,15 +24,11 @@ def plot_Delta_J_2d(x, y, Delta, J, ax=None):
         ax = plt.gca()
 
     m = pcolormesh_complex(x, y, Delta, ax=ax)
-    I = J.real
-    xc = (x[1:] + x[:-1]) / 2
-    yc = (y[1:] + y[:-1]) / 2
-    Ix = I[0, :-1, :-1, 0].T
-    Iy = I[1, :-1, :-1, 0].T
+
     lw = np.hypot(Ix, Iy)
     lw = 6 * (lw / max(1.0, lw.max())) ** 0.5
+    ax.streamplot(xc, yc, Ix.T, Iy.T, color=(0.5, 0.5, 0.5), linewidth=lw.T)
 
-    ax.streamplot(xc, yc, Ix, Iy, color=(0.5, 0.5, 0.5), linewidth=lw)
     ax.set_xlabel(r"$x$")
     ax.set_ylabel(r"$y$")
     plt.colorbar(m.phase, label=r"$\varphi\,/ \pi$", ax=ax)

usadelndsoc/solver.py

@@ -937,6 +937,14 @@ def _with_workers(func):
     return wrapper
 
 
+def singlet(x):
+    return (x[..., 0, 0] + x[..., 1, 1]) / 2
+
+
+def singlet_m(x):
+    return x[..., None, None] * S_0
+
+
 @_with_workers
 def _self_consistency(
     solver,
@@ -967,12 +975,12 @@ def _self_consistency(
     T = float(T)
     T_c0 = np.asarray(T_c0)
 
-    if T_c0.shape + (2, 2) != solver.Delta.shape:
+    if T_c0.shape != solver.Delta.shape[:-2]:
         raise ValueError("T_c0 must have compatible shape with Delta")
 
     Delta_mask = (solver.mask == MASK_NONE) & (T_c0 > 0) & (T < T_c0)
 
-    z = solver.Delta[Delta_mask].copy().ravel().view(float)
+    z = singlet(solver.Delta[Delta_mask]).ravel().view(float)
     if constraint_fun is not None:
         constraint_x = np.asarray(constraint_x)
         z = np.r_[constraint_x.astype(float), z]
@@ -1010,13 +1018,13 @@ def _self_consistency(
                     return v.copy()
 
             cx = z[: constraint_x.size].reshape(constraint_x.shape)
-            Delta1 = solver.Delta.copy()
-            Delta1[Delta_mask] = z[constraint_x.size :].view(complex).reshape(-1, 2, 2)
+            Delta1 = singlet(solver.Delta)
+            Delta1[Delta_mask] = z[constraint_x.size :].view(complex).ravel()
 
             if constraint_fun is not None:
                 res_cons = constraint_fun(cx, Delta1).ravel()
                 # Constraint function updates terminal phases
-                solver.Delta[...] = Delta1
+                solver.Delta[...] = singlet_m(Delta1)
 
             res, J, m = _self_consistent_Delta_f(
                 solver, Delta1, T, T_c0, workers=workers, **solver_kw
@@ -1088,10 +1096,8 @@ def _self_consistency(
             y = np.linspace(-solver.Ly / 2, solver.Ly / 2, solver.shape[1])
 
             plt.clf()
-            ddd = solver.Delta.copy()
-            ddd[Delta_mask] = z[constraint_x.size :].view(complex).reshape(-1, 2, 2)
-
-            ddd = ddd[..., 0, 0]
+            ddd = singlet(solver.Delta)
+            ddd[Delta_mask] = z[constraint_x.size :].view(complex).ravel()
 
             if plot == "circle" or min(ddd.shape) == 1:
                 ddd = ddd.squeeze()
@@ -1113,7 +1119,7 @@ def _self_consistency(
         success = False
 
     Delta = solver.Delta.copy()
-    Delta[Delta_mask] = z[constraint_x.size :].view(complex).reshape(-1, 2, 2)
+    Delta[Delta_mask] = singlet_m(z[constraint_x.size :].view(complex).ravel())
     cx = z[: constraint_x.size]
 
     if success and (np.isnan(J).any() or np.isnan(Delta).any()):
@@ -1140,7 +1146,7 @@ def _self_consistent_Delta_f(
     w = w[len(w) // 2 :][::-1]
     a = a[len(a) // 2 :][::-1]
 
-    rtot = np.zeros(solver.Delta.shape, dtype=complex)
+    rtot = np.zeros(tuple(solver.shape), dtype=complex)
     Jtot = np.zeros(tuple(solver.shape) + (4, 4, 4), dtype=complex)
 
     mask1 = solver.mask == MASK_NONE
@@ -1148,7 +1154,7 @@ def _self_consistent_Delta_f(
     rtot[mask1 & ~mask2] = Delta[mask1 & ~mask2]
     mask = mask1 & mask2
 
-    solver.Delta[mask] = Delta[mask]
+    solver.Delta[mask] = singlet_m(Delta[mask])
 
     if workers is not None:
         jobs = []
@@ -1167,7 +1173,7 @@ def _self_consistent_Delta_f(
         rtot[mask] += rtotx[mask]
         Jtot += Jtotx
 
-    rtot[mask] -= np.log(T_c0[mask, None, None] / T) * solver.Delta[mask]
+    rtot[mask] -= np.log(T_c0[mask] / T) * singlet(solver.Delta[mask])
 
     return rtot, Jtot, mask
 
@@ -1200,7 +1206,7 @@ def _mp_one(args, solver=None, solver_kw=None):
 
         for wx, ax in zip(w, a):
             res = solver.solve(omega=wx, **solver_kw)
-            r = solver.Delta.A / abs(wx) - res.F
+            r = singlet(solver.Delta) / abs(wx) - singlet(res.F)
             rtot += (2 * np.pi * ax) * r
             Jtot += (-2j * np.pi * ax) * res.J
     finally: