Preprocessing#

This module contains utilities for preprocessing molecular structures and data.

Core Preprocessing Functions#

Preprocessing pipeline for PDB structures.

Provides a pipeline for preprocessing PDB structures: 1. parse_entry() - Load structure and fetch entry-level metadata like ligand fit-to-density 2. process_assemblies() - Filter atoms, resolve clashes, create Assembly records 3. process_pn_units() - Extract PN units and contacts 4. process_interfaces() - Create Interface records from contacts 5. _add_closest_chain_info() - Add closest_20/50_pn_unit_iids to assemblies (uses outputs from steps 1-4)

class atomworks.ml.preprocessing.preprocess.PreprocessConfig(contact_distance: float = 5.0, clash_distance: float = 1.0, ignore_residues: tuple[str, ...] = (), polymer_pn_unit_limit: int = 1000, build_assembly: str = 'all', add_missing_atoms: bool = True, remove_waters: bool = True, remove_ccds: tuple[str, ...] = ('SO4', 'GOL', 'EDO', 'PO4', 'ACT', 'PEG', 'DMS', 'TRS', 'PGE', 'PG4', 'FMT', 'EPE', 'MPD', 'MES', 'CD', 'IOD'), fix_ligands_at_symmetry_centers: bool = True, fix_arginines: bool = True, convert_mse_to_met: bool = True, hydrogen_policy: Literal['remove', 'infer', 'keep'] = 'remove', ligand_scores: tuple[str, ...] = ('RSCC', 'RSR', 'completeness', 'intermolecular_clashes', 'is_best_instance', 'ranking_model_fit', 'ranking_model_geometry'), count_af3_tokens: bool = True, sample_large_assembly_interfaces: bool = True, large_assembly_threshold: int = 20, interface_cutoff_distance: float = 15.0)[source]#

Bases: object

Configuration for preprocessing pipeline.

All parameters have sensible defaults. Create custom configs by passing only the parameters you want to change.

add_missing_atoms: bool = True#
build_assembly: str = 'all'#
clash_distance: float = 1.0#
contact_distance: float = 5.0#
convert_mse_to_met: bool = True#
count_af3_tokens: bool = True#
fix_arginines: bool = True#
fix_ligands_at_symmetry_centers: bool = True#
hydrogen_policy: Literal['remove', 'infer', 'keep'] = 'remove'#
ignore_residues: tuple[str, ...] = ()#
interface_cutoff_distance: float = 15.0#
large_assembly_threshold: int = 20#
ligand_scores: tuple[str, ...] = ('RSCC', 'RSR', 'completeness', 'intermolecular_clashes', 'is_best_instance', 'ranking_model_fit', 'ranking_model_geometry')#
polymer_pn_unit_limit: int = 1000#
remove_ccds: tuple[str, ...] = ('SO4', 'GOL', 'EDO', 'PO4', 'ACT', 'PEG', 'DMS', 'TRS', 'PGE', 'PG4', 'FMT', 'EPE', 'MPD', 'MES', 'CD', 'IOD')#
remove_waters: bool = True#
sample_large_assembly_interfaces: bool = True#
atomworks.ml.preprocessing.preprocess.parse_entry(path: Path, config: PreprocessConfig) → dict[str, Any][source]#

Load structure and fetch ligand validity scores.

Parameters:

  • path – Path to structure file (.cif, .cif.gz, .pdb, .pdb.gz).

  • config – Preprocessing configuration.

Returns:

Parser dict with ‘ligand_validity_scores’ key added.

atomworks.ml.preprocessing.preprocess.preprocess(path: Path, config: PreprocessConfig | None = None) → tuple[list[Assembly], list[PNUnit], list[Interface]][source]#

Preprocess a structure file into normalized records.

This is the main entry point that runs all 5 stages: 1. parse_entry() - Load and prepare structure 2. process_assemblies() - Filter and create Assembly records 3. process_pn_units() - Extract PN units and contacts 4. process_interfaces() - Create Interface records 5. _add_closest_chain_info() - Add closest chain info for large assemblies

Parameters:

  • path – Path to structure file (.cif, .cif.gz, .pdb, .pdb.gz).

  • config – Preprocessing configuration. Uses defaults if None.

Returns:

Tuple of (assemblies, pn_units, interfaces).

atomworks.ml.preprocessing.preprocess.process_assemblies(parser_dict: dict[str, Any], config: PreprocessConfig) → tuple[list[Assembly], dict[str, AtomArray]][source]#

Filter atoms, resolve clashes, and create Assembly records.

Parameters:

  • parser_dict – Output from parse_entry().

  • config – Preprocessing configuration.

Returns:

Tuple of (Assembly list, dict mapping assembly_id to filtered AtomArray).

atomworks.ml.preprocessing.preprocess.process_interfaces(pn_units: list[PNUnit], contact_info: list[ContactInfo]) → list[Interface][source]#

Create Interface records from contacts.

Parameters:

  • pn_units – PNUnit records from process_pn_units().

  • contact_info – ContactInfo records from process_pn_units().

Returns:

List of Interface records.

atomworks.ml.preprocessing.preprocess.process_pn_units(assemblies: list[Assembly], assembly_atoms: dict[str, AtomArray], parser_dict: dict[str, Any], config: PreprocessConfig) → tuple[list[PNUnit], list[ContactInfo]][source]#

Extract PN units and compute contacts.

Parameters:

  • assemblies – Assembly records from process_assemblies().

  • assembly_atoms – Filtered AtomArrays from process_assemblies().

  • parser_dict – Original parser dict (for chain_info, ligand_info, ligand_validity_scores).

  • config – Preprocessing configuration.

Returns:

Tuple of (PNUnit list, ContactInfo list).

atomworks.ml.preprocessing.preprocess.save_records(out_dir: Path, pdb_id: str, assemblies: list[Assembly], pn_units: list[PNUnit], interfaces: list[Interface], exclude: set[str] | None = None) → None[source]#

Save preprocessing result to parquet files.

Creates:

out_dir/assemblies/{pdb_id}.parquet out_dir/pn_units/{pdb_id}.parquet out_dir/interfaces/{pdb_id}.parquet

Parameters:

  • out_dir – Base output directory.

  • pdb_id – PDB identifier for file naming.

  • assemblies – List of Assembly records.

  • pn_units – List of PNUnit records.

  • interfaces – List of Interface records.

  • exclude – Optional set of field names to exclude from output.

