Coverage for dynasor/core/reciprocal.py: 100%

1import numpy as np 

2 

3from dynasor.core.rho_j_q_numba import _rho_j_q_exp 

4from dynasor.core.rho_j_q_numba import _rho_j_q_poly 

5from dynasor.core.rho_j_q_numba import _rho_q_incoherent_exp 

6from dynasor.core.rho_j_q_numba import _rho_q_incoherent_poly 

7from dynasor.core.rho_j_q_numba import _rho_q_exp 

8from dynasor.core.rho_j_q_numba import _rho_q_poly 

9from dynasor.logging_tools import warn_if_numba_threading_layer_is_slow 

10 

11# alpha is the argument in exp(i * alpha), i.e. alpha = q . r. For too large alpha the 

12# polynomial kernels become less accurate, so above this limit we fall back to the more 

13# stable exp kernels. The reduction only breaks down at 6.6e6 (see _PIO2_1 in 

14# rho_j_q_numba), so 1e5 leaves a wide margin and is far above any ordinary cell. 

15_MAX_ALPHA_FOR_POLYNOMIAL_KERNEL = 1e5 

16 

17# Number of phases, Nx * Nq, below which the polynomial path is not worth taking. 

18# Small systems and few q-points therefore use the exp kernels. 

19_MIN_NUMBER_OF_PHASES_FOR_POLYNOMIAL_KERNEL = 10_000 

20 

21 

22def _use_poly(x, q): 

23 """Whether the polynomial kernels are both valid and worth using for this call.""" 

24 if (x.size // 3) * len(q) < _MIN_NUMBER_OF_PHASES_FOR_POLYNOMIAL_KERNEL: 

25 return False 

26 # max(abs(.)) without building a temporary; float() also makes integer input safe 

27 x_max = max(float(x.max()), -float(x.min())) 

28 q_max = max(float(q.max()), -float(q.min())) 

29 return 3 * x_max * q_max <= _MAX_ALPHA_FOR_POLYNOMIAL_KERNEL 

30 

31 

32def calc_rho_q(x, q): 

33 """Calculate rho(q) of particle coordinates x. 

34 

35 Calls one of the numba kernels to calculate the particle density in q-space. 

36 Particle coordinates and q-space points of interest are 

37 passed as input via x and q, respectively. 

38 """ 

39 

40 assert x.shape[1] == 3 

41 assert q.shape[1] == 3 

42 

43 Nq = len(q) 

44 

45 rho_q = np.zeros(Nq, dtype=np.complex128) 

46 

47 if _use_poly(x, q): 

48 x = np.ascontiguousarray(x.T) 

49 _rho_q_poly(x, q, rho_q) 

50 else: 

51 x = np.ascontiguousarray(x) 

52 _rho_q_exp(x, q, rho_q) 

53 

54 warn_if_numba_threading_layer_is_slow() 

55 

56 return rho_q 

57 

58 

59def calc_rho_j_q(x, v, q): 

60 """As calc_rho_q, but calculate also velocities in q-space. 

61 """ 

62 assert x.shape == v.shape 

63 

64 assert x.shape[1] == 3 

65 assert v.shape[1] == 3 

66 assert q.shape[1] == 3 

67 

68 Nq = len(q) 

69 

70 rho_q = np.zeros(Nq, dtype=np.complex128) 

71 j_q = np.zeros((Nq, 3), dtype=np.complex128) 

72 

73 if _use_poly(x, q): 

74 x = np.ascontiguousarray(x.T) 

75 v = np.ascontiguousarray(v.T) 

76 _rho_j_q_poly(x, v, q, rho_q, j_q) 

77 else: 

78 x = np.ascontiguousarray(x) 

79 v = np.ascontiguousarray(v) 

80 _rho_j_q_exp(x, v, q, rho_q, j_q) 

81 

82 warn_if_numba_threading_layer_is_slow() 

83 

84 return rho_q, j_q 

85 

86 

87def _calc_rho_q_incoherent(x, q, out=None): 

88 """Calculate real self densities from a contiguous ``(N_time, 3, N_atoms)`` array.""" 

89 assert x.ndim == 3 

90 assert x.shape[1] == 3 

91 n_time = x.shape[0] 

92 n_qpoints, _ = q.shape 

93 if out is None: 

94 out = np.empty((n_time, n_qpoints), dtype=np.float64) 

95 else: 

96 assert out.shape == (n_time, n_qpoints) 

97 assert out.dtype == np.float64 

98 

99 x = np.ascontiguousarray(x) 

100 if _use_poly(x, q): 

101 _rho_q_incoherent_poly(x, q, out) 

102 else: 

103 _rho_q_incoherent_exp(x, q, out) 

104 

105 warn_if_numba_threading_layer_is_slow() 

106 

107 return out