solver: add some helper features

usadelndsoc/solver.py

@@ -36,12 +36,39 @@ except ImportError:
 from .core import Core, MASK_NONE, MASK_TERMINAL, MASK_VACUUM, _DeltaArray
 from .matsubara import get_matsubara_sum
 
-__all__ = ["Solver", "Result", "cpr", "tr", "MASK_NONE", "MASK_TERMINAL", "MASK_VACUUM"]
+__all__ = [
+    "Solver",
+    "Result",
+    "cpr",
+    "MASK_NONE",
+    "MASK_TERMINAL",
+    "MASK_VACUUM",
+    "tr",
+    "LEFT",
+    "UP",
+    "RIGHT",
+    "DOWN",
+    "S_x",
+    "S_y",
+    "S_z",
+    "S_0",
+]
 
 _log = logging.getLogger(__name__)
 _log_solve = logging.getLogger(__name__ + ".solve")
 
 
+S_x = np.array([[0, 1], [1, 0]])
+S_y = np.array([[0, -1j], [1j, 0]])
+S_z = np.array([[1, 0], [0, -1]])
+S_0 = np.array([[1, 0], [0, 1]])
+
+LEFT = 0
+UP = 1
+RIGHT = 2
+DOWN = 3
+
+
 def _core_property(name):
     def fget(self):
         return getattr(self._core, name)
@@ -599,6 +626,18 @@ class Result:
     def J(self):
         return self._core.grad_A(self.Phi).transpose(0, 1, 2, 4, 3) * (1 / 16)
 
+    @property
+    def J_c(self):
+        t3 = np.diag([1, 1, -1, -1])
+        return tr(self.J @ t3)
+
+    @property
+    def J_s(self):
+        jx = tr(self.J @ np.kron(s0, sx))
+        jy = tr(self.J @ np.kron(s0, sy))
+        jz = tr(self.J @ np.kron(s0, sz))
+        return np.array([jx, jy, jz]).transpose(1, 2, 3, 0)
+
     @property
     def S(self):
         return self._core.grad_Omega(self.Phi).transpose(0, 1, 3, 2)