Constants#

Constants for ML preprocessing.

class atomworks.ml.preprocessing.constants.ClashSeverity(*values)[source]#

Bases: Enum

Severity levels for atomic clashes in a structure.

MILD = 'mild'#

Any clashes (polymer or non-polymer).

MODERATE = 'moderate'#

Any polymers are clashing.

NO_CLASH = 'no-clash'#

No clashes detected.

SEVERE = 'severe'#

More than 50% of polymers are clashing.

class atomworks.ml.preprocessing.constants.DistanceThresholds[source]#

Bases: object

Centralized distance thresholds (Angstroms) for preprocessing.

CLASH: Final[float] = 1.0#

Atoms closer than this are considered clashing.

CONTACT: Final[float] = 5.0#

Distance threshold for atoms in contact.

INTERFACE: Final[float] = 15.0#

Cutoff for interface atom identification in large assembly sampling.

atomworks.ml.preprocessing.constants.PDB_IDS_THAT_CRASH_PREPROCESSING: Final[tuple[str, ...]] = ('1m4x', '1uf2', '3jan', '3jbp', '3k1q', '3zif', '4f5x', '4l3b', '4u3n', '4v5x', '5dat', '5j7v', '5jus', '5tbw', '5zap', '6b43', '6cgr', '6ftg', '6ftj', '6lgl', '6ncl', '6nhj', '6nu3', '6q1f', '6tb3', '6vyr', '6w19', '6woo', '6ylh', '6zvk', '7as5', '7bho', '7bsi', '7btb', '7bw6', '7fj1', '7lhd', '7nwh', '7qiz', '7r5j', '7r5k', '7tbi', '7tbk', '7v4t')#

PDB IDs that crash or hang the preprocessing pipeline.

Warning

These include extremely large assemblies (ribosomes, nuclear pore complexes) and complex structures (DNA origami) that exceed memory/time limits.

atomworks.ml.preprocessing.constants.PDB_IDS_TO_EXCLUDE: Final[frozenset[str]] = frozenset({'1m4x', '1uf2', '3jan', '3jbp', '3k1q', '3zif', '4f5x', '4l3b', '4u3n', '4v5x', '5dat', '5j7v', '5jus', '5tbw', '5zap', '6b43', '6cgr', '6ftg', '6ftj', '6lgl', '6ncl', '6nhj', '6nu3', '6q1f', '6tb3', '6vyr', '6w19', '6woo', '6ylh', '6zvk', '7as5', '7bho', '7bsi', '7btb', '7bw6', '7fj1', '7lhd', '7nwh', '7qiz', '7r5j', '7r5k', '7tbi', '7tbk', '7v4t'})#

PDB IDs to exclude from preprocessing (only those that crash/hang).

