💊 EquiPocket-3D: Geometric Deep Learning for Drug Discovery

SE(3)-Equivariant 3D Graph Neural Network & Point Transformer for AlphaFold Drug Binding Pocket Detection

---

🧬 Overview

EquiPocket-3D is an open-source geometric deep learning engine designed to predict druggable small-molecule binding pockets and interaction surfaces directly from AlphaFold and PDB 3D protein structures.

In rational structure-based drug design, locating druggable binding cavities on target macromolecules is the critical first step before virtual screening and generative ligand docking. EquiPocket-3D implements a Pure PyTorch SE(3)-Equivariant Graph Neural Network (EGNN) architecture with multi-task prediction heads:

1. Per-Residue Druggability Classification: Invariant scoring $\hat{p}_i \in [0, 1]$ identifying residues lining active binding sites.
2. 3D Pocket Centroid Regression: Equivariant attention pooling predicting 3D spatial center coordinates $(\hat{x}, \hat{y}, \hat{z}) \in \mathbb{R}^3$.
3. Pocket Bounding Sphere Radius: Invariant regression of cavity interaction radius $\hat{r} \in \mathbb{R}^+$ in Angstroms.

                     ┌─────────────────────────────────────────┐
                     │   AlphaFold / PDB 3D Protein Structure  │
                     │         (Atomic Coordinates x,y,z)      │
                     └────────────────────┬────────────────────┘
                                          │
                                          ▼
                     ┌─────────────────────────────────────────┐
                     │       SE(3)-Equivariant Backbone        │
                     │    (Multi-Layer EGNN Message Passing)   │
                     └────────────┬────────────────────────────┘
                                  │
                  ┌───────────────┴───────────────┐
                  │                               │
                  ▼                               ▼
     ┌────────────────────────┐      ┌────────────────────────┐
     │ Invariant Druggability │      │ Equivariant 3D Pocket  │
     │   Classification Head  │      │  Centroid & Radius Head│
     │      p_i in [0, 1]     │      │   (x, y, z, r in Å)    │
     └────────────┬───────────┘      └────────────┬───────────┘
                  │                               │
                  └───────────────┬───────────────┘
                                  ▼
                     ┌─────────────────────────────────────────┐
                     │   Interactive 3D Visualizer Dashboard   │
                     │    (3Dmol.js Spinning Ribbon & Cavity)  │
                     └─────────────────────────────────────────┘

---

🚀 Key Innovations

• Strict SE(3)-Equivariance & Invariance: Guarantees that if a protein rotates or translates in 3D space, predicted pocket coordinates rotate identically ($\hat{C}(Rx + t) = R \hat{C}(x) + t$), while predicted druggability scores remain strictly invariant.
• Direct AlphaFold API Integration: Fetches predicted structures on the fly by UniProt Accession ID (e.g. P00533 for EGFR, P01116 for KRAS, P24941 for CDK2).
• Geometric Bio-Loss: Formulates a combined 3D Chamfer Distance coordinate loss, Radius Huber regularization, and Druggability Focal Loss to conquer severe surface vs. pocket class imbalance.
• Interactive 3D Web Dashboard: Built with Gradio and embedded 3Dmol.js for real-time 3D protein ribbon rendering and glowing pocket cavity sphere visualization.

---

🔬 Mathematical Formulation: SE(3)-Equivariance in PyTorch

Given a 3D protein graph with $N$ amino acid residues where $x_i \in \mathbb{R}^3$ are 3D spatial coordinates and $h_i \in \mathbb{R}^d$ are invariant node features:

1. Invariant Edge Messages

$$m_{ij} = \phi_m\left(h_i^l, h_j^l, |x_i^l - x_j^l|^2\right)$$

2. Equivariant Coordinate Updates

$$x_i^{l+1} = x_i^l + \sum_{j \in \mathcal{N}(i)} \frac{x_i^l - x_j^l}{|x_i^l - x_j^l| + \epsilon} \cdot \phi_x(m_{ij})$$

3. Invariant Feature Updates

$$h_i^{l+1} = \text{LayerNorm}\left(h_i^l + \phi_h\left(h_i^l, \sum_{j \in \mathcal{N}(i)} m_{ij}\right)\right)$$

4. Equivariant Pocket Centroid Head

$$\hat{C}_k = \sum_{i=1}^N \alpha_{ki} x_i, \quad \text{where } \alpha_{ki} = \text{softmax}_i\left(\frac{q_k^\top k_i}{\sqrt{d}}\right)$$

---

⚡ Quick Start & Installation

# Clone repository
git clone https://github.com/margot-bonilla/equipocket_3d.git
cd equipocket_3d

# Set up virtual environment and install package
make env
make install

---

🧪 Automated Testing

Run the automated mathematical proof and integration test suite:

make test

---

🖥️ Interactive 3D Visualizer Demo

Launch the browser-based 3D drug discovery dashboard:

make demo
# or: python app.py

---

📜 License & Citation

Distributed under the GNU Affero General Public License v3.0 (AGPL-3.0).