DOCSEARCHbyAlgolia
Submit Docs

chai-lab

repository·main·Indexed 24 days ago

https://github.com/chaidiscovery/chai-lab

Tools for Chai-1, a multi-modal foundation model for molecular structure prediction of proteins, small molecules, DNA, RNA, and glycosylations. The library provides a Python API and CLI for running inference, managing Multiple Sequence Alignments (MSAs) via .aligned.pqt files, and specifying structural constraints through covalent bond restraints and inter-chain contact files.

Tokens
3.9K
Snippets
7
Records
25
Agent score
84%

What's inside chai-lab

  1. How to provide MSAs to Chai-1

    main

    Chai-1 supports Multiple Sequence Alignments (MSAs) in two ways:

    1. Local Files: Provide an aligned.pqt file. This format is similar to a3m but includes additional metadata columns for source databases and sequence pairing keys.
    2. Automatic Generation: Use the --use-msa-server flag to automatically generate MSAs via the ColabFold MMseqs2 server.

  2. How to provide custom templates to Chai-1

    main

    Template loading is a two-step process:

    1. Read an m8 file: This file provides a table of template hits to load.
    2. Load hits: The model either downloads the identifier from RCSB or looks for a local file in the directory defined by CHAI_TEMPLATE_CIF_FOLDER.

    Note: Local files must follow the naming convention: $CHAI_TEMPLATE_CIF_FOLDER/identifier.cif.gz.

  3. Use experimental restraints and covalent bonds

    main
    Chai-1 allows users to guide the folding process using experimental restraints. These can specify inter-chain contacts or covalent bonds at various resolutions. This is useful for specifying branched ligands or known structural constraints.

  4. Specify multi-ring glycans using abbreviated syntax

    main

    Chai-1 uses an abbreviated syntax within the fasta record to define bonds within glycans. This syntax builds the glycan outward from a root sugar.

    Syntax: ROOT_CCD(connection_syntax NEXT_CCD)

    • The X-Y syntax (e.g., 4-1) draws a bond between the atom specified by X in the preceding ring and the atom specified by Y in the subsequent ring.

    Linear Chain Example:

    >glycan|4-NAG-in-a-linear-chain
NAG(4-1 NAG(4-1 NAG(4-1 NAG)))

    Branched Glycan Example:

    >glycan|branched-glycan
NAG(4-1 NAG(4-1 BMA(3-1 MAN)(6-1 MAN)))

    In the branched example, a root NAG is connected to a NAG, which is connected to a BMA. The BMA then branches to two MAN rings via 3-1 and 6-1 connections.

    To connect these complex glycans to a protein, use the standard restraints file format to link a protein residue/atom to the root glycan's atom (e.g., @C1).

    NAG(4-1 NAG(4-1 BMA(3-1 MAN)(6-1 MAN)))

  5. Use the .aligned.pqt file format for MSA information

    main

    To improve Chai-1 performance beyond 'single-sequence mode', you can provide evolutionary information via Multiple Sequence Alignments (MSAs) using the .aligned.pqt file format. This format is a Parquet-encoded dataframe that acts as an augmented a3m file, allowing you to track multiple databases and specify how sequences across different chains should be paired.

    Each .aligned.pqt file must contain exactly one query sequence and all its associated alignment hits. The file must include the following four columns:

    • sequence: The alignment hits in a3m format.
    • source_database: The database origin. Valid values are uniprot, uniref90, bfd_uniclust, mgnify, or query. The query key must appear exactly once as the first row. If using a custom database, uniref90 is a recommended catchall.
    • pairing_key: A string used to pair alignments across different chains in a complex (similar to AlphaFold-Multimer). Sequences with the same pairing_key will have their MSAs matched. This is typically a species identifier.
    • comment: A free-text string for human readability (ignored by the model).

    Note: The source_database is featurized as an input to the model, so its value affects the prediction.

    |