atomworks.ml.preprocessing.constants.PDB_IDS_WITH_UNPHYSICAL_BONDS: Final[tuple[str, ...]] = ('185d', '193d', '1a0r', '1ay3', '1b4d', '1b6j', '1by5', '1c0t', '1c0u', '1c1b', '1can', '1cao', '1ci7', '1cpi', '1dtq', '1dtt', '1eva', '1evd', '1fav', '1fwm', '1gcm', '1jaa', '1jc8', '1jcp', '1jgc', '1l9w', '1niu', '1nma', '1nmb', '1pfe', '1pw8', '1qtn', '1qur', '1u8g', '1v4f', '1v6q', '1v7h', '1vs2', '1vtg', '1xhb', '1xql', '1xvk', '1xvn', '1xvr', '1zet', '2adw', '2bvr', '2bvs', '2c15', '2c3u', '2ddc', '2et0', '2fni', '2h77', '2j32', '2j58', '2jhg', '2llj', '2n4n', '2n6h', '2n6i', '2nbl', '2nmp', '2vlc', '2wdy', '2xe4', '2xjh', '2xji', '2xqt', '2yct', '3ah8', '3awo', '3bo7', '3bp2', '3f4z', '3go3', '3k87', '3kch', '3mgn', '3nuv', '3q9g', '3u4w', '3u51', '4auo', '4b3q', '4bp9', '4bqd', '4bxe', '4cak', '4ce5', '4chx', '4fa5', '4fa9', '4h0t', '4j5t', '4ks6', '4ktx', '4l1s', '4lnm', '4nec', '4oy5', '4rca', '4v15', '4w9f', '5aa2', '5bmm', '5bqc', '5c55', '5cs2', '5ctv', '5e83', '5eq7', '5eq8', '5f1w', '5g21', '5h1b', '5hkj', '5hwn', '5lgr', '5lvs', '5m3z', '5ol2', '5yty', '5ytz', '5zur', '6bpu', '6bpw', '6cdh', '6eik', '6euv', '6fbt', '6fcr', '6fcs', '6gmq', '6gmx', '6ibt', '6jze', '6m9z', '6skj', '6te5', '6tg4', '6th7', '6thq', '6uf4', '6uf7', '6uf8', '6uf9', '6vju', '6wve', '6x8k', '6x8l', '6xte', '6xu2', '6xu3', '6yfy', '6zuk', '7ajz', '7ak4', '7axp', '7axs', '7azd', '7bgq', '7dq0', '7dq8', '7f0c', '7f0f', '7jr4', '7oov', '7oow', '7opw', '7owm', '7pfz', '7pox', '7q1l', '7q9s', '7qgv', '7qtv', '7r0y', '7rnp', '7s50', '7ubz', '7v58', '7w3u', '7x6r', '7x97', '7x9f', '7xdj', '7z0l', '7z86', '7z9r', '7zss', '7zwh', '8a99', '8agy', '8ai7', '8an9', '8ao0', '8brn', '8cep', '8cx5', '8jug', '8sa8', '8sa9', '8sab', '8sng', '8snj', '8su6', '1b23', '1bcs', '1bjg', '1bwc', '1bx5', '1c2q', '1c2w', '1c4c', '1cfm', '1chw', '1d61', '1dd4', '1ddi', '1e1f', '1e3d', '1e7u', '1efd', '1f8c', '1f8d', '1fp4', '1fq5', '1fq8', '1g20', '1g21', '1g8j', '1ga5', '1gx8', '1h2j', '1hcy', '1hkx', '1hl9', '1i94', '1i95', '1i96', '1i97', '1j11', '1j12', '1j1a', '1jhi', '1jpc', '1kk4', '1kmk', '1ksj', '1l0l', '1m21', '1mbu', '1mbx', '1md2', '1n1h', '1n7d', '1nte', '1o7t', '1oao', '1pwf', '1q2c', '1q6k', '1q9x', '1qbv', '1qh1', '1qlb', '1qni', '1qnl', '1qnu', '1r6o', '1r6q', '1rpg', '1rqd', '1sgc', '1siw', '1sp4', '1sqb', '1sqq', '1sqv', '1sqx', '1t3p', '1th3', '1thc', '1tlf', '1twn', '1uwy', '1vje', '1vl3', '1vqw', '1w3l', '1w4r', '1w9m', '1wck', '1wdr', '1wvm', '1xv6', '1y4z', '1y6f', '1yfg', '1z7l', '2b7r', '2b7s', '2bjk', '2bvl', '2c3p', '2c5c', '2c6c', '2c6f', '2cch', '2cju', '2da8', '2dc7', '2dc8', '2dca', '2dcb', '2dcd', '2e1a', '2ef5', '2f1o', '2fwi', '2fyu', '2hhl', '2iwk', '2ja6', '2ja7', '2ja8', '2jji', '2m5b', '2min', '2nqx', '2o7a', '2olb', '2osm', '2p8x', '2rl7', '2uwn', '2uyn', '2uyv', '2uzs', '2v8f', '2v93', '2vzm', '2w31', '2wpu', '2x8j', '2xcr', '2xrz', '2xsj', '2y6e', '2yaj', '2yiv', '2ylj', '2ylo', '2zur', '3af7', '3bbf', '3bhv', '3c0v', '3c9e', '3cgb', '3d1l', '3d5x', '3dt0', '3en9', '3enh', '3esc', '3esd', '3f4i', '3h8l', '3hhm', '3hr3', '3i39', '3kku', '3kvs', '3l0o', '3lya', '3ni3', '3np2', '3ott', '3pa0', '3pot', '3pyk', '3q9h', '3q9i', '3q9j', '3rpd', '3rwr', '3rzs', '3sd3', '3sz0', '3szf', '3t0k', '3t2z', '3t4g', '3t5g', '3tuz', '3tvk', '3up9', '3usl', '3usm', '3uso', '3vah', '3w6a', '3ze9', '3zp9', '3zx0', '486d', '4a7n', '4a93', '4ag7', '4b09', '4b0n', '4c3a', '4c4m', '4car', '4cft', '4cpq', '4cps', '4cty', '4ctz', '4cu0', '4cu1', '4cul', '4cum', '4cun', '4cv5', '4cvg', '4cwv', '4cww', '4cwx', '4cwy', '4cwz', '4cx0', '4cx1', '4cx2', '4cyg', '4d34', '4d35', '4d36', '4d37', '4d38', '4d3a', '4daw', '4dlp', '4e0l', '4e0m', '4e0n', '4e0o', '4e54', '4e5z', '4ef1', '4ef2', '4eqw', '4eu9', '4fy0', '4g5f', '4gjs', '4got', '4i9b', '4ivh', '4j7v', '4j8w', '4j9u', '4jm0', '4js6', '4jsa', '4jsk', '4jsl', '4jsm', '4jss', '4k0k', '4k3f', '4k5h', '4k5i', '4k5j', '4k5k', '4kcp', '4kcq', '4kcr', '4kcs', '4kky', '4luw', '4mrr', '4n4k', '4n4l', '4n4m', '4n7k', '4n7l', '4n9r', '4nhp', '4nhq', '4nmm', '4ntp', '4ntr', '4nw8', '4nw9', '4o1w', '4ote', '4oxd', '4p4v', '4p4w', '4p4x', '4p4y', '4p4z', '4pa5', '4q5t', '4q8d', '4qgn', '4qo5', '4qp1', '4r6c', '4r7u', '4rlo', '4tt5', '4ttt', '4u6w', '4uaz', '4ubb', '4udk', '4uh7', '4uh8', '4uh9', '4uha', '4upq', '4upr', '4ups', '4upt', '4wc8', '4wna', '4x0s', '4xtd', '4xti', '4z12', '4z13', '4zc0', '4zjx', '4zqw', '5a4m', '5adj', '5adk', '5adl', '5adm', '5adn', '5brv', '5ci8', '5d51', '5dal', '5ddl', '5e9r', '5ecx', '5eul', '5f1t', '5fez', '5ff0', '5ff3', '5ff4', '5fj2', '5fj3', '5fle', '5fvy', '5fvz', '5giy', '5hes', '5hpp', '5ieg', '5ik9', '5iwd', '5jiw', '5kfp', '5lh6', '5lms', '5lpy', '5mb1', '5mdj', '5mdk', '5mdl', '5mf2', '5msg', '5mwx', '5n5r', '5ny6', '5nz6', '5oar', '5ob6', '5ob7', '5ob8', '5oeg', '5sur', '5sus', '5sut', '5suu', '5tpc', '5uhr', '5uod', '5v4g', '5v4h', '5v63', '5v64', '5v65', '5viv', '5vu9', '5vv7', '5vv8', '5vv9', '5vva', '5vvg', '5vvn', '5w4h', '5w4i', '5w4j', '5wnq', '5wnr', '5wns', '5wp6', '5x7h', '5zzf', '6b1r', '6b41', '6bo1', '6bo2', '6c6k', '6cdk', '6cg3', '6cg4', '6cg5', '6dr4', '6dr5', '6dr6', '6e1a', '6e5i', '6egw', '6ejw', '6erm', '6evj', '6f2f', '6f2h', '6f2i', '6f2j', '6f2m', '6frm', '6frn', '6gmi', '6gob', '6h8t', '6hf6', '6hf7', '6hkq', '6hmy', '6hsn', '6hso', '6hy6', '6i24', '6i25', '6ic8', '6jzm', '6kf5', '6ks1', '6lzb', '6m97', '6m9g', '6ma8', '6mik', '6mu5', '6nys', '6o7m', '6o9t', '6o9v', '6ofy', '6otl', '6pl9', '6pla', '6plb', '6pv2', '6q8j', '6qcv', '6qcw', '6qcx', '6qfu', '6qfv', '6qfx', '6qjs', '6qqi', '6qqk', '6qsa', '6qsh', '6qss', '6qxc', '6r2r', '6r93', '6rcf', '6rke', '6rqa', '6rr7', '6ruk', '6rve', '6s46', '6s4q', '6s50', '6s6w', '6s9z', '6swm', '6sym', '6tb2', '6tho', '6tvk', '6tvy', '6twc', '6u7y', '6udr', '6vfb', '6vjr', '6vu1', '6vu4', '6vvj', '6wgo', '6wnh', '6wxm', '6wxz', '6xuv', '6ycs', '6yh0', '6ykh', '6yo0', '6yxg', '6z4i', '6z7r', '6z8j', '6zre', '7a0d', '7a6q', '7abr', '7b0g', '7b3r', '7b97', '7bay', '7bpa', '7bpf', '7bpg', '7dix', '7efi', '7fbs', '7g96', '7jqr', '7jqs', '7jqt', '7jqu', '7jrh', '7jxn', '7jxo', '7jyy', '7jz0', '7ksv', '7ktm', '7kvv', '7l6r', '7l6t', '7lok', '7m4t', '7ne0', '7nh5', '7nwq', '7o1s', '7o25', '7o55', '7o6j', '7odg', '7odh', '7onm', '7onq', '7onv', '7otv', '7ph4', '7ph7', '7pr3', '7pwb', '7qdj', '7qhr', '7qpw', '7qpy', '7qpz', '7qy3', '7rnb', '7rsr', '7rtz', '7t4h', '7tne', '7tpn', '7tpo', '7tq0', '7tq9', '7tqc', '7tqe', '7tqf', '7tqh', '7tqi', '7tqj', '7tqk', '7tqw', '7u1b', '7uyn', '7uyo', '7uyp', '7vcn', '7vdt', '7wz9', '7xkj', '7y44', '7yd7', '7yd8', '7yeb', '7yec', '7yee', '7yx0', '7z60', '7z6a', '7z6d', '7zcy', '7zkn', '7zu8', '7zx5', '7zxl', '7zz3', '8a33', '8a5s', '8a6t', '8ab2', '8ai0', '8aqx', '8aqy', '8b7g', '8bqt', '8by0', '8cbz', '8crz', '8cs0', '8cs2', '8ec9', '8eca', '8es1', '8f4y', '8fo0', '8fxi', '8hyi', '8jeb', '8omy', '8on3', '8oqf', '8oqg', '8p4e', '8qm3', '8qml', '8w6p')#

