Merge pull request #60 from SSCHAcode/fast_w_update

mesonepigreco · web-flow · commit a8394dbcd848 · 2022-05-20T18:15:07.000+02:00
Now we removed a lot of diagonalization from the code, this should speedup by a significative amount the code.
diff --git a/Modules/Ensemble.py b/Modules/Ensemble.py
@@ -136,6 +136,10 @@ def __init__(self, dyn0, T0, supercell = None, **kwargs):
         self.stresses = []
         self.xats = []
         self.units = UNITS_DEFAULT
+        self.w_0 = None
+        self.pols_0 = None
+        self.current_w = None
+        self.current_pols = None
         
         self.sscha_energies = []
         self.sscha_forces = []
@@ -221,6 +225,12 @@ def __setattr__(self, name, value):
         else:
             super(Ensemble, self).__setattr__(name, value)
 
+        if name == "dyn_0":
+            self.w_0, self.pols_0 = value.DiagonalizeSupercell()
+            self.current_dyn = value.Copy()
+            self.current_w = self.w_0.copy()
+            self.current_pols = self.pols_0.copy()
+
 
     def convert_units(self, new_units):
         """
@@ -1297,14 +1307,40 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         
         self.current_T = newT
 
-            
+        # Check if the dynamical matrix has changed
+        changed_dyn = np.max([np.max(np.abs(self.current_dyn.dynmats[i] - new_dynamical_matrix.dynmats[i])) for i in range(len(self.current_dyn.q_tot))])
+        print("DYN CHANGED BY:", changed_dyn)
+        changed_dyn = changed_dyn > 1e-30
+
+        # Prepare the new displacements
+        super_struct0 = self.dyn_0.structure.generate_supercell(self.supercell)
+        super_structure = new_dynamical_matrix.structure.generate_supercell(self.supercell)
+        old_disps = np.zeros(np.shape(self.u_disps), dtype = np.double)
+        Nat_sc = super_structure.N_atoms
+
+        self.u_disps[:,:] = self.xats.reshape((self.N, 3*Nat_sc)) - np.tile(super_structure.coords.ravel(), (self.N,1))
+        old_disps[:,:] = self.xats.reshape((self.N, 3*Nat_sc)) - np.tile(super_struct0.coords.ravel(), (self.N,1))
+
+        if changed_dyn:
+            w_new, pols = new_dynamical_matrix.DiagonalizeSupercell()#new_super_dyn.DyagDinQ(0)
+            self.current_w = w_new.copy()
+            self.current_pols = pols.copy()
+        else:
+            w_new = self.current_w.copy()
+            pols  = self.current_pols.copy()
+            #w_new, pols = new_dynamical_matrix.DiagonalizeSupercell()#new_super_dyn.DyagDinQ(0)
+            #self.current_w = w_new.copy()
+            #self.current_pols = pols.copy()
+        # Update sscha energies and forces
+        self.sscha_energies[:], self.sscha_forces[:,:,:] = new_dynamical_matrix.get_energy_forces(None, displacement = self.u_disps, w_pols = (w_new, pols))
+
         
         t1 = time.time()
         # Get the frequencies of the original dynamical matrix
-        super_struct0 = self.dyn_0.structure.generate_supercell(self.supercell)
         #super_dyn = self.dyn_0.GenerateSupercellDyn(self.supercell)
 
-        w_original, pols_original = self.dyn_0.DiagonalizeSupercell()#super_dyn.DyagDinQ(0)
+        w_original = self.w_0.copy()
+        pols_original = self.pols_0.copy()
         
         # Exclude translations
         if not self.ignore_small_w:
@@ -1321,10 +1357,8 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         old_a = SCHAModules.thermodynamic.w_to_a(w, self.T0)
         
         # Now do the same for the new dynamical matrix
-        super_structure = new_dynamical_matrix.structure.generate_supercell(self.supercell)
         #new_super_dyn = new_dynamical_matrix.GenerateSupercellDyn(self.supercell)
         
-        w_new, pols = new_dynamical_matrix.DiagonalizeSupercell()#new_super_dyn.DyagDinQ(0)
 
         if not self.ignore_small_w:
             trans_mask = CC.Methods.get_translations(pols, super_structure.get_masses_array())
@@ -1341,11 +1375,11 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         if violating_sum_rule:
             ERR_MSG = """
 ERROR WHILE UPDATING THE WEIGHTS
