Support ChiQ in JO/cthyb-seg solver

src/dcore/dcore_bse.py

@@ -49,7 +49,7 @@ def _compare_str_list(list1, list2):
 
 
 def calc_g2_in_impurity_model(solver_name, solver_params, mpirun_command, basis_rot, Umat, gf_struct, beta, n_iw,
-                              Sigma_iw, Gloc_iw, num_wb, num_wf, only_chiloc, ish, freqs=None):
+                              Sigma_iw, Gloc_iw, num_wb, num_wf, save_chiloc, only_chiloc, ish, freqs=None):
     """
 
     Calculate G2 in an impurity model
@@ -71,6 +71,9 @@ def calc_g2_in_impurity_model(solver_name, solver_params, mpirun_command, basis_
     s_params = copy.deepcopy(solver_params)
     s_params['random_seed_offset'] = 1000 * ish
 
+    s_params['save_chiloc'] = save_chiloc
+    s_params['only_chiloc'] = only_chiloc
+
     work_dir_org = os.getcwd()
     work_dir = 'work/imp_shell' + str(ish) + '_bse'
     if not os.path.isdir(work_dir):
@@ -83,7 +86,7 @@ def calc_g2_in_impurity_model(solver_name, solver_params, mpirun_command, basis_
     # Solve the model
     rot = impurity_solvers.compute_basis_rot(basis_rot, sol)
     if flag_box:
-        xloc, chiloc = sol.calc_Xloc_ph(rot, mpirun_command, num_wf, num_wb, s_params, only_chiloc)
+        xloc, chiloc = sol.calc_Xloc_ph(rot, mpirun_command, num_wf, num_wb, s_params)
     else:
         xloc, chiloc = sol.calc_Xloc_ph_sparse(rot, mpirun_command, freqs, num_wb, s_params)
 
@@ -167,6 +170,56 @@ def g(_i, _j, _w):
                     data[j] -= g(i2, i1, wf1) * g(i3, i4, wf2)
 
 
+def calc_X0_loc(gimp, spin_names, num_wb, num_wf, freqs=None):
+    """
+    calculate X0_loc from G_imp
+        X0_loc[(i1, i2, i3, i4)][wb, wf1, wf2] = - G[(i3, i1)][wf1] * G[(i2, i4)][wf2+wb]
+
+    if freqs is None: box frequency sampling
+    else: sparse sampling
+
+    """
+
+    # g_ij[(i, j)] = data[iw]
+    g_ij = {}
+    for isp, sp in enumerate(spin_names):
+        # data[iw, orb1, orb2]
+        norb = gimp[sp].data.shape[1]
+        assert norb == gimp[sp].data.shape[2]
+        for o1, o2 in product(list(range(norb)), repeat=2):
+            i = o1 + isp*norb
+            j = o2 + isp*norb
+            g_ij[(i, j)] = gimp[sp].data[:, o1, o2]
+
+    assert g_ij[(0, 0)].shape[0] % 2 == 0
+    w0 = g_ij[(0, 0)].shape[0] // 2
+
+    # reduce Matsubara frequencies
+    for key, g in g_ij.items():
+        g_ij[key] = g[w0-num_wf:w0+num_wf+num_wb]
+
+    # get maximum index
+    max_ij = numpy.max(numpy.array([key for key in g_ij.keys()]))
+
+    chi0_loc = {}
+    for i1, i2, i3, i4 in product(range(max_ij+1), repeat=4):
+        # skip if g_ij has no data
+        if (not (i3, i1) in g_ij) or (not (i2, i4) in g_ij):
+            continue
+
+        if freqs is None:
+            # box sampling
+            chi0 = numpy.zeros((num_wb, 2*num_wf), dtype=numpy.complex128)
+            for wb in range(num_wb):
+                chi0[wb, :] = - g_ij[(i3, i1)][0:2*num_wf] * g_ij[(i2, i4)][wb:wb+2*num_wf]
+            chi0_loc[(i1, i2, i3, i4)] = chi0
+        else:
+            # sparse sampling
+            raise NotImplementedError
+
+    return chi0_loc
+
+
 def gen_sparse_freqs_fix_boson(boson_freq, Lambda, sv_cutoff):
     """
     Generate sparse frequency points (wf1, wf2) for fixed bosonic freq wb
