Code cleanup

Author: gmatteo

abipy/dfpt/qha_2D.py

@@ -43,13 +43,9 @@ def from_json_file(cls,
         For the meaning of the other arguments see from_gsr_ddb_paths.
         """
         data = mjson_load(filepath)
-
-        bo_strains_ac = [data["strains_a"], data["strains_c"]]
-        phdos_strains_ac = [data["strains_a"], data["strains_c"]]
-
         return cls.from_gsr_ddb_paths(nqsmall_or_qppa,
                                       data["gsr_relax_paths"], data["ddb_relax_paths"],
-                                      bo_strains_ac, phdos_strains_ac,
+                                      data["bo_strains_ac"], data["phdos_strains_ac"],
                                       anaget_kwargs=anaget_kwargs, smearing_ev=smearing_ev, verbose=verbose)
 
     @classmethod
@@ -203,8 +199,8 @@ def plot_energies(self, ax=None, **kwargs) -> Figure:
         ax, fig, plt = get_ax_fig_plt(ax, figsize=(10, 8))
         ax = fig.add_subplot(111, projection='3d')  # Create a 3D subplot
 
-        a0 =self.lattice_a[:,0]
-        c0 =self.lattice_c[0,:]
+        a0 = self.lattice_a[:,0]
+        c0 = self.lattice_c[0,:]
 
         X, Y = np.meshgrid(c0, a0)
 
@@ -216,7 +212,7 @@ def plot_energies(self, ax=None, **kwargs) -> Figure:
 
         initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
         xy_init = np.array(initial_guess)
-        min_x0,min_y0,min_energy= self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
+        min_x0, min_y0, min_energy= self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
 
         x_new = np.linspace(min(self.lattice_a[:,0]), max(self.lattice_a[:,0]), 100)
         y_new = np.linspace(min(self.lattice_c[0,:]), max(self.lattice_c[0,:]), 100)
@@ -227,10 +223,10 @@ def plot_energies(self, ax=None, **kwargs) -> Figure:
         ax.plot_surface(x_grid, y_grid, energy_interp, cmap='viridis', alpha=0.6)
 
         # Set labels
-        ax.set_xlabel('Lattice Parameter C (Å)')
-        ax.set_ylabel('Lattice Parameter A (Å)')
+        ax.set_xlabel('Lattice parameter C (Å)')
+        ax.set_ylabel('Lattice parameter A (Å)')
         ax.set_zlabel('Energy (eV)')
-        ax.set_title('Energy Surface in 3D')
+        ax.set_title('BO Energy Surface in 3D')
 
         return fig
 
@@ -266,7 +262,7 @@ def find_minimum(self, f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.0
         return xy[0], xy[1], min_energy
 
     @add_fig_kwargs
-    def plot_free_energies(self, tstart=800 , tstop=0 ,num=5, ax=None, **kwargs) -> Figure:
+    def plot_free_energies(self, tstart=800, tstop=0, num=5, ax=None, **kwargs) -> Figure:
         """
         Plot free energy as a function of temperature in a 3D plot.
 
@@ -286,8 +282,8 @@ def plot_free_energies(self, tstart=800 , tstop=0 ,num=5, ax=None, **kwargs) ->
 
             X, Y = np.meshgrid(self.lattice_c[0,:], self.lattice_a[:,0])
             for  e in ( tot_en.T ):
-                ax.plot_surface(X, Y ,e, cmap='viridis', alpha=0.7)
-                ax.plot_wireframe(X, Y ,e, cmap='viridis')
+                ax.plot_surface(X, Y, e, cmap='viridis', alpha=0.7)
+                ax.plot_wireframe(X, Y, e, cmap='viridis')
 
             min_x = np.zeros(num)
             min_y = np.zeros(num)
@@ -336,23 +332,23 @@ def plot_free_energies(self, tstart=800 , tstop=0 ,num=5, ax=None, **kwargs) ->
             initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
             xy_init = np.array(initial_guess)
             for j, e in enumerate(tot_en2.T):
-                f_interp = RectBivariateSpline(a,c, e , kx=4, ky=4)
-                min_x[j],min_y[j],min_tot_en2[j]= self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
+                f_interp = RectBivariateSpline(a, c, e, kx=4, ky=4)
+                min_x[j],min_y[j],min_tot_en2[j] = self.find_minimum(f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
 
             X, Y = np.meshgrid(c, a)
             for e in tot_en2.T:
                 ax.plot_wireframe(X, Y, e, cmap='viridis')
-                ax.plot_surface(X, Y ,e, cmap='viridis', alpha=0.7)
+                ax.plot_surface(X, Y, e, cmap='viridis', alpha=0.7)
 
-            ax.scatter(min_y,min_x,min_tot_en2, color='c', s=100)
-            ax.plot(min_y,min_x,min_tot_en2, color='c')
+            ax.scatter(min_y, min_x, min_tot_en2, color='c', s=100)
+            ax.plot(min_y, min_x, min_tot_en2, color='c')
 
         ax.scatter(self.lattice_c[0,self.iy0], self.lattice_a[self.ix0,0], self.energies[self.ix0, self.iy0], color='red', s=100)
 
         ax.set_xlabel('C')
         ax.set_ylabel('A')
-        ax.set_zlabel('Energy (eV)')
-        #ax.set_title('Energies as a 3D Plot')
+        ax.set_zlabel('Free energy (eV)')
+        #ax.set_title('Free energies as a 3D Plot')
         plt.savefig("energy.pdf", format="pdf", bbox_inches="tight")
 
         return fig
@@ -403,11 +399,11 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
 
         elif (len(self.lattice_a_from_phdos)==3 or len(self.lattice_c_from_phdos)==3):
 
-            dF_dA  = np.zeros( num)
-            dF_dC  = np.zeros( num)
-            d2F_dA2= np.zeros( num)
-            d2F_dC2= np.zeros( num)
-            d2F_dAdC = np.zeros( num)
+            dF_dA = np.zeros(num)
+            dF_dC = np.zeros(num)
+            d2F_dA2 = np.zeros(num)
+            d2F_dC2 = np.zeros(num)
+            d2F_dAdC = np.zeros(num)
             a0 = self.lattice_a_from_phdos[1,1]
             c0 = self.lattice_c_from_phdos[1,1]
             da = self.lattice_a_from_phdos[0,1]-self.lattice_a_from_phdos[1,1]
@@ -436,7 +432,7 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
             A0 = self.lattice_a[self.ix0,self.iy0]
             C0 = self.lattice_c[self.ix0,self.iy0]
             scale= self.volumes[self.ix0,self.iy0]/A0**2/C0
-            min_v=min_x**2*min_y*scale
+            min_v = min_x**2*min_y*scale
 
             dt = tmesh[1] - tmesh[0]
             alpha_a = (min_x[2:] - min_x[:-2]) / (2 * dt) / min_x[1:-1]
@@ -447,9 +443,10 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
             #ax.plot(tmesh[1:-1], alpha_v, linestyle='--', color='darkorange', label=r"$\alpha_v$ E$\infty$Vib2")
 
         # Save the data
-        data_to_save = np.column_stack((tmesh[1:-1],alpha_v,alpha_a,alpha_c))
-        columns= [ '#Tmesh', 'alpha_v' ,  'alpha_a', 'alpha_c']
+        data_to_save = np.column_stack((tmesh[1:-1], alpha_v, alpha_a, alpha_c))
+        columns = ['#Tmesh', 'alpha_v', 'alpha_a', 'alpha_c']
         file_path = 'thermal-expansion_data.txt'
+        print(f"Writing thermal expansion data to: {file_path}"
         np.savetxt(file_path, data_to_save, fmt='%4.6e', delimiter='\t\t',  header='\t\t\t'.join(columns), comments='')
 
         ax.grid(True)
@@ -533,7 +530,7 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure
 
             A0 = self.lattice_a[self.ix0,self.iy0]
             C0 = self.lattice_c[self.ix0,self.iy0]
-            scale= self.volumes[self.ix0,self.iy0]/A0**2/C0
+            scale = self.volumes[self.ix0,self.iy0]/A0**2/C0
             min_volumes = min_x**2 * min_y * scale
 
             axs[0].plot(tmesh, min_x, color='c', label=r"$a$ (E$\infty$Vib2)", linewidth=2)
@@ -571,9 +568,10 @@ def get_vib_free_energies(self, tstart: float, tstop: float, num: int) -> np.nda
         Return: A 3D array of vibrational free energies of shape (num_c, num_a, num_temp)
         """
         f = np.zeros((len(self.lattice_c_from_phdos[0]), len(self.lattice_a_from_phdos[:, 0]), num))
+
         for i in range(len(self.lattice_a_from_phdos[:, 0])):
             for j in range(len(self.lattice_c_from_phdos[0])):
-                dos = self.phdoses[i][j]
-                if dos is not None:
-                    f[j][i] = dos.get_free_energy(tstart, tstop, num).values
+                phdos = self.phdoses[i][j]
+                if phdos is not None:
+                    f[j][i] = phdos.get_free_energy(tstart, tstop, num).values
         return f

abipy/flowtk/flows.py

@@ -753,7 +753,7 @@ def _iflat_tasks_wti(self, status=None, op="==", nids=None, with_wti=True):
                         else:
                             yield task
 
-    def abivalidate_inputs(self) -> tuple:
+    def abivalidate_inputs(self, verbose: int = 1) -> tuple:
         """
         Run ABINIT in dry mode to validate all the inputs of the flow.
 
@@ -777,7 +777,11 @@ def abivalidate_inputs(self) -> tuple:
         isok, tuples = True, []
         for task in self.iflat_tasks():
             t = task.input.abivalidate()
-            if t.retcode != 0: isok = False
+            if t.retcode != 0:
+                isok = False
+                if verbose:
+                    print(t.stderr_file.read())
+
             tuples.append(t)
 
         return isok, tuples

abipy/flowtk/qha_2d.py

@@ -26,8 +26,8 @@ class Qha2dFlow(Flow):
     def from_scf_input(cls,
                        workdir: PathLike,
                        scf_input: AbinitInput,
-                       bo_strains_ac,
-                       phdos_strains_ac,
+                       bo_strains_ac: list[list],
+                       phdos_strains_ac: list[list]:,
                        ngqpt,
                        with_becs: bool,
                        with_quad: bool,
@@ -67,10 +67,10 @@ def finalize(self):
         data = {"bo_strains_ac": work.bo_strains_ac, "phdos_strains_ac": work.phdos_strains_ac}
 
         # Build list of strings with path to the relevant output files ordered by V.
-        data["gsr_relax_paths"] = [task.gsr_path for task in work.relax_tasks_deformed]
+        data["gsr_relax_paths"] = [task.gsr_path for task in work.relax_tasks_strained]
 
         gsr_relax_entries, gsr_relax_volumes = [], []
-        for task in work.relax_tasks_deformed:
+        for task in work.relax_tasks_strained:
             with task.open_gsr() as gsr:
                 gsr_relax_entries.append(dict(
                     volume=gsr.structure.volume,
@@ -120,9 +120,9 @@ def from_scf_input(cls,
         # Save attributes in work
         work.initial_scf_input = scf_input
 
+        # Make sure ench row is a numpy array.
         work.bo_strains_ac = bo_strains_ac
         work.phdos_strains_ac = phdos_strains_ac
-        # Make sure ench row is a numpy array.
         for i in range(2):
             work.bo_strains_ac[i] = np.array(bo_strains_ac[i])
             work.phdos_strains_ac[i] = np.array(phdos_strains_ac[i])
@@ -153,23 +153,23 @@ def on_ok(self, sender):
             # Get relaxed structure and build new task for structural relaxation at fixed volume.
             relaxed_structure = sender.get_final_structure()
 
-            relax_template = self.relax_template
-            self.relax_tasks_deformed = []
+            self.relax_tasks_strained = []
 
             import itertools
             for s1, s3 in itertools.product(self.bo_strains_ac[0], self.bo_strains_ac[1]):
-                deformation_name = f"{s1=}, {s3=}"
+                strain_name = f"{s1=}, {s3=}"
                 # Apply strain to the structure
                 strain_tensor = np.diag([s1, s1, s3])
                 strained_structure = relaxed_structure.apply_strain(strain_tensor, inplace=False)
                 #print("strained_structure:", strained_structure)
 
                 # Relax deformed structure with fixed unit cell.
-                task = self.register_relax_task(relax_template.new_with_structure(strained_structure, optcell=0))
+                task = self.register_relax_task(self.relax_template.new_with_structure(strained_structure, optcell=0))
+
                 task.bo_strain = (s1, s3)
                 task.in_phdos_strains = np.any(np.abs(s1 - self.phdos_strains_ac[0]) < 1e-3) and \
                                         np.any(np.abs(s3 - self.phdos_strains_ac[1]) < 1e-3)
-                self.relax_tasks_deformed.append(task)
+                self.relax_tasks_strained.append(task)
 
             self.flow.allocate(build=True)
 
@@ -185,9 +185,9 @@ def on_all_ok(self):
         # Build phonon works for the different relaxed structures.
         self.ph_works = []
 
-        for task in self.relax_tasks_deformed:
+        for task in self.relax_tasks_strained:
             s1, s3 = task.bo_strain
-            deformation_name = f"{s1=}, {s3=}"
+            strain_name = f"{s1=}, {s3=}"
 
             relaxed_structure = task.get_final_structure()
             scf_input = self.initial_scf_input.new_with_structure(relaxed_structure)
@@ -197,11 +197,10 @@ def on_all_ok(self):
                                                     with_becs=self.with_becs, ddk_tolerance=None)
 
                 # Reduce the number of files produced in the DFPT tasks to avoid possible disk quota issues.
-                prtvars = dict(prtden=0, prtpot=0)
                 for task in ph_work[1:]:
-                    task.input.set_vars(**prtvars)
+                    task.input.set_vars(prtden=0, prtpot=0)
 
-                ph_work.set_name(deformation_name)
+                ph_work.set_name(strain_name)
                 self.ph_works.append(ph_work)
                 self.flow.register_work(ph_work)
 

abipy/flowtk/vzsisa.py

@@ -171,9 +171,8 @@ def on_all_ok(self):
                                                 ddk_tolerance=None)
 
             # Reduce the number of files produced in the DFPT tasks to avoid possible disk quota issues.
-            prtvars = dict(prtden=0, prtpot=0)
             for task in ph_work[1:]:
-                task.input.set_vars(**prtvars)
+                task.input.set_vars(prtden=0, prtpot=0)
 
             self.flow.register_work(ph_work)
             self.ph_works.append(ph_work)

abipy/flowtk/zsisa.py

@@ -64,11 +64,11 @@ def finalize(self):
         data = {"eps": work.eps, "spgrp_number": work.spgrp_number}
 
         # Build list of strings with path to the relevant output files ordered by V.
-        data["gsr_relax_paths"] = [task.gsr_path for task in work.relax_tasks_deformed]
+        data["gsr_relax_paths"] = [task.gsr_path for task in work.relax_tasks_strained]
         data["strain_inds"] = work.strain_inds
 
         gsr_relax_entries, gsr_relax_volumes = [], []
-        for task in work.relax_tasks_deformed:
+        for task in work.relax_tasks_strained:
             with task.open_gsr() as gsr:
                 gsr_relax_entries.append(dict(
                     volume=gsr.structure.volume,
@@ -143,19 +143,13 @@ def on_ok(self, sender):
             # Get relaxed structure and build new task for structural relaxation at fixed volume.
             relaxed_structure = sender.get_final_structure()
 
-            self.deformed_structures_dict, self.strain_inds, self.spgrp_number = generate_deformations(relaxed_structure, self.eps)
+            self.strained_structures_dict, self.strain_inds, self.spgrp_number = generate_deformations(relaxed_structure, self.eps)
 
-            relax_template = self.relax_template
-            self.relax_tasks_deformed = []
-            for structure in self.deformed_structures_dict.values():
+            self.relax_tasks_strained = []
+            for structure in self.strained_structures_dict.values():
                 # Relax deformed structure with fixed unit cell.
-                task = self.register_relax_task(relax_template.new_with_structure(structure, optcell=0))
-                self.relax_tasks_deformed.append(task)
-
-            # Build work for elastic properties (clamped-ions)
-            # activate internal strain and piezoelectric part.
-            #from abipy.flowtk.dfpt import ElasticWork
-            #elastic_work = ElasticWork.from_scf_input(scf_input, with_relaxed_ion=True, with_piezo=True)
+                task = self.register_relax_task(self.relax_template.new_with_structure(structure, optcell=0))
+                self.relax_tasks_strained.append(task)
 
             self.flow.allocate(build=True)
 
@@ -171,25 +165,24 @@ def on_all_ok(self):
         # Build phonon works for the different relaxed structures.
         self.ph_works = []
 
-        for task, deform_name in zip(self[1:], self.deformed_structures_dict.keys(), strict=True):
+        for task, strain_name in zip(self[1:], self.strained_structures_dict.keys(), strict=True):
             relaxed_structure = task.get_final_structure()
             scf_input = self.initial_scf_input.new_with_structure(relaxed_structure)
             ph_work = PhononWork.from_scf_input(scf_input, self.ngqpt, is_ngqpt=True, tolerance=None,
                                                 with_becs=self.with_becs, ddk_tolerance=None)
 
             # Reduce the number of files produced in the DFPT tasks to avoid possible disk quota issues.
-            prtvars = dict(prtden=0, prtpot=0)
             for task in ph_work[1:]:
-                task.input.set_vars(**prtvars)
+                task.input.set_vars(prtden=0, prtpot=0)
 
-            ph_work.set_name(deform_name)
+            ph_work.set_name(strain_name)
             self.ph_works.append(ph_work)
             self.flow.register_work(ph_work)
 
             # Add task for electron DOS calculation to edos_work
             if self.edos_ngkpt is not None:
                 edos_input = scf_input.make_edos_input(self.edos_ngkpt)
-                self.edos_work.register_nscf_task(edos_input, deps={ph_work[0]: "DEN"}).set_name(deform_name)
+                self.edos_work.register_nscf_task(edos_input, deps={ph_work[0]: "DEN"}).set_name(strain_name)
 
         if self.edos_ngkpt is not None:
             self.flow.register_work(self.edos_work)
@@ -227,6 +220,9 @@ def from_relax_input(cls,
         for pressure_gpa in work.pressures_gpa:
             for temperature in work.temperatures:
                 #strtarget = zsisa.get_strtarget(temperature, pressure)
+                #converged, new_structure, strtarget = zsisa.get_new_guess(relaxed_structure,
+                #    self.temperature, self.pressure_gpa, tolerance)
+
                 new_input = relax_input.new_with_vars(strtarget=strtarget)
                 task = work.register_task(new_input, task_class=ThermalRelaxTask)
                 # Attach pressure and temperature
@@ -239,6 +235,10 @@ def on_all_ok(self):
         """
         Implement the post-processing step at the end of the Work.
         """
+        # Build work for elastic properties (clamped-ions)
+        # activate internal strain and piezoelectric part.
+        #from abipy.flowtk.dfpt import ElasticWork
+        #elastic_work = ElasticWork.from_scf_input(scf_input, with_relaxed_ion=True, with_piezo=True)
         return super().on_all_ok()
 
 
@@ -248,10 +248,15 @@ def _on_ok(self):
         results = super()._on_ok()
         zsisa = self.work.zsisa
 
+        relaxed_structure = sender.get_final_structure()
+        converged, new_structure, strtarget = zsisa.get_new_guess(relaxed_structure,
+            self.temperature, self.pressure_gpa, tolerance)
+
         # Check for convergence.
-        #if not self.collinear_done:
-        #    self.input.set_vars(strtarget=strtarget)
-        #    self.finalized = False
-        #    self.restart()
+        if not converged:
+            #self.input.set_structure(new_structure)
+            self.input.set_vars(strtarget=strtarget)
+            self.finalized = False
+            self.restart()
 
         return results