-
Notifications
You must be signed in to change notification settings - Fork 1
/
pbc.py
executable file
·188 lines (161 loc) · 7.5 KB
/
pbc.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
import warnings
import numpy as np
from pymatgen.core.periodic_table import _pt_data
from pymatgen.core.structure import Lattice
from pymatgen.core.structure import Molecule
from pymatgen.core.structure import PeriodicSite
from pymatgen.core.structure import Site
from pymatgen.core.structure import Structure
from pymatgen.util.coord import pbc_shortest_vectors
"""
2020/07/29 copied from ocelot dev branch
PBCparser: get unwrapped structure and mols
method here should not rely on ocelot schema
"""
def AtomicRadius_string(species_string):
d = _pt_data[species_string]
at_r = d.get("Atomic radius", "no data")
try:
return float(at_r)
except ValueError:
warnings.warn('unknown element {}'.format(species_string))
return 1.0
def AtomicRadius(site: Site):
d = _pt_data[site.species_string]
at_r = d.get("Atomic radius", "no data")
try:
return float(at_r)
except ValueError:
warnings.warn('unknown element {}'.format(site.species_string))
return 1.0
class PBCparser:
@staticmethod
def get_dist_and_trans(lattice: Lattice, fc1, fc2):
"""
get the shortest distance and corresponding translation vector between two frac coords
:param lattice: pmg lattic obj
:param fc1:
:param fc2:
:return:
"""
v, d2 = pbc_shortest_vectors(lattice, fc1, fc2, return_d2=True)
if len(np.array(fc1).shape) == 1:
fc = lattice.get_fractional_coords(v[0][0]) + fc1 - fc2
else:
fc = np.zeros((len(fc1), len(fc2), 3))
for i in range(len(fc1)):
for j in range(len(fc2)):
fc_vector = np.dot(v[i][j], lattice.inv_matrix)
fc[i][j] = fc_vector + fc1[i] - fc2[j]
return np.sqrt(d2), fc
@staticmethod
def unwrap(pstructure: Structure):
"""
unwrap the structure, extract isolated mols
this will modify psites *in-place*, properties will be inherited and a new property
'imol' will be written
psite with imol=x is an element of both mols[x] and unwrap_pblock_list[x]
this method is not supposed to modify siteid!
:param pstructure: periodic structure obj from pymatgen
:return: mols, unwrap_str_sorted, unwrap_pblock_list
"""
psites = pstructure.sites
cutoff_matrix = np.zeros((len(psites), len(psites)))
for i in range(len(psites)):
for j in range(i + 1, len(psites)):
cutoff = AtomicRadius(psites[i]) + AtomicRadius(psites[j])
cutoff *= 1.3
cutoff_matrix[i][j] = cutoff
cutoff_matrix[j][i] = cutoff
if all('iasym' in s.properties for s in psites):
sameiasym = True
warnings.warn('if a percolating step involves hydrogens, only percolate to sites with the same iasym')
else:
sameiasym = False
pindices = range(len(psites))
visited = []
block_list = []
# unwrap = []
unwrap_block_list = []
unwrap_pblock_list = []
while len(visited) != len(psites):
# initialization
unvisited = [idx for idx in pindices if idx not in visited]
ini_idx = unvisited[0]
block = [ini_idx]
# unwrap.append(psites[ini_idx])
unwrap_block = [Site(psites[ini_idx].species_string, psites[ini_idx].coords,
properties=psites[ini_idx].properties)]
unwrap_pblock = [psites[ini_idx]]
pointer = 0
while pointer != len(block):
outside = [idx for idx in pindices if idx not in block and idx not in visited]
for i in outside:
si = psites[i]
sj = psites[block[pointer]]
if sameiasym and (si.species_string == 'H' or sj.species_string == 'H'):
if si.properties['iasym'] != sj.properties['iasym']:
continue
distance, fctrans = PBCparser.get_dist_and_trans(pstructure.lattice,
psites[block[pointer]].frac_coords,
psites[i].frac_coords, )
fctrans = np.rint(fctrans)
cutoff = cutoff_matrix[block[pointer]][i]
if distance < cutoff:
block.append(i)
psites[i]._frac_coords += fctrans
# psites[i] = PeriodicSite(psites[i].species_string, psites[i].frac_coords + fctrans,
# pstructure.lattice, properties=deepcopy(psites[i].properties))
unwrap_block.append(
# Site(psites[i].species_string, psites[i].coords, properties=deepcopy(psites[i].properties))
Site(psites[i].species_string, psites[i].coords, properties=psites[i].properties)
)
unwrap_pblock.append(psites[i])
visited.append(block[pointer])
pointer += 1
unwrap_block_list.append(unwrap_block)
unwrap_pblock_list.append(unwrap_pblock)
block_list.append(block)
unwrap = []
for i in range(len(unwrap_block_list)):
for j in range(len(unwrap_block_list[i])):
unwrap_block_list[i][j].properties['imol'] = i
unwrap_pblock_list[i][j].properties['imol'] = i
unwrap.append(unwrap_pblock_list[i][j])
mols = [Molecule.from_sites(sites) for sites in unwrap_block_list]
# unwrap_structure = Structure.from_sites(sorted(unwrap, key=lambda x: x.species_string))
unwrap_structure = Structure.from_sites(unwrap, to_unit_cell=False)
return mols, unwrap_structure, unwrap_pblock_list
@staticmethod
def unwrap_and_squeeze(pstructure: Structure):
"""
after unwrapping, the mols can be far away from each other, this tries to translate them s.t. they stay together
:param pstructure:
:return:
"""
mols, unwrap_structure, psiteblocks = PBCparser.unwrap(pstructure)
if len(mols) > 1:
refpoint = mols[0].center_of_mass
refpoint = unwrap_structure.lattice.get_fractional_coords(refpoint)
for i in range(1, len(mols)):
varmol = mols[i]
varpoint = varmol.center_of_mass
varpoint = unwrap_structure.lattice.get_fractional_coords(varpoint)
distance, fctrans = PBCparser.get_dist_and_trans(unwrap_structure.lattice, refpoint, varpoint)
for j in range(len(psiteblocks[i])):
psiteblocks[i][j]._frac_coords += fctrans
squeeze_psites = []
squeeze_mols = []
pblock: [PeriodicSite]
for pblock in psiteblocks:
squeeze_block = []
for ps in pblock:
squeeze_psites.append(ps)
squeeze_block.append(
Site(ps.species_string, ps.coords, properties=ps.properties)) # do we need deepcopy?
mol = Molecule.from_sites(squeeze_block)
squeeze_mols.append(mol)
squeeze_unwrap_structure = Structure.from_sites(squeeze_psites)
return squeeze_mols, squeeze_unwrap_structure, psiteblocks
else:
return mols, unwrap_structure, psiteblocks