Add EquiBind-PyG: Rigid Protein–Ligand Docking Model for PyTorch Geometric (contrib)

[soowei-su]

This Pull Request adds a modular, PyG-native implementation of the rigid EquiBind docking model to the torch_geometric.contrib namespace.
The implementation targets clarity, modularity, and ease of extension, making it suitable as a teaching reference and a foundation for future docking research inside the PyG ecosystem.

---

What’s Included

1. Model Implementation

• EquiBindRigid model with:

  • Equivariant EGNN / iEGNN message passing layers
  • Keypoint generation + alignment
  • Rigid-body transformation prediction
  • Coordinate refinement
  • Batch support via PyG Batch

2. Geometry Utilities

• Differentiable Kabsch alignment (kabsch.py)
• RMSD and centroid metrics (metrics.py)

3. Loss Functions

• Modular rigid-docking loss (losses.py) combining:

  • Coordinate RMSD
  • Kabsch RMSD
  • Centroid loss
  • Optional regularization hooks

4. EGNN / iEGNN Layers

• Clean, readable implementations adapted for PyG messaging conventions.
• Designed to be reusable outside this model.

5. contrib API Exposure

All components are exported via:

torch_geometric.contrib.equibind_pyg

6. Minimal Unit Test

test_equibind_pyg.py verifies:

• Module import structure
• A forward pass with synthetic ligand/receptor graphs
• Correct output tensor shapes

7. Documentation

A lightweight README.md describing:

• Model architecture
• Usage
• Supported components
• Design constraints for a rigid-only baseline

---

Test Status

• pytest test/contrib/test_equibind_pyg.py -q → PASS
• All added files are isolated under contrib/ and do not affect core PyG functionality.

---

Notes

• This contribution intentionally implements only the rigid docking portion of EquiBind to keep scope manageable for contrib.
• The codebase is structured to allow later extensions (e.g., torsion prediction, full flexible docking).
• No external dependencies beyond PyTorch and PyG.

---

Acknowledgments

This work adapts the original EquiBind method (Stärk et al., ICML 2022) into a modern PyG idiom with a focus on readability and modularity.