| sequence | source_database | pairing_key  | comment                       |
| -------- | --------------- | ------------ | ----------------------------- |
| RKDSS... | query           |              | query sequence                |
| RKDES... | uniref90        |              | A fun sequence from uniref90  |
| RKSES... | uniprot         | Mus musculus | A mouse sequence from uniprot |
|

  6. Use automatic MSA generation via ColabFold server

    main

    If you do not have local MSA files, you can use the ColabFold server for automatic generation.

    Via CLI Use the chai-lab command with the --msa-server flag:

    chai-lab fold input.fasta output_directory --msa-server

    Via Python API Set msa_server=True when calling run_inference:

    candidates = run_inference(
    ...
    msa_server=True,
    ...
)

  7. Install Chai-1

    main

    You can install the stable version of chai_lab from PyPI or the latest development version directly from GitHub.

    System Requirements:

    • OS: Linux
    • Python: 3.10 or later
    • Hardware: GPU with CUDA and bfloat16 support.
      • Recommended: A100 80GB, H100 80GB, or L40S 48GB.
      • Minimum: A10, A30, or consumer-grade RTX 4090.

  8. Specify covalent bonds and modifications in Chai-1

    main

    Chai-1 allows you to specify covalent linkages between atoms in a folded complex using a restraints file. This is useful for modeling glycosylation or arbitrary non-canonical bonds.

    Important Constraints:

    • Intra-chain bonds: Chai-1 was not trained on intra-chain bonds (e.g., disulfides); behavior for these is unverified.
    • Modified Amino Acids: Do not use bond restraints to specify modified amino acids that have an existing CCD code. Instead, include the CCD code in parentheses directly in the protein sequence (e.g., RKDES(MSE)EES for selenomethionine).
    • Atom Naming: For non-glycan ligands, ensure atom names match the version of rdkit used in chai-lab. You can check your version with uv pip list | grep rdkit.

  9. Format restraints for Chai-1 using CSV

    main

    Chai-1 allows users to guide protein complex folding by providing "restraints" in a .csv file. These restraints specify inter-chain contacts at different resolutions.

    Restraint Types

    • contact: Specifies an interaction between two specific residues in two distinct chains. Both res_idxA and res_idxB must be provided.
    • pocket: A coarser, asymmetric restraint. It specifies that an entire chain is in contact with a specific residue in another chain. To use this, leave the res_idxA field empty.

    Column Definitions

    ColumnDescription
    restraint_idA unique identifier for the restraint.
    chainAIdentifier for the first chain (A-Z based on input order).
    res_idxAResidue and 1-based index (e.g., D4). Must match the input sequence.
    chainBIdentifier for the second chain (A-Z based on input order).
    res_idxBResidue and 1-based index (e.g., G7).
    connection_typeEither contact or pocket.
    confidenceCurrently unused (future-proofing).
    min_distance_angstromCurrently unused (future-proofing).
    max_distance_angstromThe upper bound distance for the restraint.
    commentUser-provided metadata (ignored by the model).

    Important Rules

    • Residue Indexing: Use a concatenation of the residue letter and its 1-based index (e.g., R84). The model validates that the residue at that index matches your input; a mismatch will trigger an error.
    • Chain Identifiers: Chains are assigned letters A, B, C, etc., in alphabetical order based on their order in the input file.

  10. Connect non-glycan ligands via covalent bonds

    main

    You can connect residues to non-glycan ligands using the same covalent connection type in the restraints file. For non-glycan ligands, it is recommended to use a SMILES string in the fasta file.

    Note on Leaving Atoms: When using SMILES for non-sugar ligands, specify the SMILES string without any leaving atoms (e.g., hydroxyl groups that would be lost during bond formation). While Chai-1 automatically handles leaving atoms for glycan sugar rings, it does not do this for other ligand types.

  11. Convert A3M files to .aligned.pqt via CLI or Python

    main

    You can convert standard .a3m files into the .aligned.pqt format using the provided utility functions.

    Command Line Interface Use the chai CLI to perform the conversion:

    chai a3m-to-pqt --help

    Python API Use the merge_multi_a3m_to_aligned_dataframe function located in chai_lab.data.parsing.msas.aligned_pqt.

    Warning: The default parsing logic for pairing keys is based on species annotation in UniProt files. If you need to use different pairing logic or different data sources, you should implement custom parsing logic to generate your .aligned.pqt files.

    chai a3m-to-pqt --help

  12. Configure Model Weight and Template Storage

    main

    Chai-1 uses environment variables to manage where data is stored:

    • CHAI_DOWNLOADS_DIR: Controls the location where model weights are downloaded. By default, they are stored in the package's downloads directory within site-packages.
    • CHAI_TEMPLATE_CIF_FOLDER: Specifies a directory for custom (non-RCSB) structure .cif.gz files used as templates. The files must be named $CHAI_TEMPLATE_CIF_FOLDER/identifier.cif.gz, where identifier matches the entry in your m8 template file.