Coverage for gamdpy/tools/Evaluator.py: 82%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda 

5 

6from ..simulation.get_default_compute_flags import get_default_compute_flags 

7 

8# gamdpy 

9import gamdpy as gp 

10 

11class Evaluator: 

12 """ Evaluates interactions between particles in a configuration. 

13  

14 This class can be used to evaluate interactions between particles in a configuration 

15 that are different from the ones of the Simulation class. 

16 

17 Parameters 

18 ---------- 

19 

20 configuration : rumd.Configuration 

21 The configuration for which the interactions are to be evaluated. 

22 

23 interactions : an interaction 

24 Interactions such as pair potentials, bonds, external fields, etc. 

25 

26 compute_plan : dict, optional 

27 A dictionary with the compute plan for the simulation. If None, a default compute plan is used. 

28 

29 verbose : bool, optional 

30 If True, print information about the interactions. 

31  

32 """ 

33 def __init__(self, configuration, interactions, compute_plan=None, verbose=True): 

34 

35 self.configuration = configuration 

36 

37 if compute_plan==None: 

38 self.compute_plan = gp.get_default_compute_plan(self.configuration) 

39 else: 

40 self.compute_plan = compute_plan 

41 

42 compute_flags = get_default_compute_flags() 

43 

44 # Make sure interactions is a list 

45 if type(interactions) == list: 

46 self.interactions = interactions 

47 else: 

48 self.interactions = [interactions, ] 

49 

50 self.interactions_kernel, self.interactions_params = gp.add_interactions_list( 

51 self.configuration, 

52 self.interactions, 

53 compute_plan=self.compute_plan, 

54 compute_flags=compute_flags, 

55 verbose=verbose) 

56 

57 self.evaluater_func = self.make_evaluater_func(self.configuration, self.interactions_kernel, self.compute_plan, verbose) 

58 

59 def make_evaluater_func(self, configuration, compute_interactions, compute_plan, verbose=True): 

60 D, num_part = configuration.D, configuration.N 

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

62 num_blocks = (num_part - 1) // pb + 1 

63 

64 if gridsync: 

65 # Return a kernel that evaluates interactions  

66 @cuda.jit 

67 def evaluater(vectors, scalars, ptype, r_im, sim_box, interaction_params): 

68 grid = cuda.cg.this_grid() 

69 compute_interactions(grid, vectors, scalars, ptype, sim_box, interaction_params) 

70 return 

71 return evaluater[num_blocks, (pb, tp)] 

72 

73 else: 

74 

75 # Return a Python function that evaluates interactions 

76 def evaluater(vectors, scalars, ptype, r_im, sim_box, interaction_params): 

77 compute_interactions(0, vectors, scalars, ptype, sim_box, interaction_params) 

78 return 

79 return evaluater 

80 return 

81 

82 def evaluate(self, configuration=None): 

83 

84 if configuration==None: 

85 configuration = self.configuration 

86 configuration.copy_to_device() 

87 self.evaluater_func(configuration.d_vectors, 

88 configuration.d_scalars, 

89 configuration.d_ptype, 

90 configuration.d_r_im, 

91 configuration.simbox.d_data, 

92 self.interactions_params, 

93 ) 

94 configuration.copy_to_host() 

95 return