-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathMolecularDynamics.py
115 lines (94 loc) · 5.63 KB
/
MolecularDynamics.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
class MolecularDynamics:
def __init__(self, n_systems=1, n_atoms=0, n_node=0, n_pbc=0):
self.n_systems = n_systems
self.n_atoms = n_atoms
self.n_node = n_node
self.n_pbc = max(1, n_pbc)
self.n_vecs = n_atoms + n_node
self.ctx = cl.create_some_context()
self.queue = cl.CommandQueue(self.ctx)
self.buffers = {}
self.kernels = {}
self.init_buffers()
self.init_kernels()
def init_buffers(self):
mf = cl.mem_flags
self.buffers['atoms'] = cl.Buffer(self.ctx, mf.READ_WRITE, size=self.n_systems * self.n_vecs * 4 * 4)
self.buffers['forces'] = cl.Buffer(self.ctx, mf.READ_WRITE, size=self.n_systems * self.n_vecs * 4 * 4)
self.buffers['REQs'] = cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_atoms * 4 * 4)
self.buffers['neighs'] = cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_atoms * 4 * 4)
self.buffers['neighCell'] = cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_atoms * 4 * 4)
# Add more buffers as needed
def init_kernels(self):
# Load and compile OpenCL program
with open('kernels.cl', 'r') as f:
prg = cl.Program(self.ctx, f.read()).build()
self.kernels['getMMFFf4'] = prg.getMMFFf4
self.kernels['getNonBond'] = prg.getNonBond
self.kernels['updateAtomsMMFFf4'] = prg.updateAtomsMMFFf4
# Set up kernel arguments
self.setup_kernel_args()
def setup_kernel_args(self):
# Set up arguments for each kernel
pass
def run_simulation(self, n_steps):
for _ in range(n_steps):
self.calculate_forces()
self.update_positions()
def calculate_forces(self):
# Enqueue kernels for force calculation
pass
def update_positions(self):
# Enqueue kernel for position update
pass
def setup_buffers(self):
mf = cl.mem_flags
self.buffers = {
'atoms': cl.Buffer(self.ctx, mf.READ_WRITE, size=self.n_systems * self.n_vecs * 4 * 4),
'forces': cl.Buffer(self.ctx, mf.READ_WRITE, size=self.n_systems * self.n_vecs * 4 * 4),
'velocities': cl.Buffer(self.ctx, mf.READ_WRITE, size=self.n_systems * self.n_vecs * 4 * 4),
'REQs': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_atoms * 4 * 4),
'neighs': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_atoms * 4 * 4),
'neighCell': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_atoms * 4 * 4),
'MMpars': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_node * 4 * 4),
'bLs': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_node * 4 * 4),
'bKs': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_node * 4 * 4),
'Ksp': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_node * 4 * 4),
'Kpp': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_node * 4 * 4),
'lvecs': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * 3 * 3 * 4),
'ilvecs': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * 3 * 3 * 4),
'pbcshifts': cl.Buffer(self.ctx, mf.READ_ONLY, size=self.n_systems * self.n_pbc * 4 * 4),
}
def pack_system(self, system):
# Pack system data into numpy arrays
atoms = np.array(system.apos, dtype=np.float32)
REQs = np.array(system.REQs, dtype=np.float32)
neighs = np.array(system.neighs, dtype=np.int32)
neighCell = np.array(system.neighCell, dtype=np.int32)
MMpars = np.array(system.apars, dtype=np.float32)
bLs = np.array(system.bLs, dtype=np.float32)
bKs = np.array(system.bKs, dtype=np.float32)
Ksp = np.array(system.Ksp, dtype=np.float32)
Kpp = np.array(system.Kpp, dtype=np.float32)
lvec = np.array(system.lvec, dtype=np.float32)
ilvec = np.array(system.invLvec, dtype=np.float32)
pbcshifts = np.array(system.pbc_shifts, dtype=np.float32)
return atoms, REQs, neighs, neighCell, MMpars, bLs, bKs, Ksp, Kpp, lvec, ilvec, pbcshifts
def upload_system(self, system, index):
data = self.pack_system(system)
offset = index * self.n_vecs
cl.enqueue_copy(self.queue, self.buffers['atoms'], data[0], device_offset=offset*4*4)
cl.enqueue_copy(self.queue, self.buffers['REQs'], data[1], device_offset=index*self.n_atoms*4*4)
cl.enqueue_copy(self.queue, self.buffers['neighs'], data[2], device_offset=index*self.n_atoms*4*4)
cl.enqueue_copy(self.queue, self.buffers['neighCell'], data[3], device_offset=index*self.n_atoms*4*4)
cl.enqueue_copy(self.queue, self.buffers['MMpars'], data[4], device_offset=index*self.n_node*4*4)
cl.enqueue_copy(self.queue, self.buffers['bLs'], data[5], device_offset=index*self.n_node*4*4)
cl.enqueue_copy(self.queue, self.buffers['bKs'], data[6], device_offset=index*self.n_node*4*4)
cl.enqueue_copy(self.queue, self.buffers['Ksp'], data[7], device_offset=index*self.n_node*4*4)
cl.enqueue_copy(self.queue, self.buffers['Kpp'], data[8], device_offset=index*self.n_node*4*4)
cl.enqueue_copy(self.queue, self.buffers['lvecs'], data[9], device_offset=index*3*3*4)
cl.enqueue_copy(self.queue, self.buffers['ilvecs'], data[10], device_offset=index*3*3*4)
cl.enqueue_copy(self.queue, self.buffers['pbcshifts'], data[11], device_offset=index*self.n_pbc*4*4)
def upload_all_systems(self, systems):
for i, system in enumerate(systems):
self.upload_system(system, i)