summaryrefslogtreecommitdiff
path: root/interacticle/simulation.py
diff options
context:
space:
mode:
Diffstat (limited to 'interacticle/simulation.py')
-rw-r--r--interacticle/simulation.py394
1 files changed, 394 insertions, 0 deletions
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)