Pipelines#

This module contains pipeline implementations for different molecular structure prediction tasks.

AF3 Pipeline#

atomworks.ml.pipelines.af3.build_af3_transform_pipeline(*, is_inference: bool, protein_msa_dirs: list[dict], rna_msa_dirs: list[dict], n_recycles: int = 5, crop_size: int = 384, crop_center_cutoff_distance: float = 15.0, crop_contiguous_probability: float = 0.5, crop_spatial_probability: float = 0.5, max_atoms_in_crop: int | None = None, undesired_res_names: list[str] = ['144', '15P', '1PE', '2F2', '2JC', '3HR', '3SY', '7N5', '7PE', '9JE', 'AAE', 'ABA', 'ACE', 'ACN', 'ACT', 'ACY', 'AZI', 'BAM', 'BCN', 'BCT', 'BDN', 'BEN', 'BME', 'BO3', 'BTB', 'BTC', 'BU1', 'C8E', 'CAD', 'CAQ', 'CBM', 'CCN', 'CIT', 'CL', 'CLR', 'CM', 'CMO', 'CO3', 'CPT', 'CXS', 'D10', 'DEP', 'DIO', 'DMS', 'DN', 'DOD', 'DOX', 'EDO', 'EEE', 'EGL', 'EOH', 'EOX', 'EPE', 'ETF', 'FCY', 'FJO', 'FLC', 'FMT', 'FW5', 'GOL', 'GSH', 'GTT', 'GYF', 'HED', 'IHP', 'IHS', 'IMD', 'IOD', 'IPA', 'IPH', 'LDA', 'MB3', 'MEG', 'MES', 'MLA', 'MLI', 'MOH', 'MPD', 'MRD', 'MSE', 'MYR', 'N', 'NA', 'NH2', 'NH4', 'NHE', 'NO3', 'O4B', 'OHE', 'OLA', 'OLC', 'OMB', 'OME', 'OXA', 'P6G', 'PE3', 'PE4', 'PEG', 'PEO', 'PEP', 'PG0', 'PG4', 'PGE', 'PGR', 'PLM', 'PO4', 'POL', 'POP', 'PVO', 'SAR', 'SCN', 'SEO', 'SIN', 'SO4', 'SPD', 'SPM', 'SR', 'STE', 'STO', 'STU', 'TAR', 'TBU', 'TME', 'TRS', 'UNK', 'UNL', 'UNX', 'UPL', 'URE'], conformer_generation_timeout: float = 5.0, use_element_for_atom_names_of_atomized_tokens: bool = False, max_n_template: int = 20, n_template: int = 4, template_max_seq_similarity: float = 60.0, template_min_seq_similarity: float = 10.0, template_min_length: int = 10, template_allowed_chain_types: list[ChainType] = [ChainType.POLYPEPTIDE_L, ChainType.RNA], template_distogram_bins: Tensor = tensor([3.2500, 4.5338, 5.8176, 7.1014, 8.3851, 9.6689, 10.9527, 12.2365, 13.5203, 14.8041, 16.0878, 17.3716, 18.6554, 19.9392, 21.2230, 22.5068, 23.7905, 25.0743, 26.3581, 27.6419, 28.9257, 30.2095, 31.4932, 32.7770, 34.0608, 35.3446, 36.6284, 37.9122, 39.1959, 40.4797, 41.7635, 43.0473, 44.3311, 45.6149, 46.8986, 48.1824, 49.4662, 50.7500]), template_default_token: str = '<G>', template_lookup_path: PathLike | None = None, template_base_dir: PathLike | None = None, max_msa_sequences: int = 10000, n_msa: int = 10000, dense_msa: bool = True, msa_cache_dir: PathLike | str | None = None, sigma_data: float = 16.0, diffusion_batch_size: int = 48, run_confidence_head: bool = False, return_atom_array: bool = True, pad_dna_p_skip: float = 0.0, b_factor_min: float | None = None, b_factor_max: float | None = None) → Transform[source]#

Build the AF3 pipeline with specified parameters.

This function constructs a pipeline of transforms for processing protein structures in a manner similar to AlphaFold 3. The pipeline includes steps for removing hydrogens, adding annotations, atomizing residues, cropping, adding templates, encoding features, and generating reference molecule features.

Parameters:

  • crop_size (int, optional) – The size of the crop. Defaults to 384.

  • crop_center_cutoff_distance (float, optional) – The cutoff distance for spatial cropping. Defaults to 15.0.

  • crop_contiguous_probability (float, optional) – The probability of using contiguous cropping. Defaults to 0.5.

  • crop_spatial_probability (float, optional) – The probability of using spatial cropping. Defaults to 0.5.

  • conformer_generation_timeout (float, optional) – The timeout for conformer generation in seconds. Defaults to 10.0.

Returns:

A composed pipeline of transforms.

Return type:

Transform

Raises:

  • AssertionError – If the crop probabilities do not sum to 1.0, if the crop size is not positive,

  • or if the crop center cutoff distance is not positive. –

Note

The cropping method is chosen randomly based on the provided probabilities. The pipeline includes steps for processing the structure, adding annotations, and generating features required for AF3-like predictions.

