Coverage for calorine/nep/tensor_conventions.py: 100%

1import numpy as np 

2 

3# ASE's own canonical Voigt-6 convention (see ase.stress.voigt_notation). 

4ASE_VOIGT6_ORDER = ('xx', 'yy', 'zz', 'yz', 'xz', 'xy') 

5 

6# Reduced-6 order used by the standalone `nep` executable's virial_*.out, 

7# stress_*.out, and polarizability_*.out files (training/prediction mode), 

8# and therefore by calorine.nep.io.read_structures. Established by GPUMD's 

9# main_nep/structure.cu: `reduced_index = {0, 3, 5, 3, 1, 4, 5, 4, 2}` (maps 

10# full-3x3 flat positions onto this 6-slot array; used identically for 

11# parsing both `virial=` and `pol=` extxyz fields). Also matches 

12# `get_polarizability` in src/nepy/nep.cpp. Independently confirmed 

13# empirically for stress by applying isolated shear strains and checking 

14# that each produces a response concentrated at the matching component. 

15NEP_REDUCED6_ORDER = ('xx', 'yy', 'zz', 'xy', 'yz', 'xz') 

16 

17# Raw, unreduced, row-major full-3x3 order for Born effective charges (BEC 

18# is not symmetric, so no reduced form applies), used identically by 

19# src/nepy/nep.cpp (g_bec[n1+N*0..8]), GPUMD's src/force/nep_charge.cu 

20# (identical kernel/layout), and src/measure/dump_xyz.cu (writes `bec` with 

21# no reindexing). 

22BEC_FULL9_ORDER = ('xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz') 

23 

24 

25def reduced6_to_full_3x3(values, order=NEP_REDUCED6_ORDER): 

26 """Expands an array whose last axis holds a reduced-6 symmetric tensor 

27 (in the component order given by `order`) into an array with two 

28 trailing 3x3 axes, e.g. shape (6,) -> (3,3) or (natoms,3,6) -> (natoms,3,3,3). 

29 Use ``order=ASE_VOIGT6_ORDER`` for ASE-Voigt input (equivalent to 

30 :func:`ase.stress.voigt_6_to_full_3x3_stress`, provided here for 

31 uniformity with the other conventions in this module). 

32 

33 Parameters 

34 ---------- 

35 values 

36 Array whose last axis has length 6, holding a reduced symmetric 

37 tensor in the component order given by `order`. 

38 order 

39 Component order of `values`' last axis. Defaults to 

40 :data:`NEP_REDUCED6_ORDER`. 

41 """ 

42 values = np.asarray(values) 

43 idx = {label: i for i, label in enumerate(order)} 

44 index_grid = np.array([ 

45 [idx['xx'], idx['xy'], idx['xz']], 

46 [idx['xy'], idx['yy'], idx['yz']], 

47 [idx['xz'], idx['yz'], idx['zz']], 

48 ]) 

49 return values[..., index_grid] 

50 

51 

52def reduced6_permutation(source_order, target_order=ASE_VOIGT6_ORDER): 

53 """Returns the index array that permutes a reduced-6 array from 

54 `source_order` to `target_order`. 

55 

56 Parameters 

57 ---------- 

58 source_order 

59 Component order to permute from. 

60 target_order 

61 Component order to permute to. Defaults to :data:`ASE_VOIGT6_ORDER`. 

62 """ 

63 return [source_order.index(label) for label in target_order] 

64 

65 

66def nep_reduced6_to_ase_voigt6(values): 

67 """Permutes a reduced-6 symmetric-tensor array (last axis) from `nep`'s 

68 native order (:data:`NEP_REDUCED6_ORDER`) to ASE's Voigt-6 order 

69 (:data:`ASE_VOIGT6_ORDER`). 

70 

71 Parameters 

72 ---------- 

73 values 

74 Array whose last axis has length 6, holding a reduced symmetric 

75 tensor in :data:`NEP_REDUCED6_ORDER`. 

76 """ 

77 return np.asarray(values)[..., reduced6_permutation(NEP_REDUCED6_ORDER)]