Skip to content

Commit

Permalink
qha_2d: use previous solution as initial guess in find_minimum
Browse files Browse the repository at this point in the history
  • Loading branch information
@gmatteo
gmatteo committed Dec 23, 2024
1 parent abd83d5 commit 7e6a2bb
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 71 deletions.
135 changes: 72 additions & 63 deletions abipy/dfpt/qha_2D.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
Capable of calculating anisotropic thermal expansion and lattice constants for uniaxial configurations.
Requires PHDOS.nc and DDB files for GSR calculations or _GSR.nc files.
If PHDOS.nc is available for all structures, normal interpolation for QHA will be applied.
Supports the use of six PHDOS.nc files for specific structures to employ the E_infVib2 approximation.
Supports the use of six PHDOS.nc files for specific structures to employ the EinfVib2 approximation.
"""
from __future__ import annotations

Expand All @@ -14,7 +14,7 @@

from scipy.interpolate import RectBivariateSpline #, RegularGridInterpolator
#from monty.collections import dict2namedtuple
#from monty.functools import lazy_property
from monty.functools import lazy_property
from abipy.tools.plotting import add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt
from abipy.tools.typing import Figure
from abipy.tools.serialization import HasPickleIO, mjson_load
Expand Down Expand Up @@ -188,6 +188,20 @@ def __init__(self, structures, phdoses, energies, structures_from_phdos,
# Find index of minimum energy
self.min_energy_idx = np.unravel_index(np.nanargmin(self.energies), self.energies.shape)

@lazy_property
def use_qha(self) -> bool:
"""True if we are in full QHA_2D mode."""
return len(self.lattice_a_from_phdos) == len(self.lattice_a) and len(self.lattice_c_from_phdos) == len(self.lattice_c)

@lazy_property
def use_einfvib2(self) -> bool:
return len(self.lattice_a_from_phdos) == 3 and len(self.lattice_c_from_phdos) == 3

def get_initial_guess_ac(self) -> np.array:
"""Return the initial guess for (a, c):"""
initial_guess = [1.005*self.lattice_a[self.ix0, 0], 1.005 * self.lattice_c[0,self.iy0]]
return np.array(initial_guess)

@add_fig_kwargs
def plot_energies(self, ax=None, **kwargs) -> Figure:
"""
Expand All @@ -210,8 +224,8 @@ def plot_energies(self, ax=None, **kwargs) -> Figure:

f_interp = RectBivariateSpline(a0, c0, self.energies, kx=4, ky=4)

initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()

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)
Expand Down Expand Up @@ -247,13 +261,14 @@ def find_minimum(self, f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.0
xy = np.array(xy_init)
dx = dy = 0.001

for _ in range(max_iter):
for it in range(max_iter):
grad = [
(f_interp(xy[0] + dx, xy[1]) - f_interp(xy[0] - dx, xy[1])) / (2 * dx),
(f_interp(xy[0], xy[1] + dy) - f_interp(xy[0], xy[1] - dy)) / (2 * dy),
]
xy -= step_size * np.ravel(grad)
if np.linalg.norm(grad) < tol:
#print(f"Converged after {it} iterations with {tol=}")
break
else:
raise RuntimeError(f"Could not reach {tol=} after {max_iter=}")
Expand All @@ -277,37 +292,31 @@ def plot_free_energies(self, tstart=800, tstop=0, num=5, ax=None, **kwargs) -> F
a0 = self.lattice_a[:,0]
c0 = self.lattice_c[0,:]

if (len(self.lattice_a_from_phdos)==len(self.lattice_a) or len(self.lattice_c_from_phdos)==len(self.lattice_c)):
if self.use_qha:
tot_en = self.energies[np.newaxis, :].T + ph_energies + self.volumes[np.newaxis, :].T * self.pressure / abu.eVA3_GPa

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')

min_x = np.zeros(num)
min_y = np.zeros(num)
min_tot_en = np.zeros(num)

initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()
min_x, min_y, min_tot_en = np.zeros(num), np.zeros(num), np.zeros(num)
for j,e in enumerate(tot_en.T):
f_interp = RectBivariateSpline(a0, c0, e, kx=4, ky=4)
min_x[j],min_y[j],min_tot_en[j]= self.find_minimum(f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
min_x[j], min_y[j], min_tot_en[j]= self.find_minimum(f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
xy_init = min_x[j], min_y[j]

ax.scatter(min_y,min_x,min_tot_en, color='c', s=100)
ax.plot(min_y,min_x,min_tot_en, color='c')

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)
elif self.use_einfvib2:
a0 = self.lattice_a[1,1]
c0 = self.lattice_c[1,1]
da = self.lattice_a[0,1]-self.lattice_a[1,1]
dc = self.lattice_c[1,0]-self.lattice_c[1,1]
dF_dA, dF_dC, d2F_dA2, d2F_dC2, d2F_dAdC = (np.zeros(num) for _ in range(5))

for i, e in enumerate(ph_energies.T):
dF_dA[i]=(e[0,1]-e[2,1])/(2*da)
dF_dC[i]=(e[1,0]-e[1,2])/(2*dc)
Expand All @@ -320,20 +329,17 @@ def plot_free_energies(self, tstart=800, tstop=0, num=5, ax=None, **kwargs) -> F
tot_en2 = tot_en2+ (self.lattice_c[np.newaxis, :].T - c0)*dF_dC + 0.5*(self.lattice_c[np.newaxis, :].T - c0)**2*d2F_dC2
tot_en2 = tot_en2+ (self.lattice_c[np.newaxis, :].T - c0)*(self.lattice_a[np.newaxis, :].T - a0)*d2F_dAdC

min_x = np.zeros(num)
min_y = np.zeros(num)
min_tot_en2 = np.zeros(num)

a = self.lattice_a[:,0]
c = self.lattice_c[0,:]
a_phdos = self.lattice_a[:,0]
c_phdos = self.lattice_c[0,:]

initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()
min_x, min_y, min_tot_en2 = np.zeros(num), np.zeros(num), np.zeros(num)
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)
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)
xy_init = min_x[j], min_y[j]

