Support complex-Hermitian Green's functions in the multi-orbital solver

README.md

@@ -25,7 +25,10 @@ The names of parameters are case-insensitive.
     - Filename storing G(τ)
     - When `num_flavor > 1`, "{filein_g}.{a}_{b}" are read for G_{ab}, where a,b = 0,1,...,num_flavor-1
 - column: int
-    - Index of column storing G(τ) (0-origin)
+    - Index of column storing Re G(τ) (0-origin)
+- column_imag: int (optional)
+    - Index of column storing Im G(τ) (0-origin)
+    - If omitted, Im G(τ) is treated as 0 (real input, backward compatible)
 - beta: float
     - Inverse temperature β
 - max_omega: float

pyproject.toml

@@ -44,3 +44,6 @@ build-backend = "hatchling.build"
 
 [tool.mypy]
 files = "src"
+
+[tool.uv.sources]
+admmsolver = { git = "https://github.com/SpM-lab/admmsolver", branch = "complex-hermitian-support" }

src/spmac/main.py

@@ -49,9 +49,24 @@ def optimize(
     return loglambdas[i]
 
 
+def _load_gtau_from_file(
+    filename: str, column: int, column_imag: Optional[int] = None
+) -> NDArray:
+    """Load G(τ) from a text file.
+
+    If ``column_imag`` is given, return complex G = Re + i Im.
+    Otherwise return the real column only (imaginary part is treated as 0).
+    """
+    data = np.loadtxt(filename)
+    g = data[:, column]
+    if column_imag is not None:
+        g = g + 1j * data[:, column_imag]
+    return g
+
+
 def run(
-    params: Dict[str, Any], Gtau: NDArray[np.float64]
-) -> Tuple[SolverBase, NDArray[np.float64], float]:
+    params: Dict[str, Any], Gtau: NDArray
+) -> Tuple[SolverBase, NDArray, float]:
     params = dict_with_lowerkey(params)
     verbose = params.get("verbose", True)
     n_trials = params.get("num_trials", 10)
@@ -103,18 +118,21 @@ def main():
 
     nflavor = params.get("num_flavor", 1)
     column: int = params["column"]
+    column_imag: Optional[int] = params.get("column_imag")
+    print(f"column: {column}, column_imag: {column_imag}")
     if nflavor == 1:
-        Gtau = np.loadtxt(params["filein_g"])[:, column]
+        Gtau = _load_gtau_from_file(params["filein_g"], column, column_imag)
     else:
-        Gtau = np.zeros((1, nflavor, nflavor))
+        dtype = np.complex128 if column_imag is not None else np.float64
+        Gtau = np.zeros((1, nflavor, nflavor), dtype=dtype)
         ntau = -1
         for i in range(nflavor):
             for j in range(nflavor):
                 filename = params["filein_g"] + f".{i}_{j}"
-                Gt = np.loadtxt(filename)[:, column]
+                Gt = _load_gtau_from_file(filename, column, column_imag)
                 if ntau < 0:
                     ntau = len(Gt)
-                    Gtau = np.zeros((ntau, nflavor, nflavor))
+                    Gtau = np.zeros((ntau, nflavor, nflavor), dtype=dtype)
                 Gtau[:, i, j] = Gt[:]
 
     solver, rho_l, best_loglambda = run(params, Gtau)

src/spmac/solver/multiple.py

@@ -229,8 +229,10 @@ def callback():
             print("  Not converged")
         return opt
 
-    def predict_rho(self, rho_l) -> NDArray[np.float64]:
-        r_l = rho_l.reshape((-1, self.nflavor, self.nflavor))
+    def predict_rho(self, rho_l) -> NDArray:
+        r_l = np.asarray(rho_l).reshape((-1, self.nflavor, self.nflavor))
+        # Keep the dtype of rho_l: complex off-diagonal rho_ab(omega) must not be
+        # truncated to its real part (needed for complex-Hermitian / SOC data).
         if self.use_sparse_ir:
             return np.einsum("lab,lw->wab", r_l, self.basis.v(self.ws))
         else:
@@ -242,9 +244,11 @@ def predict_Gtau(self, rho_l, idx=None):
         r_l = rho_l.reshape((-1, self.nflavor, self.nflavor))
         if self.use_sparse_ir:
             sampler = sparse_ir.TauSampling(self.basis, self.ts[idx])