-    
+
 Error, one dynamical matrix does not satisfy the acoustic sum rule.
-       If this problem arises on a sscha run, 
-       it may be due to a gradient that violates the sum rule.
-       Please, be sure you are not using a custom gradient function.
+    If this problem arises on a sscha run, 
+    it may be due to a gradient that violates the sum rule.
+    Please, be sure you are not using a custom gradient function.
 
 DETAILS OF ERROR:
     Number of translatinal modes in the original dyn = {}
@@ -1360,22 +1394,15 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         w = np.array(w/2, dtype = np.float64)
         new_a = SCHAModules.thermodynamic.w_to_a(w, newT)
         
-        Nat_sc = super_structure.N_atoms
         
         # Get the new displacements in the supercell
         t3 = time.time()
-        old_disps = np.zeros(np.shape(self.u_disps), dtype = np.double)
-
-        self.u_disps[:,:] = self.xats.reshape((self.N, 3*Nat_sc)) - np.tile(super_structure.coords.ravel(), (self.N,1))
-        old_disps[:,:] = self.xats.reshape((self.N, 3*Nat_sc)) - np.tile(super_struct0.coords.ravel(), (self.N,1))
-
         # for i in range(self.N):
         #     self.u_disps[i, :] = (self.xats[i, :, :] - super_structure.coords).reshape( 3*Nat_sc )
             
         #     old_disps[i,:] = (self.xats[i, :, :] - super_dyn.structure.coords).reshape( 3*Nat_sc )
             
         #     # TODO: this method recomputes the displacements, it is useless since we already have them in self.u_disps
-        self.sscha_energies[:], self.sscha_forces[:,:,:] = new_dynamical_matrix.get_energy_forces(None, displacement = self.u_disps, w_pols = (w_new, pols))
 
         t4 = time.time()
 
@@ -1422,6 +1449,8 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
             rho_tmp[i] *= np.exp(-0.5 * (v_new - v_old) )
         # Lets try to use this one
         self.rho = rho_tmp
+    
+
 
         #print("\n".join(["%8d) %16.8f" % (i+1, r) for i, r in enumerate(self.rho)]))
         
@@ -1620,7 +1649,7 @@ def get_free_energy(self, return_error = False):
                 The free energy in the current dynamical matrix and at the ensemble temperature
         """
         
-        free_energy = self.current_dyn.GetHarmonicFreeEnergy(self.current_T)
+        free_energy = self.current_dyn.GetHarmonicFreeEnergy(self.current_T, w_pols = (self.current_w, self.current_pols))
         
         # We got the F_0 
         # Now we can compute the free energy difference
diff --git a/Modules/Minimizer.py b/Modules/Minimizer.py
@@ -240,7 +240,9 @@ def run_step(self, gradient, kl_new):
                 self.step *= self.decrement_step
 
                 if self.verbose:
-                    print("Step too large, reducing to {}".format(self.step))
+                    print("Step too large (scalar = {} | kl_ratio = {}), reducing to {}".format(scalar, kl_ratio, self.step))
+                    #print("Direction: ", self.direction)
+                    #print("Gradient: ", gradient)
             else:
                 # The step is good, therefore next step perform a new direction
                 self.new_direction = True
diff --git a/Modules/SchaMinimizer.py b/Modules/SchaMinimizer.py
@@ -24,6 +24,8 @@
 It is possible to use it to perform the anharmonic minimization 
 """
 
+from inspect import signature
+
 #import Ensemble
 import numpy as np
 import difflib
@@ -415,16 +417,19 @@ def minimization_step(self, custom_function_gradient = None):
             struct_grad, struct_grad_err =  self.ensemble.get_average_forces(True)
             #print "SHAPE:", np.shape(struct_grad)
             struct_grad_reshaped = - struct_grad.reshape( (3 * self.dyn.structure.N_atoms))
-            t2 = time.time()
-            
-            print ("Time elapsed to compute the structure gradient:", t2 - t1, "s")
             
             
             # Preconditionate the gradient for the wyckoff minimization
             if self.precond_wyck:
