From 75ef8db0d058158c10951a4184186f8c6cc27acc Mon Sep 17 00:00:00 2001 From: Adrian Kummerlaender Date: Sat, 27 Mar 2021 22:35:43 +0100 Subject: Initial public commit of this basic MD code Simulation of interacting particles --- interacticle/simulation.py | 394 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 394 insertions(+) create mode 100644 interacticle/simulation.py (limited to 'interacticle/simulation.py') diff --git a/interacticle/simulation.py b/interacticle/simulation.py new file mode 100644 index 0000000..6da1dd1 --- /dev/null +++ b/interacticle/simulation.py @@ -0,0 +1,394 @@ +import numpy as np +import scipy.linalg as sp + +import pyopencl as cl +mf = cl.mem_flags +from pyopencl.tools import get_gl_sharing_context_properties + +import OpenGL.GL as gl +from OpenGL.arrays import vbo + +from string import Template + +from collections import namedtuple, defaultdict +from itertools import product + +Atom = namedtuple('Atom', 'x y z vx vy vz mass') + +Bond = namedtuple('Bond', 'i j kb k0') +Angle = namedtuple('Angle', 'i j k theta0 r0') +Torsion = namedtuple('Torsion', 'i j k l phi0 r0') + +LennardJones = namedtuple('LennardJones', 'mass sigma epsilon') +Coulomb = namedtuple('Coulomb', 'mass charge') + +def rotationMatrix(axis, theta): + return sp.expm(np.cross(np.eye(3), axis/sp.norm(axis)*np.radians(theta))) + +class Molecule: + def __init__(self): + self.atoms = [ ] + self.connections = [ ] + +class MoleculeCollection: + def __init__(self): + self.index = 0 + self.current_molecule = 0 + self.indices_of_molecule = defaultdict(lambda: [ ]) + + self.atoms = [ ] + + self.molecules = [ ] + + self.neighborhood_step = 100 + self.statistics_step = 100 + + self.max_lennard_jones = 100 + self.max_coulomb = 100 + + self.lennard_jones = defaultdict(lambda: None) + self.coulomb = defaultdict(lambda: None) + + self.bonds = defaultdict(lambda: [ ]) + self.angles = defaultdict(lambda: [ ]) + self.torsions = defaultdict(lambda: [ ]) + + self.cutoff = 0.242*2.5 + self.skin = 0.1 + + self.tau = 0.0005 + self.domain_size = 2 + self.target_temperature = None + + def add(self, molecule, origin=(0,0,0), rotations=None): + for i, atom in enumerate(molecule.atoms): + x, y, z = atom.x, atom.y, atom.z + if rotations: + for rotation in rotations: + x, y, z = np.dot(rotationMatrix(*rotation), [x, y, z]) + self.atoms.append(Atom(origin[0]+x, origin[1]+y, origin[2]+z, atom.vy, atom.vy, atom.vz, atom.mass)) + self.indices_of_molecule[self.current_molecule].append(self.index+i) + self.molecules.append(self.current_molecule) + + for c in molecule.connections: + if isinstance(c, Bond): + self.bonds[self.index+c.i].append(Bond(self.index+c.i, self.index+c.j, c.kb, c.k0)) + self.bonds[self.index+c.j].append(Bond(self.index+c.j, self.index+c.i, c.kb, c.k0)) + elif isinstance(c, Angle): + tmp = Angle(self.index+c.i, self.index+c.j, self.index+c.k, + c.theta0, c.r0) + self.angles[self.index+c.i].append(tmp) + self.angles[self.index+c.j].append(tmp) + self.angles[self.index+c.k].append(tmp) + elif isinstance(c, Torsion): + tmp = Torsion(self.index+c.i, self.index+c.j, self.index+c.k, self.index+c.l, + c.phi0, c.r0) + self.torsions[self.index+c.i].append(tmp) + self.torsions[self.index+c.j].append(tmp) + self.torsions[self.index+c.k].append(tmp) + self.torsions[self.index+c.l].append(tmp) + else: + raise Exception('Unknown connection type') + + self.index = len(self.atoms) + self.current_molecule = self.current_molecule + 1 + + def potential(self, config): + if isinstance(config, LennardJones): + self.lennard_jones[config.mass] = config + elif isinstance(config, Coulomb): + self.coulomb[config.mass] = config + + def serialize(self): + n_atoms = len(self.atoms) + atom_pos, atom_vel = np.zeros((n_atoms,4), dtype=np.float32), np.zeros((n_atoms,4), dtype=np.float32) + for i, atom in enumerate(self.atoms): + atom_pos[i,:] = [atom.x, atom.y, atom.z, atom.mass] + atom_vel[i,:] = [atom.vx, atom.vy, atom.vz, 0] + + molecules = np.array(self.molecules, dtype=np.uint32) + + max_bonds = max(map(lambda x: len(x[1]), self.bonds.items()), default=0) + n_bonds = n_atoms * max_bonds + bond_count, bond_indices, bond_kb, bond_k0 = np.zeros((n_atoms,1), dtype=np.uint32), np.zeros((n_bonds,1), dtype=np.uint32), np.zeros((n_bonds,1), dtype=np.float32), np.zeros((n_bonds,1), dtype=np.float32) + for i, _ in enumerate(self.atoms): + bond_count[i] = len(self.bonds[i]) + curr = i * max_bonds + for j, bond in enumerate(self.bonds[i]): + bond_indices[curr+j] = bond.j + bond_kb[curr+j] = bond.kb + bond_k0[curr+j] = bond.k0 + + max_angles = max(map(lambda x: len(x[1]), self.angles.items()), default=0) + n_angles = n_atoms * max_angles + angle_count, angle_indices, angle_theta0, angle_r0 = np.zeros((n_atoms,1), dtype=np.uint32), np.zeros((3*n_angles,1), dtype=np.uint32), np.zeros((n_angles,1), dtype=np.float32), np.zeros((n_angles,1), dtype=np.float32) + for i, _ in enumerate(self.atoms): + angle_count[i] = len(self.angles[i]) + curr = i * max_angles + for j, angle in enumerate(self.angles[i]): + angle_indices[0*n_angles+curr+j] = angle.i + angle_indices[1*n_angles+curr+j] = angle.j + angle_indices[2*n_angles+curr+j] = angle.k + angle_theta0[curr+j] = np.radians(angle.theta0) + angle_r0[curr+j] = angle.r0 + + max_torsions = max(map(lambda x: len(x[1]), self.torsions.items()), default=0) + n_torsions = n_atoms * max_torsions + torsion_count, torsion_indices, torsion_phi0, torsion_r0 = np.zeros((n_atoms,1), dtype=np.uint32), np.zeros((4*n_torsions,1), dtype=np.uint32), np.zeros((n_torsions,1), dtype=np.float32), np.zeros((n_torsions,1), dtype=np.float32) + curr = 0 + for i, _ in enumerate(self.atoms): + torsion_count[i] = len(self.torsions[i]) + curr = i * max_torsions + for j, torsion in enumerate(self.torsions[i]): + torsion_indices[0*n_torsions+curr+j] = torsion.i + torsion_indices[1*n_torsions+curr+j] = torsion.j + torsion_indices[2*n_torsions+curr+j] = torsion.k + torsion_indices[3*n_torsions+curr+j] = torsion.l + torsion_phi0[curr+j] = np.radians(torsion.phi0) + torsion_r0[curr+j] = torsion.r0 + + return atom_pos, atom_vel, molecules, bond_count, bond_indices, bond_kb, bond_k0, angle_count, angle_indices, angle_theta0, angle_r0, torsion_count, torsion_indices, torsion_phi0, torsion_r0 + +def build_kernel(domain_size, n_atoms, n_bonds, n_angles, n_torsions, max_bonds, max_angles, max_torsions, max_lj, max_coulomb, lennard_jones, coulomb): + lj_sigma_expr = " + ".join(map(lambda lj: f"(iMass == {lj.mass}) * {lj.sigma}", lennard_jones)) + lj_epsilon_expr = " + ".join(map(lambda lj: f"(iMass == {lj.mass}) * {lj.epsilon}", lennard_jones)) + coulomb_charge_expr = " + ".join(map(lambda c: f"(iMass == {c.mass}) * {c.charge}", coulomb)) + + with open('interacticle/kernel.cl', 'r') as kernel_src: + return Template(kernel_src.read()).substitute( + domain_size = domain_size, + n_atoms = n_atoms, + n_bonds = n_bonds, + n_angles = n_angles, + n_torsions = n_torsions, + max_lj = max_lj, + max_coulomb = max_coulomb, + max_bonds = max_bonds, + max_angles = max_angles, + max_torsions = max_torsions, + lj_sigma_expr = lj_sigma_expr if lj_sigma_expr else "0", + lj_epsilon_expr = lj_epsilon_expr if lj_epsilon_expr else "0", + coulomb_charge_expr = coulomb_charge_expr if coulomb_charge_expr else "0") + +class Simulation: + def __init__(self, setup, opengl = False): + self.domain_size = np.float32(setup.domain_size) + self.cutoff = np.float32(setup.cutoff) + self.skin = np.float32(setup.skin) + self.adaptive_skin = setup.skin + self.neighborhood_step = setup.neighborhood_step + self.statistics_step = setup.statistics_step + self.tau = np.float32(setup.tau) + self.target_temperature = setup.target_temperature + self.verbose = True + + self.np_atom_pos, self.np_atom_vel, self.np_molecules, self.np_bond_count, self.np_bond_indices, self.np_bond_kb, self.np_bond_k0, self.np_angle_count, self.np_angle_indices, self.np_angle_theta0, self.np_angle_r0, self.np_torsion_count, self.np_torsion_indices, self.np_torsion_phi0, self.np_torsion_r0 = setup.serialize() + + self.n_molecules = np.max(self.np_molecules)+1 + self.n_atoms = self.np_atom_pos.shape[0] + self.n_bonds = self.np_bond_kb.shape[0] + self.n_angles = self.np_angle_r0.shape[0] + self.n_torsions = self.np_torsion_r0.shape[0] + + self.max_lj = setup.max_lennard_jones + self.max_coulomb = setup.max_coulomb + self.max_bonds = max(self.np_bond_count)[0] + self.max_angles = max(self.np_angle_count)[0] + self.max_torsions = max(self.np_torsion_count)[0] + + self.n_lj = self.n_atoms * setup.max_lennard_jones + self.np_lj_count, self.np_lj_indices, self.np_lj_sigma, self.np_lj_epsilon = np.zeros((self.n_atoms,1), dtype=np.uint32), np.zeros((self.n_lj,1), dtype=np.uint32), np.zeros((self.n_lj,1), dtype=np.float32), np.zeros((self.n_lj,1), dtype=np.float32) + + self.n_coulomb = self.n_atoms * setup.max_coulomb + self.np_coulomb_count, self.np_coulomb_indices, self.np_coulomb_charge = np.zeros((self.n_atoms,1), dtype=np.uint32), np.zeros((self.n_coulomb,1), dtype=np.uint32), np.zeros((self.n_coulomb,1), dtype=np.float32) + + self.step = 0 + self.opengl = opengl + + self.kernel_src = build_kernel(setup.domain_size, self.n_atoms, self.n_bonds, self.n_angles, self.n_torsions, self.max_bonds, self.max_angles, self.max_torsions, self.max_lj, self.max_coulomb, setup.lennard_jones.values(), setup.coulomb.values()) + + def setup(self): + self.platform = cl.get_platforms()[0] + if self.opengl: + self.context = cl.Context( + properties=[(cl.context_properties.PLATFORM, self.platform)] + get_gl_sharing_context_properties()) + else: + self.context = cl.Context( + properties=[(cl.context_properties.PLATFORM, self.platform)]) + self.queue = cl.CommandQueue(self.context) + self.program = cl.Program(self.context, self.kernel_src).build( + '-cl-single-precision-constant') + + self.cl_molecules = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_molecules) + + self.np_force = np.zeros((self.n_atoms,4), dtype=np.float32) + + self.cl_force_prev = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_force) + self.cl_force_curr = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_force) + + self.cl_velocity = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_atom_vel) + + self.np_lj_shift = np.zeros((self.n_lj,4), dtype=np.float32) + + self.cl_lj_count = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_lj_count) + self.cl_lj_indices = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_lj_indices) + self.cl_lj_sigma = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_lj_sigma) + self.cl_lj_epsilon = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_lj_epsilon) + self.cl_lj_shift = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_lj_shift) + + self.np_coulomb_shift = np.zeros((self.n_coulomb,4), dtype=np.float32) + + self.cl_coulomb_count = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_coulomb_count) + self.cl_coulomb_indices = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_coulomb_indices) + self.cl_coulomb_charge = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_coulomb_charge) + self.cl_coulomb_shift = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_coulomb_shift) + + if self.n_bonds > 0: + self.cl_bond_count = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_bond_count) + self.cl_bond_indices = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_bond_indices) + self.cl_bond_kb = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_bond_kb) + self.cl_bond_k0 = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_bond_k0) + + if self.n_angles > 0: + self.cl_angle_count = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_angle_count) + self.cl_angle_indices = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_angle_indices) + self.cl_angle_theta0 = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_angle_theta0) + self.cl_angle_r0 = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_angle_r0) + + if self.n_torsions > 0: + self.cl_torsion_count = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_torsion_count) + self.cl_torsion_indices = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_torsion_indices) + self.cl_torsion_phi0 = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_torsion_phi0) + self.cl_torsion_r0 = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_torsion_r0) + + if self.opengl: + self.gl_position = vbo.VBO(data=self.np_atom_pos, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER) + self.gl_position.bind() + self.cl_position = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_position)) + cl.enqueue_acquire_gl_objects(self.queue, [self.cl_position]) + else: + self.cl_position = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_atom_pos) + + self.kernel_args = (self.queue, (self.n_atoms,), None, self.cl_position) + + if self.target_temperature and not callable(self.target_temperature): + self.randomize_velocities(self.target_temperature) + + def compute_intramolecular_forces(self): + if self.n_bonds > 0 and self.n_angles > 0 and self.n_torsions > 0: + self.program.compute_intramolecular( + *self.kernel_args, + self.cl_force_curr, + self.cl_bond_count, self.cl_bond_indices, self.cl_bond_kb, self.cl_bond_k0, + self.cl_angle_count, self.cl_angle_indices, self.cl_angle_theta0, self.cl_angle_r0, + self.cl_torsion_count, self.cl_torsion_indices, self.cl_torsion_phi0, self.cl_torsion_r0) + else: + if self.n_bonds > 0: + self.program.compute_bonds( + *self.kernel_args, + self.cl_force_curr, + self.cl_bond_count, self.cl_bond_indices, self.cl_bond_kb, self.cl_bond_k0) + if self.n_angles > 0: + self.program.compute_angles( + *self.kernel_args, + self.cl_force_curr, + self.cl_angle_count, self.cl_angle_indices, self.cl_angle_theta0, self.cl_angle_r0) + if self.n_torsions > 0: + self.program.compute_torsions( + *self.kernel_args, + self.cl_force_curr, + self.cl_torsion_count, self.cl_torsion_indices, self.cl_torsion_phi0, self.cl_torsion_r0) + + def compute_intermolecular_forces(self): + self.program.compute_lennard_jones( + *self.kernel_args, + self.cl_force_curr, + self.cl_lj_count, self.cl_lj_indices, self.cl_lj_sigma, self.cl_lj_epsilon, self.cl_lj_shift, + self.cutoff) + self.program.compute_coulomb( + *self.kernel_args, + self.cl_force_curr, + self.cl_coulomb_count, self.cl_coulomb_indices, self.cl_coulomb_charge, self.cl_coulomb_shift, + self.cutoff) + + def update_neighborhoods(self): + self.program.wrap_molecules(*self.kernel_args, self.cl_molecules) + self.program.update_potential_neighborhoods( + *self.kernel_args, + self.cl_molecules, + self.cl_lj_count, self.cl_lj_indices, self.cl_lj_sigma, self.cl_lj_epsilon, self.cl_lj_shift, + self.cl_coulomb_count, self.cl_coulomb_indices, self.cl_coulomb_charge, self.cl_coulomb_shift, + self.cutoff, self.skin).wait() + + def evolve(self): + if self.opengl: + cl.enqueue_acquire_gl_objects(self.queue, [self.cl_position]) + + if self.step % self.statistics_step == 0: + self.read_velocities() + temperature, rmsv = self.get_temperature(), self.get_rms_velocity() + if self.target_temperature and temperature > 0: + if callable(self.target_temperature): + target = self.target_temperature(self.step, temperature) + if target: + self.scale_velocities(np.sqrt(target/temperature)) + else: + self.scale_velocities(np.sqrt(self.target_temperature/temperature)) + if self.verbose: + print(f"t={self.step*self.tau:.2f} T={temperature:.0f}, V_rms={rmsv:.0f}") + + if self.step % self.neighborhood_step == 0: + if self.adaptive_skin: + skin_bound = np.max(np.linalg.norm(self.np_atom_vel[:,0:3], axis=1)) * self.tau * self.neighborhood_step + self.skin = np.float32(max(2*skin_bound, self.adaptive_skin)) + self.update_neighborhoods() + if self.verbose: + min_n, max_n, mean_n = self.get_neighborhood_characteristics() + print(f"t={self.step*self.tau:.2f} minN={min_n}, maxN={max_n}, meanN={mean_n}, skin={self.skin:.3f}") + + self.program.evolve_x(*self.kernel_args, self.cl_velocity, self.cl_force_prev, self.cl_force_curr, self.tau) + if self.n_bonds > 0 or self.n_angles > 0 or self.n_torsions > 0: + self.compute_intramolecular_forces() + self.compute_intermolecular_forces() + self.program.evolve_v(*self.kernel_args, self.cl_velocity, self.cl_force_prev, self.cl_force_curr, self.tau).wait() + self.step = self.step + 1 + + def read_velocities(self): + cl.enqueue_copy(self.queue, self.np_atom_vel, self.cl_velocity).wait() + + def get_neighborhood_characteristics(self): + cl.enqueue_copy(self.queue, self.np_lj_count, self.cl_lj_count).wait() + return np.min(self.np_lj_count), np.max(self.np_lj_count), int(np.mean(self.np_lj_count)) + + def get_velocity_norms(self): + return np.linalg.norm(self.np_atom_vel[:,0:3], axis=1) + + def get_temperature(self): + velocity_norms = np.square(np.linalg.norm(self.np_atom_vel[:,0:3], axis=1)) # nm^2 ps^-2 + masses = self.np_atom_pos[:,3:4].flatten() # u + energy = np.sum(velocity_norms * masses * 0.5) # u nm^2 ps^-2 + kb = 0.00831446262 # Boltzmann's constant [kJ mol^-1 K^-1] + return energy * (2 / (3*self.n_atoms*kb)) # K + + def get_rms_velocity(self): + velocity_norms = (np.linalg.norm(self.np_atom_vel[:,0:3], axis=1)) # nm^2 ps^-2 + return np.sqrt(np.mean(velocity_norms) * 1e6) + + def randomize_velocities(self, temperature): + kb = 1.38064852e-23 # Boltzmann's constant [J/K] + masses = self.np_atom_pos[:,3:4].flatten() * 1.66053907e-27 # Mass in [kg] + self.np_atom_vel[:,0] = np.random.normal(0, np.sqrt(kb*temperature/masses), self.n_atoms) * 0.001 + self.np_atom_vel[:,1] = np.random.normal(0, np.sqrt(kb*temperature/masses), self.n_atoms) * 0.001 + self.np_atom_vel[:,2] = np.random.normal(0, np.sqrt(kb*temperature/masses), self.n_atoms) * 0.001 + cl.enqueue_copy(self.queue, self.cl_velocity, self.np_atom_vel).wait() + + def scale_velocities(self, scale): + self.np_atom_vel = self.np_atom_vel * scale; + cl.enqueue_copy(self.queue, self.cl_velocity, self.np_atom_vel).wait() + + def gl_draw_particles(self): + gl.glEnableClientState(gl.GL_VERTEX_ARRAY) + self.gl_position.bind() + gl.glVertexPointer(4, gl.GL_FLOAT, 0, self.gl_position) + gl.glDrawArrays(gl.GL_POINTS, 0, self.n_atoms) + gl.glDisableClientState(gl.GL_VERTEX_ARRAY) -- cgit v1.2.3