Coverage for gamdpy/interactions/relaxtemp.py: 64%

1import math 

2 

3import numba 

4import numpy as np 

5from numba import cuda 

6 

7from .make_fixed_interactions import make_fixed_interactions # tether is an example of 'fixed' interactions 

8 

9# Abstract Base Class and type annotation 

10from .interaction import Interaction 

11from gamdpy import Configuration 

12 

13class Relaxtemp(Interaction): 

14 """ Thermostat using simple kinetic temperature relaxation of each particles.  

15  

16 Parameters 

17 ---------- 

18 (a) A list of particle indices to be thermostated and associated list of relaxation time  

19 (b) A list of particle types to be thermostated and associated list of relaxation times 

20 

21 See examples/poiseuille.py 

22 """ 

23 

24 def __init__(self): 

25 self.indices_set = False 

26 

27 

28 def set_relaxation_from_lists(self, particle_indices, temperature, relax_times): 

29 

30 ntau, npart, ntemp = len(tau), len(pindices), len(temperature) 

31 

32 if ntau != npart or ntau != ntemp or npart != ntemp: 

33 raise ValueError( 

34 "Each particle must have exactly one relax time and temperature - arrays must be same length") 

35 

36 indices, params = [], [] 

37 for n in range(npart): 

38 indices.append([n, particle_indices[n]]) 

39 params.append([temperature(n), relax_times[n]]) 

40 

41 self.relax_params = np.array(params, dtype=np.float32) 

42 self.indices = np.array(indices, dtype=np.int32) 

43 

44 self.indices_set = True 

45 

46 

47 def set_relaxation_from_types(self, particle_types, temperature, relax_times, configuration): 

48 

49 ntypes, ntau, ntemp = len(particle_types), len(relax_times), len(temperature) 

50 

51 if ntypes != ntau or ntypes != ntemp or ntemp != ntau: 

52 raise ValueError("Each type must have exactly one relax time - arrays must be same length") 

53 

54 indices, params = [], [] 

55 counter = 0 

56 for n in range(configuration.N): 

57 for m in range(ntypes): 

58 if configuration.ptype[n] == particle_types[m]: 

59 indices.append([counter, n]) 

60 params.append([temperature[m], relax_times[m]]) 

61 counter = counter + 1 

62 break 

63 

64 self.relax_params = np.array(params, dtype=np.float32) 

65 self.indices = np.array(indices, dtype=np.int32) 

66 

67 self.indices_set = True 

68 

69 def get_params(self, configuration: Configuration, compute_plan: dict, verbose=False) -> tuple: 

70 

71 if self.indices_set == False: 

72 raise ValueError("Indices not defined") 

73 

74 

75 self.d_pindices = cuda.to_device(self.indices) 

76 self.d_relax_params = cuda.to_device(self.relax_params); 

77 

78 return (self.d_pindices, self.d_relax_params) 

79 

80 def get_kernel(self, configuration: Configuration, compute_plan: dict, compute_flags: dict[str,bool], verbose=False): 

81 # Unpack parameters from configuration and compute_plan 

82 D, N = configuration.D, configuration.N 

83 pb, tp, gridsync, UtilizeNIII = [compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']] 

84 num_blocks = (N - 1) // pb + 1 

85 

86 # Not sure if compute_flags is relevant here?? NB, Nov 2024 

87 

88 # Get indices values (instead of dictonary entries)  

89 v_id = configuration.vectors.indices['v'] 

90 m_id = configuration.sid['m'] 

91 

92 def relaxtemp_calculator(vectors, scalars, ptype, sim_box, indices, values): 

93 v = vectors[v_id][indices[1]] 

94 m = scalars[indices[1]][m_id] 

95 Tdesired = values[indices[0]][0] 

96 tau = values[indices[0]][1] 

97 

98 Tparticle = m / numba.float32(3.0) * (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]) 

99 

100 one = numba.float32(1.0) 

101 fac = math.sqrt(one + tau * (Tdesired / Tparticle - one)) 

102 

103 for k in range(D): 

104 v[k] = v[k] * fac 

105 

106 return 

107 

108 return make_fixed_interactions(configuration, relaxtemp_calculator, compute_plan, verbose=False)