Add Open Babel Rex module support

docs/modules.rst

@@ -5,6 +5,7 @@ Rush Modules
 
    rush.exess
    rush.nnxtb
+   rush.openbabel
    rush.prepare_protein
    rush.prepare_complex
    rush.auto3d

docs/rush.openbabel.rst

@@ -0,0 +1,7 @@
+Open Babel
+==========
+
+.. automodule:: rush.openbabel
+   :members:
+   :show-inheritance:
+   :undoc-members:

src/rush/client.py

@@ -95,6 +95,7 @@ def _get_project_id() -> str:
         "exess_qmmm_rex": "github:talo/tengu-exess/61b1874f8df65a083e9170082250473fd8e46978#exess_qmmm_rex",
         "mmseqs2_rex": "github:talo/tengu-colabfold/749a096d082efdac3ac13de4aaa98aee3347d79d#mmseqs2_rex",
         "nnxtb_rex": "github:talo/tengu-nnxtb/4e733660264d38faab5d23eadc41ca86fd6ff97a#nnxtb_rex",
+        "openbabel_openbabel_protonate_rex": "github:talo/openbabel_rex/9f75cf354e62561f15bfa0d7728f52f7b0812f01#openbabel_openbabel_protonate_rex",
         "pbsa_rex": "github:talo/pbsa-cuda/f8b1c357fddfebf7e0c51a84f8d4e70958440c00#pbsa_rex",
         "prepare_protein_rex": "github:talo/tengu-prepare-protein/911d9fb69c269f8783b7cce2de3e87ee9333fb08#prepare_protein_rex",
     }

src/rush/openbabel.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+"""
+Open Babel module helpers for the Rush Python client.
+"""
+
+import sys
+from pathlib import Path
+from string import Template
+
+from gql.transport.exceptions import TransportQueryError
+
+from .client import (
+    RunOpts,
+    RunSpec,
+    _get_project_id,
+    _submit_rex,
+    collect_run,
+    upload_object,
+)
+
+
+def openbabel_openbabel_protonate_rex(
+    input_json: Path | str,
+    config_rex: str,
+    run_spec: RunSpec = RunSpec(target="Bullet"),
+    run_opts: RunOpts = RunOpts(),
+    collect=False,
+):
+    """
+    Protonate a ligand with Open Babel using TRC JSON input.
+
+    Args:
+        input_json: Path to the TRC JSON input object.
+        config_rex: Rex expression for the config struct, e.g.
+            "openbabel_openbabel_protonate_rex::ProtonateConfig { ph = Some 7.4, babel_libdir = None, babel_datadir = None }".
+        run_spec: Rush compute resources to request.
+        run_opts: Rush run metadata.
+        collect: Whether to wait for completion and return outputs.
+    """
+
+    input_vobj = upload_object(input_json)
+
+    rex = Template(
+        """let
+  obj_j = λ j →
+    VirtualObject { path = j, format = ObjectFormat::json, size = 0 },
+  run = λ input →
+    openbabel_openbabel_protonate_rex_s
+      ($run_spec)
+      ($config_rex)
+      (obj_j input)
+in
+  run "$input_vobj_path"
+"""
+    ).substitute(
+        run_spec=run_spec._to_rex(),
+        config_rex=config_rex,
+        input_vobj_path=input_vobj["path"],
+    )
+
+    try:
+        run_id = _submit_rex(_get_project_id(), rex, run_opts)
+        if collect:
+            return collect_run(run_id)
+        return run_id
+    except TransportQueryError as e:
+        if e.errors:
+            print("Error:", file=sys.stderr)
+            for error in e.errors:
+                print(f"  {error['message']}", file=sys.stderr)

tests/data/openbabel/input.json

@@ -0,0 +1,62 @@
+{
+  "topology": {
+    "schema_version": "0.2.0",
+    "symbols": [
+      "O",
+      "H",
+      "H"
+    ],
+    "geometry": [
+      0.0,
+      0.0,
+      0.0,
+      0.7586,
+      0.0,
+      0.5043,
+      -0.7586,
+      0.0,
+      0.5043
+    ],
+    "labels": null,
+    "partial_charges": null,
+    "formal_charges": null,
+    "connectivity": null,
+    "stereochemistry": null,
+    "velocities": null,
+    "fragments": null,
+    "fragment_formal_charges": null,
+    "fragment_partial_charges": null,
+    "fragment_multiplicities": null
+  },
+  "residues": {
+    "residues": [
+      [
+        0,
+        1,
+        2
+      ]
+    ],
+    "seqs": [
+      "LIG"
+    ],
+    "seq_ns": [
+      1
+    ],
+    "insertion_codes": [
+      ""
+    ],
+    "labeled": null,
+    "labels": null
+  },
+  "chains": {
+    "chains": [
+      [
+        0
+      ]
+    ],
+    "alpha_helices": null,
+    "beta_sheets": null,
+    "labeled": null,
+    "labels": null
+  }
+}

tests/test_onboarding_openbabel.py

@@ -0,0 +1,50 @@
+import os
+import time
+from pathlib import Path
+
+import pytest
+from gql.transport.exceptions import TransportConnectionFailed
+
+
+def test_openbabel_wrapper_import():
+    import rush.openbabel as openbabel
+
+    assert hasattr(openbabel, "openbabel_openbabel_protonate_rex")
+
+
+def test_openbabel_module_lock():
+    from rush import client
+
+    assert "openbabel_openbabel_protonate_rex" in client.MODULE_LOCK
+
+
+@pytest.mark.skipif(
+    not os.getenv("RUSH_TOKEN") or not os.getenv("RUSH_PROJECT"),
+    reason="RUSH_TOKEN and RUSH_PROJECT are required for live runs",
+)
+def test_openbabel_protonate_run():
+    from rush import openbabel
+    from rush.client import RunError, RunSpec
+
+    input_json = Path(__file__).resolve().parent / "data" / "openbabel" / "input.json"
+    config_rex = (
+        "openbabel_openbabel_protonate_rex::ProtonateConfig "
+        "{ ph = Some 7.4, babel_libdir = None, babel_datadir = None }"
+    )
+
+    result = None
+    for attempt in range(3):
+        try:
+            result = openbabel.openbabel_openbabel_protonate_rex(
+                input_json,
+                config_rex,
+                run_spec=RunSpec(target="Bullet"),
+                collect=True,
+            )
+            break
+        except TransportConnectionFailed:
+            if attempt == 2:
+                raise
+            time.sleep(2)
+
+    assert not isinstance(result, RunError)