PDB IDs with bond geometries that violate basic plausibility assumptions.

atomworks.ml.preprocessing.constants.PDB_REGEX: Final[Pattern[str]] = re.compile('^[0-9A-Za-z]{4}$')#

Regex pattern for valid PDB IDs (4 alphanumeric characters).

atomworks.ml.preprocessing.constants.PREPROCESSING_SUPPORTED_CHAIN_TYPES: Final[list[ChainType]] = [ChainType.NON_POLYMER, ChainType.POLYPEPTIDE_L, ChainType.POLYPEPTIDE_D, ChainType.DNA, ChainType.RNA, ChainType.BRANCHED, ChainType.DNA_RNA_HYBRID]#

Chain types supported during preprocessing.

atomworks.ml.preprocessing.constants.PREPROCESSING_SUPPORTED_CHAIN_TYPES_INTS: Final[list[int]] = [8, 6, 5, 3, 7, 10, 4]#

Integer values of chain types supported during preprocessing.

atomworks.ml.preprocessing.constants.TRAINING_SUPPORTED_CHAIN_TYPES: Final[list[ChainType]] = [ChainType.NON_POLYMER, ChainType.POLYPEPTIDE_L, ChainType.DNA, ChainType.RNA, ChainType.BRANCHED]#

Chain types supported during training.

atomworks.ml.preprocessing.constants.TRAINING_SUPPORTED_CHAIN_TYPES_INTS: Final[list[int]] = [8, 6, 3, 7, 10]#

Integer values of chain types supported during training.

Utilities#

MSA#

Note that the following functions can be called via the command line. See Multiple Sequence Alignment in AtomWorks for more details.

Finding#

Find existing MSA files and their locations.

atomworks.ml.preprocessing.msa.finding.find_msas(sequences: list[str], msa_dirs: list[dict] | list[PathLike] | None = None) → tuple[list[str], dict[str, Path]][source]#

Find existing MSA files for sequences and return missing sequences with MSA path mapping.

Parameters:

  • sequences – Protein sequences to find MSAs for.

  • msa_dirs – Directories to search. Accepts: - None: Uses PROTEIN_MSA_DIRS env var - list[PathLike]: Simple format, auto-detects metadata - list[dict]: Pre-computed format with ‘dir’, ‘extension’, ‘directory_depth’ keys

Returns:

  • missing_sequences: List of sequences without existing MSA files

  • sequence_to_msa_path: Dict mapping sequences to their MSA file paths

Return type:

Tuple of (missing_sequences, sequence_to_msa_path) where

Examples

Find MSAs with default settings (uses PROTEIN_MSA_DIRS env var):

sequences = ["MKKKEVE...", "MSYIWRQ..."]
missing, found_paths = find_msas(sequences)

Find MSAs with explicit directories:

dirs = get_msa_dirs("MY_MSA_DIRS")
missing, found_paths = find_msas(sequences, msa_dirs=dirs)
atomworks.ml.preprocessing.msa.finding.find_paired_msas(sequences: list[str], msa_dirs: list[dict] | list[PathLike] | None = None) → tuple[list[str], dict[str, Path]][source]#

Find the paired MSA files for one complex’s full unique sequence set.

Mirrors find_msas()’s shape, but for a paired/complex query: sequences is treated as the complete set of unique chains in ONE complex – the same unit organize_paired_msas/make_msas_colabfold_server operate on – and complex_id is derived from them automatically via get_complex_id, not something the caller computes or passes separately.

Parameters:

  • sequences – The full unique-sequence set of one complex.

  • msa_dirs – Directories to search, each containing paired MSA output (e.g. make_msas_colabfold_server’s <output_dir>/paired). Accepts the same formats as find_msas().

Returns:

Tuple of (missing_sequences, sequence_to_msa_path), same shape as find_msas().

