basisset.py 6.75 KB
Newer Older
1
from abc import ABC, abstractmethod
2
3
4
from collections import OrderedDict
import numpy as np
from typing import Tuple, List
5
from pint import UnitRegistry
6

7
8
from nomad.parsing.legacy import Backend
from nomad.metainfo import Section
9
from nomad.utils import RestrictedDict
10

11
ureg = UnitRegistry()
12

13
14

def get_basis_set(context, backend: Backend, logger) -> RestrictedDict:
15
16
17
18
19
20
21
22
23
24
25
26
    """Decide which type of basis set settings are applicable to the entry and
    return a corresponding settings as a RestrictedDict.

    Args:
        context: The calculation context.
        backend: Backend from which values are extracted.
        logger: Shared logger.

    Returns:
        RestrictedDict or None: Returns the extracted settings as a
        RestrictedDict. If no suitable basis set settings could be identified,
        returns None.
27
    """
28
    settings: BasisSet = None
29
    program_name = backend.entry_archive.section_run[0].program_name
30
    if program_name == "exciting":
31
        settings = BasisSetExciting(context, backend, logger)
32
    elif program_name == "FHI-aims":
33
        settings = BasisSetFHIAims(context, backend, logger)
34
35
    else:
        return None
36

37
    return settings.to_dict()
38
39


40
class BasisSet(ABC):
41
42
43
    """Abstract base class for basis set settings. The idea is to create
    subclasses that inherit this class and hierarchically add new mandatory and
    optional settings with the setup()-function.
44
45
46
47
    """
    def __init__(self, context, backend, logger):
        """
        """
48
49
50
51
        self._ctx = context
        self._backend = backend
        self._logger = logger
        mandatory, optional = self.setup()
52
        self.settings = RestrictedDict(mandatory, optional, forbidden_values=[None])
53

54
    @abstractmethod
55
56
57
    def to_dict(self) -> RestrictedDict:
        """Used to extract basis set settings from the backend and returning
        them as a RestrictedDict.
58
        """
59
        pass
60
61

    @abstractmethod
62
63
    def setup(self) -> Tuple:
        """Used to define a list of mandatory and optional settings for a
64
        subclass.
65

66
67
68
        Returns:
            Should return a tuple of two lists: the first one defining
            mandatory keys and the second one defining optional keys.
69
        """
70
71
72
        mandatory: List = []
        optional: List = []
        return mandatory, optional
73
74


75
class BasisSetFHIAims(BasisSet):
76
    """Basis set settings for 'FHI-Aims' (code-dependent).
77
    """
78
    def setup(self) -> Tuple:
79
        # Get previously defined values from superclass
80
        mandatory, optional = super().setup()
81
82
83
84

        # Add new values
        mandatory += ["fhiaims_basis"]

85
        return mandatory, optional
86

87
    def to_dict(self):
88
        # Get basis set settings for each species
89
        aims_bs = self._ctx.representative_method.x_fhi_aims_section_controlIn_basis_set
90
91
92
93
94
95
96
97
98
99
100
101
        if aims_bs is not None:
            bs_by_species = {}
            for this_aims_bs in aims_bs:
                this_bs_dict = self._values_to_dict(this_aims_bs, level=2)
                this_species = this_aims_bs['x_fhi_aims_controlIn_species_name'][0]
                bs_by_species[this_species] = this_bs_dict

            # Sort alphabetically by species label
            if bs_by_species:
                basis = OrderedDict()
                for k in sorted(bs_by_species.keys()):
                    basis[k] = bs_by_species[k]
102
103
104
                self.settings["fhiaims_basis"] = basis

        return self.settings
105
106
107
108
109
110

    @classmethod
    def _values_to_dict(cls, data, level=0):
        result = None
        if data is None:
            return None
111
        elif isinstance(data, (Section, dict)):
112
113
114
115
116
117
118
119
120
            result = OrderedDict()
            for k in sorted(cls._filtered_section_keys(data)):
                v = data.get(k, None)
                result[k] = cls._values_to_dict(v, level=level + 1)
        elif isinstance(data, (list)):
            result = []
            for k in range(len(data)):
                v = data[k]
                result.append(cls._values_to_dict(v, level=level + 1))
121
        elif isinstance(data, (np.ndarray)):
122
123
124
125
126
            result = data.tolist()
        else:
            result = data
        return result

127
128
129
130
131
132
133
134
135
136
137
138
    @classmethod
    def _filtered_section_keys(cls, section):
        for k in section.keys():
            # skip JSON-specific keys
            if k == '_gIndex':
                continue
            if k == '_name':
                continue
            else:
                # json values and subsections
                yield k

139

140
class BasisSetExciting(BasisSet):
141
    """Basis set settings for 'Exciting' (code-dependent).
142
    """
143
    def setup(self) -> Tuple:
144
        # Get previously defined values from superclass
145
        mandatory, optional = super().setup()
146
147

        # Add new values
148
149
150
151
152
153
154
        mandatory += [
            "muffin_tin_settings",
            "rgkmax",
            "gkmax",
            "lo",
            "lmaxapw",
        ]
155

156
157
        return mandatory, optional

158
    def to_dict(self):
159
160
161
        """Special case of basis set settings for Exciting code. See list at:
        https://gitlab.mpcdf.mpg.de/nomad-lab/encyclopedia-general/wikis/FHI-visit-preparation
        """
162
        # Add the muffin-tin settings for each species ordered alphabetically by atom label
163
        try:
164
165
            groups = self._ctx.representative_system.x_exciting_section_atoms_group
            groups = sorted(groups, key=lambda group: group.x_exciting_geometry_atom_labels)
166
            muffin_tin_settings = OrderedDict()
167
            for group in groups:
168
                label = group.x_exciting_geometry_atom_labels
169
                try:
170
                    muffin_tin_settings["{}_muffin_tin_radius".format(label)] = "%.6f" % (group.x_exciting_muffin_tin_radius.to(ureg.angstrom).magnitude)
171
                except KeyError:
172
                    muffin_tin_settings["{}_muffin_tin_radius".format(label)] = None
173
                try:
174
                    muffin_tin_settings["{}_muffin_tin_points".format(label)] = "%d" % group.x_exciting_muffin_tin_points
175
                except KeyError:
176
177
                    muffin_tin_settings["{}_muffin_tin_points".format(label)] = None
            self.settings["muffin_tin_settings"] = muffin_tin_settings
178
179
        except KeyError:
            pass
180

181
        # Other important method settings
182
        system = self._ctx.representative_system
183
        try:
184
            self.settings['rgkmax'] = "%.6f" % (system.x_exciting_rgkmax.magnitude)
185
186
        except KeyError:
            pass
187
        try:
188
            self.settings['gkmax'] = "%.6f" % (1e-10 * system.x_exciting_gkmax.magnitude)
189
190
        except KeyError:
            pass
191
        try:
192
            self.settings['lo'] = "%d" % (system.x_exciting_lo)
193
194
        except KeyError:
            pass
195
        try:
196
            self.settings['lmaxapw'] = "%d" % (system.x_exciting_lmaxapw)
197
198
        except KeyError:
            pass
199
200

        return self.settings