Add run_qha_2d.py example

gmatteo · gmatteo · commit 36fe56ab0519 · 2024-12-22T19:03:59.000+01:00
diff --git a/abipy/dfpt/deformation_utils.py b/abipy/dfpt/deformation_utils.py
@@ -51,7 +51,7 @@ def _add(name, new_rprim, i, j, k, l, m, n) -> None:
         structures_new[name] = new_structure
         strain_inds.append([i, j, k, l, m, n])
 
-    i, j, k, l, m, n = 6 * [None]
+    i, j, k, l, m, n = 6 * [0]
 
     if 1 <= spgrp_number <= 2:
         disp=[[1,1,1,1,1,1],  [0,1,1,1,1,1],  [2,1,1,1,1,1],  [1,0,1,1,1,1],  [1,2,1,1,1,1],  [1,1,0,1,1,1],
@@ -193,5 +193,5 @@ def _add(name, new_rprim, i, j, k, l, m, n) -> None:
     else:
         raise ValueError(f"Invalid {spgrp_number=}")
 
-    return structures_new, np.array(strain_inds, dtype=object), spgrp_number
+    return structures_new, np.array(strain_inds, dtype=int), spgrp_number
 
diff --git a/abipy/dfpt/qha_2D.py b/abipy/dfpt/qha_2D.py
@@ -12,44 +12,96 @@
 import numpy as np
 import abipy.core.abinit_units as abu
 
+from scipy.interpolate import RectBivariateSpline #, RegularGridInterpolator
 #from monty.collections import dict2namedtuple
 #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
 from abipy.electrons.gsr import GsrFile
 from abipy.dfpt.ddb import DdbFile
 from abipy.dfpt.phonons import PhdosFile # PhononBandsPlotter, PhononDos,
-from scipy.interpolate import RectBivariateSpline #, RegularGridInterpolator
+from abipy.dfpt.vzsisa import anaget_phdoses_with_gauss
+
 
 
-class QHA_2D:
+class QHA_2D(HasPickleIO):
     """
     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):
+    @classmethod
+    def from_json_file(cls,
+                       filepath: PathLike,
+                       nqsmall_or_qppa: int,
+                       anaget_kwargs: dict | None = None,
+                       smearing_ev: float | None = None,
+                       verbose: int = 0) -> Vzsisa:
+        """
+        Build an instance from a json file `filepath` typically produced by an Abipy flow.
+        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_files(ddb_files, phdos_paths_2D, gsr_file="DDB")
+        return cls.from_gsr_ddb_paths(nqsmall_or_qppa,
+                                      data["gsr_relax_paths"], data["ddb_relax_paths"],
+                                      bo_strains_ac, phdos_strains_ac,
+                                      anaget_kwargs=anaget_kwargs, smearing_ev=smearing_ev, verbose=verbose)
 
     @classmethod
-    def from_files(cls, gsr_paths_2D, phdos_paths_2D, gsr_file="GSR.nc") -> QHA_2D:
+    def from_gsr_ddb_paths(cls,
+                           nqsmall_or_qppa: int,
+                           gsr_paths,
+                           ddb_paths,
+                           bo_strains_ac,
+                           phdos_strains_ac,
+                           anaget_kwargs: dict | None = None,
+                           smearing_ev: float | None = None,
+                           verbose: int = 0) -> QHA_2D:
+        """
+        Creates an instance from a list of GSR files and a list of DDB files.
+        This is a simplified interface that computes the PHDOS.nc files automatically
+        from the DDB files by invoking anaddb
+
+        Args:
+            nqsmall_or_qppa: Define the q-mesh for the computation of the PHDOS.
+                if > 0, it is interpreted as nqsmall
+                if < 0, it is interpreted as qppa.
+            gsr_paths: list of paths to GSR files.
+            ddb_paths: list of paths to DDB files.
+            bo_strains_ac: List of strains for the a and the c lattice vector.
+            phdos_strains_ac: List of strains for the a and the c lattice vector.
+            anaget_kwargs: dict with extra arguments passed to anaget_phdoses_with_gauss.
+            smearing_ev: Gaussian smearing in eV.
+            verbose: Verbosity level.
+        """
+        phdos_paths, phbands_paths = anaget_phdoses_with_gauss(nqsmall_or_qppa, smearing_ev, ddb_paths, anaget_kwargs, verbose)
+
+        new = cls.from_files(ddb_files, phdos_paths_2D, bo_strains_ac, phdos_strains_ac, gsr_file="GSR.nc")
+        #new.pickle_dump(workdir, basename=None)
+        return new
+
+    @classmethod
+    def from_files(cls, gsr_paths_2D, phdos_paths_2D, bo_strains_ac, phdos_strains_ac, 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.
+            bo_strains_ac: List of strains for the a and the c lattice vector.
+            phdos_strains_ac: List of strains for the a and the c lattice vector.
         """
         energies, structures, phdoses , structures_from_phdos = [], [], [],[]
 
+        #shape = (len(strains_a), len(strains_c))
+        #gsr_paths_2d = np.reshape(gsr_paths_2D, shape)
+        #phdos_paths_2d = np.reshape(phdos_paths_2D, shape)
+
         if gsr_file == "GSR.nc":
             # Process GSR files
             for row in gsr_paths_2D:
@@ -98,22 +150,29 @@ def from_files(cls, gsr_paths_2D, phdos_paths_2D, gsr_file="GSR.nc") -> QHA_2D:
             phdoses.append(row_doses)
             structures_from_phdos.append(row_structures)
 
-        return cls(structures, phdoses, energies , structures_from_phdos)
+        return cls(structures, phdoses, energies, structures_from_phdos, bo_strains_ac, phdos_strains_ac)
 
     def __init__(self, structures, phdoses, energies, structures_from_phdos,
+                 bo_strains_ac, phdos_strains_ac,
                  eos_name: str='vinet', pressure: float=0.0):
         """
         Args:
             structures (list): List of structures at different volumes.
             phdoses: List of density of states (DOS) data for phonon calculations.
             energies (list): SCF energies for the structures in eV.
+            bo_strains_ac: List of strains for the a and the c lattice vector.
+            phdos_strains_ac: List of strains for the a and the c lattice vector.
             eos_name (str): Expression used to fit the energies (e.g., 'vinet').
             pressure (float): External pressure in GPa to include in p*V term.
         """
         self.phdoses = phdoses
         self.structures = structures
         self.structures_from_phdos = structures_from_phdos
         self.energies = np.array(energies, dtype=np.float64)
+
+        self.bo_strains_ac = bo_strains_ac
+        self.phdos_strains_ac = phdos_strains_ac
+
         self.eos_name = eos_name
         self.pressure = pressure
         self.volumes = np.array([[s.volume if s else np.nan for s in row] for row in structures])
@@ -509,7 +568,7 @@ def get_vib_free_energies(self, tstart: float, tstop: float, num: int) -> np.nda
             tstop: Stop temperature.
             num: Number of temperature points.
 
-        Return: A 3D array of vibrational free energies.
+        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])):
diff --git a/abipy/dfpt/qha_general_stress.py b/abipy/dfpt/qha_general_stress.py
@@ -46,20 +46,21 @@ def from_json_file(cls,
             verbose: Verbosity level
         """
         data = mjson_load(filepath)
-        ddb_paths = data["ddb_relax_paths"]
-        gsr_path = data["gsr_relax_paths"]
+        ddb_relax_paths = data["ddb_relax_paths"]
+        gsr_relax_paths = data["gsr_relax_paths"]
 
         phdos_paths, phbands_paths = anaget_phdoses_with_gauss(nqsmall_or_qppa, smearing_ev,
-                                                               ddb_paths, anaget_kwargs, verbose)
+                                                               ddb_relax_paths, anaget_kwargs, verbose)
 
         # Create a 6D array with shape (3, 3, 3, 3, 3, 3) initialized with None.
         gsr_paths_6d = np.full((3, 3, 3, 3, 3, 3), None, dtype=object)
         phdos_paths_6d = np.full((3, 3, 3, 3, 3, 3), None, dtype=object)
 
-        cell6_inds = np.array(data["cell6_inds"], dtype=np.int)
-        for gsr_path, phdos_path, cell_ind in zip(gsr_paths, phdos_paths, cell6_inds, strict=True):
-            gsr_paths_6d[cell_ind] = gsr_path
-            phdos_paths_6d[cell_ind] = phdos_path
+        strain_inds = data["strain_inds"]
+        #print(strain_inds)
+        for gsr_path, phdos_path, inds in zip(gsr_relax_paths, phdos_paths, strain_inds, strict=True):
+            gsr_paths_6d[inds] = gsr_path
+            phdos_paths_6d[inds] = phdos_path
 
         new = cls.from_files(gsr_paths_6d, phdos_paths_6d, gsr_guess, model='ZSISA')
         #new.pickle_dump(workdir, basename=None)
diff --git a/abipy/dfpt/tests/test_deformation_utils.py b/abipy/dfpt/tests/test_deformation_utils.py
@@ -0,0 +1,20 @@
+"""Tests for deformation_utils module"""
+import numpy as np
+import abipy.data as abidata
+
+from abipy.core.testing import AbipyTest
+from abipy.dfpt.deformation_utils import generate_deformations
+
+
+class DeformationUtilsTest(AbipyTest):
+
+    def test_generate_deformations(self):
+        """Testing generate_deformations"""
+        eps = 0.005
+        si = self.get_structure("Si")
+        structures_dict, strain_inds, spgrp_number = generate_deformations(si, eps)
+        assert len(structures_dict) == 3
+        assert spgrp_number == 227
+        assert strain_inds.shape == (3, 6)
+        self.assert_equal(strain_inds[:, 0], [0, 1, 2])
+        assert np.all(strain_inds[1:, 1] == 0)
diff --git a/abipy/dfpt/tests/test_qha_2d.py b/abipy/dfpt/tests/test_qha_2d.py
@@ -1,5 +1,6 @@
 """Tests for QHA_2D"""
 import os
+import numpy as np
 import abipy.data as abidata
 
 from abipy.dfpt.qha_2D import QHA_2D
@@ -8,20 +9,25 @@
 
 class Qha2dTest(AbipyTest):
 
-    def test_qha_2d(self):
+    def test_zsisa_approximation(self):
 
-        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 = [[os.path.join(root, f"scale_{s1}_{s3}/out_GSR_DDB") for s3 in strains_c] for s1 in strains_a]
-        dos_paths = [[os.path.join(root, f"scale_{s1}_{s3}/out_PHDOS.nc") for s3 in strains_c1] for s1 in strains_a1]
+        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]
+        dos_paths = [[os.path.join(root, f"scale_{s1}_{s3}/out_PHDOS.nc") for s3 in phdos_strains_c] for s1 in phdos_strains_a]
+
+        bo_strains_ac = [bo_strains_a, bo_strains_c]
+        bo_strains_ac = (np.array(bo_strains_ac) - 1000) / 100
+        phdos_strains_ac = [phdos_strains_a, phdos_strains_c]
+        phdos_strains_ac = (np.array(phdos_strains_ac) - 1000) / 100
 
