Push get_free_energy_hessian_dev() and environment

Modules/Ensemble.py

@@ -6,6 +6,8 @@
 import time
 #from scipy.special import tanh, sinh, cosh
 
+import sscha.Classify as Classify
+import sscha.qClassify as qClassify
 
 """
 This is part of the program python-sscha
@@ -250,6 +252,19 @@ def __init__(self, dyn0, T0, supercell = None, **kwargs):
         # A flag that memorize if the ensemble has also the stresses
         self.has_stress = True
 
+        # Store data for calculations of 3FC elements:
+        self.a = np.zeros( (Nsc * 3), dtype = np.double)
+        self.new_pol = np.zeros( (Nsc, Nsc * 3, 3), dtype = np.double)
+        # If the element is sym:
+        self.ur = np.zeros( (self.N, Nsc * 3)) 
+        self.upsilon = np.zeros( (Nsc*3, Nsc * 3))
+        # Symmetry data for sym 3FC:
+        self.nsym = 0
+        self.s_cart = np.zeros( (3, 3, 48) , dtype = np.float64, order = "F")
+        self.s_inv_cart = np.zeros( (3, 3, 48) , dtype = np.float64, order = "F")
+        self.irt = np.zeros( (48, Nsc), dtype = np.intc, order = "F")
+        self.translations_irt = np.zeros( (Nsc, np.prod(self.supercell)), dtype = np.intc, order = "F") 
+
         # A flag for each configuration that check if it possess a force and a stress
         self.force_computed = None
         self.stress_computed = None
@@ -3831,7 +3846,254 @@ def get_free_energy_hessian(self, include_v4 = False, get_full_hessian = True, v
             return dyn_hessian, d3* 2.0 # Ha to Ry
         return dyn_hessian
 
+    def get_free_energy_hessian_dev(self, include_v4 = False, do_scf = True, get_full_hessian = True, verbose = False):
+        """
+        Dev function.
+
+        GET THE FREE ENERGY ODD CORRECTION DEV
+        ======================================
 
