Feature: Add number_hydrogens= keyword argument to capping ILP

Author: bertrand-caron

helpers/capping.py

@@ -11,6 +11,7 @@
 def get_best_capped_molecule_with_ILP(
     molecule: 'Molecule',
     net_charge: Optional[int] = None,
+    number_hydrogens: Optional[int] = None,
     enforce_octet_rule: bool = True,
     allow_radicals: bool = False,
     debug: Optional[TextIO] = None,
@@ -21,6 +22,7 @@ def get_best_capped_molecule_with_ILP(
     Args:
         ``molecule``: Molecule to be capped. Some of its atoms should have the ``capped`` attribute set to False.
         ``net_charge``: (Optional) Constraint the total net charge for the capped molecule.
+        ``number_hydrogens``: (Optional) Constraint the total number of hydrogens for the capped molecule (including hydrogens already present in the molecule).
         ``enforce_octet_rule``: (Optional) Constraint organic elements (H, C, N, O) to satisfy the octet rule.
         ``allow_radicals``: (Optional) Allow unpaired non-bonded electrons.
         ``debug``: (Optional) Print very detailed debugging information
@@ -94,6 +96,8 @@ def possible_capping_strategies_for_atom(atom: Atom) -> List[Capping_Strategy]:
 
     ELECTRON_MULTIPLIER = (2 if not allow_radicals else 1)
 
+    hydrogens_before_capping = len([1 for atom in molecule.atoms.values() if atom.element == 'H'])
+
     counter_charges = {}
     fragment_switches, fragment_scores, fragment_H_scores = {}, {}, {}
     capping_atoms_for = {}
@@ -187,7 +191,8 @@ def possible_capping_strategies_for_atom(atom: Atom) -> List[Capping_Strategy]:
     ]
 
     H_size_objective = MAX(lpSum([F_i * fragment_H_scores[uncapped_atom_id, i] for ((uncapped_atom_id, i), F_i) in fragment_switches.items()]))
-    if sum([fragment_H_scores[uncapped_atom_id, i] for ((uncapped_atom_id, i), F_i) in fragment_switches.items()]) != 0:
+    has_non_null_H_size_objective = (sum([fragment_H_scores[uncapped_atom_id, i] for ((uncapped_atom_id, i), F_i) in fragment_switches.items()]) != 0)
+    if has_non_null_H_size_objective:
         OBJECTIVES.append(H_size_objective)
 
     total_size_objective = MIN(lpSum([F_i * fragment_scores[uncapped_atom_id, i] for ((uncapped_atom_id, i), F_i) in fragment_switches.items()]))
@@ -202,6 +207,17 @@ def possible_capping_strategies_for_atom(atom: Atom) -> List[Capping_Strategy]:
     if net_charge is not None:
         problem += (lpSum(charges.values()) == net_charge, 'Known net charge')
 
+    if number_hydrogens is not None:
+        problem += (
+            lpSum(
+                [
+                    F_i * fragment_H_scores[uncapped_atom_id, i]
+                    for ((uncapped_atom_id, i), F_i) in fragment_switches.items()
+                ],
+            ) + hydrogens_before_capping == number_hydrogens,
+            'Total number of hydrogens: {0}'.format(number_hydrogens)
+        )
+
     for atom in molecule.atoms.values():
         problem += (
             charges[atom.index]

test_peptide_capping.py

@@ -5,6 +5,7 @@
 import unittest
 from sys import stdout
 from os.path import basename
+from typing import Optional
 
 from fragment_capping.helpers.molecule import Atom, Molecule, molecule_from_pdb_str
 from fragment_capping.helpers.compare import compare_capped_molecules
@@ -20,7 +21,7 @@ def element_for_line(line: str) -> str:
 def atom_name_for_line(line: str) -> str:
     return line[13:16].replace(' ', '')
 
-def test_capping_peptide(peptide_pdb_filepath: str, peptide_net_charge: int) -> None:
+def test_capping_peptide(peptide_pdb_filepath: str, peptide_net_charge: int, peptide_number_hydrogens: Optional[int] = None) -> None:
     peptide_name, _ = basename(peptide_pdb_filepath).split('.')
 
     with open(peptide_pdb_filepath) as fh:
@@ -72,6 +73,7 @@ def enforce_valence_aromatic_rings() -> None:
     new_molecule = input_molecule.get_best_capped_molecule_with_ILP(
         enforce_octet_rule=True,
         net_charge=old_charge,
+        number_hydrogens=peptide_number_hydrogens,
     )
     print((datetime.now() - now).total_seconds())
     new_formula = new_molecule.formula(charge=True)
@@ -100,13 +102,13 @@ def enforce_valence_aromatic_rings() -> None:
 
 class Test_Peptide_Capping(unittest.TestCase):
     def test_capping_AAA(self):
-        test_capping_peptide('pdbs/AAA.pdb', 0)
+        test_capping_peptide('pdbs/AAA.pdb', 0, None)
 
     def test_capping_ARNDCEQGHILKMFPSTWYV(self):
-        test_capping_peptide('pdbs/ARNDCEQGHILKMFPSTWYV.pdb', 0)
+        test_capping_peptide('pdbs/ARNDCEQGHILKMFPSTWYV.pdb', 0, None)
 
     def test_capping_2OVN_all(self) -> None:
-        test_capping_peptide('pdbs/2OVN_with_connects.pdb', -4)
+        test_capping_peptide('pdbs/2OVN_with_connects.pdb', -4, 144)
 
 if __name__ == '__main__':
     unittest.main(warnings='ignore', verbosity=2)