usadelndsoc: expose alpha/D in Solver, add convenience feature

usadelndsoc/core.py

@@ -132,6 +132,14 @@ class Core:
     def Delta(self):
         return _DeltaArray(self.Omega)
 
+    @property
+    def Ux(self):
+        return self.U[:, :, 0, :, :]
+
+    @property
+    def Uy(self):
+        return self.U[:, :, 1, :, :]
+
 
 class _DeltaArray:
     def __init__(self, Omega):

usadelndsoc/solver.py

@@ -36,7 +36,7 @@ except ImportError:
 from .core import Core, MASK_NONE, MASK_TERMINAL, MASK_VACUUM, _DeltaArray
 from .matsubara import get_matsubara_sum
 
-__all__ = ["Solver", "Result", "cpr", "MASK_NONE", "MASK_TERMINAL", "MASK_VACUUM"]
+__all__ = ["Solver", "Result", "cpr", "tr", "MASK_NONE", "MASK_TERMINAL", "MASK_VACUUM"]
 
 _log = logging.getLogger(__name__)
 _log_solve = logging.getLogger(__name__ + ".solve")
@@ -96,8 +96,12 @@ class Solver:
     y = _core_property("y")
     mask = _core_property("mask")
     Omega = _core_property("Omega")
+    alpha = _core_property("alpha")
+    D = _core_property("D")
     Delta = _core_property("Delta")
     U = _core_property("U")
+    Ux = _core_property("Ux")
+    Uy = _core_property("Uy")
 
     Phi = _array_property("_Phi")
     J_tot = _array_property("_J_tot")
@@ -407,7 +411,7 @@ class Solver:
         I = np.eye(2)
         s = np.sign(self._core.omega)
 
-        G = s * np.linalg.solve(I + g @ gt, I - g @ gt)
+        G = np.linalg.solve(I + g @ gt, I - g @ gt)
         if G[:, :, 0, 0].real.min() < -0.05 or G[:, :, 1, 1].real.min() < -0.05:
             return False
         return True
@@ -1081,8 +1085,8 @@ def cpr(
                 _log.info(f"CPR full cycle obtained: stop (pi reverse cross)")
                 break
             t3 = np.diag([1, 1, -1, -1])
-            I1 = (state.Js[-1] @ t3).trace(axis1=-2, axis2=-1)
-            I2 = (state.Js[-2] @ t3).trace(axis1=-2, axis2=-1)
+            I1 = tr(state.Js[-1] @ t3)
+            I2 = tr(state.Js[-2] @ t3)
             if (
                 state.phis[-2] < np.pi
                 and state.phis[-1] >= np.pi
@@ -1166,10 +1170,17 @@ def cpr(
         state.Js.append(solver.J_tot.copy())
 
         t3 = np.diag([1, 1, -1, -1])
-        J_avg = np.mean((solver.J_tot[:, :, 0].real @ t3).trace(axis1=-2, axis2=-1))
+        J_avg = np.mean(tr(solver.J_tot[:, :, 0].real @ t3))
         _log.info(f"CPR #{j}: phi = {phi}, J = {J_avg}")
 
         if filename is not None:
             state.save(filename)
 
     return np.asarray(state.phis), np.asarray(state.Deltas), np.asarray(state.Js)
+
+
+def tr(M):
+    """
+    Trace over last two dimensions
+    """
+    return np.trace(M, axis1=-2, axis2=-1)