diff --git a/openff/interchange/components/mdconfig.py b/openff/interchange/components/mdconfig.py index ce4c39662..23ae72547 100644 --- a/openff/interchange/components/mdconfig.py +++ b/openff/interchange/components/mdconfig.py @@ -221,7 +221,7 @@ def write_lammps_input( def _get_coeffs_of_constrained_bonds_and_angles( interchange: "Interchange", - ) -> tuple[set[int], set[int]]: + ) -> tuple[set[int], set[int], set[int]]: """ Get coefficients of bonds and angles that appear to be constrained. @@ -233,7 +233,7 @@ def _get_coeffs_of_constrained_bonds_and_angles( """ constraint_styles = {key.associated_handler for key in interchange["Constraints"].potentials} - if len(constraint_styles.difference({"Bonds", "Angles"})) > 0: + if len(constraint_styles.difference({"Bonds", "Angles", "Constraints"})) > 0: raise NotImplementedError( "Found unsupported constraints case in LAMMPS input writer.", ) @@ -246,6 +246,10 @@ def _get_coeffs_of_constrained_bonds_and_angles( key.id for key in interchange["Constraints"].potentials if key.associated_handler == "Angles" } + constraint_ids = { + key.id for key in interchange["Constraints"].potentials if key.associated_handler == "Constraints" + } + return ( { key @@ -261,12 +265,20 @@ def _get_coeffs_of_constrained_bonds_and_angles( ).items() if val.id in constrained_angle_smirks }, + { + key + for key, val in dict( + enumerate(interchange["Constraints"].potentials), + ).items() + if val.id in constraint_ids + }, ) # zero-indexed here ( constrained_bond_coeffs, constrained_angle_coeffs, + constraint_coeffs, ) = _get_coeffs_of_constrained_bonds_and_angles(interchange) # Construct the input file in memory so nothing is written to disk if an @@ -360,7 +372,7 @@ def _get_coeffs_of_constrained_bonds_and_angles( "thermo_style custom ebond eangle edihed eimp epair evdwl ecoul elong etail pe\n\n", ) - if len(constrained_bond_coeffs.union(constrained_angle_coeffs)) > 0: + if len(constrained_bond_coeffs.union(constrained_angle_coeffs).union(constraint_coeffs)) > 0: # https://docs.lammps.org/fix_shake.html # TODO: Apply fix to just a group (sub-group)? lmp.write( @@ -377,6 +389,11 @@ def _get_coeffs_of_constrained_bonds_and_angles( f"a {' '.join([str(val + 1) for val in constrained_angle_coeffs])}", ) + if constraint_coeffs: + lmp.write( + f"a {' '.join([str(val + 1) for val in constraint_coeffs])}", + ) + lmp.write("\n") if self.coul_method == _PME: