Fast wpols

Modules/Ensemble.py

@@ -34,6 +34,7 @@
 import sscha.Parallel as Parallel
 from sscha.Parallel import pprint as print
 
+import difflib
 
 import SCHAModules
 #import sscha_HP_odd
@@ -98,7 +99,7 @@
 class Ensemble:
     __debug_index__ = 0
 
-    def __init__(self, dyn0, T0, supercell = None):
+    def __init__(self, dyn0, T0, supercell = None, **kwargs):
         """
         PREPARE THE ENSEMBLE
         ====================
@@ -115,6 +116,7 @@ def __init__(self, dyn0, T0, supercell = None):
                 The temperature used to generate the ensemble.
             supercell: optional, list of int
                 The supercell dimension. If not provided, it will be determined by dyn0
+            **kwargs : any other attribute of the ensemble
         """
 
         # N is the number of element in the ensemble
@@ -178,6 +180,39 @@ def __init__(self, dyn0, T0, supercell = None):
         # Get the extra quantities
         self.extra_quantities = {}
 
+
+        # Setup the attribute control
+        self.__total_attributes__ = [item for item in self.__dict__.keys()]
+        self.fixed_attributes = True # This must be the last attribute to be setted
+
+
+        # Setup any other keyword given in input (raising the error if not already defined)
+        for key in kwargs:
+            self.__setattr__(key, kwargs[key])
+
+
+    def __setattr__(self, name, value):
+        """
+        This method is used to set an attribute.
+        It will raise an exception if the attribute does not exists (with a suggestion of similar entries)
+        """
+
+
+        if "fixed_attributes" in self.__dict__:
+            if name in self.__total_attributes__:
+                super(Ensemble, self).__setattr__(name, value)
+            elif self.fixed_attributes:
+                similar_objects = str( difflib.get_close_matches(name, self.__total_attributes__))
+                ERROR_MSG = """
+        Error, the attribute '{}' is not a member of '{}'.
+        Suggested similar attributes: {} ?
+        """.format(name, type(self).__name__,  similar_objects)
+
+                raise AttributeError(ERROR_MSG)
+        else:
+            super(Ensemble, self).__setattr__(name, value)
+
+
     def convert_units(self, new_units):
         """
         CONVERT ALL THE VARIABLE IN A COHERENT UNIT OF MEASUREMENTS
@@ -1178,15 +1213,15 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         super_struct0 = self.dyn_0.structure.generate_supercell(self.supercell)
         #super_dyn = self.dyn_0.GenerateSupercellDyn(self.supercell)
 
-        w, pols = self.dyn_0.DiagonalizeSupercell()#super_dyn.DyagDinQ(0)
+        w_original, pols_original = self.dyn_0.DiagonalizeSupercell()#super_dyn.DyagDinQ(0)
 
         # Exclude translations
         if not self.ignore_small_w:
-            trans_original = CC.Methods.get_translations(pols, super_struct0.get_masses_array()) 
+            trans_original = CC.Methods.get_translations(pols_original, super_struct0.get_masses_array()) 
         else:
-            trans_original = np.abs(w) < CC.Phonons.__EPSILON_W__
+            trans_original = np.abs(w_original) < CC.Phonons.__EPSILON_W__
 
-        w = w[~trans_original]
+        w = w_original[~trans_original]
 
         # Convert from Ry to Ha and in fortran double precision
         w = np.array(w/2, dtype = np.float64)
@@ -1198,12 +1233,12 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         super_structure = new_dynamical_matrix.structure.generate_supercell(self.supercell)
         #new_super_dyn = new_dynamical_matrix.GenerateSupercellDyn(self.supercell)
 
-        w, pols = new_dynamical_matrix.DiagonalizeSupercell()#new_super_dyn.DyagDinQ(0)
+        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())
         else:
-            trans_mask = np.abs(w) < CC.Phonons.__EPSILON_W__
+            trans_mask = np.abs(w_new) < CC.Phonons.__EPSILON_W__
 
 
         # Check if the new dynamical matrix satisfies the sum rule
@@ -1230,7 +1265,7 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
             print(ERR_MSG)
             raise ValueError(ERR_MSG)
 
-        w = w[~trans_mask]
+        w= w_new[~trans_mask]
         w = np.array(w/2, dtype = np.float64)
         new_a = SCHAModules.thermodynamic.w_to_a(w, newT)
 
@@ -1249,7 +1284,7 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
         #     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)
+        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()
 
@@ -1271,8 +1306,8 @@ def update_weights(self, new_dynamical_matrix, newT, update_q = False):
 
 
         # Get the covariance matrices of the ensemble
-        ups_new = np.real(new_dynamical_matrix.GetUpsilonMatrix(self.current_T))
-        ups_old = np.real(self.dyn_0.GetUpsilonMatrix(self.T0))
+        ups_new = np.real(new_dynamical_matrix.GetUpsilonMatrix(self.current_T, w_pols = (w_new, pols)))
+        ups_old = np.real(self.dyn_0.GetUpsilonMatrix(self.T0, w_pols = (w_original, pols_original)))
 
         # Get the normalization ratio
         #norm = np.sqrt(np.abs(np.linalg.det(ups_new) / np.linalg.det(ups_old))) 
@@ -1755,11 +1790,22 @@ def get_preconditioned_gradient(self, subtract_sscha = True, return_error = Fals
         if verbose:
             print (" [GRADIENT] Time to call the fortran code:", t2 - t1, "s")
 
-        # Perform the fourier transform
-        q_grad = CC.Phonons.GetDynQFromFCSupercell(grad, np.array(self.current_dyn.q_tot),
-                                                   self.current_dyn.structure, supercell_dyn.structure)
-        q_grad_err = CC.Phonons.GetDynQFromFCSupercell(grad_err, np.array(self.current_dyn.q_tot),
-                                                       self.current_dyn.structure, supercell_dyn.structure)
+        # If we are at gamma, we can skip this part
+        # Which makes the code faster
+        if np.prod(self.dyn_0.GetSupercell()) > 1:
+
+            # Perform the fourier transform
+            q_grad = CC.Phonons.GetDynQFromFCSupercell(grad, np.array(self.current_dyn.q_tot),
+                                                    self.current_dyn.structure, supercell_dyn.structure)
+            q_grad_err = CC.Phonons.GetDynQFromFCSupercell(grad_err, np.array(self.current_dyn.q_tot),
+                                                        self.current_dyn.structure, supercell_dyn.structure)
+        else:
+            nat3, _ = grad.shape
+            q_grad = np.zeros( (1, nat3, nat3), dtype = np.double)
+            q_grad_err = np.zeros_like(q_grad)
+            q_grad[0, :, :] = grad 
+            q_grad_err[0, :, :] = grad_err
+
         t1 = time.time()
         if verbose:
             print (" [GRADIENT] Time to get back in fourier space:", t1 - t2, "s")
@@ -3014,7 +3060,7 @@ def get_energy_forces(self, ase_calculator, compute_stress = True, stress_numeri
 
             # Print the status
             if Parallel.am_i_the_master() and verbose:
-                print ("Computing configuration %d out of %d" % (i+1, stop))
+                print ("Computing configuration %d out of %d (nat = %d)" % (i+1, stop, struct.N_atoms))
                 sys.stdout.flush()
 
             # Avoid for errors
@@ -3023,27 +3069,28 @@ def get_energy_forces(self, ase_calculator, compute_stress = True, stress_numeri
             while run:
                 try:
                     energy = atms.get_total_energy() / Rydberg # eV => Ry
+                    # Get energy, forces (and stress)
+                    energy = atms.get_total_energy() / Rydberg # eV => Ry
+                    forces_ = atms.get_forces() / Rydberg # eV / A => Ry / A
+                    if compute_stress:
+                        if not stress_numerical:
+                            stress[9*i0 : 9*i0 + 9] = -atms.get_stress(False).reshape(9) * Bohr**3 / Rydberg  # ev/A^3 => Ry/bohr
+                        else:
+                            stress[9*i0 : 9*i0 + 9] = -ase_calculator.calculate_numerical_stress(atms, voigt = False).ravel()* Bohr**3 / Rydberg 
+
+                    # Copy into the ensemble array
+                    energies[i0] = energy
+                    forces[nat3*i0 : nat3*i0 + nat3] = forces_.reshape( nat3 )
                     run = False
                 except:
                     print ("Rerun the job %d" % i)
                     count_fails += 1
                     if count_fails >= 5:
                         run = False
-                        sys.stderr.write("Error in the ASE calculator for more than 5 times\n")
+                        struct.save_scf("error_struct.scf")
+                        sys.stderr.write("Error in the ASE calculator for more than 5 times\n     while computing 'error_struct.scf'")
                         raise
 
-            # Get energy, forces (and stress)
-            energy = atms.get_total_energy() / Rydberg # eV => Ry
-            forces_ = atms.get_forces() / Rydberg # eV / A => Ry / A
-            if compute_stress:
-                if not stress_numerical:
-                    stress[9*i0 : 9*i0 + 9] = -atms.get_stress(False).reshape(9) * Bohr**3 / Rydberg  # ev/A^3 => Ry/bohr
-                else:
-                    stress[9*i0 : 9*i0 + 9] = -ase_calculator.calculate_numerical_stress(atms, voigt = False).ravel()* Bohr**3 / Rydberg 
-
-            # Copy into the ensemble array
-            energies[i0] = energy
-            forces[nat3*i0 : nat3*i0 + nat3] = forces_.reshape( nat3 )
 
 
             i0 += 1

Modules/Relax.py

@@ -47,7 +47,7 @@
 class SSCHA(object):
 
     def __init__(self, minimizer = None, ase_calculator=None, N_configs=1, max_pop = 20, 
-                 save_ensemble = False, cluster = None):
+                 save_ensemble = False, cluster = None, **kwargs):
         """
         This module initialize the relaxer. It may perform
         constant volume or pressure relaxation using fully anharmonic potentials.
