basisset.py 6.73 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 nomad.units import ureg
6

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

11
12

def get_basis_set(context, backend: Backend, logger) -> RestrictedDict:
13
14
15
16
17
18
19
20
21
22
23
24
    """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.
25
    """
26
    settings: BasisSet = None
27
    program_name = backend.entry_archive.section_run[0].program_name
28
    if program_name == "exciting":
29
        settings = BasisSetExciting(context, backend, logger)
30
    elif program_name == "FHI-aims":
31
        settings = BasisSetFHIAims(context, backend, logger)
32
33
    else:
        return None
34

35
    return settings.to_dict()
36
37


38
class BasisSet(ABC):
39
40
41
    """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.
42
43
44
45
    """
    def __init__(self, context, backend, logger):
        """
        """
46
47
48
49
        self._ctx = context
        self._backend = backend
        self._logger = logger
        mandatory, optional = self.setup()
50
        self.settings = RestrictedDict(mandatory, optional, forbidden_values=[None])
51

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

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

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


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

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

83
        return mandatory, optional
84

85
    def to_dict(self):
86
        # Get basis set settings for each species
87
        aims_bs = self._ctx.representative_method.x_fhi_aims_section_controlIn_basis_set
88
89
90
91
92
93
94
95
96
97
98
99
        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]
100
101
102
                self.settings["fhiaims_basis"] = basis

        return self.settings
103
104
105
106
107
108

    @classmethod
    def _values_to_dict(cls, data, level=0):
        result = None
        if data is None:
            return None
109
        elif isinstance(data, (Section, dict)):
110
111
112
113
114
115
116
117
118
            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))
119
        elif isinstance(data, (np.ndarray)):
120
121
122
123
124
            result = data.tolist()
        else:
            result = data
        return result

125
126
127
128
129
130
131
132
133
134
135
136
    @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

137

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

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

154
155
        return mandatory, optional

156
    def to_dict(self):
157
158
159
        """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
        """
160
        # Add the muffin-tin settings for each species ordered alphabetically by atom label
161
        try:
162
163
            groups = self._ctx.representative_system.x_exciting_section_atoms_group
            groups = sorted(groups, key=lambda group: group.x_exciting_geometry_atom_labels)
164
            muffin_tin_settings = OrderedDict()
165
            for group in groups:
166
                label = group.x_exciting_geometry_atom_labels
167
                try:
168
                    muffin_tin_settings["{}_muffin_tin_radius".format(label)] = "%.6f" % (group.x_exciting_muffin_tin_radius.to(ureg.angstrom).magnitude)
169
                except Exception:
170
                    muffin_tin_settings["{}_muffin_tin_radius".format(label)] = None
171
                try:
172
                    muffin_tin_settings["{}_muffin_tin_points".format(label)] = "%d" % group.x_exciting_muffin_tin_points
173
                except Exception:
174
175
                    muffin_tin_settings["{}_muffin_tin_points".format(label)] = None
            self.settings["muffin_tin_settings"] = muffin_tin_settings
176
        except Exception:
177
            pass
178

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

        return self.settings