-                struct_precond = GetStructPrecond(self.ensemble.current_dyn, ignore_small_w = self.ensemble.ignore_small_w)
+                w_pols = None
+                if len(self.dyn.q_tot) == 1:
+                    w_pols = (self.ensemble.current_w, self.ensemble.current_pols)
+                struct_precond = GetStructPrecond(self.ensemble.current_dyn, ignore_small_w = self.ensemble.ignore_small_w, w_pols = w_pols)
                 struct_grad_precond = struct_precond.dot(struct_grad_reshaped)
                 struct_grad = struct_grad_precond.reshape( (self.dyn.structure.N_atoms, 3))
+            t2 = time.time()
+            
+            print ("Time elapsed to compute the structure gradient:", t2 - t1, "s")
             
             
             # Apply the symmetries to the forces
@@ -448,7 +453,35 @@ def minimization_step(self, custom_function_gradient = None):
 
         # Perform the gradient restriction
         if custom_function_gradient is not None:
-            custom_function_gradient(dyn_grad, struct_grad)    
+
+            # Check the number of parameters
+            try:
+                sig = signature(custom_function_gradient).parameters
+            except Exception as e:
+                print(e)
+
+                MSG = '''
+While inspecting the custom_function_gradient an error was rised.
+     Maybe you did not pass the minimizer a valid function?                
+'''             
+                raise ValueError(MSG)
+
+            
+            if len(sig) not in [2,3]:
+                MSG = '''
+Error, the custom_function_gradient must have either 2 or 3 arguments:
+    - dynamical_matrix_gradient
+    - structure gradient
+    - [Optional] The minimizer (self) object
+
+      The function you provided accepts {} arguments instead.
+'''.format(len(sig))
+                raise ValueError(MSG)
+
+            if len(sig) == 3:
+                custom_function_gradient(dyn_grad, struct_grad, self)
+            else:
+                custom_function_gradient(dyn_grad, struct_grad)    
             
 
         # Append the gradient modulus to the minimization info
@@ -500,22 +533,23 @@ def minimization_step(self, custom_function_gradient = None):
 
 
             # If we have imaginary frequencies, force the kl ratio to zero
+        
+            # Update the ensemble
+            try:
+                self.update()
+            except np.linalg.LinAlgError as error:
+                print("Diagonalization error:")
+                print(error)
+                print("Reducing the minimization step...")
+                new_kl_ratio = 0 # Force step reduction
+                is_diag_ok = False 
+                diag_error_counter += 1
+            
             if self.check_imaginary_frequencies():
                 print("Immaginary frequencies found! Redoing the step.")
                 new_kl_ratio = 0
                 is_diag_ok = False
                 imag_freq_counter += 1            