Reference:

AlphaFold 3 Supplementary Information

RF2AA Pipeline#

class atomworks.ml.pipelines.rf2aa.RF2AAInputs(seq: ndarray, msa: ndarray, msa_masked: ndarray, msa_full: ndarray, mask_msa: ndarray, xyz: ndarray, mask: ndarray, idx_pdb: ndarray, xyz_t: ndarray, t1d: ndarray, mask_t: ndarray, xyz_prev: ndarray, mask_prev: ndarray, same_chain: ndarray, unclamp: ndarray, negative: ndarray, atom_frames: ndarray, bond_feats: ndarray, dist_matrix: ndarray, chirals: ndarray, ch_label: ndarray, symmgp: str, task: str, example_id: str)[source]#

Bases: NamedTuple

A named tuple containing the inputs to the RF2AA model.

atom_frames: ndarray#

Alias for field number 16

bond_feats: ndarray#

Alias for field number 17

ch_label: ndarray#

Alias for field number 20

chirals: ndarray#

Alias for field number 19

dist_matrix: ndarray#

Alias for field number 18

example_id: str#

Alias for field number 23

classmethod from_dict(data: dict) → RF2AAInputs[source]#
idx_pdb: ndarray#

Alias for field number 7

mask: ndarray#

Alias for field number 6

mask_msa: ndarray#

Alias for field number 4

mask_prev: ndarray#

Alias for field number 12

mask_t: ndarray#

Alias for field number 10

msa: ndarray#

Alias for field number 1

msa_full: ndarray#

Alias for field number 3

msa_masked: ndarray#

Alias for field number 2

negative: ndarray#

Alias for field number 15

num_atoms() → int[source]#
num_res() → int[source]#
same_chain: ndarray#

Alias for field number 13

seq: ndarray#

Alias for field number 0

symmgp: str#

Alias for field number 21

t1d: ndarray#

Alias for field number 9

task: str#

Alias for field number 22

to_atom_array(symm_copy: int = 0) → AtomArray[source]#

Decode the inputs into an AtomArray for the given symm_copy.

to_dict() → dict[source]#
unclamp: ndarray#

Alias for field number 14

xyz: ndarray#

Alias for field number 5

xyz_prev: ndarray#

Alias for field number 11

xyz_t: ndarray#

Alias for field number 8

atomworks.ml.pipelines.rf2aa.build_rf2aa_transform_pipeline(*, protein_msa_dirs: list[dict], rna_msa_dirs: list[dict], n_recycles: int = 5, crop_size: int = 256, crop_center_cutoff_distance: float = 15.0, crop_spatial_probability: float = 0.5, crop_contiguous_probability: float = 0.5, unresolved_ligand_atom_limit: int | float | None = 0.1, undesired_res_names: list[str] = ['144', '15P', '1PE', '2F2', '2JC', '3HR', '3SY', '7N5', '7PE', '9JE', 'AAE', 'ABA', 'ACE', 'ACN', 'ACT', 'ACY', 'AZI', 'BAM', 'BCN', 'BCT', 'BDN', 'BEN', 'BME', 'BO3', 'BTB', 'BTC', 'BU1', 'C8E', 'CAD', 'CAQ', 'CBM', 'CCN', 'CIT', 'CL', 'CLR', 'CM', 'CMO', 'CO3', 'CPT', 'CXS', 'D10', 'DEP', 'DIO', 'DMS', 'DN', 'DOD', 'DOX', 'EDO', 'EEE', 'EGL', 'EOH', 'EOX', 'EPE', 'ETF', 'FCY', 'FJO', 'FLC', 'FMT', 'FW5', 'GOL', 'GSH', 'GTT', 'GYF', 'HED', 'IHP', 'IHS', 'IMD', 'IOD', 'IPA', 'IPH', 'LDA', 'MB3', 'MEG', 'MES', 'MLA', 'MLI', 'MOH', 'MPD', 'MRD', 'MSE', 'MYR', 'N', 'NA', 'NH2', 'NH4', 'NHE', 'NO3', 'O4B', 'OHE', 'OLA', 'OLC', 'OMB', 'OME', 'OXA', 'P6G', 'PE3', 'PE4', 'PEG', 'PEO', 'PEP', 'PG0', 'PG4', 'PGE', 'PGR', 'PLM', 'PO4', 'POL', 'POP', 'PVO', 'SAR', 'SCN', 'SEO', 'SIN', 'SO4', 'SPD', 'SPM', 'SR', 'STE', 'STO', 'STU', 'TAR', 'TBU', 'TME', 'TRS', 'UNK', 'UNL', 'UNX', 'UPL', 'URE'], res_names_to_atomize: list[str] | None = None, max_msa_sequences: int = 10000, dense_msa: bool = True, n_msa_cluster_representatives: int = 256, msa_n_extra_rows: int = 1024, msa_mask_probability: float = 0.15, msa_mask_behavior_probs: dict[str, float] = {'do_not_replace': 0.1, 'replace_with_msa_profile': 0.1, 'replace_with_random_aa': 0.1}, order_independent_atom_frame_prioritization: bool = True, polymer_token_indices: Tensor = tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]), n_template: int = 5, pick_top_templates: bool = False, template_max_seq_similarity: float = 60.0, template_min_seq_similarity: float = 10.0, template_min_length: int = 10, template_lookup_path: PathLike | None = None, template_base_dir: PathLike | None = None, max_automorphs: int = 1000, max_isomorphs: int = 1000, use_negative_interface_examples: bool = False, unclamp_loss_probability: float = 0.1, black_hole_init: bool = True, black_hole_init_noise_scale: float = 5.0, msa_cache_dir: PathLike | str | None = None, assert_rf2aa_assumptions: bool = True, convert_feats_to_rf2aa_input_tuple: bool = True, is_inference: bool = False) → Compose[source]#

