ur-whitelab / hoomd-tf

A plugin that allows the use of Tensorflow in Hoomd-Blue for GPU-accelerated ML+MD
https://hoomd-tf.readthedocs.io
MIT License
30 stars 8 forks source link

Automatic force computation for system of many particle types #333

Closed onehalfatsquared closed 3 years ago

onehalfatsquared commented 3 years ago

Hello,

I am trying to use HOOMD-TF to evaluate forces in a system of rigid subunits, composed of 3 particle types, as well as a scaffold particle. The final goal is to implement an attractive force between subunits that is only active when that subunit is attached to the scaffold particle, which will hopefully be doable with access to the particle neighbor lists. In order to build up to this, I am simply trying to implement a basic force computation between different particle types using HOOMD-TF to use in a simulation. I am unsure how to sum up the forces from these many interaction types, as all the examples I have found use a single particle type.

Attached is a minimal python script with my attempt at this and an example trajectory to perform an initialization of the system. I have implemented two functions, one which computes the potentials I am interested in using plain HOOMD, and my attempt at computing the same forces using HOOMD-TF.

My attempt involves computing masked neighbor lists for each pair of particle types that interact, and computing forces as in the examples in the documentation. I do this for each interaction type, but I cannot sum up the forces as I believe the resulting tensors have different dimensions.

Is there an easy way to get an overall force tensor from all these contributions, or should I be going about this in a different way?

ex.zip

whitead commented 3 years ago

@mehradans92 can you take a look at this? I've put the file below for ease of discussion.

From my quick look, I think you're mistakenly computing the forces per pair. You can just sum all the energies and then compute the forces once from the summed energy. The masks are already accounted for in your energy computations and so the autograd will not propagate forces through the mask. @mehradans92 please double-check this though.

hoomdtf_test.py