atomworks.ml.preprocessing.msa.finding.find_template_alignments(sequences: list[str], template_dirs: list[dict] | list[PathLike] | None = None) → tuple[list[str], dict[str, Path]][source]#

Find existing raw template-alignment (pdb70.m8) files for sequences.

Mirrors find_msas()’s shape and hash/shard convention: template alignments are looked up by the same sequence hash as MSAs, just under a different directory/extension. To download the hit structures, see template_structures.

Parameters:

  • sequences – Protein sequences to find template alignments for.

  • template_dirs –

    Directories to search. Accepts:

    • None: No directories (all sequences reported missing).

    • list[PathLike]: Auto-assumes .m8 / directory_depth=2 (matching organize_template_alignments’s defaults).

    • list[dict]: Pre-computed format with ‘dir’, ‘extension’, ‘directory_depth’ keys.

Returns:

Tuple of (missing_sequences, sequence_to_template_alignment_path), same shape as find_msas().

atomworks.ml.preprocessing.msa.finding.get_msa_depth_and_ext_from_folder(folder: Path, max_depth: int = 10) → tuple[int, MSAFileExtension][source]#

Automatically detect the shard depth and file extension of an MSA folder.

Goes down subdirectories one level at a time until finding MSA files.

Parameters:

  • folder – Top-level MSA directory to analyze.

  • max_depth – Maximum depth to search (default: 10).

Returns:

  • shard_depth: Number of subdirectory levels (0 = files directly in folder)

  • extension: MSAFileExtension enum value for the file extension found

Return type:

Tuple of (shard_depth, extension) where

Raises:

ValueError – If no MSA files are found within max_depth levels.

Examples

For structure like /msa/ab/cd/abcd123.a3m.gz:

>>> depth, ext = get_msa_depth_and_ext_from_folder(Path("/msa"))
>>> # Returns: (2, MSAFileExtension.A3M_GZ)
atomworks.ml.preprocessing.msa.finding.get_msa_dirs(env_var_name: str = 'PROTEIN_MSA_DIRS', raise_if_not_set: bool = False) → list[dict][source]#

Get MSA directories with auto-detected metadata from an environment variable.

Parameters:

  • env_var_name – Environment variable with comma-separated paths.

  • raise_if_not_set – Raise if env var is unset. Defaults to False.

Returns:

List of dicts with dir, extension, and directory_depth keys.

atomworks.ml.preprocessing.msa.finding.get_msa_dirs_from_env(raise_if_not_set: bool = False) → list[dict][source]#

Deprecated: Use get_msa_dirs() instead.

atomworks.ml.preprocessing.msa.finding.get_paired_msa_path(complex_id: str, sequence: str, msa_dirs: list[dict] | list[PathLike] | None = None) → Path | None[source]#

Retrieve the path to a chain’s paired MSA file within a given complex.

Companion to get_msa_path (in atomworks.ml.transforms.msa._msa_loading_utils) for paired MSAs, which need two keys – which complex, and which chain within it – rather than a sequence hash alone.

Parameters:

  • complex_id – Identifier for the complex, from get_complex_id.

  • sequence – The chain’s protein sequence, within that complex.

  • msa_dirs – Directories to search, each containing one (possibly sharded) subdirectory per complex_id. Plain paths assume the writer’s default layout (see _normalize_paired_msa_dirs()); dicts use the same format as find_msas(), where directory_depth shards complex_id (default 1, matching make_msas_colabfold_server’s default), not the chain hash – chain files within a complex directory are always flat.

Returns:

The MSA file path if found, else None.

atomworks.ml.preprocessing.msa.finding.sequence_has_msa(sequence: str, msa_dirs: list[dict] | list[PathLike] | None = None) → bool[source]#

Check if a sequence has an existing MSA file in any directory.

Parameters:

  • sequence – Protein sequence to check.

  • msa_dirs – Directories to search. Accepts: - None: Uses PROTEIN_MSA_DIRS env var - list[PathLike]: Simple format, auto-detects metadata - list[dict]: Pre-computed format with ‘dir’, ‘extension’, ‘directory_depth’ keys

Returns:

True if MSA exists, False otherwise.

Filtering#

MSA filtering for AtomWorks.

Handles filtering MSA files to reduce sequence count and redundancy using HHfilter.

class atomworks.ml.preprocessing.msa.filtering.HHFilterConfig(max_sequences: int = 10000, max_identity_percent: float = 90.0, min_coverage_percent: float = 50.0)[source]#

Bases: object

Configuration for HHfilter MSA filtering.

max_identity_percent: float = 90.0#
max_sequences: int = 10000#
min_coverage_percent: float = 50.0#
class atomworks.ml.preprocessing.msa.filtering.MSAFilterConfig(input_extension: ~atomworks.enums.MSAFileExtension | str = MSAFileExtension.A3M_GZ, output_extension: ~atomworks.enums.MSAFileExtension | str = MSAFileExtension.A3M_GZ, hhfilter: ~atomworks.ml.preprocessing.msa.filtering.HHFilterConfig = <factory>, num_workers: int | None = None)[source]#

Bases: object

Configuration for MSA file filtering.

hhfilter: HHFilterConfig#
input_extension: MSAFileExtension | str = '.a3m.gz'#
num_workers: int | None = None#
output_extension: MSAFileExtension | str = '.a3m.gz'#
atomworks.ml.preprocessing.msa.filtering.count_sequences_in_msa(file_path: Path) → int[source]#

Count sequences in MSA file by counting lines starting with ‘>’.

Uses bash commands for faster counting of large files.

Parameters:

file_path – Path to the MSA file.

Returns:

Number of sequences in the MSA.

Raises:

  • FileNotFoundError – If file does not exist.

  • subprocess.CalledProcessError – If bash command fails.

atomworks.ml.preprocessing.msa.filtering.filter_msas(input_dir: PathLike, output_dir: PathLike | None = None, config: MSAFilterConfig | None = None) → None[source]#

Filter MSA files using HHfilter to reduce sequence count and redundancy.

Parameters:

  • input_dir – Source directory containing MSA files, recursively searched.

  • output_dir – Destination directory for filtered files. If None, performs in-place filtering.

  • config – MSA filtering configuration. If None, uses defaults.

Raises:

FileNotFoundError – If no MSA files are found or directories don’t exist.

Examples

Filter MSA files in-place (modifying original files):

