Skip to content

Commit

Permalink
Move classmethods before __init__ to improve readability
Browse files Browse the repository at this point in the history
  • Loading branch information
@gmatteo
gmatteo committed Dec 15, 2024
1 parent 2c53c9a commit 861c9b8
Show file tree
Hide file tree
Showing 3 changed files with 274 additions and 261 deletions.
199 changes: 98 additions & 101 deletions abipy/dfpt/qha_2D.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
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.
"""
from __future__ import annotations

import os
import abc
import numpy as np
Expand All @@ -19,13 +21,87 @@
from abipy.dfpt.phonons import PhdosFile # PhononBandsPlotter, PhononDos,
from scipy.interpolate import RectBivariateSpline #, RegularGridInterpolator


class QHA_2D:
"""
Quasi-Harmonic Approximation (QHA) analysis in 2D.
Provides methods for calculating and visualizing energy, free energy, and thermal expansion.
"""

#@classmethod
#def from_ddb_files(cls, ddb_files):
# for row in ddb_files:
# for gp in row:
# if os.path.exists(gp):

# return cls.from_files(ddb_files, phdos_paths_2D, gsr_file="DDB")

@classmethod
def from_files(cls, gsr_paths_2D, phdos_paths_2D, gsr_file="GSR.nc") -> QHA_2D:
"""
Creates an instance of QHA from 2D lists of GSR and PHDOS files.
Args:
gsr_paths_2D: 2D list of paths to GSR files.
phdos_paths_2D: 2D list of paths to PHDOS.nc files.
Returns:
A new instance of QHA.
"""
energies, structures, doses , structures_from_phdos = [], [], [],[]

if gsr_file == "GSR.nc":
# Process GSR files
for row in gsr_paths_2D:
row_energies, row_structures = [], []
for gp in row:
if os.path.exists(gp):
with GsrFile.from_file(gp) as g:
row_energies.append(g.energy)
row_structures.append(g.structure)
else:
row_energies.append(None)
row_structures.append(None)
energies.append(row_energies)
structures.append(row_structures)

elif gsr_file == "DDB":
# Process DDB files
for row in gsr_paths_2D:
row_energies, row_structures = [], []
for gp in row:
if os.path.exists(gp):
with DdbFile.from_file(gp) as g:
row_energies.append(g.total_energy)
row_structures.append(g.structure)
else:
row_energies.append(None)
row_structures.append(None)
energies.append(row_energies)
structures.append(row_structures)

else:
raise ValueError(f"Invalid {gsr_file=}")

# Process PHDOS files
for row in phdos_paths_2D:
row_doses , row_structures = [],[]
for path in row:
if os.path.exists(path):
with PhdosFile(path) as p:
row_doses.append(p.phdos)
row_structures.append(p.structure)
else:
row_doses.append(None)
row_structures.append(None)

doses.append(row_doses)
structures_from_phdos.append(row_structures)

return cls(structures, doses, energies , structures_from_phdos)

def __init__(self, structures, doses, energies, structures_from_phdos,
eos_name: str='vinet', pressure: float=0):
eos_name: str='vinet', pressure: float=0.0):
"""
Args:
structures (list): List of structures at different volumes.
Expand Down Expand Up @@ -65,28 +141,24 @@ def plot_energies(self, ax=None, **kwargs) -> Figure:
Args:
ax: Matplotlib axis for the plot. If None, creates a new figure.
"""
a0 = self.lattice_a[:, 0]
c0 = self.lattice_c[0, :]

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)
X, Y = np.meshgrid(c0, a0)

# Plot the surface
ax.plot_wireframe(X, Y, self.energies, cmap='viridis')
ax.scatter(self.lattice_c[0,self.iy0], self.lattice_a[self.ix0,0], self.energies[self.ix0, self.iy0], color='red', s=100)

f_interp = RectBivariateSpline(a0,c0, self.energies, kx=4, ky=4)
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)
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)
x_grid, y_grid = np.meshgrid(y_new, x_new)
Expand All @@ -103,7 +175,7 @@ def plot_energies(self, ax=None, **kwargs) -> Figure:

return fig

def find_minimum(self, f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01):
def find_minimum(self, f_interp, xy_init, tol=1e-6, max_iter=1000, step_size=0.01) -> tuple:
"""
Gradient descent to find the minimum of the interpolated energy surface.
Expand Down Expand Up @@ -147,10 +219,10 @@ def plot_free_energies(self, tstart=800 , tstop=0 ,num=5, ax=None, **kwargs) ->

tmesh = np.linspace(tstart, tstop, num)
ph_energies = self.get_vib_free_energies(tstart, tstop, num)
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)):
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)):
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])
Expand All @@ -165,8 +237,8 @@ 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_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)
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)