+        This function computes the third and fourth order corrections classifying the
+        polarization vector by wave-vector q. We rewrote the hessian from Bianco's paper
+        into smaller objects. The expression will be soon uploaded.
+
+        To reduce the required RAM, only the symmetry independent third and fourth order
+        force constants are saved in memory, computing the rest on the fly.
+
+        Parameters
+        ----------
+            include_v4 : bool
+                If True we include the fourth order force constant matrix.
+                This requires a lot of memory
+            get_full_hessian : bool
+                If True the full hessian matrix is returned, if false, only the correction to
+                the SSCHA dynamical matrix is returned.
+            verbose : bool
+                If true, the third order force constant tensor is written in output [Ha/bohr^3 units].
+                This can be used to interpolate the result on a bigger mesh with cellconstructor.
+
+        Returns
+        -------
+            phi_sc : Phonons()
+                The dynamical matrix of the free energy hessian in (Ry/bohr^2)
+        """
+
+        self.convert_units(UNITS_HARTREE)
+        super_structure = self.current_dyn.structure.generate_supercell(self.supercell)
+        dyn_supercell = self.current_dyn.GenerateSupercellDyn(self.supercell)
+        nr = np.prod(self.supercell)
+
+        nat_sc = dyn_supercell.structure.N_atoms
+        n_modes = nat_sc*3
+
+        mapping, rot_cart, map_uc, map_tr, T_list, T_list_frac = Classify.map_singlet(self.current_dyn, verbose = verbose)
+
+        orbit2a, orbit2s, norbit, indep_elem, n_indep_elem, tensor = Classify.recognize_doublet(self.current_dyn, mapping, map_uc, verbose = verbose)
+
+        orbit3a, orbit3s, norbit3, indep_3fc_elem, n_indep_3fc_elem, kernel_3fc, rot_3fc, mapping_triplet = Classify.recognize_triplet(self.current_dyn, mapping, map_uc, verbose=verbose)
+
+        nref2 = orbit2a.shape[0]
+        indep_fc = np.zeros((nref2, max(n_indep_elem)), dtype=np.float64)
+
+        super_structure = self.current_dyn.structure.generate_supercell(self.supercell)
+        w, pols = self.current_dyn.DiagonalizeSupercell()
+
+        n_modes = len(w)
+        nat_sc = int(np.shape(pols)[0] / 3)
+        nat = self.dyn_0.structure.N_atoms
+
+        # Get the translational modes
+        if not self.ignore_small_w:
+            trans = CC.Methods.get_translations(pols, super_structure.get_masses_array())
+        else:
+            trans = np.abs(w) < CC.Phonons.__EPSILON_W__
+
+
+        # Get the atomic types
+        ityp = super_structure.get_ityp() + 1 #Py to Fortran indexing
+        n_typ = len(self.current_dyn.structure.masses)
+
+        amass = np.zeros(n_typ, dtype = np.double)
+
+        for at_type in self.current_dyn.structure.masses:
+            index = ityp[self.current_dyn.structure.atoms.index(at_type)] - 1
+            amass[index] = self.current_dyn.structure.masses[at_type]
+
+        # Get the forces and conver in the correct units
+        f = (self.forces - self.sscha_forces)# * Bohr
+        u = self.u_disps.reshape((self.N, nat_sc, 3), order = "C") #/ Bohr
+
+        log_err = "err_yesrho"
+        self.a = SCHAModules.thermodynamic.w_to_a(w, self.current_T)
+        # Get the polarization vectors in the correct format
+        for i in range(nat_sc):
+            for j in range(n_modes):
+                self.new_pol[i, j, :] = pols[3*i : 3*(i+1), j]
+
+        #Calculating rotated displacements and upsilon matrix"
+        self.ur, self.upsilon = SCHAModules.get_ur_upsilon_matrices(self.a, self.new_pol, trans, amass, ityp, u)
+        #Obtaining symmetry data
+        qe_sym = CC.symmetries.QE_Symmetry(super_structure)
+        qe_sym.SetupFromSPGLIB() 
+
+        # SpaceGroup symmetries data
+        self.nsym, self.s_cart, self.s_inv_cart, self.irt = qe_sym.QE_nsym, qe_sym.QE_s_cart, qe_sym.QE_s_inv_cart, qe_sym.QE_irt
+        self.s_cart = np.asfortranarray(self.s_cart)
+        self.s_inv_cart = np.asfortranarray(self.s_inv_cart)
+
+        #Translations
+        self.translations_irt = qe_sym.QE_translations_irt
+        #self.prepare_sym_third_fc = False
+
+        wq = np.empty(nat*3, dtype=np.float64)
+        tmp_pol_vecs = np.empty([nat*3, nat*3], dtype=np.complex128)
+        norm_pol_vecs = np.empty([nat*3, nat*3], dtype=np.complex128)
+        pol_vecs = np.empty([self.supercell[0]*self.supercell[1]*self.supercell[2], nat*3, nat_sc*3], dtype=np.complex128)
+        l = np.empty([self.supercell[0]*self.supercell[1]*self.supercell[2], nat*3, nat_sc*3], dtype=np.complex128)
+        wq = np.empty([l.shape[0], l.shape[1]], dtype=np.float64)
+        aq = np.empty([l.shape[0], l.shape[1]], dtype=np.float64)
+
+        rcell = CC.Methods.get_reciprocal_vectors(self.current_dyn.structure.unit_cell)
+        q_list = np.empty([l.shape[0], 3], dtype=np.float64)
+        q_list_cart = np.empty([l.shape[0], 3], dtype=np.float64)
+        qi=0
+        for qstar in self.current_dyn.q_stars:
+            for q in qstar:
+                q_cryst = np.round(CC.Methods.cart_to_cryst(rcell, q)/__A_TO_BOHR__,6)
+                wq[qi], tmp_pol_vecs = self.current_dyn.DyagDinQ(iq=qi)
+                tmp_pol_vecs = np.transpose(tmp_pol_vecs) # First index mode. Second incex atom, alpha.
+                aq[qi] = SCHAModules.thermodynamic.w_to_a(wq[qi], self.current_T)
+                for mode in range(nat*3):
+                    for atom in range(nat):
+                        norm_pol_vecs[mode,atom*3:(atom+1)*3] = tmp_pol_vecs[mode,atom*3:(atom+1)*3]*aq[qi, mode]/np.sqrt(amass[ityp[atom]-1])
+                for mode in range(nat*3):
+                    for Tx in range(self.supercell[0]):
+                        for Ty in range(self.supercell[1]):
+                            for Tz in range(self.supercell[2]):
+                                T = np.array([Tx, Ty, Tz])
+                                index = Tx*self.supercell[2]*self.supercell[1]+Ty*self.supercell[2]+Tz
+                                l[qi,mode,index*nat*3:(index+1)*nat*3] = norm_pol_vecs[mode]*np.exp(2j*np.pi*np.matmul(q_cryst,T))/np.sqrt(self.supercell[0]*self.supercell[1]*self.supercell[2], dtype=np.complex128)
+                                pol_vecs[qi,mode,index*nat*3:(index+1)*nat*3] = tmp_pol_vecs[mode]*np.exp(2j*np.pi*np.matmul(q_cryst,T))/np.sqrt(self.supercell[0]*self.supercell[1]*self.supercell[2], dtype=np.complex128)
+                q_list[qi] = q_cryst
+                q_list_cart[qi] = q
+                qi+=1
+
+        Nq = l.shape[0]
+        # Get translational modes
+        transq = np.zeros(wq.shape, dtype=bool)
+        for qi in range(Nq):
+            if not self.ignore_small_w:
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore")
+                    transq[qi,:] = CC.Methods.get_translations(np.transpose(pol_vecs[qi,:,:nat*3]), super_structure.get_masses_array()[:nat])
+            else:
+                transq[qi,:] = np.abs(w) < CC.Phonons.__EPSILON_W__
+            for mode in range(nat*3):
+                if transq[qi, mode]:
+                    l[qi,mode] = np.zeros(nat_sc*3, dtype=np.complex128)
+                    pol_vecs[qi,mode] = np.zeros(nat_sc*3, dtype=np.complex128)
+
+        # Imposing eps(q)=[eps(-q)]* time-reversal criteria")
+
+        ############################## Start impose TRS ##############################
+        mappingq, orbitq1a, orbitq1s, its_zb = qClassify.map_singlet(q_list_cart, q_list, rcell*__A_TO_BOHR__, rot_cart)
+        k = 0
+        trs_qlist = np.empty(q_list.shape, dtype=np.float64)
+        trs_qlist_cart = np.empty(q_list.shape, dtype=np.float64)
+        trs_polvecs = np.empty(l.shape, dtype=np.complex128)
+        trs_l = np.empty(l.shape, dtype=np.complex128)
+        trs_aq = np.empty(aq.shape, dtype=np.float64)
+        trs_wq = np.empty(wq.shape, dtype=np.float64)
+
+        for qi, q in enumerate(q_list):
+            if its_zb[qi] == 0:
+                trs_qlist[k] = q
+                trs_qlist_cart[k] = CC.Methods.cryst_to_cart(rcell*__A_TO_BOHR__, q)
+                trs_polvecs[k] = pol_vecs[qi]
+                trs_l[k] = l[qi]
+                trs_aq[k] = aq[qi]
+                trs_wq[k] = wq[qi]
+                k+=1
+            elif its_zb[qi] == 1:
+                trs_qlist[k] = q
+                trs_qlist_cart[k] = CC.Methods.cryst_to_cart(rcell*__A_TO_BOHR__, q)
+                trs_polvecs[k] = pol_vecs[qi]
+                trs_l[k] = l[qi]
+                trs_aq[k] = aq[qi]
+                trs_wq[k] = wq[qi]
+                k+=1
+                trs_qlist[k] = (-1)*q
+                trs_qlist_cart[k] = CC.Methods.cryst_to_cart(rcell*__A_TO_BOHR__, (-1)*q)
+                trs_polvecs[k] = np.conjugate(pol_vecs[qi])
+                trs_l[k] = np.conjugate(l[qi])
+                trs_aq[k] = aq[qi]
+                trs_wq[k] = wq[qi]
+                k+=1
+        ############################## End impose TRS ##############################
+
+        # Redo the classification imposing the TRS criteria.
+        mappingq, orbitq1a, orbitq1s, its_zb = qClassify.map_singlet(trs_qlist_cart, trs_qlist, rcell*__A_TO_BOHR__, rot_cart)
+        refq2, refq2o, norbitq2, nrefq2 = qClassify.recognize_doublet(trs_qlist, mappingq)
+
+        mod = self.supercell
+        nat = self.current_dyn.structure.N_atoms
+        n_modes = self.current_dyn.structure.N_atoms*mod[0]*mod[1]*mod[2]*3
+
+        ref_3fc = SCHAModules.module_hess.get_ref3fc(nat, orbit3a, indep_3fc_elem, n_indep_3fc_elem, kernel_3fc, rot_3fc, self.ur, self.upsilon, f, self.rho, log_err, self.s_inv_cart, self.irt, self.translations_irt, verbose)
+
+        vs_red = np.empty([nrefq2,nat*3,nat*3,nat*3], dtype=np.complex128)
+        vs_red = SCHAModules.module_hess.get_ref_vsq(refq2,trs_l,rot_3fc,ref_3fc,mapping_triplet,verbose)
+        trs_gq, daq = SCHAModules.get_gq(trs_aq, trs_wq, transq, self.current_T)
+        indep_fc = SCHAModules.module_hess.get_indep2fc(vs_red, refq2, refq2o, norbitq2, orbit2a, n_indep_elem, indep_elem, rot_cart, mapping, map_uc, map_tr, T_list, trs_qlist, trs_gq, verbose)
+
+        if include_v4:
+            refq4, refq4o, norbitq4, nrefq4 = qClassify.recognize_quadruplet(trs_qlist,mappingq,verbose)
+
+            orbit4t, orbit4o, norbit_4, indep_4fc_elem, n_indep_4fc_elem, kernel_4fc, rot_4fc, mapping_quadruplet = Classify.recognize_quadruplet(self.current_dyn, mapping, map_uc, verbose)
+
+            ref_4fc = SCHAModules.module_hess.get_ref4fc(orbit4t, indep_4fc_elem, n_indep_4fc_elem, kernel_4fc, rot_4fc, self.ur, self.upsilon, f, self.rho, log_err, self.s_inv_cart, self.irt, self.translations_irt, verbose)
+
+            ws_red = np.zeros([nrefq4,nat*3,nat*3,nat*3,nat*3], dtype=np.complex128)
+            ws_red = SCHAModules.module_hess.get_ref_wsq(refq4,trs_l,rot_4fc,ref_4fc,mapping_quadruplet,verbose)
+
+            degs = qClassify.find_degeneracies(trs_wq)
+            Pmn = qClassify.construct_Pmn(mapping, orbitq1a, orbitq1s, trs_polvecs, rot_cart)
+
+            v_red, ref_3fc = SCHAModules.module_hess.get_v3_red(nat, norbit3, orbit3a, orbit3s, indep_3fc_elem, n_indep_3fc_elem, kernel_3fc, rot_3fc, self.ur, self.upsilon, f, self.rho, log_err, self.s_inv_cart, self.irt, self.translations_irt)
+            vs = SCHAModules.module_hess.get_all_vsq(trs_l, v_red, map_uc)
+
+            if do_scf:
+                ws_red_scf = SCHAModules.module_hess.get_scf_wsq(ws_red, trs_gq, refq4, refq4o, norbitq4, Pmn, degs, verbose)
+                indep_fc4 = SCHAModules.module_hess.get_indep2fc_v4(vs, ws_red_scf, refq4, refq4o, norbitq4, orbit2a, n_indep_elem, indep_elem, trs_gq, Pmn, degs, mapping, rot_cart, verbose)
+            else:
+                indep_fc4 = SCHAModules.module_hess.get_indep2fc_v4(vs, ws_red, refq4, refq4o, norbitq4, orbit2a, n_indep_elem, indep_elem, trs_gq, Pmn, degs, mapping, rot_cart, verbose)
+            indep_fc += indep_fc4
+        phi_sc_odd = np.zeros((n_modes, n_modes), dtype = np.double)
+        for ref2 in range(nref2):
+            for i in range(norbit[ref2]):
+                nat1, nat2 = orbit2a[ref2,i,:]
+                fc9 = np.dot(tensor[ref2,i,:, :n_indep_elem[ref2]], indep_fc[ref2, :n_indep_elem[ref2]])
+                for alpha in range(3):
+                    for beta in range(3):
+                        index = 3*alpha+beta 
+                        #Apply translation sym:
+                        for r in range(nr):
+                            # Translated Single atomic-cartesian index
+                            # Fortran to Py: -1
+                            phi_sc_odd[3*(self.translations_irt[nat1,r]-1)+alpha, 3*(self.translations_irt[nat2,r]-1)+beta] = fc9[index]
+        dynq_odd = CC.Phonons.GetDynQFromFCSupercell(phi_sc_odd, np.array(self.current_dyn.q_tot), 
+                                            self.current_dyn.structure, super_structure)
+        self.convert_units(UNITS_DEFAULT)
+        dynq_odd *= 2 # Ha/bohr^2 -> Ry/bohr^2
+
+        # Generate the Phonon structure by including the odd correction
+        dyn_hessian = self.current_dyn.Copy()
+        for iq in range(len(self.current_dyn.q_tot)):
+            if get_full_hessian:
+                dyn_hessian.dynmats[iq] = self.current_dyn.dynmats[iq] + dynq_odd[iq, :, :]
+            else:
+                dyn_hessian.dynmats[iq] = dynq_odd[iq, :, :]
+        return dyn_hessian
 
     def compute_ensemble(self, calculator, compute_stress = True, stress_numerical = False,
                          cluster = None, verbose = True, timer=None):