update with master

perses/app/relative_point_mutation_setup.py

@@ -22,19 +22,10 @@
 import tempfile
 
 ENERGY_THRESHOLD = 1e-2
-temperature = 298 * unit.kelvin
+temperature = 300 * unit.kelvin
 kT = kB * temperature
 beta = 1.0/kT
 ring_amino_acids = ['TYR', 'PHE', 'TRP', 'PRO', 'HIS']
-CL_CHARGE = unit.Quantity(value=-1.0, unit=unit.elementary_charge)
-CL_SIGMA = unit.Quantity(value=0.4477656957373345, unit=unit.nanometer)
-CL_EPSILON = unit.Quantity(value=0.14891274399999999, unit=unit.kilojoule_per_mole)
-NA_CHARGE = unit.Quantity(value=1.0, unit=unit.elementary_charge)
-NA_SIGMA = unit.Quantity(value=0.2439280690268249, unit=unit.nanometer)
-NA_EPSILON = unit.Quantity(value=0.3658460312, unit=unit.kilojoule_per_mole)
-O_CHARGE = unit.Quantity(value=-0.834, unit=unit.elementary_charge)
-H_CHARGE = unit.Quantity(value=0.417, unit=unit.elementary_charge)
-
 
 # Set up logger
 import logging
@@ -598,8 +589,8 @@ def __init__(self,
                  apo_box_dimensions=None,
                  flatten_torsions=False,
                  flatten_exceptions=False,
-                 vanilla=True,
-                 repartitioned=True,
+                 generate_unmodified_hybrid_topology_factory=True,
+                 generate_rest_capable_hybrid_topology_factory=False,
                  debug_dir=None,
                  **kwargs):
         """
@@ -651,9 +642,9 @@ def __init__(self,
                 in the htf, flatten torsions involving unique new atoms at lambda = 0 and unique old atoms are lambda = 1
             flatten_exceptions : bool, default False
                 in the htf, flatten exceptions involving unique new atoms at lambda = 0 and unique old atoms at lambda = 1
-            vanilla : bool, default True
+            generate_unmodified_hybrid_topology_factory : bool, default True
                 whether to generate a vanilla HybridTopologyFactory
-            repartitioned : bool, default True
+            generate_rest_capable_hybrid_topology_factory : bool, default False
                 whether to generate a RepartitionedHybridTopologyFactory
             debug_dir : str, default None
                 if specified, debug output files will be saved here
@@ -734,134 +725,47 @@ def __init__(self,
 
             topology_proposal, new_positions = point_mutation_engine.propose(sys, top, pos, tleap_prefix, is_complex, debug_dir)
 
+            # Fix naked charges in old and new systems
+            old_topology_atom_map = {atom.index: atom.residue.name for atom in topology_proposal.old_topology.atoms()}
+            new_topology_atom_map = {atom.index: atom.residue.name for atom in topology_proposal.new_topology.atoms()}
+            for i, system in enumerate([topology_proposal.old_system, topology_proposal.new_system]):
+                force_dict = {i.__class__.__name__: i for i in system.getForces()}
+                atom_map = old_topology_atom_map if i == 0 else new_topology_atom_map
+                if 'NonbondedForce' in [k for k in force_dict.keys()]:
+                    nb_force = force_dict['NonbondedForce']
+                    for idx in range(nb_force.getNumParticles()):
+                        if atom_map[idx] in ['HOH', 'WAT']: # Do not add naked charge fix to water hydrogens
+                            continue
+                        charge, sigma, epsilon = nb_force.getParticleParameters(idx)
+                        if sigma == 0*unit.nanometer:
+                            sigma = 0.06*unit.nanometer
+                            nb_force.setParticleParameters(idx, charge, sigma, epsilon)
+                        if epsilon == 0*unit.kilojoule_per_mole:
+                            epsilon = 0.0001*unit.kilojoule_per_mole
+                            nb_force.setParticleParameters(idx, charge, sigma, epsilon)
+
             # Check for charge change...
-            charge_diff = point_mutation_engine._get_charge_difference(current_resname = topology_proposal._old_topology.residue_topology.name,
-                                                                       new_resname = topology_proposal._new_topology.residue_topology.name)
+            charge_diff = point_mutation_engine._get_charge_difference(current_resname=topology_proposal.old_topology.residue_topology.name,
+                                                                       new_resname=topology_proposal.new_topology.residue_topology.name)
             _logger.info(f"charge diff: {charge_diff}")
             if charge_diff != 0:
-                new_water_indices_to_ionize = point_mutation_engine.get_water_indices(charge_diff, new_positions, topology_proposal._new_topology, radius=0.8)
+                new_water_indices_to_ionize = point_mutation_engine.get_water_indices(charge_diff, new_positions, topology_proposal.new_topology, radius=0.8)
                 _logger.info(f"new water indices to ionize {new_water_indices_to_ionize}")
-                PointMutationExecutorRBD._modify_new_system(new_water_indices_to_ionize, topology_proposal._new_system, charge_diff)
-                PointMutationExecutorRBD._modify_atom_classes(new_water_indices_to_ionize, topology_proposal)
+                self._get_ion_and_water_parameters(topology_proposal.old_system, topology_proposal.old_topology)
+                self._transform_waters_into_ions(new_water_indices_to_ionize, topology_proposal.new_system, charge_diff)
+                PointMutationExecutor._modify_atom_classes(new_water_indices_to_ionize, topology_proposal)
 
-            factories = []
-            if vanilla:
+            if generate_unmodified_hybrid_topology_factory:
                 repartitioned_endstate = None
                 self.generate_htf(HybridTopologyFactory, topology_proposal, pos, new_positions, flatten_exceptions, flatten_torsions, repartitioned_endstate, is_complex)
-            if repartitioned:
-                for repartitioned_endstate in [0, 1]:
+            if generate_rest_capable_hybrid_topology_factory:
+                 for repartitioned_endstate in [0, 1]:
                     self.generate_htf(RepartitionedHybridTopologyFactory, topology_proposal, pos, new_positions, flatten_exceptions, flatten_torsions, repartitioned_endstate, is_complex)
 
+
         if is_temp:
             shutil.rmtree(debug_dir)
 
-    def generate_htf(self, factory, topology_proposal, old_positions, new_positions, flatten_exceptions, flatten_torsions, repartitioned_endstate, is_complex):
-        htf = factory(topology_proposal=topology_proposal,
-                                      current_positions=old_positions,
-                                      new_positions=new_positions,
-                                      use_dispersion_correction=False,
-                                      functions=None,
-                                      softcore_alpha=None,
-                                      bond_softening_constant=1.0,
-                                      angle_softening_constant=1.0,
-                                      soften_only_new=False,
-                                      neglected_new_angle_terms=[],
-                                      neglected_old_angle_terms=[],
-                                      softcore_LJ_v2=True,
-                                      softcore_electrostatics=True,
-                                      softcore_LJ_v2_alpha=0.85,
-                                      softcore_electrostatics_alpha=0.3,
-                                      softcore_sigma_Q=1.0,
-                                      interpolate_old_and_new_14s=flatten_exceptions,
-                                      omitted_terms=None,
-                                      endstate=repartitioned_endstate,
-                                      flatten_torsions=flatten_torsions)
-        if is_complex:
-            if factory == HybridTopologyFactory:
-                self.complex_htf = htf
-            elif factory == RepartitionedHybridTopologyFactory:
-                if repartitioned_endstate == 0:
-                    self.complex_rhtf_0 = htf
-                elif repartitioned_endstate == 1:
-                    self.complex_rhtf_1 = htf
-        else:
-            if factory == HybridTopologyFactory:
-                self.apo_htf = htf
-            elif factory == RepartitionedHybridTopologyFactory:
-                if repartitioned_endstate == 0:
-                    self.apo_rhtf_0 = htf
-                elif repartitioned_endstate == 1:
-                    self.apo_rhtf_1 = htf
-
-    def get_complex_rhtf_0(self):
-        return self.complex_rhtf_0
-
-    def get_apo_rhtf_0(self):
-        return self.apo_rhtf_0
-
-    def get_complex_rhtf_1(self):
-        return self.complex_rhtf_1
-
-    def get_apo_rhtf_1(self):
-        return self.apo_rhtf_1
-
-    def _modify_new_system(water_atoms, system, charge_diff):
-        """
-        given a system and an array of ints (corresponding to atoms to turn into ions), modify the nonbonded particle parameters in the system such that the Os are turned into the ion of interest and the charges of the Hs are zeroed.
-        Parameters
-        ----------
-        water_atoms : np.array(int)
-            integers corresponding to particle indices to neutralize
-        system : simtk.openmm.System
-            system to modify
-        charge_diff : int
-            the charge difference between the old_system - new_system
-        Returns
-        -------
-        modify system in place
-        """
-        # Determine which ion to turn the water into
-        if charge_diff < 0: # Turn water into Cl-
-            ion_charge, ion_sigma, ion_epsilon = CL_CHARGE, CL_SIGMA, CL_EPSILON
-        elif charge_diff > 0: # Turn water into Na+
-            ion_charge, ion_sigma, ion_epsilon = NA_CHARGE, NA_SIGMA, NA_EPSILON
-
-        # Scale the nonbonded terms of the water atoms
-        force_dict = {i.__class__.__name__: i for i in system.getForces()}
-        if 'NonbondedForce' in [i for i in force_dict.keys()]:
-            nbf = force_dict['NonbondedForce']
-            for idx in water_atoms:
-                idx = int(idx)
-                charge, sigma, epsilon = nbf.getParticleParameters(idx)
-                if charge == O_CHARGE:
-                    nbf.setParticleParameters(idx, ion_charge, ion_sigma, ion_epsilon)
-                elif charge == H_CHARGE:
-                    nbf.setParticleParameters(idx, charge*0.0, sigma, epsilon)
-                else:
-                    raise Exception(f"Trying to modify an atom that is not part of a water residue. Atom index: {idx}")
-
-    @staticmethod
-    def _modify_atom_classes(water_atoms, topology_proposal):
-        """
-        Modifies:
-        - topology proposal._core_new_to_old_atom_map - add the ion(s) to neutralize
-        - topology_proposal._new_environment_atoms - remove the ion(s) to neutralize
-        - topology_proposal._old_environment_atoms - remove the ion(s) to neutralize
-        
-        Parameters
-        ----------
-        water_atoms : np.array(int)
-            integers corresponding to particle indices to turn into ions
-        topology_proposal : perses.rjmc.TopologyProposal
-            topology_proposal to modify
- 
-        """
-        for new_index in water_atoms:
-            old_index = topology_proposal._new_to_old_atom_map[new_index]
-            topology_proposal._core_new_to_old_atom_map[new_index] = old_index
-            topology_proposal._new_environment_atoms.remove(new_index)
-            topology_proposal._old_environment_atoms.remove(old_index)    
-
     def _correct_topology(self, original_topology, is_apo=True):
 
         """