X, Y = np.meshgrid(c, a)
for e in tot_en2.T:
Expand All @@ -343,6 +349,9 @@ def plot_free_energies(self, tstart=800, tstop=0, num=5, ax=None, **kwargs) -> F
ax.scatter(min_y, min_x, min_tot_en2, color='c', s=100)
ax.plot(min_y, min_x, min_tot_en2, color='c')

else:
raise RuntimeError("Invalid branch")

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')
Expand All @@ -368,20 +377,19 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
tmesh = np.linspace(tstart, tstop, num)
ph_energies = self.get_vib_free_energies(tstart, tstop, num)
ax, fig, plt = get_ax_fig_plt(ax, figsize=(10, 8)) # Ensure a valid plot axis
min_x, min_y = np.zeros(num), np.zeros(num)
min_tot_energy = np.zeros(num)
min_x, min_y, min_tot_energy = np.zeros(num), np.zeros(num), np.zeros(num)

if (len(self.lattice_a_from_phdos)==len(self.lattice_a) or len(self.lattice_c_from_phdos)==len(self.lattice_c)):
if self.use_qha:
tot_energies = self.energies[np.newaxis, :].T + ph_energies+ self.volumes[np.newaxis, :].T * self.pressure / abu.eVA3_GPa

# Initial guess for minimization
initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()

# Perform minimization for each temperature
for j, energy in enumerate(tot_energies.T):
f_interp = RectBivariateSpline(self.lattice_a[:, 0], self.lattice_c[0, :], energy, kx=4, ky=4)
min_x[j],min_y[j],min_tot_energy[j]= self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
min_x[j], min_y[j], min_tot_energy[j] = self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
xy_init = min_x[j], min_y[j]

# Calculate thermal expansion coefficients
A0, C0 = self.lattice_a[self.ix0, self.iy0], self.lattice_c[self.ix0, self.iy0]
Expand All @@ -397,17 +405,14 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
ax.plot(tmesh[1:-1], alpha_c, color='r', label=r"$\alpha_c$ (QHA)", linewidth=2)
#ax.plot(tmesh[1:-1], alpha_v, color='purple', label=r"$\alpha_v$ (QHA)", linewidth=2)

elif (len(self.lattice_a_from_phdos)==3 or len(self.lattice_c_from_phdos)==3):
elif self.use_einfvib2:

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]
dc = self.lattice_c_from_phdos[1,0]-self.lattice_c_from_phdos[1,1]

