Merge branch 'fast_parallel_gradient' into timer

Author: mesonepigreco

Modules/Ensemble.py

@@ -1895,6 +1895,50 @@ def get_fc_from_self_consistency(self, subtract_sscha = False, return_error = Fa
 
         return new_phi
 
+    def get_preconditioned_gradient_parallel(self, *args, **kwargs):
+        """
+        Compute the gradient using multiprocessing.
+        For documentation, see get_preconditioned_gradient
+        """
+
+        print("force length:", len(self.force_computed))
+
+        def work_function(argument):
+            print("force length [inside]:", len(self.force_computed))
+            ensemble_start_config, ensemble_end_config = argument
+            mask = np.zeros(self.N, dtype = bool)
+            mask[ensemble_start_config : ensemble_end_config] = True
+            print("mask length:", np.sum(mask.astype(int))," total N:", len(mask))
+            new_ensemble = self.split(mask)
+
+            gradient, _ = new_ensemble.get_preconditioned_gradient(*args, **kwargs)
+
+            av_ensemble = np.sum(new_ensemble.rho)
+            
+            return gradient * av_ensemble
+
+        nproc = CC.Settings.GetNProc()
+        list_of_inputs = []
+        index = 0
+        for i in range(nproc):
+            start_config = index
+            index += self.N // nproc
+            if i < self.N % nproc:
+                index += 1
+            end_config = index
+            
+            list_of_inputs.append( (start_config, end_config) )
+
+        gradient = CC.Settings.GoParallel(work_function, list_of_inputs, "+")
+        gradient /= np.sum(self.rho)
+
+        return gradient, np.zeros_like(gradient) + 1
+
+        
+
+        
+    
+
     def get_preconditioned_gradient(self, subtract_sscha = True, return_error = False,
                                     use_ups_supercell = True, preconditioned = 1,
                                     fast_grad = False, verbose = True, timer=None):
@@ -1910,7 +1954,7 @@ def get_preconditioned_gradient(self, subtract_sscha = True, return_error = Fals
 
         .. math::
 
-            \\Phi_{ab} = \\sum_c \\upsilon_{ac} \\left< u_c f_a\\right>_{\\Phi}
+            \\Phi_{ab} = \\sum_c \\Upsilon_{ac} \\left< u_c f_a\\right>_{\\Phi}
 
         The previous equation is true only if the :math:`\\Phi` matrix is the solution
         of the SCHA theory. Here :math:`\\vec u` are the displacements of the configurations
@@ -3109,9 +3153,12 @@ def compute_ensemble(self, calculator, compute_stress = True, stress_numerical =
         else:
             computing_ensemble.get_energy_forces(calculator, compute_stress, stress_numerical, verbose = verbose)
 
+        print("CE BEFORE MERGE:", len(self.force_computed))
+
         if should_i_merge:
             # Remove the noncomputed ensemble from here, and merge
             self.merge(computing_ensemble)
+        print("CE AFTER MERGE:", len(self.force_computed))
 
         print('ENSEMBLE ALL PROPERTIES:', self.all_properties)
 
@@ -3177,6 +3224,11 @@ def split(self, split_mask):
         N = np.sum(split_mask.astype(int))
         ens = Ensemble(self.dyn_0, self.T0, self.dyn_0.GetSupercell())
         ens.init_from_structures(structs)
+        print("Split length:", len(split_mask))
+        print("original force legnth:", len(self.force_computed))
+        print("Expected length:", len(ens.force_computed))
+        print("Splitting force length:", len(self.force_computed[split_mask]))
+
         ens.force_computed[:] = self.force_computed[split_mask]
         ens.stress_computed[:] = self.stress_computed[split_mask]
         ens.energies[:] = self.energies[split_mask]
@@ -3199,7 +3251,7 @@ def remove_noncomputed(self):
         It may be used to run a minimization even if the ensemble was not completely calculated.
         """
 
-        good_mask = self.force_computed
+        good_mask = np.copy(self.force_computed)
         if self.has_stress:
             good_mask = good_mask & self.stress_computed
 
@@ -3211,6 +3263,8 @@ def remove_noncomputed(self):
         self.sscha_energies = self.sscha_energies[good_mask]
         self.xats = self.xats[good_mask, :, :]
         self.u_disps = self.u_disps[good_mask, :]
+        self.force_computed = self.force_computed[good_mask]
+        self.stress_computed = self.stress_computed[good_mask]
 
         self.structures = [self.structures[x] for x in np.arange(len(good_mask))[good_mask]]
         self.all_properties = [self.all_properties[x] for x in np.arange(len(good_mask))[good_mask]]
@@ -3263,6 +3317,7 @@ def get_energy_forces(self, ase_calculator, compute_stress = True, stress_numeri
 
         # Setup the calculator for each structure
         parallel = False
+        print("Force computed shape:", len(self.force_computed))
         if __MPI__:
             comm = MPI.COMM_WORLD
             size = comm.Get_size()
@@ -3419,6 +3474,7 @@ def get_energy_forces(self, ase_calculator, compute_stress = True, stress_numeri
         # Reshape the arrays
         self.forces[:, :, :] = np.reshape(total_forces, (N_rand, self.current_dyn.structure.N_atoms*np.prod(self.supercell), 3), order = "C")
         self.force_computed[:] = True
+        print("Force computed shape:", len(self.force_computed))
 
         if compute_stress:
             self.stresses[:,:,:] = np.reshape(total_stress, (N_rand, 3, 3), order = "C")

Modules/Relax.py

@@ -576,15 +576,20 @@ def vc_relax(self, target_press = 0, static_bulk_modulus = 100,
                                                  cluster = self.cluster)
                 #self.minim.ensemble.get_energy_forces(self.calc, True, stress_numerical = stress_numerical)
 
+                print("RELAX force length:", len(self.minim.ensemble.force_computed))
+                
                 if ensemble_loc is not None and self.save_ensemble:
                     self.minim.ensemble.save_bin(ensemble_loc, pop)
+                print("RELAX force length:", len(self.minim.ensemble.force_computed))
 
             self.minim.population = pop
             self.minim.init(delete_previous_data = False)
 
+            print("RELAX force length:", len(self.minim.ensemble.force_computed))
             self.minim.run(custom_function_pre = self.__cfpre__,
                            custom_function_post = self.__cfpost__,
                            custom_function_gradient = self.__cfg__)
+            
 
 
             self.minim.finalize()

Modules/SchaMinimizer.py

@@ -313,6 +313,7 @@ def minimization_step(self, custom_function_gradient = None, timer=None):
                 It takes as input the two gradient (the dynamical matrix one and the structure one), and
                 modifies them (or does some I/O on it).
         """
+        print("Minimization step, force computed:", len(self.ensemble.force_computed))
 
         # Setup the symmetries
         t1 = time.time()
@@ -337,14 +338,14 @@ def minimization_step(self, custom_function_gradient = None, timer=None):
         if self.minim_dyn:
             if self.precond_dyn:
                 if timer is not None:
-                    dyn_grad, err = timer.execute_timed_function(self.ensemble.get_preconditioned_gradient, True, True, preconditioned=1)
+                    dyn_grad, err = timer.execute_timed_function(self.ensemble.get_preconditioned_gradient_parallel, True, True, preconditioned=1)
                 else:
-                    dyn_grad, err = self.ensemble.get_preconditioned_gradient(True, True, preconditioned=1)
+                    dyn_grad, err = self.ensemble.get_preconditioned_gradient_parallel(True, True, preconditioned=1)
             else:
                 if timer is not None:
-                    dyn_grad, err = timer.execute_timed_function(self.ensemble.get_preconditioned_gradient, True, True, preconditioned=0)
+                    dyn_grad, err = timer.execute_timed_function(self.ensemble.get_preconditioned_gradient_parallel, True, True, preconditioned=0)
                 else:
-                    dyn_grad, err = self.ensemble.get_preconditioned_gradient(True, True, preconditioned=0)
+                    dyn_grad, err = self.ensemble.get_preconditioned_gradient_parallel(True, True, preconditioned=0)
         else:
             dyn_grad = np.zeros( (len(self.dyn.q_tot), 3 * self.dyn.structure.N_atoms, 3 * self.dyn.structure.N_atoms), dtype = np.complex128)
             err = np.zeros_like(dyn_grad)

setup.py

@@ -9,6 +9,7 @@
 mpi_compile_args = []
 mpi_link_args = []
 
+
 # If true, do not check for parallel
 avoid_parallel_test = True