perses/rjmc/topology_proposal.py

@@ -2146,7 +2146,7 @@ def propose(self,
 
         # index_to_new_residues : dict, key : int (index) , value : str (three letter name of proposed residue)
         _logger.debug(f"\tconstructing atom map for TopologyProposal...")
-        atom_map, old_res_to_oemol_map, new_res_to_oemol_map, local_atom_map_stereo_sidechain, current_oemol_sidechain, proposed_oemol_sidechain, old_oemol_res_copy, new_oemol_res_copy  = self._construct_atom_map(residue_map, old_topology, index_to_new_residues, new_topology, ignore_sidechain_atoms=True)
+        atom_map, old_res_to_oemol_map, new_res_to_oemol_map, local_atom_map_stereo_sidechain, current_oemol_sidechain, proposed_oemol_sidechain, old_oemol_res_copy, new_oemol_res_copy  = self._construct_atom_map(residue_map, old_topology, index_to_new_residues, new_topology, ignore_sidechain_atoms=False)
         _logger.debug(f"\tadding indices of the 'C' backbone atom in the next residue and the 'N' atom in the previous")
         _logger.debug(f"\t{list(index_to_new_residues.keys())[0]}")
         extra_atom_map = self._find_adjacent_residue_atoms(old_topology, new_topology, list(index_to_new_residues.keys())[0])