@@ -338,6 +391,9 @@ def decompose_index(index):
     def save_xloc(self, xloc_ijkl, icrsh):
         self._save_common(xloc_ijkl, icrsh, 'X_loc')
 
+    def save_x0loc(self, x0loc_ijkl, icrsh):
+        self._save_common(x0loc_ijkl, icrsh, 'X0_loc')
+
     def save_chiloc(self, chiloc_ijkl, icrsh):
         self._save_common(chiloc_ijkl, icrsh, 'chi_loc_in')
 
@@ -424,6 +480,7 @@ def _calc_bse_x0q(self):
         params['div'] = lattice_model.nkdiv()
         params['bse_h5_out_file'] = os.path.abspath(self._params['bse']['h5_output_file'])
         params['use_temp_file'] = self._params['bse']['use_temp_file']
+        params['flag_save_X0_loc'] = self._params['bse']['choice_X0loc'] == 'qsum'
         if self._params['bse']['X0q_qpoints_saved'] == 'quadrant':
             params['X0q_qpoints_saved'] = 'quadrant'
         else:
@@ -507,12 +564,25 @@ def calc_num_flavors(_ish):
                                                               self._sh_quant[ish].Sigma_iw, Gloc_iw_sh[ish],
                                                               self._params['bse']['num_wb'],
                                                               self._params['bse']['num_wf'],
+                                                              self._params['bse']['save_chiloc'],
                                                               self._params['bse']['calc_only_chiloc'],
                                                               ish, freqs=freqs)
 
             if x_loc is not None:
                 subtract_disconnected(x_loc, g_imp, self.spin_block_names, freqs=freqs)
 
+            if self._params['bse']['choice_X0loc'] == 'imp':
+                print("calc_X0_loc start")
+                params = dict(
+                    spin_names = self.spin_block_names,
+                    num_wf = self._params['bse']['num_wf'],
+                    num_wb = self._params['bse']['num_wb'],
+                )
+                x0_loc = calc_X0_loc(g_imp, **params, freqs=freqs)
+                print("calc_X0_loc end")
+            else:
+                x0_loc = None
+
             # Open HDF5 file to improve performance. Close manually.
             bse.h5bse.open('a')
 
@@ -525,15 +595,22 @@ def calc_num_flavors(_ish):
                     # chi_loc
                     if chi_loc is not None:
                         bse.save_chiloc(chi_loc, icrsh=icrsh)
+                    # X0_loc
+                    if x0_loc is not None:
+                        bse.save_x0loc(x0_loc, icrsh=icrsh)
 
             bse.h5bse.close()
 
     def calc_bse(self):
         """
         Compute data for BSE
         """
-        self._calc_bse_x0q()
-        self._calc_bse_xloc()
+        self._calc_bse_x0q()  # X_0(q)
+        self._calc_bse_xloc()  # X_{loc}
+
+        # X_{0,loc} is computed in
+        #   _calc_bse_x0q()  if calc_X0loc_from_X0q = True (default)
+        #   _calc_bse_xloc() if                     = False
 
 
     def fit_Xloc(self, save_only):

src/dcore/impurity_solvers/alps_cthyb.py

@@ -125,7 +125,7 @@ def solve(self, rot, mpirun_command, params_kw):
 
         self._solve_impl(rot, mpirun_command, None, params_kw)
 
-    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw, only_chiloc):
+    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw):
         raise RuntimeError("calc_Xloc_ph is not implemented!")
 
     def calc_G2loc_ph_sparse(self, rot, mpirun_command, wsample_ph, params_kw):

