Hubbard: add useful utility functions

docs/source/conf.py

@@ -50,6 +50,7 @@
     'python': ('https://docs.python.org/3.8', None),
     'aiida': ('https://aiida.readthedocs.io/projects/aiida-core/en/latest/', None),
     'aiida_pseudo': ('http://aiida-pseudo.readthedocs.io/en/latest/', None),
+    'pymatgen': ('https://pymatgen.org/', None),
 }
 
 # Settings for the `autoapi.extenstion` automatically generating API docs

src/aiida_quantumespresso/data/hubbard_structure.py

@@ -120,9 +120,11 @@ def append_hubbard_parameter(
         parameters = HubbardParameters.from_tuple(hp_tuple)
         hubbard = self.hubbard
 
-        if parameters not in hubbard.parameters:
-            hubbard.parameters.append(parameters)
-            self.hubbard = hubbard
+        if parameters in hubbard.parameters:
+            hubbard.parameters.remove(parameters)
+
+        hubbard.parameters.append(parameters)
+        self.hubbard = hubbard
 
     def pop_hubbard_parameters(self, index: int):
         """Pop Hubbard parameters in the list.

src/aiida_quantumespresso/utils/hubbard.py

@@ -3,9 +3,10 @@
 # pylint: disable=no-name-in-module
 from itertools import product
 import os
-from typing import List, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 from aiida.orm import StructureData
+import numpy as np
 
 from aiida_quantumespresso.common.hubbard import Hubbard
 from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData
@@ -251,8 +252,6 @@ def get_hubbard_for_supercell(self, supercell: StructureData, thr: float = 1e-3)
 
         :returns: a new ``HubbbardStructureData`` with all the mapped Hubbard parameters
         """
-        import numpy as np
-
         uc_pymat = self.hubbard_structure.get_pymatgen_structure()
         sc_pymat = supercell.get_pymatgen_structure()
         uc_positions = uc_pymat.cart_coords  # positions in Cartesian coordinates
@@ -313,6 +312,243 @@ def get_hubbard_for_supercell(self, supercell: StructureData, thr: float = 1e-3)
 
         return HubbardStructureData.from_structure(structure=supercell, hubbard=new_hubbard)
 
+    def get_interacting_pairs(self) -> Dict[str, List[str]]:
+        """Return tuple of kind name interaction pairs.
+
+        :returns: dictionary of onsite kinds with a list of V kinds
+        """
+        pairs_dict = dict()
+        sites = self.hubbard_structure.sites
+        parameters = self.hubbard_structure.hubbard.parameters
+
+        onsite_indices = [p.atom_index for p in parameters if p.atom_index == p.neighbour_index]
+
+        for index in onsite_indices:
+            name = sites[index].kind_name
+            if name not in pairs_dict:
+                pairs_dict[name] = []
+
+        for parameter in parameters:
+            onsite_name = sites[parameter.atom_index].kind_name
+            neigh_name = sites[parameter.neighbour_index].kind_name
+
+            if onsite_name in pairs_dict and onsite_name != neigh_name:
+                if not neigh_name in pairs_dict[onsite_name]:
+                    pairs_dict[onsite_name].append(neigh_name)
+
+        return pairs_dict
+
+    def get_pairs_radius(
+        self,
+        onsite_index: int,
+        neighbours_names: List[str],
+        number_of_neighbours: int,
+        radius_max: float = 7.0,
+        thr: float = 1.0e-2,
+    ) -> Tuple[float, float]:
+        """Return the minimum and maximum radius of the first neighbours of the onsite site.
+
+        :param onsite_index: index in the structure of the onsite Hubbard atom
+        :param neighbours_names: kind names of the neighbours
+        :param number_of_neighbours: number of neighbours coming to select
+        :param radius_max: maximum radius (in Angstrom) to use for looking for neighbours
+        :param thr: threshold (in Angstrom) for defining the shells
+        :return: (radius min +thr, radius max -thr) defining the shells containing only the first neighbours
+        """
+        rmin = 0
+        pymat = self.hubbard_structure.get_pymatgen_structure()
+        _, neigh_indices, _, distances = pymat.get_neighbor_list(sites=[pymat[onsite_index]], r=radius_max)
+
+        sort = np.argsort(distances)
+        neigh_indices = neigh_indices[sort]
+        distances = distances[sort]
+
+        count = 0
+        for i in range(len(neigh_indices)):  # pylint: disable=consider-using-enumerate
+            index = i
+            if self.hubbard_structure.sites[neigh_indices[i]].kind_name in neighbours_names:
+                rmin = max(rmin, distances[i])
+                count += 1
+            if count == number_of_neighbours:
+                break
+
+        return rmin + thr, distances[index + 1] - thr
+
+    def get_intersites_radius(
+        self,
+        nn_finder: str = 'crystal',
+        nn_inputs: Union[Dict, None] = None,
+        radius_max: float = 7.0,
+        thr: float = 1.0e-2,
+        **_,
+    ) -> float:
+        """Return the radius (in Angstrom) for intersites from nearest neighbour finders.
+
+        It peforms a nearest neighbour analysis (via pymatgen modules) to find the first inersite
+        neighbours for all the onsite atoms. A radius is returned which can be used to
+        run an ``hp.x`` calculation. Such radius defines a shell including only the first
+        neighbours of each onsite Hubbard atom.
+
+        :param nn_finder: string defining the nearest neighbour finder; options are:
+            * `crystal`: use :class:`pymatgen.analysis.local_env.CrystalNN`
+            * `voronoi`: use :class:`pymatgen.analysis.local_env.VoronoiNN`
+        :param nn_inputs: inputs for the nearest neighbours finder; when None, standard inputs
+            are used to find geometric first neighbours (recommended)
+        :param radius_max: max radius where to look for neighbouring atoms, in Angstrom
+        :param thr: threshold (in Angstrom) for defining the shells
+        :return: radius defining the shell containing only the first neighbours
+        """
+        import warnings
+
+        from pymatgen.analysis.local_env import CrystalNN, VoronoiNN
+
+        rmin, rmax = 0.0, radius_max
+
+        if nn_finder not in ['crystal', 'voronoi']:
+            raise ValueError('`nn_finder` must be either `cyrstal` or `voronoi`')
+        if nn_inputs is None:
+            if nn_finder == 'crystal':
+                nn_inputs = {'distance_cutoffs': None, 'x_diff_weight': 0, 'porous_adjustment': False}
+            if nn_finder == 'voronoi':
+                nn_inputs = {'tol': 0.1, 'cutoff': radius_max}
+
+        voronoi = CrystalNN(**nn_inputs) if nn_finder == 'crystal' else VoronoiNN(**nn_inputs)  # pylint: disable=unexpected-keyword-arg
+
+        sites = self.hubbard_structure.sites
+        name_to_specie = {kind.name: kind.symbol for kind in self.hubbard_structure.kinds}
+        pymat = self.hubbard_structure.get_pymatgen_structure()
+        pairs = self.get_interacting_pairs()
+
+        for i, site in enumerate(sites):
+            if site.kind_name in pairs:
+                neigh_species = voronoi.get_cn_dict(pymat, i)  # e.g. {'O': 4, 'S': 2, ...}
+                number_of_neighs = 0
+
+                for neigh_name in pairs[site.kind_name]:
+                    specie = name_to_specie[neigh_name]
+
+                    if specie in neigh_species:
+                        number_of_neighs += neigh_species[specie]
+                        neigh_species.pop(specie)  # avoid 'duplicating' same specie but different (kind) name
+
+                rmin_, rmax_ = self.get_pairs_radius(i, pairs[site.kind_name], number_of_neighs, radius_max, thr)
+
+                rmin = max(rmin_, rmin)  # we want the largest to include them all
+                rmax = min(rmax_, rmax)  # we want the smallest to check whether such radius exist
+
+        if rmin > rmax:
+            warnings.warn('A common radius seems to not exist! Try lowering `thr`.')
+
+        return min(rmin, rmax)
+
+
+def initialize_hubbard_parameters(
+    structure: StructureData,
+    pairs: Dict[str, Tuple[str, float, float, Dict[str, str]]],
+    nn_finder: str = 'crystal',
+    nn_inputs: Union[Dict, None] = None,
+    fold: bool = True,
+    standardize: bool = False,
+    radius_max: float = 7.0,
+    thr: float = 1e-5,
+    **_,
+) -> HubbardStructureData:
+    """Initialize the on-site and intersite parameters using nearest neighbour finders.
+
+    It peforms a nearest neighbour analysis (via pymatgen modules) to find the first inersite
+    neighbours for all the onsite atoms. Only the atoms in the `pair`, and that are nearest
+    neighbours from the analysis, are initialized.
+
+    :param structure: a StructureData instance
+    :param pairs: dictionary of the kind
+        {onsite name: [onsite manifold, onsite value, intersites value, {neighbour name: neighbour manifold}], ...}
+        For example: {'Fe': ['3d', 5.0, 1.0, {'O':'2p', 'O1':'1s', 'Se':'4p'}]}
+    :param nn_finder: string defining the nearest neighbour finder; options are:
+        * `crystal`: use :class:`pymatgen.analysis.local_env.CrystalNN`
+        * `voronoi`: use :class:`pymatgen.analysis.local_env.VoronoiNN`
+    :param nn_inputs: inputs for the nearest neighbours finder; when None, standard inputs
+        are used to find geometric first neighbours (recommended)
+    :param fold: whether to fold in within the cell the atoms
+    :param standardize: whether to standardize the atoms and the cell via spglib (symmetry analysis)
+    :param radius_max: max radius where to look for neighbouring atoms, in Angstrom
+    :param thr: threshold to refold the atoms with crystal coordinates close to 1.0
+    :return: HubbardStructureData with initialized Hubbard parameters
+    """
+    from aiida.tools.data import spglib_tuple_to_structure, structure_to_spglib_tuple
+    from pymatgen.analysis.local_env import CrystalNN, VoronoiNN
+    from spglib import standardize_cell
+
+    if nn_finder not in ['crystal', 'voronoi']:
+        raise ValueError('`nn_finder` must be either `cyrstal` or `voronoi`')
+
+    if nn_inputs is None:
+        if nn_finder == 'crystal':
+            nn_inputs = {'distance_cutoffs': None, 'x_diff_weight': 0, 'porous_adjustment': False}
+        if nn_finder == 'voronoi':
+            nn_inputs = {'tol': 0.1, 'cutoff': radius_max}
+
+    voronoi = CrystalNN(**nn_inputs) if nn_finder == 'crystal' else VoronoiNN(**nn_inputs)  # pylint: disable=unexpected-keyword-arg
+
+    if not standardize and not fold:
+        hubbard_structure = HubbardStructureData.from_structure(structure=structure)
+
+    if fold:
+        cell, kind_info, kinds = structure_to_spglib_tuple(structure)
+        cell = list(cell)
+        cell[1] = cell[1] % 1.0
+        cell[1] = np.where(np.abs(cell[1] - 1.0) < thr, 0.0, cell[1])
+
+        folded_struccture = spglib_tuple_to_structure(cell, kind_info, kinds)
+        hubbard_structure = HubbardStructureData.from_structure(structure=folded_struccture)
+
+    if standardize:
+        cell, kind_info, kinds = structure_to_spglib_tuple(structure)
+        new_cell = standardize_cell(cell)
+        new_structure = spglib_tuple_to_structure(new_cell, kind_info, kinds)
+        hubbard_structure = HubbardStructureData.from_structure(structure=new_structure)
+
+    sites = hubbard_structure.sites
+    name_to_specie = {kind.name: kind.symbol for kind in hubbard_structure.kinds}
+    pymat = hubbard_structure.get_pymatgen_structure()
+
+    for i, site in enumerate(sites):  # pylint: disable=too-many-nested-blocks
+        if site.kind_name in pairs:
+            neigh_species = voronoi.get_cn_dict(pymat, i)  # e.g. {'O': 4, 'S': 2, ...}
+            onsite = pairs[site.kind_name]
+            hubbard_structure.append_hubbard_parameter(i, onsite[0], i, onsite[0], onsite[1])
+
+            for neigh_name in onsite[3]:
+                try:  # we want an API that allows for specifying neighbour that are not present
+                    specie = name_to_specie[neigh_name]
+
+                    count = 0
+                    if specie in neigh_species:
+                        _, neigh_indices, images, distances = pymat.get_neighbor_list(sites=[pymat[i]], r=radius_max)
+
+                        sort = np.argsort(distances)
+                        neigh_indices = neigh_indices[sort]
+                        images = images[sort]
+
+                        for index, image in zip(neigh_indices, images):
+                            if pymat[index].specie.name == specie:
+                                hubbard_structure.append_hubbard_parameter(
+                                    atom_index=i,
+                                    atom_manifold=onsite[0],
+                                    neighbour_index=int(index), # otherwise the validator complains
+                                    neighbour_manifold=onsite[3][neigh_name],
+                                    value=onsite[2],
+                                    translation=np.array(image, dtype=np.int64).tolist(),
+                                    hubbard_type='V',
+                                )
+                                count += 1
+
+                            if count >= neigh_species[specie]:
+                                break
+                except KeyError:
+                    pass
+
+    return hubbard_structure
+
 
 def get_supercell_atomic_index(index: int, num_sites: int, translation: List[Tuple[int, int, int]]) -> int:
     """Return the atomic index in 3x3x3 supercell.