_protonate_fallback improvement with formal_charge

[speleo3]

I see that you recently added a protonate command and a _protonate_fallback implementation in commit 5456a67. That's really useful, thanks!

However, I'm skeptical about the _protonate_fallback implementation, which uses valence and overall looks a bit too complicated to me. In my understanding, h_add results are best controlled by setting formal_charge first.

How about an implementation like this?

from pymol import cmd

_TITRATABLE_PKA_FORMAL_CHARGE = {
    # (resn, name): (pKa, charge below, charge above)

    ('ARG', 'NE'):  (12.48, 0, 0),
    ('ARG', 'NH1'): (12.48, 1, 0),
    ('ARG', 'NH2'): (12.48, 0, 0),

    ('LYS', 'NZ'):  (10.53, 1, 0),

    ('TYR', 'OH'):  (10.07, 0, -1),
    ('CYS', 'SG'):   (8.18, 0, -1),

    ('HIE', 'ND1'):  (6.00, 1, 0),
    ('HIP', 'ND1'):  (6.00, 1, 0),
    ('HIS', 'ND1'):  (6.00, 1, 0),

    ('HIE', 'NE2'):  (6.00, 0, 0),
    ('HIP', 'NE2'):  (6.00, 0, 0),
    ('HIS', 'NE2'):  (6.00, 0, 0),
 
    ('GLU', 'OE1'):  (4.25, 0, 0),
    ('GLU', 'OE2'):  (4.25, 0, -1),

    ('ASP', 'OD1'):  (3.65, 0, 0),
    ('ASP', 'OD2'):  (3.65, 0, -1),
}



def _get_formal_charge(resn: str, name: str, pH: float, fc_fallback: float) -> float:
    value = _TITRATABLE_PKA_FORMAL_CHARGE.get((resn, name))
    if value is None:
        return fc_fallback
    pka, fc_below, fc_above = value
    return fc_below if pH < pka else fc_above


@cmd.extendaa(cmd.auto_arg[0]["h_add"])
def protonate_fc(selection: str = "all", pH: float = 7.4, *,
                 state=-1, quiet=1, _self=cmd):
    _self.remove(f"({selection}) & hydro")
    _self.alter(
        selection,
        f"formal_charge = _get_formal_charge(resn, name, {float(pH)}, formal_charge)",
        space={"_get_formal_charge": _get_formal_charge},
        quiet=quiet)
    _self.h_add(selection, state=state, quiet=quiet)


def test_protonate_fc():
    cmd.reinitialize()
    cmd.fab("EHK", "m1")
    protonate_fc("m1", pH=3.5)
    assert cmd.count_atoms("hydro") == 30
    protonate_fc("m1", pH=5.9)
    assert cmd.count_atoms("hydro") == 29
    protonate_fc("m1", pH=6.1)
    assert cmd.count_atoms("hydro") == 28
    protonate_fc("m1", pH=13.0)
    assert cmd.count_atoms("hydro") == 27

[JarrettSJohnson]

That's fair. I should have taken a closer look at ObjectMoleculeVerifyChemistry. Feel free to submit a PR.

[speleo3]

Thanks Jarrett. I'll test this solution a bit more before opening a PR.