```py import os import sys import math import numpy as np import hoomd import hoomd.md as md import gsd import gsd.hoomd import tensorflow as tf import hoomd.htf as htf ''' Pseudoatom types A: Attractor (subunit-subunit attraction) T: Top of capsomer (repulsive with eachother) B: Bottom of capsomer (repulsive with top) C: Center of the rigid capsomers S: Scaffold Particle (attractive to bottoms) ''' class DodecModel(htf.SimModel): def setup(self, type_dict_): self.type_dict = type_dict_ def compute(self, nlist, positions): #after debugging, set these via input file subunit_attraction = 6.0 scaffold_attraction = 6.0 rho = 2.5 #compute potentials one by one #Top - Top LJ Repulsion top_sigma = 2.1 top_cut = top_sigma T_index = self.type_dict['T'] #get just the T-T pairs ttMask = htf.masked_nlist(nlist, positions[:,3], T_index, T_index) #compute sigma/r, if less than sig/cut, set to sig/cut tt_rinv = top_sigma * htf.nlist_rinv(ttMask) tt_rinv = tf.math.maximum(tt_rinv, tf.constant([top_sigma / top_cut])) #compute pairwise Lennard Jones repulsion ptt_energy = 4.0 / 2.0 * subunit_attraction * ((tt_rinv)**12 - 1) #shifted tt_energy = tf.reduce_sum(input_tensor=ptt_energy, axis = 1) forces = htf.compute_nlist_forces(ttMask, tt_energy) #forces = htf.compute_nlist_forces(nlist, tt_energy) #Top - Bottom LJ Repulsion bot_sigma = 1.8 bot_cut = bot_sigma B_index = self.type_dict['B'] #get the T-B pairs tbMask = htf.masked_nlist(nlist, positions[:,3], T_index, B_index) #compute sigma/r, if less than sig/cut, set to sig/cut tb_rinv = bot_sigma * htf.nlist_rinv(tbMask) tb_rinv = tf.math.maximum(tb_rinv, tf.constant([bot_sigma / bot_cut])) #compute pairwise Lennard Jones repulsion ptb_energy = 4.0 / 2.0 * subunit_attraction * ((tb_rinv)**12 - 1) #shifted tb_energy = tf.reduce_sum(input_tensor=ptb_energy, axis = 1) forces += htf.compute_nlist_forces(tbMask, tb_energy) #forces += htf.compute_nlist_forces(nlist, tb_energy) #Scaffold - Attractor LJ Repulsion sa_sigma = 1.0 sa_cut = sa_sigma S_index = self.type_dict['S'] A_index = self.type_dict['A'] #get the S-A pairs saMask = htf.masked_nlist(nlist, positions[:,3], S_index, A_index) #compute sigma/r, if less than sig/cut, set to sig/cut sa_rinv = sa_sigma * htf.nlist_rinv(saMask) sa_rinv = tf.math.maximum(sa_rinv, tf.constant([sa_sigma / sa_cut])) #compute pairwise Lennard Jones repulsion psa_energy = 4.0 / 2.0 * 2.0 * ((sa_rinv)**12 - 1) #shifted sa_energy = tf.reduce_sum(input_tensor=psa_energy, axis = 1) forces += htf.compute_nlist_forces(saMask, sa_energy) #forces += htf.compute_nlist_forces(nlist, sa_energy) #Attractor - Attractor Morse Attraction aa_r0 = 0.2 aa_cut = 2.0 aa_alpha = rho / aa_r0 #evaluate the potential at the cutoff y = np.exp(-aa_alpha*(aa_cut-aa_r0)) u0 = subunit_attraction * (y*y - 2*y) #get the A-A pairs aaMask = htf.masked_nlist(nlist, positions[:,3], A_index, A_index) #get distances r, take min with the cutoff aa_r = tf.norm(aaMask[:,:,:3], axis=2) aa_r = tf.math.minimum(aa_r, tf.constant([aa_cut])) #evaluate the morse potential y_aa = tf.math.exp(-aa_alpha*(aa_r-tf.constant([aa_r0]))) paa_energy = subunit_attraction * (tf.math.multiply(y_aa,y_aa) - 2 * y_aa) - u0 aa_energy = tf.reduce_sum(input_tensor=paa_energy, axis = 1) forces += htf.compute_nlist_forces(aaMask, aa_energy) #forces += htf.compute_nlist_forces(nlist, aa_energy) #Bottom - Scaffold Morse Attraction bs_r0 = 0.9 bs_cut = 2.5 #evaluate potential at cutoff y = np.exp(-aa_alpha*(bs_cut-bs_r0)) u0 = subunit_attraction * (y*y - 2*y) #get the B-S pairs bsMask = htf.masked_nlist(nlist, positions[:,3], S_index, B_index) #get distances r, take min with the cutoff bs_r = tf.norm(bsMask[:,:,:3], axis=2) bs_r = tf.math.minimum(bs_r, tf.constant([bs_cut])) #evaluate the morse potential y_bs = tf.math.exp(-aa_alpha*(bs_r-tf.constant([bs_r0]))) pbs_energy = scaffold_attraction * (tf.math.multiply(y_bs,y_bs) - 2 * y_bs) - u0 bs_energy = tf.reduce_sum(input_tensor=pbs_energy, axis = 1) forces += htf.compute_nlist_forces(bsMask, bs_energy) #forces += htf.compute_nlist_forces(nlist, bs_energy) return forces def create_potential(subunit_attraction, scaffold_attraction, rho): #create the potential corresponding to each interaction in the system #attractive morse between A sites and between S and B sites #repulsive LJ between tops to induce curvature #return the neighbor list and potentials, in case we need to modify these # Top (T) interactions top_sigma = 2.1 top_cut = top_sigma top_repulsion = subunit_attraction / 1.0 # Bottom (B) interactions bot_sigma = 1.8 bot_cut = bot_sigma # Scaffold (S) interactions s_r0 = 0.9 s_cut = 2.5 #construct neighbor list nl = md.nlist.cell(check_period=1) # Cell-based neighbor list nl.reset_exclusions(exclusions=['body']) # Lennard-Jones potential: 4 * epsilon [(sigma/r)**12 - alpha * (sigma/r)**6] lj = md.pair.lj(r_cut=top_cut, nlist=nl) #A shift is applied to the entire potential such that it is 0 at the cutoff lj.set_params(mode="shift") #set for each pair of particles #on interactions lj.pair_coeff.set('T', 'T', alpha=0, epsilon=top_repulsion, r_cut=top_cut, sigma=top_sigma) lj.pair_coeff.set('B', 'T', alpha=0, epsilon=top_repulsion, r_cut=bot_cut, sigma=bot_sigma) #off interactions lj.pair_coeff.set(['C', 'S'], ['C', 'A', 'T', 'B', 'S'], alpha=0, epsilon=0, r_cut=0, sigma=0) lj.pair_coeff.set(['B'], ['A','B'], alpha=0, epsilon=0, r_cut=0, sigma=0) lj.pair_coeff.set(['A'], ['A','T'], alpha=0, epsilon=0, r_cut=0, sigma=0) #to avoid overlap with the scaffold lj.pair_coeff.set('S', 'A', alpha=0, epsilon=2, r_cut=1.0, sigma=1.0) # Morse potential: D0[exp(-2 alpha(r-r0)) - 2 exp(-alpha(r-r0))] r0 = 0.2 alpha = rho / r0 morse = md.pair.morse(r_cut=2.0, nlist=nl) morse.set_params(mode="shift") #set for each pair of particles #off interactions morse.pair_coeff.set(['C', 'A', 'T', 'B', 'S'], ['A', 'C', 'T', 'B', 'S'], D0=0, alpha=0, r0=0, r_cut=0.0) #on interactions morse.pair_coeff.set('A', 'A', D0=subunit_attraction, alpha=alpha, r0=r0, r_cut=2.0) morse.pair_coeff.set('S', ['B'], D0=scaffold_attraction, alpha=alpha, r0=s_r0, r_cut=s_cut) #return neighbor list and potentials return nl, morse, lj def simulate(): #set up and perform the simulation # Set simulation parameters num_capsomers = 38 rho = 2.5 subunit_attraction = 6.0 scaffold_attraction = 6.0 box_size = 9.0 ts = 1e-3 Tf = 5 animation_time = 0.5 seed = 1 #init hoomd hoomd.context.initialize("--notice-level=2") #initialize the system using a snapshot try: system = hoomd.init.read_gsd('traj0.gsd', frame=0) except: print("initialization file not found.") raise #make a type dictionary - associates type by string to an int type_dict = {} types = system.particles.types for i in range(len(types)): type_dict[types[i]] = i #construct the potentials using hoomd #nl, morse, lj = create_potential(subunit_attraction, scaffold_attraction, rho) #construct the potential using hoomd-tf NN = 15 model = DodecModel(NN, type_dict_ = type_dict, output_forces = True) tfcompute = htf.tfcompute(model) #create neighborlist with body exclusions, attach to hoomdtf nlist = md.nlist.cell(check_period=1) # Cell-based neighbor list nlist.reset_exclusions(exclusions=['body']) tfcompute.attach(nlist, r_cut=2.5) # integrate rigid and non rigid structures together group_rigid = hoomd.group.rigid_center() group_nonrigid = hoomd.group.nonrigid() group_integrate = hoomd.group.union('integrate', group_rigid, group_nonrigid) #set the integration mode - langevin integrator_mode = md.integrate.mode_standard(dt=ts) md.integrate.langevin(group=group_integrate, kT=1.0, seed=seed, dscale=True) #md.integrate.brownian(group=group_integrate, kT=1.0, seed=seed, dscale=True) #set the final time and run num_steps = int(Tf / ts) hoomd.run(num_steps, profile=True) return if __name__ == "__main__": #do a simulation simulate() ```

