diff --git a/examples/brd_ghost_distribution/build_minimize_heat_noMLF.py b/examples/brd_ghost_distribution/build_minimize_heat_noMLF.py index 91f2fdd..8b80172 100644 --- a/examples/brd_ghost_distribution/build_minimize_heat_noMLF.py +++ b/examples/brd_ghost_distribution/build_minimize_heat_noMLF.py @@ -18,6 +18,8 @@ from flexibletopology.utils.reporters import H5Reporter from flexibletopology.utils.openmmutils import read_params, getParameters, setParameters +from flexibletopology.forces.static import build_protein_restraint_force +from flexibletopology.attr_minimizer import AttrMinimizer import sys from contforceplugin import ContForce @@ -85,7 +87,7 @@ pos_arr = np.array(crd.positions.value_in_unit(unit.nanometers)) # MD simulations settings -TEMPERATURES = [10, 20, 50, 100, 150, 200, 250, 300] +TEMPERATURES = [10, 10, 20, 50, 100, 150, 200, 250, 300] FRICTION_COEFFICIENT = 1/unit.picosecond TIMESTEP = 0.002*unit.picoseconds NUM_STEPS = [10000 for t in TEMPERATURES] @@ -97,7 +99,7 @@ # system building values WIDTH = 0.3 # nm # MIN_DIST = 0.04 # nm -MIN_DIST = 0.05 # nm +MIN_DIST = 0.2 # nm # force group values ghostghost_group = 29 @@ -136,8 +138,10 @@ #_______________BUILD SYSTEM & SIMULATION OBJECT_______________# bs_idxs = pdb_file.topology.select(BS_SELECTION_STRING) + prot_idxs = pdb_file.topology.select('protein') + bb_idxs = pdb_file.topology.select('protein and backbone') - print("n_ghosts is ",n_ghosts) + print(f"Building a system with {n_ghosts} ghost particles..") n_system = pdb_file.n_atoms gst_idxs = list(range(n_system, n_system+n_ghosts)) @@ -148,23 +152,19 @@ cont_force_idxs = gst_idxs + SYSTEM_CONT_FORCE_IDXS con_force.addBond(cont_force_idxs, len(cont_force_idxs), 0.25, 10000) - BUILD_UTILS = SystemBuild(psf=psf, crd=crd, pdb=pdb_file, n_ghosts=n_ghosts, + BUILD_UTILS = SystemBuild(psf=psf, pos=pos_arr, pdb=pdb_file, n_ghosts=n_ghosts, toppar_str=TOPPAR_STR, inputs_path=INPUTS_PATH, width=WIDTH, binding_site_idxs=bs_idxs, - min_dist=MIN_DIST, + min_dist=MIN_DIST, gg_min_dist=MIN_DIST, gg_group=ghostghost_group, gg_nb_group=ghostghost_nb_group, sg_group=systemghost_group, ghost_mass=GHOST_MASS, attr_bounds=BOUNDS, contForce=con_force) - print('Building the system..') - system, initial_signals, n_ghosts, psf_top, crd_pos, _ = BUILD_UTILS.build_system_forces() + system, initial_signals, n_ghosts, psf_top, pos, _ = BUILD_UTILS.build_system_forces() pressure = 1.0*unit.atmospheres temp = 300*unit.kelvin - barostat = omm.MonteCarloBarostat(pressure, temp) - system.addForce(barostat) - - print('System built') + barostat = omm.MonteCarloBarostat(pressure, temp) # add this to the system after heating step 0 coeffs = {'lambda': rest_coeff, 'charge': rest_coeff, @@ -178,7 +178,7 @@ simulation = omma.Simulation(psf_top, system, integrator, platform, prop) - simulation.context.setPositions(crd_pos) + simulation.context.setPositions(pos) # add reporters if not osp.exists(OUTPUTS_PATH): @@ -188,26 +188,31 @@ omma.PDBFile.writeFile(psf.topology, pre_min_positions, open(osp.join(OUTPUTS_PATH,'struct_before_min.pdb'), 'w')) #_______________MINIMIZE_______________# - print('Running minimization') + print('Running minimization..') print('Before min: E=', simulation.context.getState(getEnergy=True).getPotentialEnergy()) + + attr_min = AttrMinimizer(simulation, n_ghosts, BOUNDS) + print('old attr:',attr_min.get_attributes()) + print('new_attr:',attr_min.attr_minimize()) + print('After attr min: E=', simulation.context.getState(getEnergy=True).getPotentialEnergy()) + begin = time.time() simulation.minimizeEnergy(tolerance=TOL ,maxIterations=MAXITR) end = time.time() print('After min: E=', simulation.context.getState(getEnergy=True).getPotentialEnergy()) # save a PDB of the minimized positions - latest_state = simulation.context.getState(getPositions=True) + latest_state = simulation.context.getState(getPositions=True,enforcePeriodicBox=True) latest_par = getParameters(simulation, n_ghosts) omma.PDBFile.writeFile(psf_top, latest_state.getPositions(), open(osp.join(OUTPUTS_PATH,f'minimized_pos.pdb'),'w')) - print("Minimization Ends") print(f"Minimization run time = {np.round(end - begin, 3)}s") - print('Heating system') - begin = time.time() - + print('Heating system..') + heat_begin = time.time() for temp_idx, TEMP in enumerate(TEMPERATURES): + heat_step_begin = time.time() H5REPORTER_FILE = osp.join(OUTPUTS_PATH,f'traj{temp_idx}.h5') integrator = CustomHybridIntegratorRestrictedChargeVariance(n_ghosts, TEMP*unit.kelvin, FRICTION_COEFFICIENT, @@ -215,7 +220,17 @@ # _______________HEAT SYSTEM_______________ # - + + if temp_idx == 0: # restrain all system atoms + box_sizes = np.zeros((3)) + box_sizes[0] = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(unit.nanometers) + box_sizes[1] = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(unit.nanometers) + box_sizes[2] = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(unit.nanometers) + + positions = simulation.context.getState(getPositions=True,enforcePeriodicBox=True).getPositions() + prot_restr_force = build_protein_restraint_force(positions,prot_idxs,bb_idxs,box_sizes) + prot_restr_force_idx = system.addForce(prot_restr_force) + simulation = omma.Simulation(psf_top, system, integrator, platform, prop) simulation.context.setState(latest_state) @@ -245,18 +260,21 @@ latest_par = getParameters(simulation, n_ghosts) latest_state = simulation.context.getState(getPositions=True) - end = time.time() - print(f"Heating run time = {np.round(end - begin, 3)}s") - print('Done heating system; dcd saved') - + if temp_idx == 0: + system.removeForce(prot_restr_force_idx) + system.addForce(barostat) + heat_step_end = time.time() + print(f"Heating step {temp_idx} completed. Run time = {np.round(heat_step_end - heat_step_begin, 3)}s") + + heat_end = time.time() + print(f"Total heating run time = {np.round(heat_end - heat_begin, 3)}s") #_______________SAVE SIM INPUTS_______________# # save the system, topology, simulation object, positions, and parameters - print(" n_ghosts", n_ghosts) - print('Saving simulation input files') + print('Saving simulation input files..') with open(osp.join(OUTPUTS_PATH, 'system.pkl'), 'wb') as new_file: pkl.dump(system, new_file) diff --git a/examples/brd_ghost_distribution/system_build_util.py b/examples/brd_ghost_distribution/system_build_util.py index 92563a8..7b37ce1 100644 --- a/examples/brd_ghost_distribution/system_build_util.py +++ b/examples/brd_ghost_distribution/system_build_util.py @@ -22,13 +22,13 @@ class SystemBuild(object): A class that generates a ready-to-equilibrate OpenMM system object. """ - def __init__(self, psf=None, crd=None, pdb=None, target_pkl=None, n_ghosts=None, toppar_str=None, inputs_path=None, + def __init__(self, psf=None, pos=None, pdb=None, target_pkl=None, n_ghosts=None, toppar_str=None, inputs_path=None, ani_model=None, width=None, binding_site_idxs=None, min_dist=0.15, gg_min_dist=0.05, sf_weights=None, gg_group=None, gg_nb_group=None, mlforce_group=None, sg_group=None, mlforce_scale=None, ghost_mass=None,attr_bounds=None,assignFreq=None,rmax_delta=None, rest_k=None, contForce=None): self.psf = psf - self.crd = crd + self.pos = pos self.pdb = pdb self.target_pkl = target_pkl self.toppar_str = toppar_str @@ -151,21 +151,19 @@ def build_system_forces(self): init_attr = gen_init_attr(self.n_ghosts, self.attr_bounds,total_charge=0,init_lambda=1.0) com_bs = self.generate_COM(self.binding_site_idxs, self.pdb) - pos_arr = np.array(self.crd.positions.value_in_unit(unit.nanometers)) - pdb_pos = np.array([pos_arr]) - init_positions = gen_init_pos(self.n_ghosts, com_bs, self.width, pdb_pos, self.min_dist, self.gg_min_dist) + init_positions = gen_init_pos(self.n_ghosts, com_bs, self.width, self.pos, self.min_dist, self.gg_min_dist) # calculating box length - box_lengths = pos_arr.max(axis=0) - pos_arr.min(axis=0) + box_lengths = self.pos.max(axis=0) - self.pos.min(axis=0) self.psf.setBox(box_lengths[0] * unit.nanometers, box_lengths[1] * unit.nanometers, box_lengths[2] * unit.nanometers) params = read_params(self.toppar_str, self.inputs_path) - n_part_system = len(self.crd.positions) - self.crd.positions.extend(unit.quantity.Quantity(init_positions, - unit.nanometers)) + n_part_system = len(self.pos) + self.pos = np.append(self.pos, init_positions, axis=0) + system = self.psf.createSystem(params, nonbondedMethod=omma.forcefield.CutoffPeriodic, nonbondedCutoff=1*unit.nanometers, @@ -201,9 +199,10 @@ def build_system_forces(self): n_part_system=n_part_system, group_num=self.gg_nb_group, initial_sigmas=init_attr['sigma'], + initial_charges=init_attr['charge'], nb_exclusion_list=exclusion_list) - return system, init_attr, self.n_ghosts, new_psf.topology, self.crd.positions, target_features + return system, init_attr, self.n_ghosts, new_psf.topology, self.pos, target_features diff --git a/src/flexibletopology/attr_minimizer.py b/src/flexibletopology/attr_minimizer.py new file mode 100644 index 0000000..18ea125 --- /dev/null +++ b/src/flexibletopology/attr_minimizer.py @@ -0,0 +1,102 @@ +from scipy.optimize import Bounds, LinearConstraint, NonlinearConstraint, minimize +from openmm import unit +import numpy as np + +class AttrMinimizer(object): + """ + A class to minimize ghost particle simulations according to their attributes, + which are saved within an OpenMM context as global variables. + + This is valuable as the standard simulation.minimize function does not + minimize the attributes, only the positions. + + Internally, the dependent variables are represented by x, which stores the + attributes in the following order: + + x[0] - charge[0] + x[1] - sigma[0] + x[2] - epsilon[0] + x[3] - lambda[0] + x[4] - charge[1] + ... + + """ + def __init__(self, simulation, n_ghosts, bounds, charge_sum=0.0, max_charge_var=0.08): + self.simulation = simulation + + self.n_ghosts = n_ghosts + + lower_bounds = [bounds['charge'][0],bounds['sigma'][0],bounds['epsilon'][0],bounds['lambda'][0]] * n_ghosts + upper_bounds = [bounds['charge'][1],bounds['sigma'][1],bounds['epsilon'][1],bounds['lambda'][1]] * n_ghosts + self.bounds = Bounds(lower_bounds,upper_bounds) + + self.constraints = [] + if charge_sum is not None: + sum_charge_vector = [1,0,0,0] * n_ghosts + self.constraints.append(LinearConstraint([sum_charge_vector],[charge_sum],[charge_sum])) + + if max_charge_var is not None: + def cons_f(x): + sum_sq = 0.0 + for i in range(n_ghosts): + sum_sq += x[4*i]*x[4*i] + sum_sq /= n_ghosts + return [sum_sq] + def cons_J(x): + jvec = [] + for i in range(n_ghosts): + for j in range(4): + if j == 0: + jvec.append(2*x[4*i]/n_ghosts) + else: + jvec.append(0) + return [jvec] + + self.constraints.append(NonlinearConstraint(cons_f, 0, max_charge_var*max_charge_var, jac=cons_J, hess='2-point')) + + def get_attributes(self): + pars = self.simulation.context.getParameters() + + x = np.zeros((4*self.n_ghosts)) + for i in range(self.n_ghosts): + x[4*i] = pars[f'charge_g{i}'] + x[4*i + 1] = pars[f'sigma_g{i}'] + x[4*i + 2] = pars[f'epsilon_g{i}'] + x[4*i + 3] = pars[f'lambda_g{i}'] + + return x + + def set_attributes(self, x): + for i in range(self.n_ghosts): + self.simulation.context.setParameter(f'charge_g{i}',x[4*i]) + self.simulation.context.setParameter(f'sigma_g{i}',x[4*i + 1]) + self.simulation.context.setParameter(f'epsilon_g{i}',x[4*i + 2]) + self.simulation.context.setParameter(f'lambda_g{i}',x[4*i + 3]) + + return + + def _energy(self, x): + self.set_attributes(x) + return self.simulation.context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole) + + def _energy_derivs(self, x): + self.set_attributes(x) + par_derivs = self.simulation.context.getState(getParameterDerivatives=True).getEnergyParameterDerivatives() + + v = np.zeros((4*self.n_ghosts)) + for i in range(self.n_ghosts): + v[4*i] = par_derivs[f'charge_g{i}'] + v[4*i + 1] = par_derivs[f'sigma_g{i}'] + v[4*i + 2] = par_derivs[f'epsilon_g{i}'] + v[4*i + 3] = par_derivs[f'lambda_g{i}'] + + return v + + def attr_minimize(self,method='trust-constr', options={'verbose':1}): + res = minimize(self._energy, self.get_attributes(), method=method, jac=self._energy_derivs, hess='2-point', + constraints=self.constraints, options=options, bounds=self.bounds) + + self.set_attributes(res.x) + + return res.x + diff --git a/src/flexibletopology/forces/attributes.py b/src/flexibletopology/forces/attributes.py new file mode 100644 index 0000000..e64182f --- /dev/null +++ b/src/flexibletopology/forces/attributes.py @@ -0,0 +1,44 @@ +import openmm as omm + +def add_attr_force(system, + n_ghosts=None, + group_num=None, + initial_attr=None): + + eps_terms = '' + charge_terms = '' + fsig_eqns = '' + for gh_idx in range(n_ghosts): + eps_terms += f'A*fsig_{gh_idx}*(epsilon_g{gh_idx}-0.25)^2' + charge_terms += f'B*fsig_{gh_idx}*(charge_g{gh_idx}-0.25)^2' + fsig_eqns += f'fsig_{gh_idx} = 1/(1+exp(100*(sigma_g{gh_idx}-0.3))); ' + if gh_idx < n_ghosts-1: + eps_terms += ' + ' + charge_terms += ' + ' + energy_function = eps_terms + ' + ' + charge_terms + '; ' + fsig_eqns + + attr_force = omm.CustomCVForce(energy_function) + + attr_force.addGlobalParameter('A', 941.5) + attr_force.addGlobalParameter('B', 275.3) + + for gh_idx in range(n_ghosts): + # set to initial values + attr_force.addGlobalParameter(f'charge_g{gh_idx}', initial_attr['charge'][gh_idx]) + attr_force.addGlobalParameter(f'sigma_g{gh_idx}', initial_attr['sigma'][gh_idx]) + attr_force.addGlobalParameter(f'epsilon_g{gh_idx}', initial_attr['epsilon'][gh_idx]) + attr_force.addGlobalParameter(f'lambda_g{gh_idx}', initial_attr['lambda'][gh_idx]) + + # adding the del(signal)s [needed in the integrator] + attr_force.addEnergyParameterDerivative(f'charge_g{gh_idx}') + attr_force.addEnergyParameterDerivative(f'sigma_g{gh_idx}') + attr_force.addEnergyParameterDerivative(f'epsilon_g{gh_idx}') + attr_force.addEnergyParameterDerivative(f'lambda_g{gh_idx}') + + # set force parameters + attr_force.setForceGroup(group_num) + system.addForce(attr_force) + + return system + + diff --git a/src/flexibletopology/forces/dynamic.py b/src/flexibletopology/forces/dynamic.py index 1c5c47b..4d1c9f0 100644 --- a/src/flexibletopology/forces/dynamic.py +++ b/src/flexibletopology/forces/dynamic.py @@ -1,15 +1,18 @@ import openmm as omm -def add_gs_force(system, - n_ghosts=None, - n_part_system=None, - group_num=None, - initial_attr=None, - sys_attr=None, - nb_exclusion_list=None): - - for gh_idx in range(n_ghosts): - energy_function = f'lambda_g{gh_idx}*(4.0*epstot*(sor12-sor6)+138.935456*q1*q2*charge_g{gh_idx}/r);' +def add_gs_forces(system, + system_atom_idxs=None, + ghost_idxs=None, + group_num=None, + initial_attr=None, + sys_attr=None, + nb_exclusion_list=None, + elec_scale=1.0): + + n_part_system = len(sys_attr['charge']) + k_fac = 138.935456*elec_scale + for gh_idx, gh_at_idx in enumerate(ghost_idxs): + energy_function = f'lambda_g{gh_idx}*(4.0*epstot*(sor12-sor6)+{k_fac}*q1*q2*charge_g{gh_idx}/r);' energy_function += 'sor12 = sor6^2; sor6 = (sigtot/r)^6;' energy_function += f'sigtot = 0.5*(sigma1+sigma2+sigma_g{gh_idx}); epstot = sqrt(epsilon1*epsilon2*epsilon_g{gh_idx})' gs_force = omm.CustomNonbondedForce(energy_function) @@ -39,13 +42,13 @@ def add_gs_force(system, # for each force term you need to add ALL the particles even # though we only use one of them! - for p_idx in range(n_ghosts): + for p_idx in ghost_idxs: gs_force.addParticle( [1.0, 0.0, 1.0]) # add ghosts using neutral parameters that won't affect the force at all # interaction between ghost and system - gs_force.addInteractionGroup(set(range(n_part_system)), - set([n_part_system + gh_idx])) + gs_force.addInteractionGroup(set(system_atom_idxs), + set([gh_at_idx])) for j in range(len(nb_exclusion_list)): @@ -59,13 +62,15 @@ def add_gs_force(system, return system -def add_gg_nb_force(system, +def add_gg_LJ_force(system, n_ghosts=None, n_part_system=None, group_num=None, initial_sigmas=None, initial_charges=None, nb_exclusion_list=None): + + # treats the inter-ghost particle interactions as normal Lennard-Jones energy_function = f'4.0*(sor12-sor6) + 138.935456*q1*q2/r; ' energy_function += 'sor12 = sor6^2; sor6 = (sigtot/r)^6; ' @@ -136,4 +141,76 @@ def add_gg_nb_force(system, system.addForce(gg_force) return system + +def add_gg_nb_force(system, + n_ghosts=None, + n_part_system=None, + group_num=None, + gg_min_dist=0.095, + initial_charges=None, + nb_exclusion_list=None, + weak_elec_scale=1.0, + repulsive_only=False): + + # treats the inter-ghost particle interactions as normal Lennard-Jones + + k_fac = 138.935456*weak_elec_scale + if repulsive_only: + energy_function = f'4.0*(sor12-sor6) + step(q1*q2)*{k_fac}*q1*q2/r; ' + else: + energy_function = f'4.0*(sor12-sor6) + {k_fac}*q1*q2/r; ' + energy_function += 'sor12 = sor6^2; sor6 = (sig/r)^6; ' + + q_term1 = 'q1 = ' + q_term2 = 'q2 = ' + for i in range(n_ghosts): + q_term1 += f'charge_g{i}*is_par{i}_1' + q_term2 += f'charge_g{i}*is_par{i}_2' + if i < n_ghosts-1: + q_term1 += ' + ' + q_term2 += ' + ' + else: + q_term1 += '; ' + q_term2 += '; ' + + energy_function += q_term1 + q_term2 + + gg_force = omm.CustomNonbondedForce(energy_function) + + for i in range(n_ghosts): + gg_force.addPerParticleParameter(f'is_par{i}_') + gg_force.addGlobalParameter(f'sig', gg_min_dist) + gg_force.addGlobalParameter(f'charge_g{i}', initial_charges[i]) + + # make the zero indicator vector for all system atoms + zero_is_par = [0 for i in range(n_ghosts)] + + # adding the systems params to the force + for p_idx in range(n_part_system): + gg_force.addParticle(zero_is_par) + + # add all the ghost particles + for p_idx in range(n_ghosts): + ghost_is_par = [0 for i in range(n_ghosts)] + ghost_is_par[p_idx] = 1 + + gg_force.addParticle(ghost_is_par) + # adding the del(signal)s [needed in the integrator] + gg_force.addEnergyParameterDerivative(f'charge_g{p_idx}') + + # only compute interactions between ghosts + gg_force.addInteractionGroup(set(range(n_part_system,n_part_system + n_ghosts)), + set(range(n_part_system,n_part_system + n_ghosts))) + + for j in range(len(nb_exclusion_list)): + gg_force.addExclusion(nb_exclusion_list[j][0], nb_exclusion_list[j][1]) + + # set force parameters + gg_force.setForceGroup(group_num) + gg_force.setNonbondedMethod(gg_force.CutoffPeriodic) + gg_force.setCutoffDistance(1.0) + system.addForce(gg_force) + + return system + diff --git a/src/flexibletopology/forces/static.py b/src/flexibletopology/forces/static.py index 4369e8c..47b2aa5 100644 --- a/src/flexibletopology/forces/static.py +++ b/src/flexibletopology/forces/static.py @@ -1,4 +1,5 @@ import openmm as omm +import openmm.unit as unit import numpy as np def add_ghosts_to_nb_forces(system, n_ghosts, n_part_system): @@ -69,3 +70,87 @@ def _modify_other_cnb_force(force, n_ghosts): set(range(n_part_system))) return force + +def build_protein_restraint_force(positions,prot_idxs,bb_idxs,box): + + posresPROT = omm.CustomExternalForce('f*(px^2+py^2+pz^2); \ + px=min(dx, boxlx-dx); \ + py=min(dy, boxly-dy); \ + pz=min(dz, boxlz-dz); \ + dx=abs(x-x0); \ + dy=abs(y-y0); \ + dz=abs(z-z0);') + posresPROT.addGlobalParameter('boxlx',box[0]) + posresPROT.addGlobalParameter('boxly',box[1]) + posresPROT.addGlobalParameter('boxlz',box[2]) + posresPROT.addPerParticleParameter('f') + posresPROT.addPerParticleParameter('x0') + posresPROT.addPerParticleParameter('y0') + posresPROT.addPerParticleParameter('z0') + + for at_idx in prot_idxs: + if at_idx in bb_idxs: + f = 400. + else: + f = 40. + xpos = positions[at_idx].value_in_unit(unit.nanometers)[0] + ypos = positions[at_idx].value_in_unit(unit.nanometers)[1] + zpos = positions[at_idx].value_in_unit(unit.nanometers)[2] + posresPROT.addParticle(at_idx, [f, xpos, ypos, zpos]) + + return posresPROT + +def add_pars_to_force(force, initial_attr): + n_ghosts = len(initial_attr['charge']) + for gh_idx in range(n_ghosts): + force.addGlobalParameter(f'charge_g{gh_idx}', initial_attr['charge'][gh_idx]) + force.addGlobalParameter(f'sigma_g{gh_idx}', initial_attr['sigma'][gh_idx]) + force.addGlobalParameter(f'epsilon_g{gh_idx}', initial_attr['epsilon'][gh_idx]) + force.addGlobalParameter(f'lambda_g{gh_idx}', initial_attr['lambda'][gh_idx]) + + # adding the del(signal)s [needed in the integrator] + force.addEnergyParameterDerivative(f'charge_g{gh_idx}') + force.addEnergyParameterDerivative(f'sigma_g{gh_idx}') + force.addEnergyParameterDerivative(f'epsilon_g{gh_idx}') + force.addEnergyParameterDerivative(f'lambda_g{gh_idx}') + + return force + +def add_custom_cbf_com(system, group_num, ghost_particle_idxs, center_of_mass, initial_attr): + + cbf = omm.CustomCentroidBondForce(1, "0.5*k*step(d - d0)*(d - d0)^2; d = sqrt((x1-com_x)^2 + (y1-com_y)^2 + (z1-com_z)^2)") + cbf.addGlobalParameter('k', 1000) + cbf.addGlobalParameter('d0', 0.9) + + cbf.addGlobalParameter('com_x', center_of_mass[0]) + cbf.addGlobalParameter('com_y', center_of_mass[1]) + cbf.addGlobalParameter('com_z', center_of_mass[2]) + + for gh_idx in range(len(ghost_particle_idxs)): + gh_grp_idx = cbf.addGroup([ghost_particle_idxs[gh_idx]]) + cbf.addBond([gh_grp_idx]) + + cbf.setForceGroup(group_num) + + cbf = add_pars_to_force(cbf, initial_attr) + system.addForce(cbf) + + return system + +def add_custom_cbf(system, group_num, ghost_particle_idxs, anchor_idxs, initial_attr): + + cbf = omm.CustomCentroidBondForce(2, "0.5*k*step(distance(g1,g2) - d0)*(distance(g1,g2) - d0)^2") + cbf.addGlobalParameter('k', 1000) + cbf.addGlobalParameter('d0', 0.9) + + anchor_grp_idx = cbf.addGroup(anchor_idxs) + for gh_idx in range(len(ghost_particle_idxs)): + gh_grp_idx = cbf.addGroup([ghost_particle_idxs[gh_idx]]) + cbf.addBond([anchor_grp_idx, gh_grp_idx]) + + cbf.setForceGroup(group_num) + + cbf = add_pars_to_force(cbf, initial_attr) + system.addForce(cbf) + + return system diff --git a/src/flexibletopology/utils/initialize.py b/src/flexibletopology/utils/initialize.py index 0208024..6ccf086 100644 --- a/src/flexibletopology/utils/initialize.py +++ b/src/flexibletopology/utils/initialize.py @@ -18,28 +18,55 @@ def gen_init_attr(n_ghosts, attr_bounds,total_charge=None,init_lambda=None): return initial_attr -def gen_init_pos(n_ghosts, COM_BS, WIDTH, pdb_pos, min_dist=0.1, gg_min_dist=0.1): +def gen_init_pos(n_ghosts, COM_BS, WIDTH, pdb_pos, min_dist=0.1, gg_min_dist=0.1, gg_max_dist=0.25, first_width=0.3, max_trials=10000, max_attempts=10): - rand_positions = [] - n_attempts = 0 - - while len(rand_positions) < 1: + for n_attempts in range(max_attempts): + n_trials_this_attempt = 0 + rand_positions = [] + + while len(rand_positions) < 1: + if n_trials_this_attempt > max_trials: + break + n_trials_this_attempt += 1 + r_pos = np.random.uniform(low=-first_width, high=first_width,size=(1, 3)) + r_pos = r_pos+COM_BS + if len(pdb_pos) > 0: + # check if this ghost particle is too close to any system atoms + dists = np.linalg.norm(pdb_pos - r_pos, axis=1) + if np.all(dists > min_dist): + rand_positions.append(r_pos) + else: + rand_positions.append(r_pos) + + while len(rand_positions) < n_ghosts: + if n_trials_this_attempt > max_trials: + break + n_trials_this_attempt += 1 - r_pos = np.random.uniform(low=-WIDTH, high=WIDTH,size=(1, 3)) - r_pos = r_pos+COM_BS - dists = np.linalg.norm(pdb_pos - r_pos, axis=1) - if np.all(dists > min_dist): - rand_positions.append(r_pos) + r_pos = np.random.uniform(low=-WIDTH, high=WIDTH,size=(1, 3)) + r_pos = r_pos+COM_BS - while len(rand_positions) < n_ghosts: + # get distances to other ghost particles + dists_gho = np.linalg.norm(np.concatenate(rand_positions) - r_pos, axis=1) - r_pos = np.random.uniform(low=-WIDTH, high=WIDTH,size=(1, 3)) - r_pos = r_pos+COM_BS + # check if this ghost particle is close enough to an existing particle + if dists_gho.min() < gg_max_dist: + too_close = False + if len(pdb_pos) > 0: + # check if this ghost particle is too close to any system atoms + dists_pdb = np.linalg.norm(pdb_pos - r_pos, axis=1) + if dists_pdb.min() < min_dist: + too_close = True + + # check if this ghost particle is too close to any other ghost atoms + if dists_gho.min() < gg_min_dist: + too_close = True - dists_pdb = np.linalg.norm(pdb_pos - r_pos, axis=1) - dists_gho = np.linalg.norm(np.concatenate(rand_positions) - r_pos, axis=1) + if not too_close: + rand_positions.append(r_pos) - if np.all(dists_pdb > min_dist) and np.all(dists_gho > gg_min_dist): - rand_positions.append(r_pos) + if len(rand_positions) == n_ghosts: + return np.concatenate(rand_positions) - return np.concatenate(rand_positions) + raise ValueError("Error initializing ghost particles! Use a smaller number of particles, or a smaller value of min_dist and/or gg_min_dist.") + diff --git a/src/flexibletopology/utils/openmmutils.py b/src/flexibletopology/utils/openmmutils.py index b6fc9f4..de1f4a5 100644 --- a/src/flexibletopology/utils/openmmutils.py +++ b/src/flexibletopology/utils/openmmutils.py @@ -1,10 +1,24 @@ import os.path as osp import numpy as np -from openmm import unit + import openmm.app as omma +import openmm as omm +from openmm import unit EP_CONVERT= -0.2390057 +def get_platform(plat_name): + if plat_name == 'CUDA': + print("Using CUDA platform..") + platform = omm.Platform.getPlatformByName('CUDA') + prop = dict(CudaPrecision='single') + + else: + print("Using Reference platform..") + prop = {} + platform = omm.Platform.getPlatformByName('Reference') + return platform, prop + def read_params(filename, parfiles_path): extlist = ['rtf', 'prm', 'str'] @@ -23,7 +37,7 @@ def read_params(filename, parfiles_path): return params -def getParameters(sim, n_ghosts): +def getParameters(sim, n_ghosts, get_assignment=True): pars = sim.context.getParameters() par_dict = {} @@ -31,13 +45,14 @@ def getParameters(sim, n_ghosts): par_dict['charge'] = np.array([pars[f'charge_g{i}'] for i in range(n_ghosts)]) par_dict['sigma'] = np.array([pars[f'sigma_g{i}'] for i in range(n_ghosts)]) par_dict['epsilon'] = np.array([pars[f'epsilon_g{i}'] for i in range(n_ghosts)]) - par_dict['assignment'] = np.array([pars[f'assignment_g{i}'] for i in range(n_ghosts)]) + if get_assignment: + par_dict['assignment'] = np.array([pars[f'assignment_g{i}'] for i in range(n_ghosts)]) return par_dict def setParameters(sim, par_dict): n_ghosts = len(par_dict['lambda']) - for attr in ['lambda','charge','sigma','epsilon','assignment']: + for attr in par_dict.keys(): for i in range(n_ghosts): sim.context.setParameter(f'{attr}_g{i}',par_dict[attr][i]) @@ -105,17 +120,17 @@ def nb_params_from_charmm_psf(psf): return params -def add_ghosts_to_system(system, psf, n_ghosts, ghost_mass): - psf_ghost_chain = psf.topology.addChain(id='G') - psf_ghost_res = psf.topology.addResidue('ghosts', - psf_ghost_chain) +def add_ghosts_to_system(system, top, n_ghosts, ghost_mass): + ghost_chain = top.addChain(id='G') + ghost_res = top.addResidue('ghosts', + ghost_chain) # adding ghost particles to the system for i in range(n_ghosts): system.addParticle(ghost_mass) - psf.topology.addAtom(f'G{i}', - omma.Element.getBySymbol('Ar'), - psf_ghost_res, - f'G{i}') + top.addAtom(f'G{i}', + omma.Element.getBySymbol('Ar'), + ghost_res, + f'G{i}') - return system, psf + return system, top diff --git a/src/flexibletopology/utils/utils.py b/src/flexibletopology/utils/utils.py index 5286050..056cbb9 100644 --- a/src/flexibletopology/utils/utils.py +++ b/src/flexibletopology/utils/utils.py @@ -84,4 +84,3 @@ def save_ani_model(platform='cpu', except: print("Can not save the model") - diff --git a/vis_utils/flextop.tcl b/vis_utils/flextop.tcl index 976d2bb..5638266 100644 --- a/vis_utils/flextop.tcl +++ b/vis_utils/flextop.tcl @@ -153,7 +153,9 @@ proc ft_set_attr_state {mol frame} { mol scaleminmax $mol $rep_idx 0.25 0.75000 # set the sigscale using sigma (as a fraction of the Argon sigma radius: 0.34 nm) - set sigscale [expr $sigma/0.34] + set sigfac 0.8 + set sigscale [expr $sigfac*$sigma/0.34] + mol modstyle $rep_idx $mol VDW $sigscale 27.000000 # use lambda to determine whether to render in opaque (lambda > 0.7), transparent (0.3 < lambda < 0.7), or ghost (lambda < 0.3) @@ -216,6 +218,21 @@ proc ft_modmaterial {material} { } } +proc ft_modradius {radius} { + # temporarily sets all ft radii to a constant value + global ft_idxs + global ft_firstrep + + set mol [molinfo top] + # use the ft_firstrep offset to determine the index of the representation to alter + set natoms [llength $ft_idxs($mol)] + + for {set i 0} {$i < $natoms} {incr i} { + set rep_idx [expr $i + $ft_firstrep] + mol modstyle $rep_idx $mol VDW $radius 27.000000 + } +} + proc ft_goto {frame} { set mol [molinfo top] if {$frame >= 0} {