-            ret = np.zeros((len(idx), self.nflavor, self.nflavor))
+            # Keep the dtype of rho_l: complex off-diagonal G_ab(tau) must not be
+            # truncated to its real part (needed for complex-Hermitian / SOC data).
+            ret = np.zeros((len(idx), self.nflavor, self.nflavor), dtype=r_l.dtype)
             for i, j in itertools.product(range(self.nflavor), repeat=2):
-                ret[:, i, j] = np.real(sampler.evaluate(self.basis.s * r_l[:, i, j]))
+                ret[:, i, j] = sampler.evaluate(self.basis.s * r_l[:, i, j])
             return ret
         else:
             return np.einsum("lab, l, li -> iab", r_l, self.s, self.u[:,idx])
@@ -277,7 +281,8 @@ def write_Gtau(self, ts, Gts, outdir: pathlib.Path):
                     f.write(str(ts[i]))
                     for gt in Gts:
                         y = gt[i, ifl, jfl]
-                        f.write(f" {np.real(y)}")
+                        # real then imaginary part (imag is 0 for real data)
+                        f.write(f" {np.real(y)} {np.imag(y)}")
                     f.write("\n")
 
     def write_rhol(
@@ -304,7 +309,7 @@ def write_rhol(
                 if loglambda is not None:
                     f.write(f"# log_lambda = {loglambda}\n")
                 for i, r in enumerate(r_l[:, ifl, jfl]):
-                    f.write(f"{i} {np.real(r)}\n")
+                    f.write(f"{i} {np.real(r)} {np.imag(r)}\n")
 
     def write_rho(
         self,
@@ -335,4 +340,5 @@ def write_rho(
                 if loglambda is not None:
                     f.write(f"# log_lambda = {loglambda}\n")
                 for w, r in zip(ws, rs[:, ifl, jfl]):
-                    f.write(f"{w} {np.real(r)/dw}\n")
+                    # real then imaginary part (imag is 0 for real data)
+                    f.write(f"{w} {np.real(r)/dw} {np.imag(r)/dw}\n")

src/spmac/solver/single.py

@@ -232,17 +232,19 @@ def callback():
             print("  Not converged")
         return opt
 
-    def predict_rho(self, rho_l: NDArray[np.float64]) -> NDArray[np.float64]:
+    def predict_rho(self, rho_l) -> NDArray:
+        # Keep the dtype of rho_l so complex spectra are not truncated to real.
         if self.use_sparse_ir:
             return rho_l @ self.basis.v(self.ws)
         else:
             return rho_l @ self.v
 
-    def predict_Gtau(self, rho_l: NDArray[np.float64], idx=None) -> NDArray[np.float64]:
+    def predict_Gtau(self, rho_l, idx=None) -> NDArray:
         if idx is None:
             idx = np.arange(self.ntau)
         if self.use_sparse_ir:
             sampler = sparse_ir.TauSampling(self.basis, self.ts[idx])
+            # Keep the dtype of rho_l so complex G(tau) is not truncated to real.
             return sampler.evaluate(self.basis.s * rho_l)
         else:
             return rho_l @ self.u[:, idx]
@@ -268,7 +270,8 @@ def write_Gtau(self, ts, Gts, outdir: pathlib.Path):
                 f.write(str(ts[i]))
                 for gt in Gts:
                     y = gt[i]
-                    f.write(f" {np.real(y)}")
+                    # real then imaginary part (imag is 0 for real data)
+                    f.write(f" {np.real(y)} {np.imag(y)}")
                 f.write("\n")
 
     def write_rhol(
@@ -292,7 +295,7 @@ def write_rhol(
             if loglambda is not None:
                 f.write(f"# log_lambda = {loglambda}\n")
             for i, r in enumerate(rs):
-                f.write(f"{i} {np.real(r)}\n")
+                f.write(f"{i} {np.real(r)} {np.imag(r)}\n")
 
     def write_rho(
         self,
@@ -321,5 +324,5 @@ def write_rho(
             if loglambda is not None:
                 f.write(f"# log_lambda = {loglambda}\n")
             for w, r in zip(ws, rs):
-                f.write(f"{w} {np.real(r)/dw}\n")
-                # f.write(f"{w} {np.real(r)}\n")
+                # real then imaginary part (imag is 0 for real data)
+                f.write(f"{w} {np.real(r)/dw} {np.imag(r)/dw}\n")