mehradans92 commented 3 years ago

I am not sure why you are computing the attraction and repulsion terms in LJ separately while changing the pair types and adding them together.

onehalfatsquared commented 3 years ago

@whitead I can't seem to sum the energies either, as the tensors have a different number of components. It is also unclear to me which neighbor list to use in the final force call when using this approach.

@mehradans92 I'll elaborate a little on what I am doing. I am using a purely repulsive interaction between particle types T-T, T-B, and A-S. I need this to be quite strong so I am using the 1/r^12 term from the Lennard Jones potential, without including the attractive 1/r^6 term. Each of the repulsions has a different strength and cutoff, so I need to consider them separately. There are also attractive interactions between A-A and B-S particle pairs, that I'm using a Morse potential to model, again with different strengths and rest-lengths.

whitead commented 3 years ago

@onehalfatsquared Just reduce the energies to be scalars for calculating the forces. I assume you're trying to save them per-particle, but for computing the forces they can be scalars.

onehalfatsquared commented 3 years ago

Thank you, that did the trick! The simulation is now starts, but the output is not matching what I am getting from the base HOOMD implementation. In fact, it is becoming unstable whereas the HOOMD version does not. Should I reasonably expect that the forces generated with both methods be the same up to some numerical tolerance? I imagine the same noise would be generated as well if I keep the same seed?

