Kdidi/ligand clean

.github/workflows/ci.yml

@@ -17,7 +17,7 @@ concurrency:
   cancel-in-progress: true
 
 env:
-  SIF: /net/software/containers/versions/modelhub/pytorch_24.10-py3.sif
+  SIF: /net/scratch/kdidi/pytorch_25.11-py3-x11.sif
   PIP_PKGS: /net/scratch/bench/ci-pip-pkgs
 
 
@@ -163,8 +163,8 @@ jobs:
           export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
           export CMAKE_PREFIX_PATH=$(python3 -c 'import torch; print(torch.utils.cmake_prefix_path)')
 
-          pytest -p no:rerunfailures -s -v --durations=25 --cov=$GITHUB_WORKSPACE/tmol.cov \
-            --junitxml=$GITHUB_WORKSPACE/testing.cuda.junit.xml -k "cuda"
+          pytest -p no:rerunfailures -s -v --durations=25 \
+            --cov=$GITHUB_WORKSPACE/tmol.cov --junitxml=$GITHUB_WORKSPACE/testing.cuda.junit.xml -k "cuda"
 
           GPU_TEST
 
@@ -213,5 +213,18 @@ jobs:
           paths: |
             testing.*.junit.xml
 
+      - name: Test Report
+        if: always()
+        uses: dorny/test-reporter@v1
+        with:
+          name: Test Results
+          path: 'testing.*.junit.xml'
+          reporter: java-junit
+          fail-on-error: false
 
-
+      - name: Upload benchmark results
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: benchmark-results
+          path: benchmark/**/*

.gitignore

@@ -127,4 +127,16 @@ venv.bak/
 
 # macos
 ._*
+# CMake / scikit-build build artifacts
+CMakeCache.txt
+CMakeFiles/
+CMakeInit.txt
+cmake_install.cmake
+build.ninja
+.ninja_deps
+.ninja_log
+.cmake/
+.skbuild-info.json
+detect_cuda_version.cc
+
 # End of https://www.gitignore.io/api/python,jupyternotebook

containers/apptainer/tmol-dev.def

@@ -31,9 +31,13 @@ IncludeCmd: yes
     ln -s /projects /mnt/projects
 
     # Update system and install essential packages
+    # libxrender1, libx11-6, libxext6 are runtime deps of openbabel-wheel plugins
     apt-get update && apt-get install -y \
         build-essential \
         git \