@@ -69,6 +69,7 @@ def __init__(self, minimizer = None, ase_calculator=None, N_configs=1, max_pop =
             cluster : Cluster.Cluster, optional
                 If different from None, the ensemble force and energy calculations
                 will be runned in the provided cluster.
+            **kwargs : any other keyword that matches an object of this structure
         """
 
         if minimizer == None:
@@ -85,7 +86,7 @@ def __init__(self, minimizer = None, ase_calculator=None, N_configs=1, max_pop =
         # If the ensemble must be saved at each iteration.
         # 
         self.save_ensemble = save_ensemble
-        self.data_dir = ""
+        self.data_dir = "data"
 
 
 
@@ -109,6 +110,11 @@ def __init__(self, minimizer = None, ase_calculator=None, N_configs=1, max_pop =
         self.__total_attributes__ = [item for item in self.__dict__.keys()]
         self.fixed_attributes = True # This must be the last attribute to be setted
 
+
+        # Setup any other keyword given in input (raising the error if not already defined)
+        for key in kwargs:
+            self.__setattr__(key, kwargs[key])
+
     def __setattr__(self, name, value):
         """
         This method is used to set an attribute.
@@ -271,8 +277,7 @@ def relax(self, restart_from_ens = False, get_stress = False,
                 with some ase calculator)
             ensemble_loc : string
                 Where the ensemble of each population is saved on the disk. If none, it will
-                use the content of self.data_dir. If also self.data_dir is None, 
-                the ensemble will not not be saved (useful to avoid disk I/O for force fields)
+                use the content of self.data_dir. It is just a way to override the variable self.data_dir
             start_pop : int, optional
                 The starting index for the population, used only for saving the ensemble and the dynamical 
                 matrix. If None, the content of self.start_pop will be used.
@@ -286,6 +291,27 @@ def relax(self, restart_from_ens = False, get_stress = False,
 
         if ensemble_loc is None:
             ensemble_loc = self.data_dir
+
+        if (not ensemble_loc) and self.save_ensemble:
+            ERR_MSG = """
+Error, you must specify where to save the ensembles.
+       this can be done either passing ensemble_loc = "path/to/dir" 
+       for the ensemble, or by setting the data_dir attribute of this object.
+"""
+            raise IOError(ERR_MSG)
+
+        if self.save_ensemble:
+            if not os.path.exists(ensemble_loc):
+                os.makedirs(ensemble_loc)
+            else:
+                if not os.isdir(ensemble_loc):
+                    ERR_MSG = """
+Error, the specified location to save the ensemble:
+       '{}' 
+       already exists and it is not a directory.
+""".format(ensemble_loc)
+                    raise IOError(ERR_MSG)
+
 
         if start_pop is None:
             start_pop = self.start_pop