Another thing I am unsure of is whether the energy is double counted when using masked neighbor lists. Using the full list, it makes sense that particles i and j would appear in eachothers neighbor list. If I use the masked list htf.masked_nlist(nlist, positions[:,3], type_i, type_j) would this only give a neighbor list for particles of type i consisting of particles of type j? If so then I should not divide the energy by 2.

whitead commented 3 years ago

@onehalfatsquared correct for the masking. Forces should be same (assuming shifting/cut-offs/nlists). We do have unit tests for LJ simulations, checking to make sure results are same between HTF and HOOMD. @mehradans92 any ideas?

mehradans92 commented 3 years ago

@onehalfatsquared The way you are summing all your forces on different particle types, I am not sure if you can take account of degeneracy by dividing the energy by 2. @whitead would it work if the forces and energies are computed per pair type interactions and then summed over to get the total energy?

onehalfatsquared commented 3 years ago

I tried to simplify the code to help narrow down what the issue might be. I have removed (set to 0) all interactions from the system except the one between A and S particles, which has a 1/r^12 repulsion. I changed the python script to run both HOOMD and HOOMDTF simulations for a short time and compare the positions of the S particle. I am finding that even in this simpler case, different forces are being generated.

I'm thinking the issue has to be one of two things. The first might be the way I am performing the shift of the potential and dealing with values outside the cutoff. I am getting the distances, r, by using the norm function on the masked nlist. I then do a safe divide to get inverse distances, sigma/r. If r is bigger than the cutoff (sigma in this case), then sigma/r < 1, so I am taking the maximum of sigma/r and 1. If the safe divide returns a 0, this maximum should return 1. Then I compute (sigma/r)^12-1 to include the shift. Is it possible that the way I am implementing this is messing with the autograd somehow?

My other concern is whether the forces are being assigned to the correct particle when I call forces = htf.compute_nlist_forces(nlist, total_energy) There are a lot of neighbor lists floating around, and this function completes without error for any neighbor list I put in. This makes me wonder if the full nlist originally passed into the compute function is what should go here.

update.zip

Edit Update: Ignore the previous edits, I was looking at the potentials instead of the force. I did find that during each time step, at most one A particle comes into contact with the S particle. Looking at the list of forces, only a single force is non-zero, which makes me think the force from S onto A is being calculated, but not the force from A onto S (or maybe vice-versa). Tomorrow morning I will try doing the same calculation with the order of types in the masked neighbor lists swapped to see if the forces are only being calculated one way.

whitead commented 3 years ago

@onehalfatsquared So, I believe there is some autograd problem we've run into with this specific issue. Judging from the note I left myself, I believe if you're working with inverse forces you should use the implementation from htf. If you would like to use your own to deal with the shifting, I recommend you use something similar to what we have. I cannot remember why (or if I ever learned) divide_no_nan breaks autograd, but it does somehow. Hopefully that addresses your issue. If not, I'll try to convert this into a unit test and troubleshoot in more depth.

image

