1"""
2The ``io.shengBTE`` module provides functions for reading and writing data
3files in shengBTE format.
4"""
5
6import numpy as np
7from itertools import product
8from collections import namedtuple
9from ..core.structures import Supercell
10import hiphive
11
12
13def read_shengBTE_fc3(filename, prim, supercell, symprec=1e-5):
14    """Reads third-order force constants file in shengBTE format.
15
16    Parameters
17    ----------
18    filename : str
19        input file name
20    prim : ase.Atoms
21        primitive cell for the force constants
22    supercell : ase.Atoms
23        supercell onto which to map force constants
24    symprec : float
25        structural symmetry tolerance
26
27    Returns
28    -------
29    fcs : ForceConstants
30        third order force constants for the specified supercell
31    """
32
33    raw_sheng = _read_raw_sheng(filename)
34
35    sheng = _raw_to_fancy(raw_sheng, prim.cell)
36
37    fcs = _sheng_to_fcs(sheng, prim, supercell, symprec)
38
39    return fcs
40
41
42def write_shengBTE_fc3(filename, fcs, prim, symprec=1e-5, cutoff=np.inf,
43                       fc_tol=1e-8):
44    """Writes third-order force constants file in shengBTE format.
45
46    Parameters
47    -----------
48    filename : str
49        input file name
50    fcs : ForceConstants
51        force constants; the supercell associated with this object must be
52        based on prim
53    prim : ase.Atoms
54        primitive configuration (must be equivalent to structure used in the
55        shengBTE calculation)
56    symprec : float
57        structural symmetry tolerance
58    cutoff : float
59        all atoms in cluster must be within this cutoff
60    fc_tol : float
61        if the absolute value of the largest entry in a force constant is less
62        than fc_tol it will not be written
63    """
64
65    sheng = _fcs_to_sheng(fcs, prim, symprec, cutoff, fc_tol)
66
67    raw_sheng = _fancy_to_raw(sheng)
68
69    _write_raw_sheng(raw_sheng, filename)
70
71
72def _read_raw_sheng(filename):
73    """ Read shengBTE fc3 file.
74
75    Parameters
76    ----------
77    filename : str
78        input file name
79
80    Returns
81    -------
82    list
83        list with shengBTE block entries, where an entry consist of
84        [i, j, k, cell_pos2, cell_pos3, fc3]
85    """
86    lines_per_fc_block = 32
87
88    fc3_data_shengBTE = []
89    with open(filename, 'r') as f:
90        lines = f.readlines()
91        num_fcs = int(lines[0])
92        lind = 2
93        for i in range(1, num_fcs+1):
94            # sanity check
95            assert int(lines[lind]) == i, (int(lines[lind]), i)
96
97            # read cell poses
98            cell_pos2 = np.array([float(fld) for fld in lines[lind+1].split()])
99            cell_pos3 = np.array([float(fld) for fld in lines[lind+2].split()])
100
101            # read basis indices
102            i, j, k = (int(fld) for fld in lines[lind+3].split())
103
104            # read fc
105            fc3_ijk = np.zeros((3, 3, 3))
106            for n in range(27):
107                x, y, z = [int(fld) - 1 for fld in lines[lind+4+n].split()[:3]]
108                fc3_ijk[x, y, z] = float(lines[lind+4+n].split()[-1])
109
110            # append entry
111            entry = [i, j, k, cell_pos2, cell_pos3, fc3_ijk]
112            fc3_data_shengBTE.append(entry)
113            lind += lines_per_fc_block
114
115    return fc3_data_shengBTE
116
117
118_ShengEntry = namedtuple('Entry', ['site_0', 'site_1', 'site_2', 'pos_1',
119                                   'pos_2', 'fc', 'offset_1', 'offset_2'])
120
121
122def _raw_to_fancy(raw_sheng, cell):
123    """
124    Converts force constants as read from shengBTE file to the namedtuple
125    format defined above (_ShengEntry).
126    """
127    sheng = []
128    for raw_entry in raw_sheng:
129        p1, p2 = raw_entry[3:5]
130        offset_1 = np.linalg.solve(cell.T, p1).round(0).astype(int)
131        offset_2 = np.linalg.solve(cell.T, p2).round(0).astype(int)
132        entry = _ShengEntry(*(i - 1 for i in raw_entry[:3]), *raw_entry[3:],
133                            offset_1, offset_2)
134        sheng.append(entry)
135    return sheng
136
137
138def _fancy_to_raw(sheng):
139    """
140    Converts force constants namedtuple format defined above (_ShengEntry) to
141    format used for writing shengBTE files.
142    """
143    raw_sheng = []
144    for entry in sheng:
145        raw_entry = list(entry[:6])
146        raw_entry[0] += 1
147        raw_entry[1] += 1
148        raw_entry[2] += 1
149        raw_sheng.append(raw_entry)
150
151    return raw_sheng
152
153
154def _write_raw_sheng(raw_sheng, filename):
155    """ See corresponding read function. """
156
157    with open(filename, 'w') as f:
158        f.write('{}\n\n'.format(len(raw_sheng)))
159
160        for index, fc3_row in enumerate(raw_sheng, start=1):
161            i, j, k, cell_pos2, cell_pos3, fc3_ijk = fc3_row
162
163            f.write('{:5d}\n'.format(index))
164
165            f.write((3*'{:14.10f} '+'\n').format(*cell_pos2))
166            f.write((3*'{:14.10f} '+'\n').format(*cell_pos3))
167            f.write((3*'{:5d}'+'\n').format(i, j, k))
168
169            for x, y, z in product(range(3), repeat=3):
170                f.write((3*' {:}').format(x+1, y+1, z+1))
171                f.write('    {:14.10f}\n'.format(fc3_ijk[x, y, z]))
172            f.write('\n')
173
174
175def _fcs_to_sheng(fcs, prim, symprec, cutoff, fc_tol):
176    """ Produces a list containing fcs in shengBTE format
177
178    prim and fcs.supercell must be aligned.
179    """
180
181    supercell = Supercell(fcs.supercell, prim, symprec)
182    assert all(fcs.supercell.pbc) and all(prim.pbc)
183
184    n_atoms = len(supercell)
185
186    D = fcs.supercell.get_all_distances(mic=False, vector=True)
187    D_mic = fcs.supercell.get_all_distances(mic=True, vector=True)
188    M = np.eye(n_atoms, dtype=bool)
189    for i in range(n_atoms):
190        for j in range(i + 1, n_atoms):
191            M[i, j] = (np.allclose(D[i, j], D_mic[i, j], atol=symprec, rtol=0)
192                       and np.linalg.norm(D[i, j]) < cutoff)
193            M[j, i] = M[i, j]
194
195    data = {}
196    for a0 in supercell:
197        for a1 in supercell:
198            if not M[a0.index, a1.index]:
199                continue
200            for a2 in supercell:
201                if not (M[a0.index, a2.index] and M[a1.index, a2.index]):
202                    continue
203
204                offset_1 = np.subtract(a1.offset, a0.offset)
205                offset_2 = np.subtract(a2.offset, a0.offset)
206
207                sites = (a0.site, a1.site, a2.site)
208
209                key = sites + tuple(offset_1) + tuple(offset_2)
210
211                ijk = (a0.index, a1.index, a2.index)
212
213                fc = fcs[ijk]
214
215                if key in data:
216                    assert np.allclose(data[key], fc, atol=fc_tol)
217                else:
218                    data[key] = fc
219
220    sheng = []
221    for k, fc in data.items():
222        if np.max(np.abs(fc)) < fc_tol:
223            continue
224        offset_1 = k[3:6]
225        pos_1 = np.dot(offset_1, prim.cell)
226        offset_2 = k[6:9]
227        pos_2 = np.dot(offset_2, prim.cell)
228        entry = _ShengEntry(*k[:3], pos_1, pos_2, fc, offset_1, offset_2)
229        sheng.append(entry)
230
231    return sheng
232
233
234def _sheng_to_fcs(sheng, prim, supercell, symprec):
235    supercell_map = Supercell(supercell, prim, symprec)
236    fc_array = np.zeros((len(supercell),) * 3 + (3,) * 3)
237    mapped_clusters = np.zeros((len(supercell),) * 3, dtype=bool)
238
239    for atom in supercell_map:
240        i = atom.index
241        for entry in sheng:
242            if atom.site != entry.site_0:
243                continue
244            j = supercell_map.index(entry.site_1, entry.offset_1 + atom.offset)
245            k = supercell_map.index(entry.site_2, entry.offset_2 + atom.offset)
246            ijk = (i, j, k)
247            if mapped_clusters[ijk]:
248                raise Exception('Ambiguous force constant.'
249                                ' Supercell is too small')
250            fc_array[ijk] = entry.fc
251            mapped_clusters[ijk] = True
252
253    fcs = hiphive.force_constants.ForceConstants.from_arrays(
254        supercell, fc3_array=fc_array)
255    return fcs
256