root / ase / io / plt.py @ 17
Historique | Voir | Annoter | Télécharger (884 octet)
1 | 1 | tkerber | import numpy as np |
---|---|---|---|
2 | 1 | tkerber | |
3 | 1 | tkerber | from ase.atoms import Atoms |
4 | 1 | tkerber | |
5 | 1 | tkerber | def write_plt(filename, atoms, data): |
6 | 1 | tkerber | if isinstance(atoms, Atoms): |
7 | 1 | tkerber | cell = atoms.get_cell() |
8 | 1 | tkerber | else:
|
9 | 1 | tkerber | cell = np.asarray(atoms, float)
|
10 | 1 | tkerber | |
11 | 1 | tkerber | if cell.ndim == 2: |
12 | 1 | tkerber | c = cell.copy() |
13 | 1 | tkerber | cell = c.diagonal() |
14 | 1 | tkerber | c.flat[::4] = 0.0 |
15 | 1 | tkerber | if c.any():
|
16 | 1 | tkerber | raise ValueError('Unit cell must be orthorhombic!') |
17 | 1 | tkerber | |
18 | 1 | tkerber | f = open(filename, 'w') |
19 | 1 | tkerber | np.array([3, 4], np.int32).tofile(f) |
20 | 1 | tkerber | |
21 | 1 | tkerber | dims = np.array(data.shape, np.int32) |
22 | 1 | tkerber | dims[::-1].tofile(f)
|
23 | 1 | tkerber | |
24 | 1 | tkerber | for n, L in zip(dims[::-1], cell[::-1]): |
25 | 1 | tkerber | if n % 2 == 0: |
26 | 1 | tkerber | d = L / n |
27 | 1 | tkerber | np.array([0.0, L - d], np.float32).tofile(f)
|
28 | 1 | tkerber | else:
|
29 | 1 | tkerber | d = L / (n + 1)
|
30 | 1 | tkerber | np.array([d, L - d], np.float32).tofile(f) |
31 | 1 | tkerber | |
32 | 1 | tkerber | if data.dtype == complex: |
33 | 1 | tkerber | data = np.abs(data) |
34 | 1 | tkerber | data.astype(np.float32).T.tofile(f) |
35 | 1 | tkerber | f.close() |