>>> from atomworks.ml.preprocessing.msa.filtering import filter_msas, MSAFilterConfig, HHFilterConfig
>>> filter_config = MSAFilterConfig(hhfilter=HHFilterConfig(max_sequences=5000, max_identity_percent=90))
>>> filter_msas("./organized_msas", config=filter_config)  # No output_dir means in-place

Filter to a new directory with different input/output extensions:

>>> filter_config = MSAFilterConfig(input_extension=MSAFileExtension.A3M, output_extension=MSAFileExtension.A3M_GZ)
>>> filter_msas("./organized_msas", "./filtered_msas", config=filter_config)
atomworks.ml.preprocessing.msa.filtering.run_hhfilter(input_file: PathLike, output_file: PathLike, maxseq: int = 10000, id: float = 90.0, cov: float = 50.0) → None[source]#

Run HHfilter on an MSA file with specified parameters.

Parameters:

  • input_file – Path to input MSA file (handles .gz compression if needed).

  • output_file – Path to output filtered file (handles .gz compression if needed).

  • maxseq – Maximum number of sequences to keep (default: 10000).

  • id – Maximum pairwise sequence identity (%) (default: 90.0).

  • cov – Minimum coverage with query (%) (default: 50.0).

References

Generating#

Generate MSAs from protein sequences using MMseqs2 or HHblits.

Supports both GPU-accelerated MMseqs2 and CPU-based HHblits backends. MSA files are automatically organized with proper hashing, sharding, and compression.

Example

from atomworks.ml.preprocessing.msa.generating import make_msas_from_csv, MSAGenerationConfig

config = MSAGenerationConfig(backend="hhblits", threads=8, max_final_sequences=5000)
make_msas_from_csv("sequences.csv", "output_msas/", sequence_column="seq", config=config)

References

  • Mirdita, M. et al. (2022). ColabFold: making protein folding accessible to all. Nature Methods, 19, 679-682.

  • ColabFold MMseqs2 Search Script - Original implementation and documentation.

class atomworks.ml.preprocessing.msa.generating.HHblitsSearchConfig(uniref30_db_path: str | None = None, bfd_db_path: str | None = None, e_values: list[float] = <factory>, bfd_e_value: float = 0.001, min_seqs_high_cov: int = 2000, min_seqs_low_cov: int = 4000, high_cov: float = 75.0, low_cov: float = 50.0, identity: float = 90.0, max_filter_seqs: int = 100000, n_iterations: int = 4, mem: int = 64, use_bfd: bool = True, mact: float = 0.35, maxfilt: int = 10000000, neffmax: float = 20.0, cov: int = 25, maxseq: int = 1000000, realign_max: int = 100000000)[source]#

Bases: object

Configuration for HHblits iterative search parameters.

Default values match the original HHblits MSA generation pipeline. Database paths are loaded from HHBLITS_UNIREF30_DB_PATH and HHBLITS_BFD_DB_PATH env vars if not provided.

Parameters:

  • mact – MAC realignment threshold for HHblits.

  • maxfilt – Maximum hits allowed through prefilter.

  • neffmax – Maximum Neff (effective sequence count) value.

  • use_bfd – Whether to fall back to BFD when UniRef30 yields insufficient sequences.

bfd_db_path: str | None = None#
bfd_e_value: float = 0.001#
cov: int = 25#
e_values: list[float]#
high_cov: float = 75.0#
identity: float = 90.0#
low_cov: float = 50.0#
mact: float = 0.35#
max_filter_seqs: int = 100000#
maxfilt: int = 10000000#
maxseq: int = 1000000#
mem: int = 64#
min_seqs_high_cov: int = 2000#
min_seqs_low_cov: int = 4000#
n_iterations: int = 4#
neffmax: float = 20.0#
realign_max: int = 100000000#
uniref30_db_path: str | None = None#
use_bfd: bool = True#
class atomworks.ml.preprocessing.msa.generating.MMseqs2SearchConfig(filter: bool = True, num_iterations: int = 3, max_seqs: int = 10000, search_eval: float = 0.1, expand_eval: float = inf, expand_max_seq_id: float = 0.95, align_eval: float = 10.0, diff: int = 3000, qsc: float = -20.0, filter_qsc: float = 0.0, filter_max_seq_id: float = 0.95, filter_min_enable: int = 1000, filter_qid: str = '0.0,0.2,0.4,0.6,0.8,1.0', max_accept: int = 1000000, prefilter_mode: int = 0, s: float = 8.0, db_load_mode: int = 2)[source]#

Bases: object

Configuration for MMseqs2 search parameters.

Default values match those used in the working ColabFold script (reference below).

Parameters:

  • filter – Whether to filter the MSA.

  • num_iterations – Number of MMseqs2 search iterations.

  • max_seqs – Maximum number of cluster centers in the MSA.

  • search_eval – Search e-value threshold.

  • expand_eval – E-value threshold for expandaln.

  • expand_max_seq_id – Maximum sequence identity for expandaln.

  • align_eval – E-value threshold for align.

  • diff – Keep at least this many sequences in each MSA block.

  • qsc – Reduce diversity using minimum score threshold.

  • filter_qsc – Filter diversity using minimum score threshold.

  • filter_max_seq_id – Maximum sequence identity for filtering.

  • filter_min_enable – Minimum number of sequences to keep in each MSA block.

  • filter_qid – Sequence identity thresholds for filtering.

  • max_accept – Maximum accepted alignments before stopping.

  • prefilter_mode – Prefiltering algorithm (0: k-mer, 1: ungapped, 2: exhaustive).

  • s – MMseqs2 sensitivity. Lower = faster but sparser MSAs.

  • db_load_mode – Database preload mode (0: auto, 1: fread, 2: mmap, 3: mmap+touch).

References

