Phonon angular mometum + PhononWork:prepgkk option (#296)

* add plot of phonon angulor momentum

* add qptopt option in Phononwork for the generation of the q-mesh

* add support for GSTORE dependency

* add prepgkk option for phonon work

* minor change

* documetation

---------

Co-authored-by: Max Mignolet <max.mignolet@student.uliege.be>
Co-authored-by: Mignolet Maxime <mmaxime@login1.discoverer.sofiatech.bg>
Co-authored-by: Max Mignolet <M.Mignolet@uliege.be>

abipy/core/abinit_units.py

@@ -119,6 +119,7 @@ def wlabel_from_units(units: str, unicode=False) -> str:
          'cm-1': r'Frequency (cm$^{-1}$)',
          'cm^-1': r'Frequency (cm$^{-1}$)',
          'thz': r'Frequency (Thz)',
+         'hbar': r'Angular momentum ($\hbar$)',
     }
     try:
         s = d[units.lower().strip()]

abipy/dfpt/phonons.py

@@ -159,6 +159,7 @@ def from_file(cls, filepath: str) -> PhononBands:
                        qpoints=qpoints,
                        phfreqs=r.read_phfreqs(),
                        phdispl_cart=r.read_phdispl_cart(),
+                       phangmom=r.read_phangmom(),
                        amu=amu,
                        non_anal_ph=non_anal_ph,
                        epsinf=epsinf, zcart=zcart,
@@ -237,15 +238,17 @@ def set_phonopy_obj_from_ananc(self, ananc, supercell_matrix, symmetrize_tensors
                                    symprec=symprec, set_masses=set_masses)
         self.phonopy_obj = ph
 
-    def __init__(self, structure, qpoints, phfreqs, phdispl_cart, non_anal_ph=None, amu=None,
-                 epsinf=None, zcart=None, linewidths=None, phonopy_obj=None):
+    def __init__(self, structure, qpoints, phfreqs, phdispl_cart, phangmom=None,
+                 non_anal_ph=None, amu=None, epsinf=None, zcart=None, linewidths=None,
+                 phonopy_obj=None):
         """
         Args:
             structure: |Structure| object.
             qpoints: |KpointList| instance.
             phfreqs: Phonon frequencies in eV.
             phdispl_cart: [nqpt, 3*natom, 3*natom] array with displacement in Cartesian coordinates in Angstrom.
                 The last dimension stores the cartesian components.
+            phangmom: Phonon angular momentum in units of hbar.
             non_anal_ph: :class:`NonAnalyticalPh` with information of the non analytical contribution
                 None if contribution is not present.
             amu: dictionary that associates the atomic species present in the structure to the values of the atomic
@@ -269,6 +272,10 @@ def __init__(self, structure, qpoints, phfreqs, phdispl_cart, non_anal_ph=None,
         # The last dimension stores the cartesian components.
         self.phdispl_cart = phdispl_cart
 
+        # phonon angular momentum in units of hbar
+        # ndarray of shape (nqpt, 3*natom, 3)
+        self.phangmom = phangmom
+
         # Handy variables used to loop.
         self.num_atoms = structure.num_sites
         self.num_branches = 3 * self.num_atoms
@@ -1319,6 +1326,66 @@ def plot_colored_matched(self, ax=None, units="eV", qlabels=None, branch_range=N
 
         return fig
 
+    @add_fig_kwargs
+    def plot_phangmom(self, ax=None, pj_dir=[0, 0, 1], units="hbar",
+                      qlabels=None, branch_range=None, colormap="rainbow",
+                      max_colors=None, **kwargs):
+        r"""
+        Plot the phonon angular momentum with different colors for each line.
+
+        Args:
+            ax: |matplotlib-Axes| or None if a new figure should be created.
+            pj_dir: direction along which the angular momentum is projected
+                in cartesian coordinates
+            units: Units for phonon plots. Only possible option: "hbar", maybe others to come ?
+                Case-insensitive.
+            qlabels: dictionary whose keys are tuples with the reduced coordinates of the q-points.
+                The values are the labels. e.g. ``qlabels = {(0.0,0.0,0.0): "$\Gamma$", (0.5,0,0): "L"}``.
+            branch_range: Tuple specifying the minimum and maximum branch_i index to plot
+                (default: all branches are plotted).
+            colormap: matplotlib colormap to determine the colors available. The colors will be chosen not in a
+                sequential order to avoid difficulties in distinguishing the lines.
+                http://matplotlib.sourceforge.net/examples/pylab_examples/show_colormaps.html
+            max_colors: maximum number of colors to be used. If max_colors < num_braches the colors will be reapeated.
+                It may be useful to better distinguish close bands when the number of branches is large.
+
+        Returns: |matplotlib-Figure|
+        """
+        # Select the band range.
+        if branch_range is None:
+            branch_range = range(self.num_branches)
+        else:
+            branch_range = range(branch_range[0], branch_range[1], 1)
+
+        ax, fig, plt = get_ax_fig_plt(ax=ax)
+
+        # Decorate the axis (e.g add ticks and labels).
+        self.decorate_ax(ax, units=units, qlabels=qlabels)
+
+        first_xx = 0
+        lines = []
+
+        if max_colors is None:
+            max_colors = len(branch_range)
+
+        colormap = plt.get_cmap(colormap)
+
+        sorted_phangmom = np.empty_like(self.phangmom)
+        for i, pam in enumerate(self.phangmom):
+            ind = self.split_matched_indices[0][i]
+            sorted_phangmom[i, :, :] = pam[ind, :]
+
+        pj_dir = np.array(pj_dir)
+        pj_dir = pj_dir / np.linalg.norm(pj_dir)
+        proj_phangmom = np.dot(sorted_phangmom, pj_dir)
+        colors = itertools.cycle(colormap(np.linspace(0, 1, max_colors)))
+        for branch_i in branch_range:
+            kwargs = dict(kwargs)
+            kwargs['color'] = next(colors)
+            lines.extend(ax.plot(proj_phangmom[:, branch_i], **kwargs))
+
+        return fig
+
     @add_fig_kwargs
     def plot_lt_character(self, units="eV", qlabels=None, ax=None, xlims=None, ylims=None, scale_size=50,
                           use_becs=True, colormap="jet", fontsize=12, **kwargs) -> Figure:
@@ -2909,6 +2976,12 @@ def read_phdispl_cart(self):
         """
         return self.read_value("phdispl_cart", cmode="c")
 
