use amu_symbol

abipy/dfpt/gruneisen.py

@@ -13,7 +13,7 @@ from monty.termcolor import cprint
 from monty.collections import AttrDict
 from monty.functools import lazy_property
 from abipy.core.kpoints import Kpath, IrredZone, KSamplingInfo
-from abipy.core.mixins import AbinitNcFile, Has_Structure, Has_ElectronBands, NotebookWriter
+from abipy.core.mixins import AbinitNcFile, Has_Structure, NotebookWriter
 from abipy.abio.inputs import AnaddbInput
 from abipy.dfpt.phonons import PhononBands, PhononBandsPlotter, PhononDos, match_eigenvectors, get_dyn_mat_eigenvec
 from abipy.dfpt.ddb import DdbFile
@@ -22,8 +22,8 @@ from abipy.tools.plotting import add_fig_kwargs, get_ax_fig_plt, get_axarray_fig
 from abipy.flowtk import AnaddbTask
 from abipy.tools.derivatives import finite_diff
 #from abipy.tools import duck
-from abipy.core.func1d import Function1D
-from pymatgen.core.units import amu_to_kg, bohr_to_ang
+from pymatgen.core.units import amu_to_kg
+from pymatgen.core.periodic_table import Element
 
 try:
     from functools import lru_cache
@@ -161,11 +161,11 @@ class GrunsNcFile(AbinitNcFile, Has_Structure, NotebookWriter):
         return len(self.structures)
 
     @lazy_property
-    def amu(self):
+    def amu_symbol(self):
         """Atomic mass units"""
         amu_list = self.reader.read_value("atomic_mass_units")
         atomic_numbers = self.reader.read_value("atomic_numbers")
-        amu = {at: a for at, a in zip(atomic_numbers, amu_list)}
+        amu = {Element.from_Z(at).symbol: a for at, a in zip(atomic_numbers, amu_list)}
         return amu
 
     def to_dataframe(self):
@@ -800,7 +800,8 @@ class GrunsNcFile(AbinitNcFile, Has_Structure, NotebookWriter):
         if match_eigv:
             eig = np.zeros_like(self.phdispl_cart_qibz)
             for i in range(self.nvols):
-                eig[:, i] = get_dyn_mat_eigenvec(self.phdispl_cart_qibz[:, i], self.structures[i], amu=self.amu)
+                eig[:, i] = get_dyn_mat_eigenvec(self.phdispl_cart_qibz[:, i], self.structures[i],
+                                                 amu_symbol=self.amu_symbol)
 
             eig = eig.transpose((1, 0, 2, 3))
         else:

abipy/dfpt/phonons.py

@@ -4057,7 +4057,7 @@ class InteratomicForceConstants(Has_Structure):
 
 
 # TODO: amu should become mandatory.
-def get_dyn_mat_eigenvec(phdispl, structure, amu=None):
+def get_dyn_mat_eigenvec(phdispl, structure, amu=None, amu_symbol=None):
     """
     Converts the phonon displacements to the orthonormal eigenvectors of the dynamical matrix.
     Small discrepancies with the original values may be expected due to the different values of the atomic masses in
@@ -4075,20 +4075,29 @@ def get_dyn_mat_eigenvec(phdispl, structure, amu=None):
             should match the q points.
         structure: |Structure| object.
         amu: dictionary that associates the atomic numbers present in the structure to the values of the atomic
-            mass units used for the calculation. If None, values from pymatgen will be used. Note that this will
-            almost always lead to inaccuracies in the conversion.
+            mass units used for the calculation. Incompatible with amu_sumbol. If None and amu_symbol is None, values
+            from pymatgen will be used.  Note that this will almost always lead to inaccuracies in the conversion.
+        amu_symbol: dictionary that associates the symbol present in the structure to the values of the atomic
+            mass units used for the calculation. Incompatible with amu. If None and amu_symbol is None, values from
+            pymatgen will be used. that this will almost always lead to inaccuracies in the conversion.
 
     Returns:
         A |numpy-array| of the same shape as phdispl containing the eigenvectors of the dynamical matrix
     """
     eigvec = np.zeros(np.shape(phdispl), dtype=np.complex)
 
-    if amu is None:
+    if amu is not None and amu_symbol is not None:
+        raise ValueError("Only one between amu and amu_symbol should be provided!")
+
+    if amu is not None:
+        amu_symbol = {Element.from_Z(n).symbol: v for n, v in amu.items()}
+
+    if amu_symbol is None:
         warnings.warn("get_dyn_mat_eigenvec has been called with amu=None. Eigenvectors may not be orthonormal.")
-        amu = {e.number: e.atomic_mass for e in structure.composition.elements}
+        amu_symbol = {e.symbol: e.atomic_mass for e in structure.composition.elements}
 
     for j, a in enumerate(structure):
-        eigvec[...,3*j:3*(j+1)] = phdispl[...,3*j:3*(j+1)] * np.sqrt(amu[a.specie.number]*abu.amu_emass) / abu.Bohr_Ang
+        eigvec[...,3*j:3*(j+1)] = phdispl[...,3*j:3*(j+1)] * np.sqrt(amu_symbol[a.specie.symbol]*abu.amu_emass) / abu.Bohr_Ang
 
     return eigvec