metainfo.py 6.23 KB
Newer Older
1
2
3
4
"""
Some playground to try the CONCEPT.md ideas.
"""

5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from typing import Dict, List, Any, Union, Type
import json
import ase.data


class Units():

    def __getattribute__(self, name):
        return name


units = Units()


class Definition():
20
    m_definition: Any = None
21
22
23
24
25
26
27
28
29
    pass


class Property(Definition):
    pass


class Quantity(Property):
    def __init__(
30
31
32
            self,
            name: str = None,
            description: str = None,
33
            parent_section: 'Section' = None,
34
35
            shape: List[Union[str, int]] = [],
            type: Union['Enum', type] = None,
36
37
38
39
            unit: str = None,
            derived: bool = False,
            repeats: bool = False,
            synonym: str = None):
40

41
42
43
44
45
46
47
48
49
50
51
52
        self.name = name
        self.parent_section = parent_section.m_definition if parent_section is not None else None
        self.derived = derived
        self.synonym = synonym

    def __get__(self, obj: 'MetainfoObject', type: Type):
        if obj is None:
            # access on cls not obj
            return self

        if self.derived:
            derive_method = getattr(obj, 'm_derive_%s' % self.name, None)
53

54
55
            if derive_method is None:
                raise KeyError('Derived quantity %s is not implemented' % self.name)
56

57
58
59
60
61
62
63
64
65
            else:
                return derive_method()

        elif self.synonym is not None:
            return getattr(obj, self.synonym)

        else:
            return obj.m_data.get(self.name, None)

66
    def __set__(self, obj: 'MetainfoObject', value: Any):
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
        obj.m_data[self.name] = value

    def __delete__(self, obj: 'MetainfoObject'):
        del obj.m_data[self.name]

    def __repr__(self):
        base = self.name
        if self.parent_section is not None:
            return '%s.%s' % (str(self.parent_section), base)
        else:
            return base


class Section(Definition):
    def __init__(
            self,
            name: str = None,
84
            parent_section=None,
85
            repeats: bool = False,
86
87
            extends=None,
            adds_to=None):
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

        self.name = name
        self.parent_section = parent_section.m_definition if parent_section is not None else None
        self.repeats = repeats

    def __get__(self, obj: 'MetainfoObject', type: Type):
        return self

    def __repr__(self):
        base = self.name
        if self.parent_section is not None:
            return '%s.%s' % (str(self.parent_section), base)
        else:
            return base


class Reference(Property):
    pass


class Enum:
    def __init__(self, values: List[Any]):
        self.values = values


class MetainfoObjectMeta(type):
    def __new__(cls, cls_name, bases, dct):
        cls.m_definition = dct.get('m_definition', None)
        for name, value in dct.items():
            if isinstance(value, Property):
                value.name = name
                value.parent_section = cls.m_definition
120

121
        cls = super().__new__(cls, cls_name, bases, dct)
122

123
124
125
126
127
128
129
        if cls.m_definition is not None:
            if cls.m_definition.name is None:
                cls.m_definition.name = cls_name

        return cls


130
class MetainfoObject(metaclass=MetainfoObjectMeta):
131
    """
132
    Base class for all
133
    """
134
    m_definition: Any = None
135
136
137
138

    def __init__(self):
        self.m_data = dict(m_defintion=self.m_definition.name)

139
    def m_create(self, section_definition: Any, *args, **kwargs) -> Any:
140
141
142
143
144
145
146
147
148
149
150
        """
        Creates a sub section of the given section definition.
        """
        section_cls = section_definition
        definition = section_definition.m_definition
        sub_section = section_cls(*args, **kwargs)
        if definition.repeats:
            self.m_data.setdefault(definition.name, []).append(sub_section)
        else:
            # TODO test overwrite
            self.m_data[definition.name] = sub_section
151

152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
        return sub_section

    def m_get_definition(self, name):
        """
        Returns the definition of the given quantity name.
        """
        descriptor = getattr(type(self), name)
        if descriptor is None:
            raise KeyError
        elif not isinstance(descriptor, Property):
            raise KeyError

        return descriptor

    def m_to_dict(self) -> Dict[str, Any]:
        """
        Returns a JSON serializable dictionary version of this section (and all subsections).
        """
        return {
            key: value.m_to_dict() if isinstance(value, MetainfoObject) else value
            for key, value in self.m_data.items()
        }

    def m_to_json(self) -> str:
        return json.dumps(self.m_to_dict(), indent=2)

    def m_validate(self) -> bool:
        """
        Validates this sections content based on section and quantity definitions.
        Can be overwritten to customize the validation.
        """
        return True

    def __repr__(self) -> str:
        return self.m_to_json()


class Run(MetainfoObject):
    m_definition = Section()


class System(MetainfoObject):
    """
    The system is ...
    """
    m_definition = Section(parent_section=Run, repeats=True)

    n_atoms = Quantity(type=int, derived=True)

    atom_labels = Quantity(shape=['n_atoms'], type=Enum(ase.data.chemical_symbols))
    """
    Atom labels are ...
    """

    atom_species = Quantity(shape=['n_atoms'], type=int, derived=True)

208
    atom_positions = Quantity(shape=['n_atoms', 3], type=float, unit=units.m)
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244

    cell = Quantity(shape=[3, 3], type=float, unit=units.m)
    lattice_vectors = Quantity(synonym='cell')

    pbc = Quantity(shape=[3], type=bool)

    def m_derive_atom_species(self) -> List[int]:
        return [ase.data.atomic_numbers[label] for label in self.atom_labels]

    def m_derive_n_atoms(self) -> int:
        return len(self.atom_labels)


class VaspSystem(MetainfoObject):
    m_definition = Section(adds_to=System)

    vasp_specific = Quantity(type=str)


run = Run()
print(run.m_definition)

system = run.m_create(System)
system.atom_labels = ['H', 'H', 'O']
system.atom_positions = [[0, 0, 0], [1, 0, 0], [0.5, 0.5, 0]]
system.cell = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
system.pbc = [False, False, False]

print(system.atom_species)
print(system.lattice_vectors)
print(system.n_atoms)

print(system.__class__.m_definition)
print(system.m_definition)
print(system.m_get_definition('atom_labels'))

245
print(run)