+    def read_phangmom(self):
+        """
+        Real array with the phonon angular momentum in units of hbar
+        """
+        return self.read_value("phangmom", default=None)
+
     def read_amu(self):
         """The atomic mass units"""
         return self.read_value("atomic_mass_units", default=None)

abipy/flowtk/utils.py

@@ -515,6 +515,7 @@ def find_1den_files(self):
     "DKDK": {},  # irddkdk is not defined.
     #"DKDE": {"getdkde": 1},
     #"DELFD": {"getdelfd": 1},
+    "GSTORE": {"getgstore_filepath": '"indata/in_GSTORE.nc"'},
 }
 
 

abipy/flowtk/works.py

@@ -1704,7 +1704,7 @@ def from_scf_task(cls, scf_task: ScfTask,
                       qpoints, is_ngqpt=False, with_becs=False,
                       with_quad=False, with_flexoe=False, with_dvdb=True,
                       tolerance=None, ddk_tolerance=None, ndivsm=0, qptopt=1,
-                      prtwf=-1, manager=None) -> PhononWork:
+                      prtwf=-1, prepgkk=0, manager=None) -> PhononWork:
         """
         Construct a `PhononWork` from a |ScfTask| object.
         The input file for phonons is automatically generated from the input of the ScfTask.
@@ -1738,6 +1738,8 @@ def from_scf_task(cls, scf_task: ScfTask,
                 to restart the DFPT task if the calculation is not converged (worst case scenario)
                 but we avoid the output of the 1-st order WFK if the calculation converged successfully.
                 Non-linear DFPT tasks should not be affected since they assume q == 0.
+            prepgkk: option to compute only the irreducible preturbations or all perturbations
+                for the chosen set of q-point. Default: 0 (irred. pret. only)
             manager: |TaskManager| object.
         """
         if not isinstance(scf_task, ScfTask):
@@ -1761,7 +1763,8 @@ def from_scf_task(cls, scf_task: ScfTask,
         for qpt in qpoints:
             is_gamma = np.sum(qpt ** 2) < 1e-12
             if with_becs and is_gamma: continue
-            multi = scf_task.input.make_ph_inputs_qpoint(qpt, tolerance=tolerance)
+            multi = scf_task.input.make_ph_inputs_qpoint(qpt, tolerance=tolerance,
+                                                         prepgkk=prepgkk)
             for ph_inp in multi:
                 # Here we set the value of prtwf for the DFPT tasks if q != Gamma.
                 if not is_gamma: ph_inp.set_vars(prtwf=prtwf)
@@ -1777,7 +1780,7 @@ def from_scf_task(cls, scf_task: ScfTask,
     def from_scf_input(cls, scf_input: AbinitInput, qpoints, is_ngqpt=False, with_becs=False,
                        with_quad=False, with_flexoe=False, with_dvdb=True, tolerance=None,
                        ddk_tolerance=None, ndivsm=0, qptopt=1,
-                       prtwf=-1, manager=None) -> PhononWork:
+                       prtwf=-1, prepgkk=0, manager=None) -> PhononWork:
         """
         Similar to `from_scf_task`, the difference is that this method requires
         an input for SCF calculation. A new |ScfTask| is created and added to the Work.
@@ -1806,7 +1809,8 @@ def from_scf_input(cls, scf_input: AbinitInput, qpoints, is_ngqpt=False, with_be
         for qpt in qpoints:
             is_gamma = np.sum(qpt ** 2) < 1e-12
             if with_becs and is_gamma: continue
-            multi = scf_task.input.make_ph_inputs_qpoint(qpt, tolerance=tolerance)
+            multi = scf_task.input.make_ph_inputs_qpoint(qpt, tolerance=tolerance,
+                                                         prepgkk=prepgkk)
             for ph_inp in multi:
                 # Here we set the value of prtwf for the DFPT tasks if q != Gamma.
                 if not is_gamma: ph_inp.set_vars(prtwf=prtwf)