Coverage for gamdpy/calculators/calculator_structure_factor.py: 76%
142 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import itertools
3import numpy as np
5import gamdpy as gp
6import numba
9class CalculatorStructureFactor:
10 """ Calculator class for the static structure factor, S(q).
11 The calculation is done for several :math:`{\\bf q}` vectors given by
13 .. math::
15 {\\bf q} = (2\\pi n_x/L_x, 2\\pi n_y/L_y, ...)
17 where :math:`n=(n_x, n_y, ...)` is a D-dimensional vector of integers and
18 :math:`L_x`, :math:`L_y`, ... are the box lengths in the :math:`x`, :math:`y`, ... directions, respectively.
19 Note that box length are assumed to be constant during the simulation (as in a NVT simulation).
20 The collective density :math:`\\rho_{\\bf q}` is calculated as
22 .. math::
24 \\rho_{\\bf q} = \\frac{1}{\\sqrt{N}} \\sum_{n} f_n \\exp(-i {\\bf q}\\cdot {\\bf r}_n)
26 where :math:`x_n` is the position of particle :math:`n`, and :math:`f_n` is the atomic form factor for that particle
27 (one by default). The normalization constant is
29 .. math::
31 N = \\sum_{n} f_n.
33 From this, the static structure factor is defined as
35 .. math::
37 S({\\bf q}) = |\\rho_{\\bf q}|^2.
39 The method :meth:`~gamdpy.calculators.CalculatorStructureFactor.update`
40 updates the structure factor with the current configuration.
41 The method :meth:`~gamdpy.calculators.CalculatorStructureFactor.read` returns the structure factor for the q vectors in the q_direction.
43 Parameters
44 ----------
46 configuration : gamdpy.Configuration
47 The configuration object to calculate the structure factor for.
49 q_max : float or None
50 The maximum value of the q vectors.
52 n_vectors : numpy.ndarray or None
53 n-vectors defining q-vectors.
54 The shape of n_vectors, if specified, must be (N, D)
55 where N is the number of q vectors and D is the number of dimensions.
56 If None, then use generate_q_vectors method.
58 atomic_form_factors : numpy.ndarray or None
59 The atomic form factors, :math:`f_n`. If None (default), then the atomic form factors are set to 1.
60 Can be given as an array of floats, one for each atom.
62 backend : str
63 The backend to use for the calculation. Either 'CPU multi core' or 'CPU single core'.
65 See also
66 --------
68 :class:`~gamdpy.CalculatorRadialDistribution`
70 """
72 BACKENDS = ['CPU multi core', 'CPU single core', 'GPU']
74 def __init__(self,
75 configuration: gp.Configuration,
76 n_vectors: np.ndarray = None,
77 atomic_form_factors: np.ndarray = None,
78 backend='CPU multi core') -> None:
79 if backend not in self.BACKENDS:
80 raise ValueError(f'Unknown backend, {backend}. The known backends are {self.BACKENDS}.')
81 self.update_count = 0
82 self.configuration = configuration
83 self.L = self.configuration.simbox.get_lengths()
85 if n_vectors is not None:
86 # n_vectors = [[0, 0, 1], [0, 0, 2], ..., [0, 1, 0], [0, 1, 1] ..., [18, 18, 18], ...]
87 self.n_vectors = np.array(n_vectors)
88 dimension_of_space = self.configuration.D
89 if self.n_vectors.shape[1] != dimension_of_space:
90 raise ValueError('n_vectors must have the same number of columns as the number of dimensions.')
91 self.q_vectors = np.array(2 * np.pi * self.n_vectors / self.L, dtype=np.float32)
92 self.q_lengths = np.linalg.norm(self.q_vectors, axis=1)
93 self.sum_S_q = np.zeros_like(self.q_lengths)
95 self.atomic_form_factors = atomic_form_factors
96 if atomic_form_factors is None:
97 number_of_atoms = self.configuration.N
98 self.atomic_form_factors = np.ones(number_of_atoms, dtype=np.float32)
100 # List for storing data
101 self.list_of_rho_q = []
102 self.list_of_rho_S_q = []
104 self.wallclock_times = []
106 # 3 first letters is CPU or GPU
107 if backend[:3] == 'CPU':
108 self._compute_rho_q = self._generate_compute_rho_q(backend)
109 if backend == 'GPU':
110 self.nthreads = 64
111 self.update_kernel = self._make_update_kernel()
112 self._compute_rho_q = self._compute_rho_q_gpu
114 def generate_q_vectors(self, q_max:float):
115 """ Generate q-vectors inside a sphere of radius q_max """
116 dimension_of_space = self.configuration.D
117 if q_max<0.0:
118 raise ValueError(f'{q_max=} must be positive')
119 n_max = int(np.ceil(q_max * max(self.L) / (2 * np.pi)))
120 n_vectors = np.array(list(itertools.product(range(n_max), repeat=dimension_of_space)), dtype=int)
121 n_vectors = n_vectors[1:] # Remove the first vector [0, 0, 0]
122 self.q_vectors = np.array(2 * np.pi * n_vectors / self.L, dtype=np.float32)
124 # Remove q_vectors where the length is greater than q_max
125 selection = np.linalg.norm(self.q_vectors, axis=1) < q_max
126 self.q_vectors = self.q_vectors[selection]
127 self.n_vectors = n_vectors[selection]
128 self.q_lengths = np.linalg.norm(self.q_vectors, axis=1)
129 self.sum_S_q = np.zeros_like(self.q_lengths)
131 def _make_update_kernel(self):
132 import math
133 from numba import cuda
134 from numba import float32 as flt
136 def kernel(r_vectors, q_vectors, form_factors, rho_q_real, rho_q_imag):
137 num = q_vectors.shape[0] # Number of q vectors
138 tid = cuda.grid(1)
140 if tid < num:
141 this_q = q_vectors[tid]
142 real, imag = flt(0.0), flt(0.0)
143 N = flt(0.0)
144 for n in range(r_vectors.shape[0]):
145 pos = r_vectors[n]
146 dot = flt(0.0)
147 for d in range(r_vectors.shape[1]):
148 dot += pos[d] * this_q[d]
149 real += form_factors[n]*math.cos(dot)
150 imag += form_factors[n]*math.sin(dot)
151 N += form_factors[n]
152 NN = flt(1.0)/math.sqrt(N)
153 rho_q_real[tid] = NN*real
154 rho_q_imag[tid] = NN*imag
155 return
156 return cuda.jit(device=0)(kernel)
158 def _compute_rho_q_gpu(self, r_vectors, q_vectors, form_factors):
159 from numba import cuda
161 # Copy data to device
162 d_r_vectors = cuda.to_device(r_vectors) #, q_vectors, form_factors, rho_q_real, rho_q_imag
163 d_q_vectors = cuda.to_device(q_vectors)
164 d_form_factors = cuda.to_device(form_factors)
165 rho_q_real = np.zeros(q_vectors.shape[0], dtype=np.float32)
166 d_rho_q_real = cuda.to_device(rho_q_real)
167 d_rho_q_imag = cuda.device_array_like(d_rho_q_real)
169 # Compute using GPU kernel
170 start = numba.cuda.event()
171 end = numba.cuda.event()
172 num = self.q_vectors.shape[0]
173 nblocks = (num // self.nthreads) + 1
174 start.record()
175 self.update_kernel[nblocks, self.nthreads](d_r_vectors, d_q_vectors, d_form_factors, d_rho_q_real, d_rho_q_imag)
176 end.record()
177 end.synchronize()
178 self.wallclock_times.append(start.elapsed_time(end))
179 rho_q_real = d_rho_q_real.copy_to_host()
180 rho_q_imag = d_rho_q_imag.copy_to_host()
182 return rho_q_real + 1j*rho_q_imag
185 @staticmethod
186 def _generate_compute_rho_q(backend):
187 if backend == 'CPU multi core': # May raise "ImportError: scipy 0.16+ is required for linear algebra" if scipy is not installed
188 def func(r_vec: np.ndarray, q_vec: np.ndarray, form_factors: np.ndarray) -> np.ndarray:
189 N = np.sum(form_factors)
190 number_of_q_vectors = q_vec.shape[0]
191 rho_q = np.zeros(number_of_q_vectors, dtype=np.complex64)
192 for i in numba.prange(number_of_q_vectors):
193 r_dot_q = np.dot(r_vec, q_vec[i])
194 rho_q[i] = np.sum(form_factors*np.exp(1j * r_dot_q))*N**-0.5
195 return rho_q
196 return numba.njit(parallel=True)(func)
197 elif backend == 'CPU single core':
198 def func(r_vec: np.ndarray, q_vec: np.ndarray, form_factors: np.ndarray) -> np.ndarray:
199 N = np.sum(form_factors)
200 r_dot_q = np.dot(r_vec, q_vec.T)
201 rho_q = np.sum(form_factors[:, np.newaxis]*np.exp(1j * r_dot_q), axis=0)*N**-0.5
202 return rho_q
203 return numba.njit(func)
205 def update(self) -> None:
206 """ Update the structure factor with the current configuration. """
207 if not np.allclose(self.L, self.configuration.simbox.get_lengths()):
208 raise ValueError('Box length has changed. Recreate the S(q) object.')
209 this_rho_q = self._compute_rho_q(self.configuration['r'], self.q_vectors, self.atomic_form_factors)
210 self.list_of_rho_q.append(this_rho_q)
211 self.list_of_rho_S_q.append(np.abs(this_rho_q)**2)
212 self.sum_S_q += np.abs(this_rho_q) ** 2
213 self.update_count += 1
215 #def read(self, bins: int | None) -> dict:
216 def read(self, bins) -> dict:
217 """ Return the structure factor S(q) for the q vectors in the q_direction.
219 Parameters
220 ----------
222 bins : int | None
223 If bins is an integer, the data is binned (ready to be plotted).
224 If bins is None, the raw S(q) data is returned.
226 Returns
227 -------
229 dict
230 A dictionary containing the q vectors, the q lengths, the structure factor S(q),
231 the collective density rho_q, and the number of q vectors in each bin. Output depends on the value of
232 the bins parameter.
233 """
234 if isinstance(bins, int):
235 q_bins = np.linspace(0, np.max(self.q_lengths), bins+1)
236 q_binned = np.zeros(bins, dtype=np.float32)
237 S_q_binned = np.zeros(bins, dtype=np.float32)
238 q_vectors_in_bin = np.zeros(bins, dtype=int)
239 for i in range(0, bins):
240 mask = (q_bins[i] <= self.q_lengths) & (self.q_lengths < q_bins[i+1])
241 if np.sum(mask) == 0: # No q vectors in this bin
242 continue
243 q_binned[i] = np.mean(self.q_lengths[mask])
244 # Use self.list_of_rho_S_q
245 S_q_binned[i] = np.mean(self.sum_S_q[mask] / self.update_count)
246 q_vectors_in_bin[i] = np.sum(mask)
247 # Remove bins with no q vectors
248 mask = q_vectors_in_bin > 0
249 q_binned = q_binned[mask]
250 S_q_binned = S_q_binned[mask]
251 q_vectors_in_bin = q_vectors_in_bin[mask]
252 return {
253 '|q|': q_binned,
254 'S(|q|)': S_q_binned,
255 'q_vectors_in_bin': q_vectors_in_bin
256 }
257 elif bins is None:
258 # Return (un-binned) raw data
259 return {
260 'q': self.q_vectors,
261 '|q|': self.q_lengths,
262 'S(q)': self.sum_S_q/self.update_count,
263 'rho_q': np.array(self.list_of_rho_q),
264 'n_vectors': self.n_vectors,
265 'atomic_form_factors': self.atomic_form_factors
266 }
267 else:
268 raise ValueError('bins must be an integer.')
270 #def save_average(self, bins: int=None, output_filename: str="sq.dat") -> None:
271 def save_average(self, bins=None, output_filename="sq.dat"):
272 """ Save average structure factors to a file. """
273 if bins is None: bins=100
274 sq_dict = self.read(bins)
275 np.savetxt(output_filename, np.c_[sq_dict['|q|'], sq_dict['S(|q|)']], header="|q| S(|q|)")