hoomd_tf.py ```py import os import sys import math import numpy as np import hoomd import hoomd.md as md import gsd import gsd.hoomd import tensorflow as tf import hoomd.htf as htf ''' Pseudoatom types A: Attractor (subunit-subunit attraction) T: Top of capsomer (repulsive with eachother) B: Bottom of capsomer (repulsive with top) C: Center of the rigid capsomers S: Scaffold Particle (attractive to bottoms) ''' class DodecModel(htf.SimModel): def setup(self, type_dict_): self.type_dict = type_dict_ def compute(self, nlist, positions): #after debugging, set these via input file subunit_attraction = 6.0 scaffold_attraction = 10.0 rho = 2.5 #compute potentials one by one # #Top - Top Repulsion # top_sigma = 2.1 # top_cut = top_sigma # T_index = self.type_dict['T'] # #get just the T-T pairs # ttMask = htf.masked_nlist(nlist, positions[:,3], T_index, T_index) # #compute sigma/r, if less than sig/cut, set to sig/cut # #tt_rinv = top_sigma * htf.nlist_rinv(ttMask) # tt_r = tf.norm(ttMask[:,:,:3], axis = 2) # tt_rinv = top_sigma * tf.math.divide_no_nan(1.0, tt_r) # tt_rinv = tf.math.maximum(tt_rinv, tf.constant([top_sigma / top_cut])) # #compute pairwise Lennard Jones repulsion # ptt_energy = 4.0 * subunit_attraction * ((tt_rinv)**12 - 1) #shifted # tt_energy = tf.reduce_sum(input_tensor=ptt_energy, axis = 1) # tt_energy = tf.reduce_sum(input_tensor=tt_energy , axis = 0) # total_energy = tt_energy # #Top - Bottom Repulsion # bot_sigma = 1.8 # bot_cut = bot_sigma # B_index = self.type_dict['B'] # #get the T-B pairs # tbMask = htf.masked_nlist(nlist, positions[:,3], T_index, B_index) # #compute sigma/r, if less than sig/cut, set to sig/cut # #tb_rinv = bot_sigma * htf.nlist_rinv(tbMask) # tb_r = tf.norm(tbMask[:,:,:3], axis = 2) # tb_rinv = bot_sigma * tf.math.divide_no_nan(1.0, tb_r) # tb_rinv = tf.math.maximum(tb_rinv, tf.constant([bot_sigma / bot_cut])) # #compute pairwise Lennard Jones repulsion # ptb_energy = 4.0 * subunit_attraction * ((tb_rinv)**12 - 1) #shifted # tb_energy = tf.reduce_sum(input_tensor=ptb_energy, axis = 1) # tb_energy = tf.reduce_sum(input_tensor=tb_energy, axis = 0) # #total_energy += tb_energy #Scaffold - Attractor Repulsion sa_sigma = 1.0 sa_cut = sa_sigma S_index = self.type_dict['S'] A_index = self.type_dict['A'] #get the S-A pairs saMask = htf.masked_nlist(nlist, positions[:,3], S_index, A_index) #use norm to get distances between particles sa_r = tf.norm(saMask[:,:,:3], axis = 2) #compute sig/r using the safe divide for empty neighbors sa_rinv = sa_sigma * tf.math.divide_no_nan(1.0, sa_r) #if sa/r is below cutoff, sig/cut=1 sa_rinv = tf.math.maximum(sa_rinv, tf.constant([sa_sigma / sa_cut])) #compute pairwise Lennard Jones repulsion psa_energy = 4.0 * 2.0 * ((sa_rinv)**12 - 1) #shifted sa_energy = tf.reduce_sum(input_tensor=psa_energy, axis = 1) sa_energy = tf.reduce_sum(input_tensor=sa_energy, axis = 0) total_energy = sa_energy # #Attractor - Attractor Attraction # aa_r0 = 0.2 # aa_cut = 2.0 # aa_alpha = rho / aa_r0 # #evaluate the potential at the cutoff # y = np.exp(-aa_alpha*(aa_cut-aa_r0)) # u0 = subunit_attraction * (y*y - 2*y) # #get the A-A pairs # aaMask = htf.masked_nlist(nlist, positions[:,3], A_index, A_index) # #get distances r, take min with the cutoff # aa_r = tf.norm(aaMask[:,:,:3], axis=2) # aa_r = tf.where( aa_r < 0.01, aa_cut * tf.ones_like( aa_r ), aa_r ) # aa_r = tf.math.minimum(aa_r, tf.constant([aa_cut])) # #evaluate the morse potential # y_aa = tf.math.exp(-aa_alpha*(aa_r-tf.constant([aa_r0]))) # paa_energy = subunit_attraction * (tf.math.multiply(y_aa,y_aa) - 2 * y_aa) - u0 # aa_energy = tf.reduce_sum(input_tensor=paa_energy, axis = 1) # aa_energy = tf.reduce_sum(input_tensor=aa_energy, axis = 0) # #total_energy += aa_energy # #Bottom - Scaffold Attraction # bs_r0 = 0.9 #tweak??? # bs_cut = 2.5 # #evaluate potential at cutoff # y = np.exp(-aa_alpha*(bs_cut-bs_r0)) # u0 = scaffold_attraction * (y*y - 2*y) # #get the B-S pairs # bsMask = htf.masked_nlist(nlist, positions[:,3], S_index, B_index) # #get distances r, take min with the cutoff # bs_r = tf.norm(bsMask[:,:,:3], axis=2) # bs_r = tf.where( bs_r < 0.01, bs_cut * tf.ones_like( bs_r ), bs_r ) # bs_r = tf.math.minimum(bs_r, tf.constant([bs_cut])) # #evaluate the morse potential # y_bs = tf.math.exp(-aa_alpha*(bs_r-tf.constant([bs_r0]))) # pbs_energy = scaffold_attraction * (tf.math.multiply(y_bs,y_bs) - 2 * y_bs) - u0 # bs_energy = tf.reduce_sum(input_tensor=pbs_energy, axis = 1) # bs_energy = tf.reduce_sum(input_tensor=bs_energy, axis = 0) # #total_energy += bs_energy #compute forces from total energy forces = htf.compute_nlist_forces(nlist, total_energy) return forces def create_potential(subunit_attraction, scaffold_attraction, rho): #create the potential corresponding to each interaction in the system #attractive morse between A sites and between S and B sites #repulsive LJ between tops to induce curvature #return the neighbor list and potentials, in case we need to modify these # Top (T) interactions top_sigma = 2.1 top_cut = top_sigma top_repulsion = subunit_attraction / 1.0 # Bottom (B) interactions bot_sigma = 1.8 bot_cut = bot_sigma # Scaffold (S) interactions s_r0 = 0.9 s_cut = 2.5 #construct neighbor list nl = md.nlist.cell(check_period=1) # Cell-based neighbor list nl.reset_exclusions(exclusions=['body']) # Lennard-Jones potential: 4 * epsilon [(sigma/r)**12 - alpha * (sigma/r)**6] lj = md.pair.lj(r_cut=top_cut, nlist=nl) #A shift is applied to the entire potential such that it is 0 at the cutoff lj.set_params(mode="shift") #set for each pair of particles #on interactions lj.pair_coeff.set('T', 'T', alpha=0, epsilon=top_repulsion*0, r_cut=top_cut, sigma=top_sigma) lj.pair_coeff.set('B', 'T', alpha=0, epsilon=top_repulsion*0, r_cut=bot_cut, sigma=bot_sigma) #off interactions lj.pair_coeff.set(['C', 'S'], ['C', 'A', 'T', 'B', 'S'], alpha=0, epsilon=0, r_cut=0, sigma=0) lj.pair_coeff.set(['B'], ['A','B'], alpha=0, epsilon=0, r_cut=0, sigma=0) lj.pair_coeff.set(['A'], ['A','T'], alpha=0, epsilon=0, r_cut=0, sigma=0) #to avoid overlap with the scaffold lj.pair_coeff.set('S', 'A', alpha=0, epsilon=2, r_cut=1.0, sigma=1.0) # Morse potential: D0[exp(-2 alpha(r-r0)) - 2 exp(-alpha(r-r0))] r0 = 0.2 alpha = rho / r0 morse = md.pair.morse(r_cut=2.0, nlist=nl) morse.set_params(mode="shift") #set for each pair of particles #off interactions morse.pair_coeff.set(['C', 'A', 'T', 'B', 'S'], ['A', 'C', 'T', 'B', 'S'], D0=0, alpha=0, r0=0, r_cut=0.0) #on interactions morse.pair_coeff.set('A', 'A', D0=subunit_attraction, alpha=alpha, r0=r0, r_cut=2.0) morse.pair_coeff.set('S', ['B'], D0=scaffold_attraction, alpha=alpha, r0=s_r0, r_cut=s_cut) morse.disable() #return neighbor list and potentials return nl, morse, lj def simulateHD(): #set up and perform the simulation with hoomd # Set simulation parameters num_capsomers = 38 rho = 2.5 subunit_attraction = 6.0 scaffold_attraction = 10.0 box_size = 9.0 ts = 1e-3 Tf = 5 animation_time = 0.5 seed = 1 #init hoomd hoomd.context.initialize("--notice-level=2") #initialize the system using a snapshot try: system = hoomd.init.read_gsd('traj0.gsd', frame=0) except: print("initialization file not found.") raise #make potential nl, morse, lj = create_potential(subunit_attraction, scaffold_attraction, rho) # integrate rigid and non rigid structures together group_rigid = hoomd.group.rigid_center() group_nonrigid = hoomd.group.nonrigid() group_integrate = hoomd.group.union('integrate', group_rigid, group_nonrigid) #set the integration mode - langevin integrator_mode = md.integrate.mode_standard(dt=ts) md.integrate.langevin(group=group_integrate, kT=1.0, seed=seed, dscale=True) #set the final time and run num_steps = int(Tf / ts) hoomd.run(num_steps) #get the position of the scaffold to test if outputs are the same N = len(system.particles) for i in range(N): if (system.particles[i].type == 'S'): scaffold = system.particles[i].position break return scaffold def simulateHDTF(): #set up and perform the simulation with hoomdTF # Set simulation parameters num_capsomers = 38 rho = 2.5 subunit_attraction = 6.0 scaffold_attraction = 10.0 box_size = 9.0 ts = 1e-3 Tf = 5 animation_time = 0.5 seed = 1 #init hoomd hoomd.context.initialize("--notice-level=2") #initialize the system using a snapshot try: system = hoomd.init.read_gsd('traj0.gsd', frame=0) except: print("initialization file not found.") raise #make a type dictionary - associates type by string to an int type_dict = {} types = system.particles.types for i in range(len(types)): type_dict[types[i]] = i #construct the potential using hoomd-tf NN = 120 model = DodecModel(NN, type_dict_ = type_dict, output_forces = True) tfcompute = htf.tfcompute(model) #create neighborlist with body exclusions, attach to hoomdtf nlist = md.nlist.cell(check_period=1) # Cell-based neighbor list nlist.reset_exclusions(exclusions=['body']) tfcompute.attach(nlist, r_cut=2.5) # integrate rigid and non rigid structures together group_rigid = hoomd.group.rigid_center() group_nonrigid = hoomd.group.nonrigid() group_integrate = hoomd.group.union('integrate', group_rigid, group_nonrigid) #set the integration mode - langevin integrator_mode = md.integrate.mode_standard(dt=ts) md.integrate.langevin(group=group_integrate, kT=1.0, seed=seed, dscale=True) #set the final time and run num_steps = int(Tf / ts) hoomd.run(num_steps) #get the position of the scaffold to test if outputs are the same N = len(system.particles) for i in range(N): if (system.particles[i].type == 'S'): scaffold = system.particles[i].position break return scaffold if __name__ == "__main__": #do a simulation ref_position = simulateHD() test_position = simulateHDTF() print("HOOMD Scaffold {} \n HOOMDTF Scaffold {}".format(ref_position, test_position)) ```
onehalfatsquared commented 3 years ago

