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

25 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-12-21 12:02 +0000

1import numpy as np 

2 

3from dynasor.core.rho_j_q_numba import rho_q as rho_q_numba 

4from dynasor.core.rho_j_q_numba import rho_j_q as rho_j_q_numba 

5 

6 

7def calc_rho_q(x, q): 

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

9 

10 Will call external function rho_q to calculate the 

11 particle density in q-space. 

12 Particle coordinates and q-space points of interest are 

13 passed as input via x and q, respectively. 

14 """ 

15 

16 assert x.shape[1] == 3 

17 assert q.shape[1] == 3 

18 

19 Nx, _ = x.shape 

20 Nq, _ = q.shape 

21 

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

23 

24 x = x.copy() # Don't ask why 

25 rho_q_numba(x, q, rho_q) 

26 

27 return rho_q 

28 

29 

30def calc_rho_j_q(x, v, q): 

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

32 """ 

33 assert x.shape == v.shape 

34 

35 assert x.shape[1] == 3 

36 assert v.shape[1] == 3 

37 assert q.shape[1] == 3 

38 

39 Nx, _ = x.shape 

40 Nq, _ = q.shape 

41 

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

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

44 

45 x = x.copy() 

46 v = v.copy() 

47 

48 rho_j_q_numba(x, v, q, rho_q, j_q) 

49 

50 return rho_q, j_q