Coverage for gamdpy/calculators/calculator_radial_distribution.py: 78%
94 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import numpy as np
2import numba
3import math
4from numba import cuda
6from ..simulation.get_default_compute_plan import get_default_compute_plan
8#############################################################
9#### Radial Distribution Function
10#############################################################
12class CalculatorRadialDistribution():
13 """ Calculator class for the radial distribution function, g(r)
15 Parameters
16 ----------
18 configuration : gamdpy.Configuration
19 The configuration object for which the radial distribution function is calculated.
21 bins : int
22 The number of bins in the radial distribution function.
24 compute_plan : dict
26 ptype : optional
27 If specified, array ptypes used to calculate g(r). If not specfified default is configuration.ptype
29 Example
30 -------
32 >>> import gamdpy as gp
33 >>> sim = gp.get_default_sim()
34 >>> calc_rdf = gp.CalculatorRadialDistribution(sim.configuration, bins=1000)
35 >>> for _ in sim.run_timeblocks():
36 ... calc_rdf.update() # Current configuration to rdf
37 >>> rdf_data = calc_rdf.read() # Read the rdf data as a dictionary
38 >>> r = rdf_data['distances'] # Pair distances
39 >>> rdf = rdf_data['rdf'] # Radial distribution function
40 """
42 def __init__(self, configuration, bins, compute_plan=None, ptype=None) -> None:
43 self.configuration = configuration
44 self.d_ptype = cuda.to_device(ptype) if ptype is not None else None
45 self.bins = bins
46 self.count = 0 # How many times have statistics been added to?
48 self.compute_plan = compute_plan
49 if self.compute_plan is None:
50 self.compute_plan = get_default_compute_plan(configuration=configuration)
52 # Allocate space for statistics
53 self.rdf_list = []
54 nptypes = int(np.max(ptype if ptype is not None else configuration.ptype)) + 1
55 self.gr_bins = np.zeros((nptypes, nptypes, self.bins), dtype=np.float64)
56 self.d_gr_bins = cuda.to_device(self.gr_bins)
57 self.host_array_zeros = np.zeros(self.d_gr_bins.shape, dtype=self.d_gr_bins.dtype)
59 # Make kernel for updating statistics
60 self.update_kernel = self.make_updater_kernel(configuration, self.compute_plan)
62 def make_updater_kernel(self, configuration, compute_plan, verbose=False):
63 # Unpack parameters from configuration and compute_plan
64 D, num_part = configuration.D, configuration.N
65 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
66 num_blocks = (num_part - 1) // pb + 1
68 # Unpack indices for scalars to be compiled into kernel
69 r_id, = [configuration.vectors.indices[key] for key in ['r', ]]
71 # Prepare user-specified functions for inclusion in kernel(s)
72 ptype_function = numba.njit(configuration.ptype_function)
73 #params_function = numba.njit(pair_potential.params_function)
74 dist_sq_function = numba.njit(configuration.simbox.get_dist_sq_function())
76 def rdf_calculator_full(vectors, sim_box, ptype, d_gr_bins):
77 """ Calculate g(r) fresh
78 Kernel configuration: [num_blocks, (pb, tp)]
79 """
81 bins = d_gr_bins.shape[2] # reading number of bins from size of the device array
82 min_box_dim = min(sim_box[:D])
83 bin_width = (min_box_dim / 2) / bins # TODO: Chose more directly!
85 my_block = cuda.blockIdx.x
86 local_id = cuda.threadIdx.x
87 global_id = my_block * pb + local_id
88 my_t = cuda.threadIdx.y
90 if global_id < num_part:
91 ptype1 = ptype[global_id]
92 for i in range(0, num_part, pb * tp):
93 for j in range(pb):
94 other_global_id = j + i + my_t * pb
95 ptype2 = ptype[other_global_id]
96 if other_global_id != global_id and other_global_id < num_part:
97 dist_sq = dist_sq_function(vectors[r_id][other_global_id], vectors[r_id][global_id],
98 sim_box)
100 # Calculate g(r)
101 if dist_sq < (min_box_dim / 2) ** 2:
102 dist = math.sqrt(dist_sq)
103 if dist < min_box_dim / 2:
104 bin_index = int(dist / bin_width)
105 cuda.atomic.add(d_gr_bins, (ptype1, ptype2, bin_index), 1)
107 return
109 return cuda.jit(device=0)(rdf_calculator_full)[num_blocks, (pb, tp)]
111 def update(self):
112 """ Update the radial distribution function with the current configuration. """
113 self.count += 1
114 self.update_kernel(self.configuration.d_vectors,
115 self.configuration.simbox.d_data,
116 self.d_ptype if self.d_ptype is not None else self.configuration.d_ptype,
117 self.d_gr_bins)
118 self.rdf_list.append(self.d_gr_bins.copy_to_host())
119 self.d_gr_bins = cuda.to_device(self.host_array_zeros)
121 def read(self):
122 """ Read the radial distribution function
124 Returns
125 -------
127 dict
128 A dictionary containing the distances and the radial distribution function.
129 """
130 bins = self.rdf_list[0].shape[2]
131 min_box_dim = np.min(self.configuration.simbox.get_lengths())
132 bin_width = (min_box_dim / 2) / bins
133 rdf_ptype = np.array(self.rdf_list)
135 # Normalize the g(r) lengths # Compute in setup and normalize om the fly
136 rho = self.configuration.N / self.configuration.simbox.get_volume()
137 for i in range(bins): # Normalize one bin/distance at a time
138 r_outer = (i + 1) * bin_width
139 r_inner = i * bin_width
140 D = self.configuration.D
141 if D % 2 == 0:
142 n = D//2
143 unit_hypersphere_volume = math.pi**n/math.factorial(n)
144 else:
145 n = (D - 1) // 2
146 unit_hypersphere_volume = ( 2**D * math.pi**n * math.factorial(n) ) / math.factorial(D)
147 shell_volume = unit_hypersphere_volume * (r_outer**D - r_inner**D)
148 expected_num = rho * shell_volume
149 rdf_ptype[:, :, :, i] /= (expected_num * self.configuration.N)
150 rdf = rdf_ptype.sum(axis=1).sum(axis=1)
151 ptype = self.d_ptype.copy_to_host() if self.d_ptype is not None else self.configuration.d_ptype.copy_to_host()
152 for j in range(rdf_ptype.shape[1]):
153 n_j = np.sum(ptype == j) / len(ptype)
154 rdf_ptype[:, j, :, :] /= n_j
155 for k in range(rdf_ptype.shape[2]):
156 n_k = np.sum(ptype == k) / len(ptype)
157 rdf_ptype[:, :, k, :] /= n_k
158 distances = np.arange(0, bins) * bin_width
159 return {'distances': distances, 'rdf': rdf, 'rdf_ptype': rdf_ptype, "ptype": ptype}
161 def save_average(self, output_filename="rdf.dat", save_ptype=False) -> None:
162 """ Save the average radial distribution function to a file
164 Parameters
165 ----------
167 output_filename : str
168 The name of the file to which the radial distribution function is saved.
170 save_ptype : bool
171 Save the type of each particle in a file name ptype_* (default False)
173 """
175 rdf_dict = self.read()
176 rdf_total = np.mean(rdf_dict['rdf'], axis=0)
177 rdf_ij = []
178 header_ij = " "
179 for i in range(rdf_dict["rdf_ptype"].shape[1]):
180 for j in range(rdf_dict["rdf_ptype"].shape[2]):
181 rdf_ij.append(np.mean(rdf_dict['rdf_ptype'][:, i, j, :], axis=0))
182 header_ij += f"g[{i}-{j}](r) "
183 np.savetxt(output_filename, np.array([rdf_dict['distances'], rdf_total, *rdf_ij]).T, header=f"r g(r) {header_ij} ptype")
184 np.savetxt(f"ptype_{output_filename}", np.array(rdf_dict["ptype"]).T, header="ptype")