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)