Skip to content

Commit

Permalink
Code cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@gmatteo
gmatteo committed Dec 23, 2024
1 parent 36fe56a commit abd83d5
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 74 deletions.
64 changes: 31 additions & 33 deletions abipy/dfpt/qha_2D.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -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)
Expand All @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
8 changes: 6 additions & 2 deletions abipy/flowtk/flows.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down
29 changes: 14 additions & 15 deletions abipy/flowtk/qha_2d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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)

Expand All @@ -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)
Expand All @@ -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)

Expand Down
3 changes: 1 addition & 2 deletions abipy/flowtk/vzsisa.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
49 changes: 27 additions & 22 deletions abipy/flowtk/zsisa.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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()


Expand All @@ -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

0 comments on commit abd83d5

Please sign in to comment.