align_eval: float = 10.0#
db_load_mode: int = 2#
diff: int = 3000#
expand_eval: float = inf#
expand_max_seq_id: float = 0.95#
filter: bool = True#
filter_max_seq_id: float = 0.95#
filter_min_enable: int = 1000#
filter_qid: str = '0.0,0.2,0.4,0.6,0.8,1.0'#
filter_qsc: float = 0.0#
max_accept: int = 1000000#
max_seqs: int = 10000#
num_iterations: int = 3#
prefilter_mode: int = 0#
qsc: float = -20.0#
s: float = 8.0#
search_eval: float = 0.1#
class atomworks.ml.preprocessing.msa.generating.MSAGenerationConfig(sharding_pattern: str = '/0:2/', output_extension: str = '.a3m.gz', use_env: bool = True, gpu: bool = False, gpu_server: bool = False, threads: int = 4, use_local_temp_dir: bool = True, max_final_sequences: int = 10000, check_existing: bool = False, existing_msa_dirs: list[~os.PathLike] | None = None, search_config: ~atomworks.ml.preprocessing.msa.generating.MMseqs2SearchConfig = <factory>, backend: str = 'mmseqs2', hhblits_search_config: ~atomworks.ml.preprocessing.msa.generating.HHblitsSearchConfig | None = None)[source]#

Bases: object

Configuration for MSA generation.

This dataclass encapsulates all user-facing configuration options for MSA generation, supporting both MMseqs2 and HHblits backends.

Parameters:

  • sharding_pattern – Directory sharding pattern for file organization.

  • output_extension – File extension and compression for output files.

  • use_env – Whether to include environmental (metagenomic) database (MMseqs2 only).

  • gpu – Whether to use GPU acceleration (MMseqs2 only).

  • gpu_server – Whether to use GPU server (requires gpu=True, MMseqs2 only).

  • threads – Number of CPU threads for search operations (used by both MMseqs2 and HHblits).

  • use_local_temp_dir – Whether to use local temporary directory for intermediate files (MMseqs2 only).

  • max_final_sequences – Maximum number of sequences in the final MSA after HHFilter.

  • check_existing – Whether to check for existing MSAs before generation.

  • existing_msa_dirs – Directories to check for existing MSAs. If None, uses PROTEIN_MSA_DIRS env var.

  • search_config – Advanced MMseqs2 search configuration (MMseqs2 only).

  • backend – MSA generation backend, either "mmseqs2" or "hhblits".

  • hhblits_search_config – HHblits search configuration (HHblits only). If None when using HHblits backend, a default config is constructed at generation time.

References

  • Mirdita, M. et al. (2022). ColabFold: making protein folding accessible to all. Nature Methods, 19, 679-682.

backend: str = 'mmseqs2'#
check_existing: bool = False#
existing_msa_dirs: list[PathLike] | None = None#
gpu: bool = False#
gpu_server: bool = False#
hhblits_search_config: HHblitsSearchConfig | None = None#
max_final_sequences: int = 10000#
output_extension: str = '.a3m.gz'#
search_config: MMseqs2SearchConfig#
sharding_pattern: str = '/0:2/'#
threads: int = 4#
use_env: bool = True#
use_local_temp_dir: bool = True#
atomworks.ml.preprocessing.msa.generating.create_fasta_with_hashed_headers(sequences: list[str], output_file: PathLike) → None[source]#

Create a FASTA file from sequence strings with SHA-256 hashed headers.

Note: For wrapped sequences with deduplication, see create_fasta_file_from_df().

Parameters:

  • sequences – List of protein sequence strings.

  • output_file – Path to the output FASTA file.

atomworks.ml.preprocessing.msa.generating.make_msas_from_csv(csv_file: PathLike, output_dir: PathLike, sequence_column: str | None = None, config: MSAGenerationConfig | None = None) → None[source]#

Generate MSAs from sequences in a CSV file.

Parameters:

  • csv_file – Path to CSV file containing protein sequences.

  • output_dir – Directory where organized MSA files will be saved.

  • sequence_column – Name of column containing sequences. If None, CSV must have exactly one column.

  • config – MSA generation configuration. If None, uses default config.

Examples

Generate MSAs from single-column CSV:

make_msas_from_csv("sequences.csv", "output_msas/")

Generate MSAs from multi-column CSV:

make_msas_from_csv("data.csv", "output_msas/", sequence_column="sequence")
atomworks.ml.preprocessing.msa.generating.make_msas_hhblits(sequences: str | list[str], output_dir: PathLike, max_final_sequences: int = 10000, sharding_pattern: str = '/0:2/', output_extension: str = '.a3m.gz', search_config: HHblitsSearchConfig | None = None, cpu: int = 4) → None[source]#

Generate MSAs from protein sequences using HHblits (CPU-only, HH-suite).

Runs iterative HHblits searches per sequence against UniRef30 (and optionally BFD), then organizes and filters the results following the same post-processing as MMseqs2.

Parameters:

  • sequences – A single protein sequence string or list of protein sequences.

  • output_dir – Path to the output directory where MSA files will be saved.

  • max_final_sequences – Maximum number of sequences in final MSAs after filtering.

  • sharding_pattern – Directory sharding pattern (e.g., “/0:2/”).

  • output_extension – Output file extension (.a3m, .a3m.gz, .a3m.zst, .afa, .afa.gz, .afa.zst).

  • search_config – HHblits search configuration. If None, uses defaults.

  • cpu – Number of CPU threads for HHblits (-cpu).

atomworks.ml.preprocessing.msa.generating.make_msas_mmseqs(sequences: str | list[str], output_dir: PathLike, gpu: bool = False, gpu_server: bool = False, use_local_temp_dir: bool = True, max_final_sequences: int = 10000, sharding_pattern: str = '/0:2/', output_extension: str = '.a3m.gz', search_config: MMseqs2SearchConfig | None = None, threads: int = 4) → None[source]#

Generate MSAs directly from protein sequences.

Parameters:

  • sequences – A single protein sequence string or list of protein sequences.

  • output_dir – Path to the output directory where MSA files will be saved.

  • gpu – Whether to use GPU acceleration.

  • gpu_server – Whether to use GPU server (requires gpu=True).

  • use_local_temp_dir – Whether to use local temporary directory for intermediate files.

  • max_final_sequences – Maximum number of sequences in final MSAs after filtering.

  • sharding_pattern – Directory sharding pattern (e.g., “/0:2/”).

  • output_extension – Output file extension (.a3m, .a3m.gz, .a3m.zst, .afa, .afa.gz, .afa.zst).

  • search_config – Advanced MMseqs2 search configuration (includes num_iterations and max_seqs).

  • threads – Number of CPU threads for MMseqs2 search.

Examples

make_msas_mmseqs(
    ["MSYIWRQLGSPTVAITLSVSTVIYVTVICPIVFIHLFGDHL...", "MKKKEVEKDDLIENASRVASCISIFLIIASTTMYIFIGLKI..."], "output_msas/"
)

