Coverage for dynasor/core/frame_processors.py: 34%

1"""Factory function that builds the callable used to populate 

2``rho_qs_dict`` (and, if requested, ``jz_qs_dict``/``jper_qs_dict``) on 

3trajectory frames, dispatched on the configured compute backend. 

4 

5Used by :mod:`dynasor.correlation_functions` to keep the public API 

6functions free of backend-selection details. 

7""" 

8import numpy as np 

9 

10from dynasor.compute_config import check_backend_is_available 

11from dynasor.logging_tools import logger 

12from dynasor.core.reciprocal import calc_rho_q, calc_rho_j_q 

13 

14 

15def get_frame_processor(q_points, backend, precision, gpu_batch_size, 

16 calculate_currents=False, q_directions=None): 

17 r"""Build the callable used to compute the reciprocal-space density 

18 (and, if requested, current) for trajectory frames. 

19 

20 Parameters 

21 ---------- 

22 q_points 

23 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in 

24 Cartesian coordinates. 

25 backend 

26 Compute backend, one of ``'numba'``, ``'torch'``, or ``'cupy'``. 

27 precision 

28 Floating point precision used by the GPU backends (``'float64'`` or 

29 ``'float32'``). 

30 gpu_batch_size 

31 Number of frames processed per GPU kernel call. Only used here for 

32 logging; the batching itself is done by 

33 :func:`_prefetch_batched_iter <dynasor.trajectory.prefetch._prefetch_batched_iter>`. 

34 calculate_currents 

35 If True, also compute the longitudinal and transverse current 

36 components. Requires velocities to be available on the frames. 

37 q_directions 

38 Unit vectors along :attr:`q_points`. Required if 

39 :attr:`calculate_currents` is True. 

40 

41 Returns 

42 ------- 

43 A tuple ``(is_batched, processor)``. If :attr:`is_batched` is True, 

44 :attr:`processor` takes a list of frames and is meant for use with 

45 :func:`_prefetch_batched_iter <dynasor.trajectory.prefetch._prefetch_batched_iter>`. 

46 Otherwise, :attr:`processor` takes a single frame and is meant for use as 

47 a :class:`WindowIterator <dynasor.trajectory.WindowIterator>` 

48 ``element_processor`` or with :func:`map`. 

49 """ 

50 check_backend_is_available(backend) 

51 

52 if backend == 'torch': 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true

53 from dynasor.core.reciprocal_torch import ( 

54 make_q_tensor, calc_rho_q_torch_batched, calc_rho_j_q_torch_batched 

55 ) 

56 q_d = make_q_tensor(q_points, device='cuda', precision=precision) 

57 logger.info(f'Using PyTorch backend (device={q_d.device}, precision={precision}, ' 

58 f'batch_size={gpu_batch_size})') 

59 

60 def fn_batch(frames): 

61 for f in frames: 

62 f.rho_qs_dict = {} 

63 if calculate_currents: 

64 f.jz_qs_dict = {} 

65 f.jper_qs_dict = {} 

66 for atom_type in frames[0].positions_by_type: 

67 x_batch = np.stack([f.positions_by_type[atom_type] for f in frames]) 

68 if calculate_currents: 

69 v_batch = np.stack([f.velocities_by_type[atom_type] for f in frames]) 

70 rho_np, j_np = calc_rho_j_q_torch_batched(x_batch, v_batch, q_d) 

71 for b, frame in enumerate(frames): 

72 frame.rho_qs_dict[atom_type] = rho_np[b] 

73 j_qs = j_np[b] 

74 jz_qs = np.sum(j_qs * q_directions, axis=1) 

75 frame.jz_qs_dict[atom_type] = jz_qs 

76 frame.jper_qs_dict[atom_type] = j_qs - jz_qs[:, None] * q_directions 

77 else: 

78 rho_np = calc_rho_q_torch_batched(x_batch, q_d) 

79 for b, frame in enumerate(frames): 

80 frame.rho_qs_dict[atom_type] = rho_np[b] 

81 return frames 

82 

83 return True, fn_batch 

84 

85 elif backend == 'cupy': 85 ↛ 86line 85 didn't jump to line 86 because the condition on line 85 was never true

86 from dynasor.core.reciprocal_cupy import ( 

87 make_q_array, calc_rho_q_cupy_batched, calc_rho_j_q_cupy_batched 

88 ) 

89 q_d = make_q_array(q_points, precision=precision) 

90 logger.info(f'Using CuPy backend (precision={precision}, batch_size={gpu_batch_size})') 

91 

92 def fn_batch(frames): 

93 for f in frames: 

94 f.rho_qs_dict = {} 

95 if calculate_currents: 

96 f.jz_qs_dict = {} 

97 f.jper_qs_dict = {} 

98 for atom_type in frames[0].positions_by_type: 

99 x_batch = np.stack([f.positions_by_type[atom_type] for f in frames]) 

100 if calculate_currents: 

101 v_batch = np.stack([f.velocities_by_type[atom_type] for f in frames]) 

102 rho_np, j_np = calc_rho_j_q_cupy_batched(x_batch, v_batch, q_d) 

103 for b, frame in enumerate(frames): 

104 frame.rho_qs_dict[atom_type] = rho_np[b] 

105 j_qs = j_np[b] 

106 jz_qs = np.sum(j_qs * q_directions, axis=1) 

107 frame.jz_qs_dict[atom_type] = jz_qs 

108 frame.jper_qs_dict[atom_type] = j_qs - jz_qs[:, None] * q_directions 

109 else: 

110 rho_np = calc_rho_q_cupy_batched(x_batch, q_d) 

111 for b, frame in enumerate(frames): 

112 frame.rho_qs_dict[atom_type] = rho_np[b] 

113 return frames 

114 

115 return True, fn_batch 

116 

117 else: # numba: original per-frame path, unchanged 

118 def f2_rho(frame): 

119 rho_qs_dict = dict() 

120 for atom_type in frame.positions_by_type.keys(): 

121 x = frame.positions_by_type[atom_type] 

122 rho_qs_dict[atom_type] = calc_rho_q(x, q_points) 

123 frame.rho_qs_dict = rho_qs_dict 

124 return frame 

125 

126 def f2_rho_and_j(frame): 

127 rho_qs_dict = dict() 

128 jz_qs_dict = dict() 

129 jper_qs_dict = dict() 

130 

131 for atom_type in frame.positions_by_type.keys(): 

132 x = frame.positions_by_type[atom_type] 

133 v = frame.velocities_by_type[atom_type] 

134 rho_qs, j_qs = calc_rho_j_q(x, v, q_points) 

135 jz_qs = np.sum(j_qs * q_directions, axis=1) 

136 jper_qs = j_qs - (jz_qs[:, None] * q_directions) 

137 

138 rho_qs_dict[atom_type] = rho_qs 

139 jz_qs_dict[atom_type] = jz_qs 

140 jper_qs_dict[atom_type] = jper_qs 

141 

142 frame.rho_qs_dict = rho_qs_dict 

143 frame.jz_qs_dict = jz_qs_dict 

144 frame.jper_qs_dict = jper_qs_dict 

145 return frame 

146 

147 return (False, f2_rho_and_j) if calculate_currents else (False, f2_rho)