/
atoms.py
320 lines (264 loc) · 12 KB
/
atoms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
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
208
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
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
"""
This file includes the definition of an Atoms class which can be used with
System
TODO
----
- Iterators for masked/unmasked atoms
- Iterators for selected/unselected atoms
"""
import numpy as np
import warnings
import os
from pyscal.attributes import AttrSetter, read_yaml, MyList
attr_docs = read_yaml(os.path.join(os.path.dirname(__file__), "data/annotations.yaml"))
class Atoms(dict, AttrSetter):
def __init__(self, atoms=None):
#self.update(atoms=atoms)
self._nreal = 0
self._nghost = 0
self._lattice_constant = None
self._lattice = None
AttrSetter.__init__(self)
if atoms is not None:
self.from_dict(atoms)
def __dir__(self):
attrs = ["natoms", "nreal", "nghost",
"ntotal", "from_dict", "iter_atoms", "apply_mask", "remove_mask",
"apply_selection", "remove_selection", "delete", "composition"]
return attrs + list(self._map_dict.keys())
def __getitem__(self, key):
if isinstance(key, slice):
return self._get_atoms(key)
elif isinstance(key, int):
return self._get_atoms(key)
else:
val = dict.__getitem__(self, key)
return val
def __setitem__(self, key, val):
dict.__setitem__(self, key, MyList(val))
def __repr__(self):
dictrepr = dict.__repr__(self)
return '%s(%s)' % (type(self).__name__, dictrepr)
def _repr_json(self):
#convert to atom base dict
disp_atoms = {f"atom {x}": self._get_atoms(x) for x in range(self.natoms)}
return disp_atoms
#def update(self, atoms):
# for k, v in dict(*args, **kwargs).items():
# self[k] = v
def __add__(self, atoms):
if not 'positions' in atoms.keys():
raise ValueError('positions is a necessary key in atoms')
nop = len(atoms["positions"])
val_length_check = np.prod([len(val)==nop for key, val in atoms.items()])
if not val_length_check:
raise ValueError("All times in the atoms dict should have same length as positions")
#now add necessary keys-ids, types, ghost
maxid = max(self["ids"])
if not 'ids' in atoms.keys():
atoms['ids'] = [maxid+x+1 for x in range(nop)]
#print(self["ids"], atoms["ids"])
else:
if len(set(atoms['ids']).intersection(set(self['ids']))):
raise ValueError("Atom id already exists, unique ID is required")
atoms['ghost'] = [False for x in range(nop)]
if not 'types' in atoms.keys():
atoms['types'] = [1 for x in range(nop)]
if not 'species' in atoms.keys():
atoms['species'] = [None for x in range(nop)]
if not 'mask_1' in atoms.keys():
atoms['mask_1'] = [False for x in range(nop)]
if not 'mask_2' in atoms.keys():
atoms['mask_2'] = [False for x in range(nop)]
if not 'condition' in atoms.keys():
atoms['condition'] = [True for x in range(nop)]
if not 'head' in atoms.keys():
atoms['head'] = [self.natoms+x for x in range(nop)]
common_keys = list(set(self.keys()).intersection(set(atoms.keys())))
_ = [self[key].extend(atoms[key]) for key in common_keys]
extra_keys_add = len(atoms.keys()) - len(common_keys)
extra_keys_exist = len(self.keys()) - len(common_keys)
if extra_keys_add > 0:
warnings.warn("Some keys present in the atoms are add are not present in existing atoms, they were ignored")
if extra_keys_exist > 0:
warnings.warn("Some keys present in the existing Atoms were not present in the atoms to add, please recalculate")
self._nreal += nop
return self
def __radd__(self, atoms):
"""
Reverse add method
"""
if ntraj == 0:
return self
else:
return self.__add__(atoms)
#def _add_attribute()
@property
def natoms(self):
return self._nreal
@property
def nreal(self):
return self._nreal
@property
def nghost(self):
return self._nghost
@property
def ntotal(self):
return self._nreal + self._nghost
def from_dict(self, atoms):
if not 'positions' in atoms.keys():
raise ValueError('positions is a necessary key in atoms')
nop = len(atoms["positions"])
if not 'ids' in atoms.keys():
atoms['ids'] = [x+1 for x in range(nop)]
atoms['ghost'] = [False for x in range(nop)]
if not 'types' in atoms.keys():
atoms['types'] = [1 for x in range(nop)]
if not 'species' in atoms.keys():
atoms['species'] = [None for x in range(nop)]
if not 'mask_1' in atoms.keys():
atoms['mask_1'] = [False for x in range(nop)]
if not 'mask_2' in atoms.keys():
atoms['mask_2'] = [False for x in range(nop)]
if not 'condition' in atoms.keys():
atoms['condition'] = [True for x in range(nop)]
if not 'head' in atoms.keys():
atoms['head'] = [x for x in range(nop)]
for key, val in atoms.items():
self[key] = MyList(val)
self._nreal = len(val)
#add attributes
mapdict = {"positions": "positions",
"ids": "ids",
"types": "types",
"species": "species",
"mask": {"primary": "mask_1", "secondary": "mask_2"},
"selection": "condition",
"condition": "condition",
"head": "head"}
#add extra keys that might be needed; non-standard ones
for key, val in atoms.items():
if key not in ["positions", "ids", "types", "species", "mask_1", "mask_2", "condition", "head"]:
mapdict[key] = key
self._add_attribute(mapdict)
def _convert_to_list(self, data):
"""
Check if the given item is a list, if not convert to a single item list
"""
if not isinstance(data, list):
data = [data]
return data
def _get_atoms(self, index):
atom_dict = {key: self._convert_to_list(self[key][index]) for key in self.keys()}
return Atoms(atom_dict)
def _delete_atoms(self, indices):
del_real = np.sum([1 for x in indices if x < self._nreal])
del_ghost = np.sum([1 for x in indices if x >= self._nreal])
for key in self.keys():
for index in sorted(indices, reverse=True):
del self[key][index]
td = len(indices)
self._nreal = int(self.nreal - del_real)
self._nghost = int(self.nghost - del_ghost)
def iter_atoms(self):
for index in range(self.nreal):
atom_dict = {key: [self[key][index]] for key in self.keys()}
yield Atoms(atom_dict)
def _generate_bool_list(self, ids=None, indices=None, condition=None, selection=False):
#necessary checks
non_nones = sum(x is not None for x in [ids, indices, condition])
if non_nones > 1:
raise ValueError("Only one of ids, indices or condition should be provided")
elif ((non_nones == 0) and (selection==False)):
warnings.warn("No conditions provided, all atoms will be included")
#generate a list of indices
if selection:
indices = [x for x in range(self.nreal) if self["condition"][x]]
elif ids is not None:
if not isinstance(ids, list):
ids = [ids]
indices = [x for x in range(len(self["ids"])) if self["ids"][x] in ids]
if len(indices) == 0:
raise ValueError("No ids found to delete")
if len(indices) != len(ids):
warnings.warn("Not all ids were found")
elif condition is not None:
indices = [x for x in range(self.nreal) if condition(self._get_atoms(x))]
elif indices is None:
indices = [x for x in range(self.nreal)]
if not isinstance(indices, list):
indices = [indices]
bool_list = [ True if x in indices else False for x in range(self.nreal)]
return bool_list
def _apply_mask(self, masks, mask_type):
if (mask_type == 'primary') or (mask_type == 'all'):
for i in range(self.ntotal):
self["mask_1"][i] = masks[self["head"][i]]
if (mask_type == 'secondary') or (mask_type == 'all'):
for i in range(self.ntotal):
self["mask_2"][i] = masks[self["head"][i]]
def apply_mask(self, mask_type="primary", ids=None, indices=None, condition=None, selection=False):
masks = self._generate_bool_list(ids=ids, indices=indices, condition=condition, selection=selection)
self._apply_mask(masks, mask_type)
def remove_mask(self, mask_type="primary", ids=None, indices=None, condition=None, selection=False):
masks = self._generate_bool_list(ids=ids, indices=indices, condition=condition, selection=selection)
masks = [not x for x in masks]
self._apply_mask(masks, mask_type)
def _apply_selection(self, condition):
for i in range(self.ntotal):
self["condition"][i] = condition[self["head"][i]]
def _validate_condition(self, condition):
if not (len(condition)==self.nreal):
raise ValueError("condition should have same length as atoms")
for c, x in enumerate(condition):
try:
x = bool(x)
condition[c] = x
except:
pass
if not isinstance(x, bool):
raise ValueError("Condition elements should be boolean")
return condition
def apply_selection(self, ids=None, indices=None, condition=None):
if isinstance(condition, list):
masks = self._validate_condition(condition)
else:
masks = self._generate_bool_list(ids=ids, indices=indices, condition=condition)
self._apply_selection(masks)
def remove_selection(self, ids=None, indices=None, condition=None):
if isinstance(condition, list):
masks = self._validate_condition(condition)
else:
masks = self._generate_bool_list(ids=ids, indices=indices, condition=condition)
masks = [not x for x in masks]
self._apply_selection(masks)
def delete(self, ids=None, indices=None, condition=None, selection=False):
#delete atoms
#reassign ids
#reassign indices
#reassign heads
masks = self._generate_bool_list(ids=ids, indices=indices, condition=condition, selection=selection)
delete_list = [masks[self["head"][x]] for x in range(self.ntotal)]
delete_ids = [x for x in range(self.ntotal) if masks[x]]
self._delete_atoms(delete_ids)
@property
def _type_dict(self):
sp = []
types, typecounts = np.unique(self["types"][:self.nreal], return_counts=True)
for t in types:
for count, tx in enumerate(self["types"][:self.nreal]):
if t==tx:
sp.append(self["species"][count])
break
return dict([x for x in zip(types, sp)])
@property
def composition(self):
if self["species"][0] is None:
typelist = self["types"][:self.nreal]
types, typecounts = np.unique(typelist, return_counts=True)
concdict = dict([(t, typecounts[c]/np.sum(typecounts)) for c, t in enumerate(types)])
else:
typelist = self["species"][:self.nreal]
types, typecounts = np.unique(typelist, return_counts=True)
concdict = {str(t): typecounts[c]/np.sum(typecounts) for c, t in enumerate(types)}
return concdict