Organizing#

MSA file organization for AtomWorks.

Prepares arbitrary directory of MSA files into an AtomWorks-compatible structure.

class atomworks.ml.preprocessing.msa.organizing.MSAOrganizationConfig(input_extension: MSAFileExtension | str = MSAFileExtension.A3M_GZ, output_extension: MSAFileExtension | str = MSAFileExtension.A3M_GZ, sharding_pattern: str | None = '/0:2/', copy_files: bool = True, num_workers: int | None = None, execution_model: str = 'process', check_existing: bool = False, existing_msa_dirs: list[PathLike] | None = None)[source]#

Bases: object

Configuration for MSA file organization and processing.

Used by organize_msas to control how existing MSA files are organized in the output directory. Controls input and output file extensions with automatic compression/decompression.

check_existing: bool = False#
copy_files: bool = True#
execution_model: str = 'process'#
existing_msa_dirs: list[PathLike] | None = None#
input_extension: MSAFileExtension | str = '.a3m.gz'#
num_workers: int | None = None#
output_extension: MSAFileExtension | str = '.a3m.gz'#
sharding_pattern: str | None = '/0:2/'#
atomworks.ml.preprocessing.msa.organizing.extract_first_sequence_from_msa(file_path: Path) → str[source]#

Extract the first sequence from an MSA file.

Parameters:

file_path – Path to the MSA file. May be compressed (.gz). A3M and FASTA formats supported.

Returns:

First sequence string (without header).

Raises:

  • ValueError – If MSA file has no sequences or invalid format.

  • FileNotFoundError – If file does not exist.

atomworks.ml.preprocessing.msa.organizing.organize_msas(input_dir: PathLike, output_dir: PathLike, config: MSAOrganizationConfig | None = None) → None[source]#

Organize an arbitrary directory of MSA files into AtomWorks-compatible structure.

After organization, files will be named by the SHA-256 hash of the first sequence, and can be loaded via the AtomWorks MSA loading and pairing Transforms.

For each MSA file found in the source directory, this function: 1. Computes the SHA-256 hash of the sequence, which will become the new filename (so we don’t store duplicates) 2. Shards files into subdirectories based on hash and directory structure 3. Automatically handles compression/decompression based on input/output extensions

Parameters:

  • input_dir – Source directory containing MSA files, recursively searched.

  • output_dir – Destination directory for organized files. Must contain sufficient space.

  • config – MSA processing configuration. If None, uses defaults. See MSAOrganizationConfig for details.

Raises:

  • ValueError – If configuration is invalid.

  • FileNotFoundError – If no MSA files are found or directories don’t exist.

Examples

Organize MSA files with default configuration:

>>> from atomworks.ml.preprocessing.msa.organization import organize_msas
>>> organize_msas("./raw_msas", "./organized_msas")

Organize with custom configuration:

>>> from atomworks.ml.preprocessing.msa.organization import organize_msas, MSAOrganizationConfig
>>> from atomworks.enums import MSAFileExtension
>>> config = MSAOrganizationConfig(
...     input_extension=MSAFileExtension.A3M, output_extension=MSAFileExtension.A3M_GZ, sharding_pattern="/0:2/"
... )
>>> organize_msas("./raw_msas", "./organized_msas", config=config)
atomworks.ml.preprocessing.msa.organizing.organize_paired_msas(input_dir: PathLike, output_dir: PathLike, config: MSAOrganizationConfig | None = None) → str[source]#

Organize raw per-chain a3m files from ONE paired/complex query into AtomWorks’ standard paired-MSA layout.

Unlike organize_msas (which treats every file in input_dir as an independent, unrelated alignment – one per unique sequence), this treats all files found in input_dir as belonging to the SAME complex: their first sequences, collectively, determine the complex_id (see get_complex_id()) – computed from file content, the same way organize_msas derives each file’s identity from content, not passed in or tracked separately by the caller.

Output layout: <output_dir>/<shard-of-complex_id>/<complex_id>/<chain_hash><extension> – flat within the complex’s own directory (a complex has few chains), but complex_id itself is sharded the same way a sequence hash would be (a large fine-tuning run can have many complexes). Discoverable via find_paired_msas() / get_paired_msa_path().

Parameters:

  • input_dir – Directory containing one raw a3m file per chain in the complex (recursively searched, same as organize_msas).

  • output_dir – Destination directory.

  • config – MSA processing configuration. sharding_pattern shards complex_id (not the files within it, which are always flat). check_existing / existing_msa_dirs / num_workers / execution_model are not used – a complex has too few files to need parallelism, and existing-MSA deduplication belongs at the sequence-submission level, before this is ever called.

Returns:

The complex_id this batch of files was organized under.

Raises:

  • FileNotFoundError – If no MSA files are found or input_dir doesn’t exist.

  • FileExistsError – If the complex’s directory already exists (complete, or partial from an interrupted run – delete it to regenerate).

atomworks.ml.preprocessing.msa.organizing.organize_template_alignments(sequence_to_content: dict[str, str], output_dir: PathLike, sharding_pattern: str | None = '/0:2/', output_extension: str = '.m8') → None[source]#

Organize raw ColabFold template-alignment hit tables (pdb70.m8 rows) into AtomWorks’ standard hash-sharded layout.

Uses the same sequence-hash/sharding convention as organize_msas, so a sequence’s template alignment and its unpaired MSA land at the corresponding shard. Unlike MSA files, template hit tables don’t self-identify their query sequence (no FASTA-style header), so the sequence must be supplied explicitly rather than extracted from file content.

Discoverable via find_template_alignments().

Parameters:

  • sequence_to_content – Maps each query sequence to its raw pdb70.m8 hit lines (tab-separated text, one hit per line). Sequences with falsy content (no hits) are skipped – no empty file is written.

  • output_dir – Destination directory for organized files.

  • sharding_pattern – Directory sharding pattern (e.g. “/0:2/”), matching organize_msas’s convention.

  • output_extension – Output file extension.

atomworks.ml.preprocessing.msa.organizing.validate_msa_file_extension(file_extension: str) → None[source]#

Validate that a file extension is supported for MSA processing.

Parameters:

file_extension – File extension to validate.

Raises:

ValueError – If file extension is not supported.

On this page
Show Source