+        libxrender1 \
+        libx11-6 \
+        libxext6 \
         && rm -rf /var/lib/apt/lists/*
 
     # Fix CUDA compat warning in Apptainer (read-only filesystem)

containers/docker/tmol-dev.Dockerfile

@@ -16,6 +16,13 @@ LABEL description="tmol - Dependencies only (for development and scoring)"
 # Create directories for bind mounts
 RUN mkdir -p /tmol_host /projects /net /squash
 
+# Install system dependencies (runtime deps of openbabel-wheel plugins)
+RUN apt-get update && apt-get install -y \
+        libxrender1 \
+        libx11-6 \
+        libxext6 \
+    && rm -rf /var/lib/apt/lists/*
+
 # Copy pyproject.toml to extract dependencies
 COPY pyproject.toml /opt/tmol_pyproject.toml
 

dev/bin/nvprof2json.py

@@ -324,9 +324,7 @@ def inspect_db(conn):
 CUPTI_ACTIVITY_KIND_STREAM
 CUPTI_ACTIVITY_KIND_SYNCHRONIZATION
 CUPTI_ACTIVITY_KIND_UNIFIED_MEMORY_COUNTER
-""".strip().split(
-        "\n"
-    )
+""".strip().split("\n")
 
     for t in tables:
         eprint(t)

pyproject.toml

@@ -57,8 +57,11 @@ dependencies = [
     "astor>=0.8,<1",
     # Interactive
     "ipython>=8.0",
-    # Biotite
+    # Structure I/O
     "biotite>=1.2.0",
+    # Ligand preparation
+    "rdkit>=2024.3",
+    "openbabel-wheel==3.1.1.22",
 ]
 
 [project.optional-dependencies]
@@ -118,7 +121,10 @@ TMOL_BUILD_TESTS = {env = "TMOL_BUILD_TESTS", default = "OFF"}
 
 # Test settings ----------------------------------------------------------------
 [tool.pytest.ini_options]
-addopts = "--benchmark-disable --benchmark-columns=ops,mean,iqr"
+addopts = "--benchmark-disable --benchmark-columns=ops,mean,iqr -m 'not slow'"
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+]
 filterwarnings = [
     "ignore:(?s).*is not compatible with the compiler Pytorch.*:",
     "ignore:(?s).*Benchmark fixture was not used.*:",

pytest

@@ -0,0 +1 @@
+bash: Apptainer: command not found

tmol/__init__.py

@@ -6,6 +6,19 @@
 import contextlib
 from importlib.metadata import PackageNotFoundError, version
 
+
+def include_paths():
+    """C++/CUDA include paths for tmol components.
+
+    Defined before other imports because JIT extension loading
+    (tmol.utility.cpp_extension) imports this during module init.
+    """
+
+    import os.path
+
+    return [os.path.abspath(os.path.dirname(__file__) + "/..")]
+
+
 from tmol._load_ext import ensure_compiled_or_jit as _ensure_compiled_or_jit
 
 # Extensions may not be built yet (e.g. during sdist creation).
@@ -58,16 +71,7 @@
 from tmol.score.constraint.utility import create_mainchain_coordinate_constraints
 from tmol.relax.fast_relax import fast_relax
 
-
 try:
     __version__ = version("tmol")
 except PackageNotFoundError:
     __version__ = "unknown version"
-
-
-def include_paths():
-    """C++/CUDA include paths for tmol components."""
-
-    import os.path
-
-    return [os.path.abspath(os.path.dirname(__file__) + "/..")]

tmol/chemical/patched_chemdb.py

@@ -40,7 +40,7 @@ def from_raw_res(self, r):
 # delete atoms, bonds, torsions, and connections involving atom
 def remove_atom(res, atom):
     res.atoms = tuple(x for x in res.atoms if x.name != atom)
-    res.bonds = tuple((x, y) for x, y in res.bonds if x != atom and y != atom)
+    res.bonds = tuple(b for b in res.bonds if b[0] != atom and b[1] != atom)
     res.torsions = tuple(
         x for x in res.torsions if atom not in [x.a.atom, x.b.atom, x.c.atom, x.d.atom]
     )
@@ -131,7 +131,8 @@ def get_modified_atoms(patch):
         added.append(i.name)
 
     # modded finds all atoms whose CONNECTIVITY or COORDINATES have changed
-    for i, j in patch.add_bonds:
+    for bond in patch.add_bonds:
+        i, j = bond[0], bond[1]
         if i[0] == "<" and i[-1] == ">" and i not in modded:
             modded.append(i)
         if j[0] == "<" and j[-1] == ">" and j not in modded:
@@ -183,7 +184,8 @@ def _validate_raw_residue_atoms(res, allatoms):
 def _validate_raw_residue_bonds(res, allatoms):
     # illegal bonds
     bad_bonds = []
-    for i, j in res.bonds:
+    for bond in res.bonds:
+        i, j = bond[0], bond[1]
         if i not in allatoms:
             bad_bonds.append((i, i, j))
         if j not in allatoms:
@@ -382,7 +384,8 @@ def _validate_patch_atom_aliases(patch, addedatoms):
 def _validate_patch_bonds(patch, added_ats_and_conns):
     # make sure all bonds are references or added atoms
     bad_bonds = []
-    for i, j in patch.add_bonds:
+    for bond in patch.add_bonds:
+        i, j = bond[0], bond[1]
         if (i[0] != "<" or i[-1] != ">") and (i not in added_ats_and_conns):
             bad_bonds.append((i, j))
     if len(bad_bonds) > 0:
@@ -487,7 +490,8 @@ def do_patch(res, variant, resgraph, patchgraph, marked):
 
         # -1. Add atoms bonded to deleted atoms to modded set
         #     needs to be done after name map
-        for i, j in res.bonds:
+        for bond in res.bonds:
+            i, j = bond[0], bond[1]
             if i in deleted and j not in deleted:
                 modded.append(j)
             if j in deleted and i not in deleted:
@@ -538,12 +542,13 @@ def do_patch(res, variant, resgraph, patchgraph, marked):
 
         # 4. add bonds
         newbonds = []
-        for i, j in variant.add_bonds:
+        for bond in variant.add_bonds:
+            i, j, btype = bond[0], bond[1], bond[2]
             if i in namemap:
                 i = namemap[i]
             if j in namemap:
                 j = namemap[j]
-            newbonds.append((i, j))
+            newbonds.append((i, j, btype))
         newres.bonds = (*newres.bonds, *newbonds)
 
         # 5. update icoors

tmol/chemical/restypes.py

@@ -1,5 +1,7 @@
+from enum import IntEnum
+
 from frozendict import frozendict
-from toolz.curried import concat, map, compose, groupby
+from toolz.curried import concat, compose, groupby
 import typing
 from typing import Mapping, Optional, NewType, Tuple, Sequence, List, Set, Union
 import attr
@@ -22,6 +24,23 @@
 ConnectionIndex = NewType("ConnectionIndex", int)
 BondCount = NewType("BondCount", int)
 
+
+class BondType(IntEnum):
+    SINGLE = 1
+    DOUBLE = 2
+    TRIPLE = 3
+    RING = 4
+    AROMATIC = 5
+
+
+BOND_TYPE_FROM_STR = {
+    "SINGLE": BondType.SINGLE,
+    "DOUBLE": BondType.DOUBLE,
+    "TRIPLE": BondType.TRIPLE,
+    "RING": BondType.RING,
+    "AROMATIC": BondType.AROMATIC,
+}
+
 # As of cattr 24.1.0, more types must be explicitly registered in order to
 # use cattr.structure. We use that here
 cattr.register_structure_hook(numpy.dtype, lambda d, _: numpy.dtype(d))
@@ -152,12 +171,30 @@ def _setup_coord_dtype(self):
     def _setup_bond_indices(self):
         bondi = compose(list, sorted, set, concat)(
             [(ai, bi), (bi, ai)]
-            for ai, bi in map(map(self.atom_to_idx.get), self.bonds)
+            for ai, bi in (
+                (self.atom_to_idx.get(b[0]), self.atom_to_idx.get(b[1]))
+                for b in self.bonds
+            )
         )
         bond_array = numpy.array(bondi, dtype=numpy.int32)
         bond_array.flags.writeable = False
         return bond_array
 
+    bond_to_type: Mapping = attr.ib()
+
+    @bond_to_type.default
+    def _setup_bond_to_type(self):
+        bt_map = {}
+        for b in self.bonds:
+            a0, a1, btype_str = b[0], b[1], b[2]
+            ai = self.atom_to_idx.get(a0)
+            bi = self.atom_to_idx.get(a1)
+            if ai is not None and bi is not None:
+                bt_int = int(BOND_TYPE_FROM_STR.get(btype_str, BondType.SINGLE))
+                bt_map[(ai, bi)] = bt_int
+                bt_map[(bi, ai)] = bt_int
+        return frozendict(bt_map)
+
     @property
     def n_conn(self):
         return len(self.connections)
@@ -184,6 +221,20 @@ def _setup_ordered_connections(self):
             [self.atom_to_idx[c.atom] for c in self.connections], dtype=numpy.int32
         )
 
+    connection_bond_types: numpy.ndarray = attr.ib()
+
+    @connection_bond_types.default
+    def _setup_connection_bond_types(self):
+        arr = numpy.array(
+            [
+                int(BOND_TYPE_FROM_STR.get(c.type, BondType.SINGLE))
+                for c in self.connections
+            ],
+            dtype=numpy.int32,
+        )
+        arr.flags.writeable = False
+        return arr
+
     # The set of "all-bonds" includes both inter- and intra- block chemical bonds
     # Each all-bond (think of "all" here as an adjective like "inter" or "intra")
     # is described as a pair (atom-ind1, unresolved-atom id) where the