Coverage for gamdpy/interactions/bonds.py: 63%
102 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
5from .make_fixed_interactions import make_fixed_interactions # bonds is an example of 'fixed' interactions
7# Abstract Base Class and type annotation
8from .interaction import Interaction
9from gamdpy import Configuration
12class Bonds(Interaction):
13 """ Fixed bond interactions between particles, such as harmonic bonds or FENE bonds.
15 Parameters
16 ----------
18 bond_potential : function
19 A function that takes the distance between two particles and the bond type as arguments and returns the potential energy, force and laplacian.
20 See :func:gamdpy.potential_functions.harmonic_bond_function for an example.
22 potential_params : list
23 A list of parameters for each bond type. Each entry is a list of parameters for a specific bond type.
25 indices : list
26 A list of lists, each containing the indices of the two particles involved in a bond and the bond
28 See Also
29 --------
31 gamdpy.harmonic_bond_function : Harmonic bond potential
33 """
34 def __init__(self, bond_potential, potential_params, indices):
35 self.bond_potential = bond_potential
36 self.potential_params = potential_params
37 self.indices = indices
39 def get_params(self, configuration: Configuration, compute_plan: dict, verbose=False) -> tuple:
40 self.N = configuration.N
41 self.potential_params_array = np.array(self.potential_params, dtype=np.float32)
42 assert len(self.potential_params_array.shape)==2 # for now...
43 self.num_types = self.potential_params_array.shape[0]
44 self.indices_array = np.array(self.indices, dtype=np.int32)
45 assert self.indices_array.shape[1] == 3 # i, j, bond_type
46 assert max(self.indices_array[:,-1]) <= self.num_types
48 if verbose:
49 print(f'Setting up bond interactions: {self.N} particles,')
50 print(f'{self.num_types} bond types, {self.indices_array.shape[0]} bonds in total.')
52 self.d_indices = cuda.to_device(self.indices_array)
53 self.d_params = cuda.to_device(self.potential_params_array)
54 return (self.d_indices, self.d_params)
57 def get_kernel(self, configuration: Configuration, compute_plan: dict, compute_flags: dict[str,bool], verbose=False):
58 # Unpack parameters from configuration and compute_plan
59 D, N = configuration.D, configuration.N
60 pb, tp, gridsync, UtilizeNIII = [compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']]
61 num_blocks = (N - 1) // pb + 1
63 compute_u = compute_flags['U']
64 compute_w = compute_flags['W']
65 compute_lap = compute_flags['lapU']
66 compute_stresses = compute_flags['stresses']
68 if verbose:
69 print('get_kernel: Bonds:')
70 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}')
71 print(f'\tNumber (virtual) particles: {num_blocks*pb}')
72 print(f'\tNumber of threads {num_blocks*pb*tp}')
73 if compute_stresses:
74 print('\tIncluding computation of stress tensor')
76 # Unpack indices for vectors and scalars to be compiled into kernel
77 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']]
79 if compute_u:
80 u_id = configuration.sid['U']
81 if compute_w:
82 w_id = configuration.sid['W']
83 if compute_lap:
84 lap_id = configuration.sid['lapU']
86 if compute_stresses:
87 sx_id = configuration.vectors.indices['sx']
88 if D > 1:
89 sy_id = configuration.vectors.indices['sy']
90 if D > 2:
91 sz_id = configuration.vectors.indices['sz']
92 if D > 3:
93 sw_id = configuration.vectors.indices['sw']
96 dist_sq_dr_function = numba.njit(configuration.simbox.get_dist_sq_dr_function())
97 bondpotential_function = numba.njit(self.bond_potential)
99 virial_factor = numba.float32( 0.5/configuration.D)
100 half = numba.float32(0.5)
102 def bond_calculator(vectors, scalars, ptype, sim_box, indices, values):
104 dr = cuda.local.array(shape=D,dtype=numba.float32)
105 dist_sq = dist_sq_dr_function(vectors[r_id][indices[1]], vectors[r_id][indices[0]], sim_box, dr)
106 u, s, umm = bondpotential_function(math.sqrt(dist_sq), values[indices[2]])
108 for k in range(D):
109 cuda.atomic.add(vectors, (f_id, indices[0], k), -dr[k]*s) # Force
110 cuda.atomic.add(vectors, (f_id, indices[1], k), +dr[k]*s)
111 if compute_w:
112 cuda.atomic.add(scalars, (indices[0], w_id), dr[k]*dr[k]*s*virial_factor) # Virial
113 cuda.atomic.add(scalars, (indices[1], w_id), dr[k]*dr[k]*s*virial_factor)
114 if compute_stresses:
115 cuda.atomic.add(vectors[sx_id], (indices[0], k), - half * s * dr[0] * dr[k])
116 cuda.atomic.add(vectors[sx_id], (indices[1], k), - half * s * dr[0] * dr[k])
117 if D > 1:
118 cuda.atomic.add(vectors[sy_id], (indices[0], k), - half * s * dr[1] * dr[k])
119 cuda.atomic.add(vectors[sy_id], (indices[1], k), - half * s * dr[1] * dr[k])
120 if D > 2:
121 cuda.atomic.add(vectors[sz_id], (indices[0], k), - half * s * dr[2] * dr[k])
122 cuda.atomic.add(vectors[sz_id], (indices[1], k), - half * s * dr[2] * dr[k])
123 if D > 3:
124 cuda.atomic.add(vectors[sw_id], (indices[0], k), - half * s * dr[3] * dr[k])
125 cuda.atomic.add(vectors[sw_id], (indices[1], k), - half * s * dr[3] * dr[k])
128 if compute_u:
129 cuda.atomic.add(scalars, (indices[0], u_id), u*numba.float32(0.5)) # Potential enerrgy
130 cuda.atomic.add(scalars, (indices[1], u_id), u*numba.float32(0.5))
131 if compute_lap:
132 lap = numba.float32(1-D)*s + umm # Laplacian
133 cuda.atomic.add(scalars, (indices[0], lap_id), lap)
134 cuda.atomic.add(scalars, (indices[1], lap_id), lap)
137 return
139 return make_fixed_interactions(configuration, bond_calculator, compute_plan, verbose=False)
141 def get_exclusions(self, configuration, max_number_exclusions=20):
142 exclusions = np.zeros((configuration.N, max_number_exclusions+1), dtype=np.int32) # last entry: number of actual exclusions
143 for i, j, _ in self.indices: # bond involving particles 'i' and 'j'. We dont care about the bond type (maybe later)
144 if exclusions[i,-1] < max_number_exclusions:
145 exclusions[i,exclusions[i,-1]] = j
146 exclusions[i,-1] += 1
148 if exclusions[j,-1] < max_number_exclusions:
149 exclusions[j,exclusions[j,-1]] = i
150 exclusions[j,-1] += 1
152 assert np.max(exclusions[:,-1]) <= max_number_exclusions
153 return exclusions