-        qha = QHA_2D.from_files(gsr_paths, dos_paths, gsr_file="DDB")
+        qha = QHA_2D.from_files(gsr_paths, dos_paths, bo_strains_ac, phdos_strains_ac, gsr_file="DDB")
 
         # Test properties and methods.
         assert qha.pressure == 0
@@ -48,8 +54,78 @@ def test_qha_2d(self):
         self.assert_almost_equal(f_mat, ref_mat)
 
         if self.has_matplotlib():
-        #if False:
             qha.plot_energies(show=False)
             qha.plot_free_energies(tstop=500, tstart=0, num=6, show=False)
             qha.plot_thermal_expansion(tstop=1000, tstart=0, num=101, show=False)
             qha.plot_lattice(tstop=1000, tstart=0, num=101, show=False)
+
+    def test_qha_2d(self):
+
+        bo_strains_a = [995, 1000, 1005, 1010, 1015]
+        bo_strains_c = [995, 1000, 1005, 1010, 1015]
+        #bo_strains_a = [1000, 1005, 1010, 1015 , 1020]
+        #bo_strains_c = [1000, 1005, 1010, 1015 , 1020]
+
+        # 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_QHA")
+
+        #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 bo_strains_c] for s1 in bo_strains_a]
+
+        bo_strains_ac = [bo_strains_a, bo_strains_c]
+        bo_strains_ac = (np.array(bo_strains_ac) - 1000) / 100
+
+        #qha = QHA_2D.from_files(gsr_paths, phdos_paths, bo_strains_ac, phdos_strains_ac, gsr_file="GSR.nc")
+        qha = QHA_2D.from_files(gsr_paths, phdos_paths, bo_strains_ac, bo_strains_ac, gsr_file="DDB")
+
+        # Test properties and methods.
+        assert qha.pressure == 0
+
+        f_mat = qha.get_vib_free_energies(0, 1000, 2)
+        #print(f_mat)
+        ref_mat = [
+            [
+                [ 0.2364734 , -1.0138777 ],
+                [ 0.23334375, -1.02191455],
+                [ 0.23021055, -1.03048734],
+                [ 0.22707712, -1.03958781],
+                [ 0.22394442, -1.04917831]
+            ],
+            [
+                [ 0.23500674, -1.0175096 ],
+                [ 0.23190046, -1.0257329 ],
+                [ 0.2287909 , -1.03448119],
+                [ 0.22568001, -1.0437223 ],
+                [ 0.22257111, -1.05344296]
+            ],
+            [
+                [ 0.23352563, -1.02131569],
+                [ 0.23044276, -1.02971896],
+                [ 0.22735644, -1.03862873],
+                [ 0.22427218, -1.04801921],
+                [ 0.2211813 , -1.05787544]
+            ],
+            [
+                [ 0.23203085, -1.02529729],
+                [ 0.22897121, -1.0338748 ],
+                [ 0.22590762, -1.04294394],
+                [ 0.22284219, -1.05248871],
+                [ 0.21977746, -1.06247964]
+            ],
+            [
+                [ 0.2305236 , -1.02944732],
+                [ 0.22748688, -1.03820052],
+                [ 0.22444611, -1.04742487],
+                [ 0.22140302, -1.05711238],
+                [ 0.21836046, -1.06723406]
+            ]
+        ]
+
+        self.assert_almost_equal(f_mat, ref_mat)
+
+        if self.has_matplotlib():
+            assert qha.plot_energies(show=False)
+            assert qha.plot_free_energies(tstop=500, tstart=0, num=6, show=False)
+            assert qha.plot_thermal_expansion(tstop=1000, tstart=0, num=101, show=False)
+            assert qha.plot_lattice(tstop=1000, tstart=00, num=101, show=False)
diff --git a/abipy/dfpt/vzsisa.py b/abipy/dfpt/vzsisa.py
@@ -119,9 +119,9 @@ def from_gsr_ddb_paths(cls,
                 if < 0, it is interpreted as qppa.
             gsr_paths: list of paths to GSR files.
             ddb_paths: list of paths to DDB files.
-            anaget_kwargs: dict with extra arguments passed to anaget_phdoses_with_gauss
-            smearing_ev: Gaussian smearing in eV
-            verbose:
+            anaget_kwargs: dict with extra arguments passed to anaget_phdoses_with_gauss.
+            smearing_ev: Gaussian smearing in eV.
+            verbose: Verbosity level.
         """
         phdos_paths, phbands_paths = anaget_phdoses_with_gauss(nqsmall_or_qppa, smearing_ev, ddb_paths, anaget_kwargs, verbose)
         new = cls.from_gsr_phdos_files(gsr_paths, phdos_paths)
diff --git a/abipy/examples/flows/run_qha_2d.py b/abipy/examples/flows/run_qha_2d.py
diff --git a/abipy/flowtk/qha_2d.py b/abipy/flowtk/qha_2d.py
diff --git a/abipy/flowtk/zsisa.py b/abipy/flowtk/zsisa.py