-            else:
-                # Update the ensemble
-                try:
-                    self.update()
-                except np.linalg.LinAlgError as error:
-                    print("Diagonalization error:")
-                    print(error)
-                    print("Reducing the minimization step...")
-                    new_kl_ratio = 0 # Force step reduction
-                    is_diag_ok = False 
-                    diag_error_counter += 1
             
             if diag_error_counter >= self.max_diag_error_counter:
                 ERROR_MSG = """
@@ -1134,7 +1168,7 @@ def run(self, verbose = 1, custom_function_pre = None, custom_function_post = No
             self.__fe__.append(np.real(fe))
             self.__fe_err__.append(np.real(err))
             
-            harm_fe = self.dyn.GetHarmonicFreeEnergy(self.ensemble.current_T) / np.prod(self.ensemble.supercell)
+            harm_fe = self.dyn.GetHarmonicFreeEnergy(self.ensemble.current_T, w_pols = (self.ensemble.current_w, self.ensemble.current_pols)) / np.prod(self.ensemble.supercell)
             anharm_fe = np.real(fe - harm_fe)
             
             # Compute the KL ratio
@@ -1295,26 +1329,35 @@ def finalize(self, verbose = 1):
                                                   self.dyn.structure.coords[i,2]))
             print ()
             print (" ==== FINAL FREQUENCIES [cm-1] ==== ")
-            super_dyn = self.dyn.GenerateSupercellDyn(self.ensemble.supercell)
-            w, pols = super_dyn.DyagDinQ(0)
-            trans = CC.Methods.get_translations(pols, super_dyn.structure.get_masses_array())
+            #super_dyn = self.dyn.GenerateSupercellDyn(self.ensemble.supercell)
+            super_struct = self.dyn.structure.generate_supercell(self.dyn.GetSupercell())
+            w = self.ensemble.current_w.copy()
+            pols = self.ensemble.current_pols.copy()
+
+            #w, pols = super_dyn.DyagDinQ(0)
+            trans = CC.Methods.get_translations(pols, super_struct.get_masses_array())
             
             for i in range(len(w)):
                 print ("Mode %5d:   freq %16.8f cm-1  | is translation? " % (i+1, w[i] * __RyToCm__), trans[i])
         
             print ()
         
             
-
     def check_imaginary_frequencies(self):
         """
         The following subroutine check if the current matrix has imaginary frequency. In this case
         the minimization is stopped.
         """
 
+        # Avoid the check if the dynamical matrix has been computed.
+        if not self.minim_dyn:
+            return False 
+
         # Get the frequencies
         #superdyn = self.dyn.GenerateSupercellDyn(self.ensemble.supercell)
-        w, pols = self.dyn.DiagonalizeSupercell()#.DyagDinQ(0)
+        w = self.ensemble.current_w.copy()
+        pols = self.ensemble.current_pols.copy()
+        #w, pols = self.dyn.DiagonalizeSupercell()#.DyagDinQ(0)
         ss = self.dyn.structure.generate_supercell(self.dyn.GetSupercell())
 
         # Get translations
@@ -1333,7 +1376,9 @@ def check_imaginary_frequencies(self):
         if w[0] < 0:
             print ("Total frequencies (excluding translations):")
             #superdyn0 = self.ensemble.dyn_0.GenerateSupercellDyn(self.ensemble.supercell)
-            wold, pold = self.ensemble.dyn_0.DiagonalizeSupercell()# superdyn0.DyagDinQ(0)
+            wold = self.ensemble.w_0.copy()
+            pold = self.ensemble.pols_0.copy()
+            #wold, pold = self.ensemble.dyn_0.DiagonalizeSupercell()# superdyn0.DyagDinQ(0)
 
             ss0 = self.ensemble.dyn_0.structure.generate_supercell(self.dyn.GetSupercell())
             
@@ -1824,7 +1869,7 @@ def ApplyFCPrecond(current_dyn, matrix, T = 0):
     
 
 
-def GetStructPrecond(current_dyn, ignore_small_w = False):
+def GetStructPrecond(current_dyn, ignore_small_w = False, w_pols = None):
     r"""
     GET THE PRECONDITIONER FOR THE STRUCTURE MINIMIZATION
     =====================================================
@@ -1845,6 +1890,8 @@ def GetStructPrecond(current_dyn, ignore_small_w = False):
     ----------
         current_dyn : Phonons()
             The current dynamical matrix
+        w_pols : 
+            If given, do not diagonalize the gradient
         
     Returns
     -------
@@ -1854,7 +1901,11 @@ def GetStructPrecond(current_dyn, ignore_small_w = False):
     """
     
     # Dyagonalize the current dynamical matrix
-    w, pols = current_dyn.DyagDinQ(0)
+    if w_pols:
+        w = w_pols[0].copy()
+        pols = w_pols[1].copy()
+    else:
+        w, pols = current_dyn.DyagDinQ(0)
     
     # Get some usefull array
     mass = current_dyn.structure.get_masses_array()
diff --git a/Modules/Utilities.py b/Modules/Utilities.py
@@ -627,6 +627,7 @@ def get_fix_rotations_CFG(dyn):
     assert np.max(np.abs(dyn.q_tot[0])) < 1e-5, "Error, the dynamical matrix must be at Gamma"
 
     nat = dyn.structure.N_atoms 
+    raise NotImplementedError("Error, not yet implemented")
     # TODO: To be ended
     
     
diff --git a/tests/test_gradient_evolution/test_cmp_gradients.py b/tests/test_gradient_evolution/test_cmp_gradients.py
diff --git a/tests/test_simple_relax/test_relax.py b/tests/test_simple_relax/test_relax.py