I've managed to get the final positions to agree to 2-3 decimal places after a short simulation, but it is still not a perfect match. It does seem that computing the forces from a masked neighbor list defined as htf.masked_nlist(nlist, positions[:,3], type_i, type_j) is only computing the forces in one direction, but I'm not sure if its type_i onto type_j or vice versa. I find that the resulting force tensor only has a single non-zero component, and no component with an equal but opposite force.

I added an additional calculation with the order of the types swapped, i.e. htf.masked_nlist(nlist, positions[:,3], type_j, type_i) and divided the energy from each calculation by 2.0. When I do this I find that the force tensor has two non-zero components, and they are indeed the same magnitude but in opposite directions. When I perform the simulation using this method, the final position of the S particle matches the HOOMD result to 2-3 decimal places.

All the above was done with my previous handling of the r_inv values. I tried using the built in nlist_rinv function as you described, and the resulting forces are slightly different. However, I made two odd observations when using this function. The first is that the deviation in the final position of the S particle is greater when using this built in function than simply using divide_no_nan. The second, which may be the cause for the first observation, is that the two non-zero forces computed using this method are no longer the same in magnitude; they only agree to 3-5 decimal places. This seems strange as I can't see why the same list of distances would give different outputs from this function...

whitead commented 3 years ago

@onehalfatsquared Excellent, I would call that a success! As you know, the tensor operations are probably slightly different on tensorflow because everything needs to be differentiable and so agreement will never be perfect.

That function, built in r_inv, has some epsilon added (as you see in source) to enable derivatives wrt parameters be numerically stable. If you're only using htf for forces wrt to position, you do not need to do this.

If you have time, would be great to hear some feedback on how to make this easier for the next person. Like things in docs that were unclear or improvements that could be made to examples.

onehalfatsquared commented 3 years ago

Ah, great! I will add in the rest of the interactions and see if I can get some structures to assemble! Thank you for the assistance!

I think the most helpful thing to do for future users would be to expand the Lennard-Jones with one particle type example to two particle types, or just add a new example if you want to keep that one simple. Showing the construction of the masked neighbor lists (in both directions), the reduction of energy to scalars, then the force computation would make the process more clear.

whitead commented 3 years ago

Great, re-open if you have more problems.