src/dcore/impurity_solvers/alps_cthyb_seg.py

@@ -23,6 +23,7 @@
 from dcore._dispatcher import *
 from ..tools import make_block_gf, launch_mpi_subprocesses, extract_H0, umat2dd, get_block_size, expand_path
 from .base import SolverBase, rotate_basis
+from dcore.program_options import parse_save_chiloc
 
 
 def to_numpy_array(g, names):
@@ -413,23 +414,20 @@ def set_blockgf_from_h5(sigma, group):
         # [(s1,o1), (s2,o2), 0]
         self.quant_to_save['nn_equal_time'] = nn_equal_time[:, :, 0]  # copy
 
-    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw, only_chiloc):
+    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw):
         """
         Compute local G2 in p-h channel
 
         For details, see SolverBase.calc_Xloc_ph
         """
 
-        if rot is not None:
-            # TODO
-            raise NotImplementedError
-
-        use_chi_loc = False
+        save_chiloc = parse_save_chiloc(params_kw['save_chiloc'], default=False)
+        only_chiloc = params_kw['only_chiloc']
 
         params_kw['cthyb.MEASURE_g2w'] = 1
         params_kw['cthyb.N_w2'] = num_wf
         params_kw['cthyb.N_W'] = num_wb
-        if use_chi_loc:
+        if save_chiloc:
             params_kw['cthyb.MEASURE_nnw'] = 1
 
         self.solve(rot, mpirun_command, params_kw)
@@ -476,7 +474,7 @@ def get_g2(_i, _j, _wb, _wf1, _wf2):
         # Save chi(wb)
         # [(s1,o1), (s2,o2), wb]
         chi_dict = None
-        if use_chi_loc:
+        if save_chiloc:
             chi_re = self._get_results("nnw_re", num_wb, orbital_symmetrize=True)
             chi_im = self._get_results("nnw_im", num_wb, orbital_symmetrize=True)
             chi_loc = chi_re + chi_im * 1.0J
@@ -487,6 +485,13 @@ def get_g2(_i, _j, _wb, _wf1, _wf2):
             for i1, i2 in product(range(2*self.n_orb), repeat=2):
                 chi_dict[(i1, i1, i2, i2)] = chi_loc[i1, i2]
 
+        # Rotate g2_dict and chi_dict back to the original basis
+        if rot is not None:
+            print("Error: 'basis_rotation' for two-particle quantities is currently not supported.", file=sys.stderr)
+            sys.exit(1)
+
+            rotate_basis(rot, self.use_spin_orbit, None, direction='backward', X_dict=g2_dict, chi_dict=chi_dict)
+
         return g2_dict, chi_dict
 
     def calc_G2loc_ph_sparse(self, rot, mpirun_command, freqs_ph, num_wb, params_kw):

src/dcore/impurity_solvers/base.py

@@ -147,7 +147,7 @@ def solve(self, rot, mpirun_command, params):
 
         # Set self.Gimp_iw, self.G_tau, self.Sigma_iw
 
-    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw, only_chiloc):
+    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw):
         """
         Compute local G2 in p-h channel
             X_loc = < c_{i1}^+ ; c_{i2} ; c_{i4}^+ ; c_{i3} >,  and
@@ -159,10 +159,12 @@ def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw, only_chil
             Number of non-negative fermionic frequencies
         num_wb: int
             Number of non-negative bosonic frequencies
-        only_chiloc: bool
-            If True, only chi_loc is computed (no Xloc).
 
         The other parameters are the same as for solve().
+        params_kw includes the following parameters in addition to the ones used in solve().
+
+            'only_chiloc': bool, if True, only chi_loc is computed (no Xloc).
+            'save_chiloc': bool, if True, chi_loc is saved.
 
         Returns
         -------
@@ -259,7 +261,7 @@ def is_Floc_computable(cls):
 
 
 
-def rotate_basis(rot, use_spin_orbit, u_matrix, Gfs=[], direction='forward'):
+def rotate_basis(rot, use_spin_orbit, u_matrix, Gfs=[], direction='forward', X_dict=None, chi_dict=None):
     """
     Rotate all Gf-like objects and U-matrix to the basis defined by rot
 