ax.scatter(min_y,min_x,min_tot_en, color='c', s=100)
ax.plot(min_y,min_x,min_tot_en, color='c')
Expand Down Expand Up @@ -237,7 +309,6 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
tstop: Stop temperature.
num: Number of temperature points.
ax: Matplotlib axis object for plotting.
**kwargs: Additional arguments for customization.
"""
tmesh = np.linspace(tstart, tstop, num)
ph_energies = self.get_vib_free_energies(tstart, tstop, num)
Expand Down Expand Up @@ -280,16 +351,15 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
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]
for i,e in enumerate(ph_energies.T):
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]
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)
d2F_dA2[i]=(e[0,1]-2*e[1,1]+e[2,1])/(da)**2
d2F_dC2[i]=(e[1,0]-2*e[1,1]+e[1,2])/(dc)**2
d2F_dAdC[i] = (e[1,1] - e[1, 0] - e[0, 1] + e[0, 0]) / ( da * dc)


tot_en2 = self.energies[np.newaxis, :].T + ph_energies[1,1] + self.volumes[np.newaxis, :].T * self.pressure / abu.eVA3_GPa
tot_en2 = tot_en2+ (self.lattice_a[np.newaxis, :].T - a0)*dF_dA + 0.5*(self.lattice_a[np.newaxis, :].T - a0)**2*d2F_dA2
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
Expand All @@ -309,7 +379,6 @@ def plot_thermal_expansion(self, tstart=800, tstop=0, num=81, ax=None, **kwargs)
scale= self.volumes[self.ix0,self.iy0]/A0**2/C0
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]
alpha_c = (min_y[2:] - min_y[:-2]) / (2 * dt) / min_y[1:-1]
Expand Down Expand Up @@ -375,15 +444,15 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure
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)
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[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]
da = self.lattice_a[0,1]-self.lattice_a[1,1]
dc = self.lattice_c[1,0]-self.lattice_c[1,1]
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 @@ -401,7 +470,7 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure
xy_init = np.array(initial_guess)
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)

A0 = self.lattice_a[self.ix0,self.iy0]
C0 = self.lattice_c[self.ix0,self.iy0]
Expand Down Expand Up @@ -431,78 +500,6 @@ def plot_lattice(self, tstart=800, tstop=0, num=81, ax=None, **kwargs) -> Figure

return fig

#@classmethod
#def from_ddb_files(cls, ddb_files):
# for row in ddb_files:
# for gp in row:
# if os.path.exists(gp):

# return cls.from_files(ddb_files, phdos_paths_2D, gsr_file="DDB")

@classmethod
def from_files(cls, gsr_paths_2D, phdos_paths_2D, gsr_file="GSR.nc"):
"""
Creates an instance of QHA from 2D lists of GSR and PHDOS files.
Args:
gsr_paths_2D: 2D list of paths to GSR files.
phdos_paths_2D: 2D list of paths to PHDOS.nc files.
Returns:
A new instance of QHA.
"""
energies, structures, doses , structures_from_phdos = [], [], [],[]

if gsr_file == "GSR.nc":
# Process GSR files
for row in gsr_paths_2D:
row_energies, row_structures = [], []
for gp in row:
if os.path.exists(gp):
with GsrFile.from_file(gp) as g:
row_energies.append(g.energy)
row_structures.append(g.structure)
else:
row_energies.append(None)
row_structures.append(None)
energies.append(row_energies)
structures.append(row_structures)

elif gsr_file == "DDB":
# Process DDB files
for row in gsr_paths_2D:
row_energies, row_structures = [], []
for gp in row:
if os.path.exists(gp):
with DdbFile.from_file(gp) as g:
row_energies.append(g.total_energy)
row_structures.append(g.structure)
else:
row_energies.append(None)
row_structures.append(None)
energies.append(row_energies)
structures.append(row_structures)

else:
raise ValueError(f"Invalid {gsr_file=}")

# Process PHDOS files
for row in phdos_paths_2D:
row_doses , row_structures = [],[]
for path in row:
if os.path.exists(path):
with PhdosFile(path) as p:
row_doses.append(p.phdos)
row_structures.append(p.structure)
else:
row_doses.append(None)
row_structures.append(None)

doses.append(row_doses)
structures_from_phdos.append(row_structures)

return cls(structures, doses, energies , structures_from_phdos)

def get_vib_free_energies(self, tstart: float, tstop: float, num: int) -> np.ndarray:
"""
Computes the vibrational free energies from phonon density of states.
Expand Down
Loading

0 comments on commit 861c9b8

Please sign in to comment.