$ show chaidiscovery/chai-lab
Chai-1 is a molecular structure prediction tool that takes FASTA sequences (proteins, DNA, RNA, ligands, glycans) and predicts 3D structures using a diffusion model. It is packaged as a Python library with a CLI (`chai-lab fold`) and a Python API (`chai_lab.chai1.run_inference`). The pipeline ingests sequences, optionally generates MSAs via ColabFold server, loads templates from RCSB, computes ESM embeddings, applies user-defined restraints (contacts, pockets, covalent bonds), tokenizes all atoms, generates a rich set of features, collates them into batches, and runs a TorchScript diffusion model to produce multiple structure samples. Outputs are ranked by confidence metrics (pTM, ipTM, pLDDT, clash score) and saved as CIF files. The codebase is well-structured with clean abstractions (feature generators, data contexts, collation), strong typing (jaxtyping, beartype), and CI quality gates. It is a hybrid of a production-grade inference tool and research code, with some copied components (rigid transformations from AlphaFold, ColabFold MSA server) and a monolithic inference entry point. The repository is actively maintained with recent commits, PyPI releases, and comprehensive examples.
The LLM identified 2 files worth highlighting for reuse:
Provides `Rigid` and `Rotation` classes for representing and manipulating rigid body transformations (rotation + translation) using quaternions and rotation matrices. Supports composition, inversion, application to points/vectors, and tensor-like indexing/stacking/cat. This is a self-contained, well-tested-style module adapted from AlphaFold/OpenFold, reusable in any 3D geometry or molecular modeling project.
Collection of generic PyTorch tensor utility functions with jaxtyping annotations: `cdist` (pairwise distances), `masked_mean`, `one_hot`, `und` (upper triangular without diagonal), `und_self`, `string_to_tensorcode`. These are reusable across any deep learning project that needs type-safe tensor manipulations, especially for structural biology or graph-like data.