@@ -268,17 +270,17 @@ def rotate_basis(rot, use_spin_orbit, u_matrix, Gfs=[], direction='forward'):
     """
 
     if direction == 'forward':
-        return _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs)
+        return _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs, X_dict, chi_dict)
     elif direction == 'backward':
         rot_conj_trans = {}
         for name, r in list(rot.items()):
             rot_conj_trans[name] = r.conjugate().transpose()
-        return _rotate_basis(rot_conj_trans, u_matrix, use_spin_orbit, Gfs)
+        return _rotate_basis(rot_conj_trans, u_matrix, use_spin_orbit, Gfs, X_dict, chi_dict)
     else:
         raise RuntimeError("Unknown direction " + direction)
 
 
-def _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs):
+def _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs, X_dict, chi_dict):
     """
     Rotate all Gf-like object and U matrix to a new local basis defined by "rot".
     :param rot: matrix
@@ -298,11 +300,52 @@ def _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs):
         for bname, gf in G:
             gf.from_L_G_R(rot[bname].transpose().conjugate(), gf, rot[bname])
 
+    n_flavors = rot_spin_full.shape[0]
+
+    if X_dict is not None:
+        shape = next(iter(X_dict.values())).shape  # (num_wb, num_wf, num_wf)
+        assert len(shape) == 3
+        array = numpy.zeros((n_flavors, n_flavors, n_flavors, n_flavors) + shape, dtype=numpy.complex128)
+        _set_from_dict(X_dict, array)
+
+        # X_{1234}(iW, iw, iw') = < c_1^+(iw) c_2(iw+iW) c_4^+(iw'+iW) c_3(iw') >
+        array = numpy.einsum("ijklxyz,im,jn,ko,lp -> mnopxyz", array,
+                    rot_spin_full, numpy.conj(rot_spin_full), numpy.conj(rot_spin_full), rot_spin_full)
+
+        X_dict.clear()
+        _set_to_dict(X_dict, array)
+
+    if chi_dict is not None:
+        shape = next(iter(chi_dict.values())).shape  # (num_wb,)
+        assert len(shape) == 1
+        array = numpy.zeros((n_flavors, n_flavors, n_flavors, n_flavors) + shape, dtype=numpy.complex128)
+        _set_from_dict(chi_dict, array)
+
+        # chi_{1234}(iW) = < c_1^+ c_2 c_4^+ c_3 >(iW)
+        array = numpy.einsum("ijklx,im,jn,ko,lp -> mnopx", array,
+                    rot_spin_full, numpy.conj(rot_spin_full), numpy.conj(rot_spin_full), rot_spin_full)
+
+        chi_dict.clear()
+        _set_to_dict(chi_dict, array)
+
     if not u_matrix is None:
         return numpy.einsum("ijkl,im,jn,ko,lp", u_matrix,
                                     numpy.conj(rot_spin_full), numpy.conj(rot_spin_full), rot_spin_full, rot_spin_full)
 
 
+def _set_from_dict(x_dict, x_array):
+    for key, val in x_dict.items():
+        i, j, k, l = key
+        x_array[i, j, k, l] = val  # val is np.array
+
+
+def _set_to_dict(x_dict, x_array):
+    n1, n2, n3, n4 = x_array.shape[:4]
+    for i, j, k, l in product(range(n1), range(n2), range(n3), range(n4)):
+        if not numpy.all(np.abs(x_array[i, j, k, l]) < 1e-8):  # if not zero matrix
+            x_dict[(i, j, k, l)] = x_array[i, j, k, l]
+
+
 class PytriqsMPISolver(SolverBase):
 
     def __init__(self, beta, gf_struct, u_mat, n_iw=1025):