dF_dA, dF_dC, d2F_dA2, d2F_dC2, d2F_dAdC = (np.zeros(num) for _ in range(5))
for i, e in enumerate(ph_energies.T):
dF_dA[i]=(e[0,1]-e[2,1])/(2*da)
dF_dC[i]=(e[1,0]-e[1,2])/(2*dc)
Expand All @@ -423,11 +428,12 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
gradient = np.zeros(2)

# Initial guess for minimization
initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()

for j, energy in enumerate(tot_en2.T):
f_interp = RectBivariateSpline(self.lattice_a[:, 0], self.lattice_c[0, :], energy, kx=4, ky=4)
min_x[j],min_y[j],min_tot_energy[j]= self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
min_x[j], min_y[j], min_tot_energy[j] = self.find_minimum(f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
xy_init = min_x[j], min_y[j]

A0 = self.lattice_a[self.ix0,self.iy0]
C0 = self.lattice_c[self.ix0,self.iy0]
Expand All @@ -442,11 +448,14 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
ax.plot(tmesh[1:-1], alpha_c, linestyle='--', color='teal', label=r"$\alpha_c$ E$\infty$Vib2")
#ax.plot(tmesh[1:-1], alpha_v, linestyle='--', color='darkorange', label=r"$\alpha_v$ E$\infty$Vib2")

else:
raise RuntimeError("Invalid branch.")

# 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']
file_path = 'thermal-expansion_data.txt'
print(f"Writing thermal expansion data to: {file_path}"
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 All @@ -471,22 +480,21 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure
"""
tmesh = np.linspace(tstart, tstop, num)
ph_energies = self.get_vib_free_energies(tstart, tstop, num)
min_x, min_y = np.zeros(num), np.zeros(num)
min_tot_energy = np.zeros(num)
min_x, min_y, min_tot_energy = np.zeros(num), np.zeros(num), np.zeros(num)
import matplotlib.pyplot as plt
fig, axs = plt.subplots(1, 3, figsize=(18, 6), sharex=True)

if (len(self.lattice_a_from_phdos)==len(self.lattice_a) or len(self.lattice_c_from_phdos)==len(self.lattice_c)):
if self.use_qha:
tot_energies = self.energies[np.newaxis, :].T + ph_energies+ self.volumes[np.newaxis, :].T * self.pressure / abu.eVA3_GPa

# Initial guess for minimization
initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()

# Perform minimization for each temperature
for j, energy in enumerate(tot_energies.T):
f_interp = RectBivariateSpline(self.lattice_a[:, 0], self.lattice_c[0, :], energy, kx=4, ky=4)
min_x[j],min_y[j],min_tot_energy[j]= self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
min_x[j], min_y[j], min_tot_energy[j] = self.find_minimum(f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
xy_init = min_x[j], min_y[j]

# Calculate thermal expansion coefficients
A0, C0 = self.lattice_a[self.ix0, self.iy0], self.lattice_c[self.ix0, self.iy0]
Expand All @@ -499,16 +507,14 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure
axs[1].plot(tmesh, min_y, color='r', label=r"$c$ (QHA)", linewidth=2)
axs[2].plot(tmesh, min_volumes, color='b', label=r"$V$ (QHA)", linewidth=2)

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)
elif self.use_einfvib2:

a0 = self.lattice_a[1,1]
c0 = self.lattice_c[1,1]
da = self.lattice_a[0,1]-self.lattice_a[1,1]
dc = self.lattice_c[1,0]-self.lattice_c[1,1]

dF_dA, dF_dC, d2F_dA2, d2F_dC2, d2F_dAdC = (np.zeros(num) for _ in range(5))
for i, e in enumerate(ph_energies.T):
dF_dA[i]=(e[0,1]-e[2,1])/(2*da)
dF_dC[i]=(e[1,0]-e[1,2])/(2*dc)
Expand All @@ -522,22 +528,26 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure
tot_en2 = tot_en2 + (self.lattice_c[np.newaxis, :].T - c0)*(self.lattice_a[np.newaxis, :].T - a0)*d2F_dAdC

# Initial guess for minimization
initial_guess = [1.005*self.lattice_a[self.ix0,0], 1.005*self.lattice_c[0,self.iy0]]
xy_init = np.array(initial_guess)
xy_init = self.get_initial_guess_ac()

for j, energy in enumerate(tot_en2.T):
f_interp = RectBivariateSpline(self.lattice_a[:, 0], self.lattice_c[0, :], energy, kx=4, ky=4)
min_x[j], min_y[j], min_tot_energy[j] = self.find_minimum( f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
min_x[j], min_y[j], min_tot_energy[j] = self.find_minimum(f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01)
xy_init = min_x[j], min_y[j]

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
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_volumes = min_x**2 * min_y * scale

axs[0].plot(tmesh, min_x, color='c', label=r"$a$ (E$\infty$Vib2)", linewidth=2)
axs[1].set_title(r"Plots of a, c, and V (E$\infty$Vib2)")
axs[1].plot(tmesh, min_y, color='r', label=r"$c$ (E$\infty$Vib2)", linewidth=2)
axs[2].plot(tmesh, min_volumes, color='b', label=r"$V$ (E$\infty$Vib2)", linewidth=2)

else:
raise RuntimeError("Invalid branch.")

axs[0].set_ylabel("a")
axs[0].legend(loc="best", shadow=True)
axs[0].grid(True)
Expand Down Expand Up @@ -571,7 +581,6 @@ def get_vib_free_energies(self, tstart: float, tstop: float, num: int) -> np.nda

for i in range(len(self.lattice_a_from_phdos[:, 0])):
for j in range(len(self.lattice_c_from_phdos[0])):
phdos = self.phdoses[i][j]
if phdos is not None:
f[j][i] = phdos.get_free_energy(tstart, tstop, num).values
if (phdos := self.phdoses[i][j]) is not None:
f[j, i] = phdos.get_free_energy(tstart, tstop, num).values
return f
25 changes: 17 additions & 8 deletions abipy/examples/plot/plot_qha_2d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,23 +7,32 @@
to compute thermodynamic properties within the quasi-harmonic approximation.
"""
import os
import numpy as np
import abipy.data as abidata

from abipy.dfpt.qha_2D import QHA_2D

strains_a = [ 995,1000, 1005, 1010, 1015 ]
strains_c = [ 995,1000, 1005, 1010, 1015 ]
strains_a1 = [ 1000, 1005, 1010 ]
strains_c1 = [ 1000, 1005, 1010 ]
bo_strains_a = [995, 1000, 1005, 1010, 1015]
bo_strains_c = [995, 1000, 1005, 1010, 1015]
phdos_strains_a = [1000, 1005, 1010]
phdos_strains_c = [1000, 1005, 1010]

# Root points to the directory in the git submodule with the output results.
root = os.path.join(abidata.dirpath, "data_v-ZSISA-QHA.git", "ZnO_ZSISA_approximation")

#gsr_paths = [[f"scale_{s1}_{s3}/out_GSR.nc" for s3 in strains_c] for s1 in strains_a]
gsr_paths = [[os.path.join(root, f"scale_{s1}_{s3}/out_GSR_DDB") for s3 in strains_c] for s1 in strains_a]
phdos_paths = [[os.path.join(root, f"scale_{s1}_{s3}/out_PHDOS.nc") for s3 in strains_c1] for s1 in strains_a1]
#gsr_paths = [[f"scale_{s1}_{s3}/out_GSR.nc" for s3 in bo_strains_c] for s1 in bo_strains_a]
gsr_paths = [[os.path.join(root, f"scale_{s1}_{s3}/out_GSR_DDB") for s3 in bo_strains_c] for s1 in bo_strains_a]
phdos_paths = [[os.path.join(root, f"scale_{s1}_{s3}/out_PHDOS.nc") for s3 in phdos_strains_c] for s1 in phdos_strains_a]

qha = QHA_2D.from_files(gsr_paths, phdos_paths, gsr_file="DDB")
bo_strains_a = (np.array(bo_strains_a) - 1000)/ 100
bo_strains_c = (np.array(bo_strains_c) - 1000)/ 100
bo_strains_ac = [bo_strains_a, bo_strains_c]

phdos_strains_a = (np.array(phdos_strains_a) - 1000)/ 100
phdos_strains_c = (np.array(phdos_strains_c) - 1000)/ 100
phdos_strains_ac = [phdos_strains_a, phdos_strains_c]

qha = QHA_2D.from_files(gsr_paths, phdos_paths, bo_strains_ac, phdos_strains_ac, gsr_file="DDB")

#%%
qha.plot_energies()
Expand Down

0 comments on commit 7e6a2bb

Please sign in to comment.