Creates a transformation pipeline for the RF2AA model, applying a series of transformations to the input data.

Parameters:

  • protein_msa_dirs (-) – The directories containing the protein MSAs and their associated file types, as a list of dictionaries. If multiple directories are provided, we will search all of them. Note that: (a) the directory structure must be flat (i.e., no subdirectories), (b) the files must be named using the SHA-256 hash of the sequence (see hash_sequence in utils/misc), and (c) order matters - we will search the directories in the order they are provided, and return the first match

  • rna_msa_dirs (-) – The directories containing the RNA MSAs and their associated file types, as a list of dictionaries. See protein_msa_dirs for directory structure details.

  • n_recycles (-) – Number of recycles for the MSA featurization. Defaults to 5.

  • crop_size (-) – Size of the crop for spatial and contiguous cropping (in number of tokens). Defaults to 384.

  • crop_center_cutoff_distance (-) – Cutoff distance for the center of the crop (in Angstroms). Defaults to 15.0.

  • crop_spatial_probability (-) – Probability of performing spatial cropping. Defaults to 0.5.

  • crop_contiguous_probability (-) – Probability of performing contiguous cropping. Defaults to 0.5.

  • unresolved_ligand_atom_limit (-) – Limit for above which a ligand is considered unresolved. many unresolved atoms has its atoms removed. If None, all atoms are kept, if between 0 and 1, the number of atoms is capped at that percentage of the crop size. If an integer >= 1, the number of unresolved atoms is capped at that number. Defaults to 0.1.

  • res_names_to_atomize (-) – List of residue names to always atomize. Note that RF2AA already atomizes all residues that are not in the encoding (i.e. that are not standard AA, RNA, DNA or special masks). Therefore only specify this if you want to deterministically atomize certain standard tokens. Defaults to None.

  • max_msa_sequences (-) – Maximum number of MSA sequences to load. Defaults to 10,000.

  • dense_msa (-) – Whether to use dense MSA pairing. Defaults to True.

  • n_msa_cluster_representatives (-) – Number of MSA cluster representatives to select. Defaults to 100.

  • msa_n_extra_rows (-) – Number of extra rows for MSA. Defaults to 100.

  • msa_mask_probability (-) – Probability of masking MSA sequences according to msa_mask_behavior_probs. Defaults to 0.15.

  • msa_mask_behavior_probs (-) – Probabilities for different MSA mask behaviors. Defaults to {“replace_with_random_aa”: 0.1, “replace_with_msa_profile”: 0.1, “do_not_replace”: 0.1}, which is the BERT style masking.

  • order_independent_atom_frame_prioritization (-) – Whether to prioritize order-independent atom frames. Defaults to True.

  • n_template (-) – Number of templates to use. Defaults to 5.

  • pick_top_templates (-) – Whether to pick the top templates if there are more than n_template. If False, the templates are selected randomly among all templates. Defaults to False.

  • template_max_seq_similarity (-) – Maximum sequence similarity cutoff for templates. Defaults to 60.0.

  • template_min_seq_similarity (-) – Minimum sequence similarity cutoff for templates. Defaults to 10.0.

  • template_min_length (-) – Minimum length cutoff for templates. Defaults to 10.

  • max_automorphs (-) – Maximum number of automorphs after which to cap small molecule ligand symmetry resolution. Defaults to 1,000.

  • max_isomorphs (-) – Maximum number of polymer isomorphs after which to cap symmetry resolution. Defaults to 1,000.

  • use_negative_interface_examples (-) – Whether to use negative interface examples. Defaults to False.

  • unclamp_loss_probability (-) – Probability of unclamping the loss during training. Defaults to 0.1.

  • black_hole_init (-) – Whether to use black hole initialization. Defaults to True.

  • black_hole_init_noise_scale (-) – Noise scale for black hole initialization. Defaults to 5.0.

  • msa_cache_dir (-) – Directory to cache the MSAs. Defaults to None.

  • assert_rf2aa_assumptions (-) – Whether to assert the RF2AA assumptions that need to be true to guarantee a successful forward & backward pass. Defaults to True.

  • convert_feats_to_rf2aa_input_tuple (-) – Whether to convert the features to the RF2AAInputs format. Defaults to True.

For more details on the parameters, see the RF2AA paper and the documentation for the respective Transforms.

Returns:

A composed transformation pipeline.

Return type:

Compose