From 6a2bcc2d1c2115c8f7081f7cf0b00107d5204ef3 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Fri, 10 Oct 2025 14:03:34 -0600 Subject: [PATCH 01/11] mmcif parsing --- aide_predict/bespoke_models/base.py | 4 +- .../utils/data_structures/structures.py | 154 ++++++++++++++---- 2 files changed, 122 insertions(+), 36 deletions(-) diff --git a/aide_predict/bespoke_models/base.py b/aide_predict/bespoke_models/base.py index e36669d..d54edb2 100644 --- a/aide_predict/bespoke_models/base.py +++ b/aide_predict/bespoke_models/base.py @@ -962,7 +962,9 @@ def _post_transform_hook(self, result, X): # Handle ragged arrays (different shapes) if self._is_ragged_array(result): if not self.pool: - raise ValueError("Pooling must be enabled to handle ragged arrays.") + # return here ... the output is not useful for direct downtream use in eg sklearn but the user may want this + logger.warning("Output is a ragged array and no pooling was specified. Returning list of arrays.") + return result else: # apply pooling to each array in the list result = [pool_func(a, axis=-2) for a in result] diff --git a/aide_predict/utils/data_structures/structures.py b/aide_predict/utils/data_structures/structures.py index fac263a..9fc7ccb 100644 --- a/aide_predict/utils/data_structures/structures.py +++ b/aide_predict/utils/data_structures/structures.py @@ -6,29 +6,29 @@ * License: MIT ''' import os -from typing import Optional, Dict, List, Union -import warnings -from dataclasses import dataclass -import glob import re +import glob import json +from dataclasses import dataclass +from typing import Optional, Dict, List +import numpy as np +from Bio.PDB import PDBParser, MMCIFParser, Structure, Chain +from Bio.PDB.DSSP import dssp_dict_from_pdb_file +from typing import Union from ..constants import AA_MAP THREE_TO_ONE_AA = {v: k for k, v in AA_MAP.items()} -import numpy as np -from Bio.PDB import PDBParser, Structure, Chain -from Bio.PDB.DSSP import dssp_dict_from_pdb_file @dataclass(eq=True) class ProteinStructure: - pdb_file: str + structure_file: str # Renamed from pdb_file to be more general chain: str = 'A' plddt_file: Optional[str] = None def __post_init__(self): - if not os.path.exists(self.pdb_file): - raise FileNotFoundError(f"PDB file not found: {self.pdb_file}") + if not os.path.exists(self.structure_file): + raise FileNotFoundError(f"Structure file not found: {self.structure_file}") if self.plddt_file and not os.path.exists(self.plddt_file): raise FileNotFoundError(f"pLDDT file not found: {self.plddt_file}") @@ -36,22 +36,64 @@ def __post_init__(self): self._sequence: Optional[str] = None self._plddt: Optional[np.ndarray] = None self._dssp: Optional[Dict[str, str]] = None + self._file_format: Optional[str] = None + + @property + def file_format(self) -> str: + """ + Determine the file format based on file extension. + + Returns: + str: 'pdb' or 'cif' + """ + if self._file_format is None: + ext = os.path.splitext(self.structure_file)[1].lower() + if ext in ['.cif', '.mmcif']: + self._file_format = 'cif' + elif ext in ['.pdb', '.ent']: + self._file_format = 'pdb' + else: + # Try to determine from content if extension is ambiguous + try: + with open(self.structure_file, 'r') as f: + first_line = f.readline().strip() + if first_line.startswith('data_'): + self._file_format = 'cif' + else: + self._file_format = 'pdb' + except: + # Default to PDB if we can't determine + self._file_format = 'pdb' + return self._file_format + + def _get_parser(self): + """ + Get the appropriate parser based on file format. + + Returns: + Parser: Either PDBParser or MMCIFParser + """ + if self.file_format == 'cif': + return MMCIFParser(QUIET=True) + else: + return PDBParser(QUIET=True) def get_sequence(self) -> str: """ - Get the amino acid sequence from the PDB file. + Get the amino acid sequence from the structure file. Returns: str: The amino acid sequence. """ if self._sequence is None: - parser = PDBParser() - structure = parser.get_structure("protein", self.pdb_file) + parser = self._get_parser() + structure = parser.get_structure("protein", self.structure_file) chain = structure[0][self.chain] self._sequence = "".join( THREE_TO_ONE_AA[residue.resname] for residue in chain - if residue.id[0] == " ") + if residue.id[0] == " " and residue.resname in THREE_TO_ONE_AA + ) return self._sequence def get_plddt(self) -> Optional[np.ndarray]: @@ -70,12 +112,20 @@ def get_plddt(self) -> Optional[np.ndarray]: def get_dssp(self) -> Dict[str, str]: """ Get the DSSP secondary structure assignments. + + Note: DSSP requires PDB format. If using CIF, consider converting first + or using alternative secondary structure assignment methods. Returns: Dict[str, str]: Dictionary of DSSP assignments. """ if self._dssp is None: - self._dssp = dssp_dict_from_pdb_file(self.pdb_file)[0] + if self.file_format == 'cif': + # DSSP typically works with PDB format + # You might want to implement CIF->PDB conversion or use alternative methods + raise NotImplementedError("DSSP analysis directly from CIF files is not yet supported. " + "Consider converting to PDB format first.") + self._dssp = dssp_dict_from_pdb_file(self.structure_file)[0] return self._dssp def validate_sequence(self, protein_sequence: str) -> bool: @@ -98,8 +148,8 @@ def get_structure(self) -> Structure: Returns: Structure: The complete protein structure. """ - parser = PDBParser() - return parser.get_structure("protein", self.pdb_file) + parser = self._get_parser() + return parser.get_structure("protein", self.structure_file) def get_chain(self) -> Chain: """ @@ -128,6 +178,7 @@ def from_af2_folder(cls, folder_path: str, chain: str = 'A') -> 'ProteinStructur This method prioritizes the top-ranked relaxed structure. If no relaxed structures are available, it selects the top-ranked unrelaxed structure. + Now supports both PDB and CIF formats. Args: folder_path (str): Path to the folder containing AlphaFold2 predictions. @@ -137,38 +188,71 @@ def from_af2_folder(cls, folder_path: str, chain: str = 'A') -> 'ProteinStructur ProteinStructure: A new ProteinStructure object. Raises: - FileNotFoundError: If no suitable PDB file is found in the folder. + FileNotFoundError: If no suitable structure file is found in the folder. """ def get_rank(filename): match = re.search(r'rank_(\d+)', filename) return int(match.group(1)) if match else float('inf') - # Search for relaxed PDB files - relaxed_pdbs = glob.glob(os.path.join(folder_path, '*relaxed*.pdb')) - if relaxed_pdbs: - pdb_file = min(relaxed_pdbs, key=get_rank) + # Search for relaxed structure files (both PDB and CIF) + relaxed_files = [] + for ext in ['*.pdb', '*.cif', '*.mmcif']: + relaxed_files.extend(glob.glob(os.path.join(folder_path, f'*relaxed*{ext}'))) + + if relaxed_files: + structure_file = min(relaxed_files, key=get_rank) else: - # If no relaxed PDFs, search for ranked PDBs - ranked_pdbs = glob.glob(os.path.join(folder_path, '*rank_*.pdb')) - if ranked_pdbs: - pdb_file = min(ranked_pdbs, key=get_rank) + # If no relaxed files, search for ranked files + ranked_files = [] + for ext in ['*.pdb', '*.cif', '*.mmcif']: + ranked_files.extend(glob.glob(os.path.join(folder_path, f'*rank*{ext}'))) + + if ranked_files: + structure_file = min(ranked_files, key=get_rank) else: - raise FileNotFoundError(f"No suitable PDB file found in {folder_path}") + raise FileNotFoundError(f"No suitable structure file found in {folder_path}") # Search for corresponding pLDDT file - plddt_file = pdb_file.replace('.pdb', '.json') - plddt_file = plddt_file.replace('_unrelaxed_', '_scores_') - plddt_file = plddt_file.replace('_relaxed_', '_scores_') + base_name = os.path.splitext(structure_file)[0] + plddt_file = base_name.replace('_unrelaxed', '_scores').replace('_relaxed', '_scores') + '.json' + if not os.path.exists(plddt_file): - raise FileNotFoundError(f"pLDDT file not found: {plddt_file}") + # Try alternative naming conventions + alt_plddt = structure_file.replace('.pdb', '.json').replace('.cif', '.json') + alt_plddt = alt_plddt.replace('_unrelaxed_', '_scores_').replace('_relaxed_', '_scores_') + if os.path.exists(alt_plddt): + plddt_file = alt_plddt + else: + plddt_file = None # Don't raise error, just set to None + + return cls(structure_file, chain, plddt_file) - return cls(pdb_file, chain, plddt_file) + # Backward compatibility properties + @property + def pdb_file(self) -> str: + """ + Backward compatibility property for pdb_file. + + Returns: + str: The structure file path. + """ + return self.structure_file + + @pdb_file.setter + def pdb_file(self, value: str) -> None: + """ + Backward compatibility setter for pdb_file. + + Args: + value (str): The structure file path. + """ + self.structure_file = value def __hash__(self) -> int: - return hash((self.pdb_file, self.chain, self.plddt_file)) + return hash((self.structure_file, self.chain, self.plddt_file)) def __repr__(self) -> str: - return f"ProteinStructure(pdb_file='{self.pdb_file}', chain='{self.chain}')" + return f"ProteinStructure(structure_file='{self.structure_file}', chain='{self.chain}', format='{self.file_format}')" class StructureMapper: @@ -204,7 +288,7 @@ def _scan_folder(self): """ for item in os.listdir(self.structure_folder): item_path = os.path.join(self.structure_folder, item) - if item.endswith('.pdb'): + if item.endswith('.pdb') or item.endswith('.cif') or item.endswith('.mmcif'): structure_id = os.path.splitext(item)[0] self.structure_map[structure_id] = ProteinStructure(item_path) elif os.path.isdir(item_path): From 98297024c5348e4c4c022bc134831345aa292bee Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Thu, 16 Oct 2025 13:00:27 -0600 Subject: [PATCH 02/11] pdb_file changed to structure_file --- aide_predict/utils/data_structures/sequences.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/aide_predict/utils/data_structures/sequences.py b/aide_predict/utils/data_structures/sequences.py index a442e1c..b915288 100644 --- a/aide_predict/utils/data_structures/sequences.py +++ b/aide_predict/utils/data_structures/sequences.py @@ -347,7 +347,7 @@ def is_in_msa(self) -> bool: return False @classmethod - def from_pdb(cls, pdb_file: str, chain: str = 'A', id: Optional[str] = None) -> 'ProteinSequence': + def from_pdb(cls, structure_file: str, chain: str = 'A', id: Optional[str] = None) -> 'ProteinSequence': """ Create a ProteinSequence from a PDB file. @@ -355,7 +355,7 @@ def from_pdb(cls, pdb_file: str, chain: str = 'A', id: Optional[str] = None) -> a ProteinSequence object with the associated structure. Args: - pdb_file (str): Path to the PDB file. + structure_file (str): Path to the PDB file. chain (str): Chain identifier to extract sequence from. Defaults to 'A'. id (Optional[str]): Identifier for the sequence. If None, uses the PDB filename. @@ -371,17 +371,17 @@ def from_pdb(cls, pdb_file: str, chain: str = 'A', id: Optional[str] = None) -> >>> print(seq) 'MAEGEITTFTALTEKFNLPPGNYKKPKLLYCSNG...' >>> print(seq.structure) - ProteinStructure(pdb_file='1abc.pdb', chain='A') + ProteinStructure(structure_file='1abc.pdb', chain='A') """ # Create structure object (this validates file existence) - structure = ProteinStructure(pdb_file=pdb_file, chain=chain) + structure = ProteinStructure(structure_file=structure_file, chain=chain) # Extract sequence from structure sequence = structure.get_sequence() # Use PDB filename as ID if none provided if id is None: - id = os.path.splitext(os.path.basename(pdb_file))[0] + id = os.path.splitext(os.path.basename(structure_file))[0] # Create sequence object with structure return cls(sequence, id=id, structure=structure) From 0894ccf4d8970eefbc2a730bea60bb725d803ced Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Mon, 24 Nov 2025 13:47:31 -0700 Subject: [PATCH 03/11] drafting method for any positions specific embedder to accept aligned sequences --- aide_predict/bespoke_models/__init__.py | 2 + aide_predict/bespoke_models/base.py | 83 +++++++- .../bespoke_models/embedders/aa_properties.py | 184 ++++++++++++++++++ aide_predict/bespoke_models/embedders/esm2.py | 57 +----- 4 files changed, 272 insertions(+), 54 deletions(-) create mode 100644 aide_predict/bespoke_models/embedders/aa_properties.py diff --git a/aide_predict/bespoke_models/__init__.py b/aide_predict/bespoke_models/__init__.py index eec5e29..61b2128 100644 --- a/aide_predict/bespoke_models/__init__.py +++ b/aide_predict/bespoke_models/__init__.py @@ -20,6 +20,7 @@ from .embedders.saprot import SaProtEmbedding from .embedders.kmer import KmerEmbedding from .embedders.ssemb import SSEmbEmbedding +from .embedders.aa_properties import AAPropertiesEmbedding TOOLS = [ @@ -40,4 +41,5 @@ SaProtEmbedding, KmerEmbedding, SSEmbEmbedding, + AAPropertiesEmbedding, ] diff --git a/aide_predict/bespoke_models/base.py b/aide_predict/bespoke_models/base.py index d54edb2..09f785d 100644 --- a/aide_predict/bespoke_models/base.py +++ b/aide_predict/bespoke_models/base.py @@ -898,7 +898,8 @@ class PositionSpecificMixin: This mixin adds functionality for handling position-specific outputs from protein models. It allows selecting specific positions to analyze, pooling across positions, and - flattening multi-dimensional outputs. + flattening multi-dimensional outputs. It can also automatically handle aligned sequences + with gaps by stripping gaps before processing and remapping embeddings back to aligned positions. Attributes: positions (Optional[List[int]]): The positions to output scores for. If None, all positions are used. @@ -908,14 +909,21 @@ class PositionSpecificMixin: - If callable: Uses the provided function for pooling - If False: No pooling is performed flatten (bool): Whether to flatten dimensions beyond the second dimension. + handle_aligned (bool): If True, automatically strip gaps before processing and remap to aligned positions. + gap_fill_value (float): Value to use for gap positions in aligned sequences (default 0.0). """ _per_position_capable: bool = True - def _init_handler(self, positions=None, pool=True, flatten=True, **kwargs): + def _init_handler(self, positions=None, pool=True, flatten=True, + handle_aligned=True, gap_fill_value=0.0, **kwargs): """Initialize position-specific attributes from kwargs.""" self.positions = positions self.pool = pool self.flatten = flatten + self.handle_aligned = handle_aligned + self.gap_fill_value = gap_fill_value + # Temporary storage for alignment mapping during transform + self._alignment_mapping = None def _is_ragged_array(self, arr): """Check if the input is a ragged array (list of arrays with different shapes).""" @@ -934,9 +942,71 @@ def _pre_transform_hook(self, X): if self.positions is not None: if not (X.aligned or len(X) == 1): raise ValueError("Input sequences must be same length / aligned for position-specific output.") + + # Handle aligned sequences if enabled + if self.handle_aligned and X.has_gaps: + # Store mapping in instance for post-transform hook + self._alignment_mapping = X.get_alignment_mapping() + # Convert mapping to have integer keys ascending from 0 + self._alignment_mapping = {i: m for i, m in enumerate(self._alignment_mapping.values())} + X = X.with_no_gaps() + + # Validate behavior is well-defined + if self.positions is None and not self.pool: + raise ValueError( + "Cannot return position-specific embeddings for sequences with gaps " + "unless positions are specified or pooling is enabled." + ) + else: + self._alignment_mapping = None return X + def _remap_to_aligned_positions(self, result, mapping, positions, fill_value): + """ + Remap embeddings from ungapped sequences back to aligned positions. + + Args: + result: List of embeddings for ungapped sequences + mapping: Dict mapping sequence index to list of aligned positions + positions: List of aligned positions to extract + fill_value: Value to use for gap positions + + Returns: + List of remapped embeddings with gaps represented by fill_value + """ + aligned_embeddings = [] + for i, emb in enumerate(result): + seq_mapping = mapping[i] + # emb shape: (1, seq_len, embedding_dim), (seq_len, embedding_dim), or (seq_len,) for 1D + # Remove batch dimension if present + if emb.ndim == 3 and emb.shape[0] == 1: + emb = emb[0] # Now (seq_len, embedding_dim) + + if emb.ndim == 1: + # Handle 1D case (e.g., single position or pooled) + emb = np.expand_dims(emb, 0) + squeeze_after = True + else: + squeeze_after = False + + # Now emb is (seq_len, embedding_dim) + seq_len = emb.shape[0] + embedding_dim = emb.shape[-1] if emb.ndim > 1 else 1 + aligned_emb = np.full((len(positions), embedding_dim), fill_value, dtype=emb.dtype) + + for j, pos in enumerate(positions): + if pos in seq_mapping: + aligned_pos = seq_mapping.index(pos) + if aligned_pos < seq_len: + aligned_emb[j] = emb[aligned_pos] + + if squeeze_after: + aligned_emb = np.squeeze(aligned_emb, axis=-1) + + aligned_embeddings.append(aligned_emb) + return aligned_embeddings + def _post_transform_hook(self, result, X): """ Process the model output to handle position selection, pooling, and flattening. @@ -950,6 +1020,15 @@ def _post_transform_hook(self, result, X): """ if result is None or len(result) == 0: return result + + # Remap to aligned positions if we have a mapping and positions were specified + if self._alignment_mapping is not None and self.positions is not None: + result = self._remap_to_aligned_positions( + result, self._alignment_mapping, self.positions, self.gap_fill_value + ) + # Clean up temporary storage + self._alignment_mapping = None + if self.pool: # get the pool function if isinstance(self.pool, str): diff --git a/aide_predict/bespoke_models/embedders/aa_properties.py b/aide_predict/bespoke_models/embedders/aa_properties.py new file mode 100644 index 0000000..3fd37d3 --- /dev/null +++ b/aide_predict/bespoke_models/embedders/aa_properties.py @@ -0,0 +1,184 @@ +# aide_predict/bespoke_models/embedders/aa_properties.py +''' +* Author: Evan Komp +* Created: 11/24/2024 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Simple amino acid property embedder for testing position-specific functionality. +''' +import numpy as np +from typing import List, Union, Optional + +from aide_predict.bespoke_models.base import ( + ProteinModelWrapper, + PositionSpecificMixin, + CanHandleAlignedSequencesMixin, + ExpectsNoFitMixin +) +from aide_predict.utils.data_structures import ProteinSequences, ProteinSequence +from aide_predict.utils.common import MessageBool + +import logging +logger = logging.getLogger(__name__) + +AVAILABLE = MessageBool(True, "AAPropertiesEmbedding is always available") + + +# Simple physicochemical properties for the 20 standard amino acids +AA_PROPERTIES = { + 'A': [1.8, 0.0, 0.0], # Alanine: hydrophobicity, charge, size + 'C': [2.5, 0.0, 0.0], # Cysteine + 'D': [-3.5, -1.0, 0.0], # Aspartic acid + 'E': [-3.5, -1.0, 0.5], # Glutamic acid + 'F': [2.8, 0.0, 1.0], # Phenylalanine + 'G': [-0.4, 0.0, -1.0], # Glycine + 'H': [-3.2, 0.5, 0.5], # Histidine + 'I': [4.5, 0.0, 0.5], # Isoleucine + 'K': [-3.9, 1.0, 0.5], # Lysine + 'L': [3.8, 0.0, 0.5], # Leucine + 'M': [1.9, 0.0, 0.5], # Methionine + 'N': [-3.5, 0.0, 0.0], # Asparagine + 'P': [-1.6, 0.0, 0.0], # Proline + 'Q': [-3.5, 0.0, 0.5], # Glutamine + 'R': [-4.5, 1.0, 1.0], # Arginine + 'S': [-0.8, 0.0, -0.5], # Serine + 'T': [-0.7, 0.0, 0.0], # Threonine + 'V': [4.2, 0.0, 0.0], # Valine + 'W': [-0.9, 0.0, 1.5], # Tryptophan + 'Y': [-1.3, 0.0, 1.0], # Tyrosine +} + + +class AAPropertiesEmbedding( + ExpectsNoFitMixin, + PositionSpecificMixin, + CanHandleAlignedSequencesMixin, + ProteinModelWrapper +): + """ + A simple amino acid property embedder for testing position-specific functionality. + + This embedder converts each amino acid to a 3-dimensional vector based on: + - Hydrophobicity (Kyte-Doolittle scale approximation) + - Charge (at physiological pH) + - Size (relative volume) + + This is a simple, fast embedder that can handle aligned sequences with gaps + and is useful for testing the PositionSpecificMixin functionality. + + Attributes: + positions (Optional[List[int]]): Specific positions to encode. If None, all positions are encoded. + pool (bool): Whether to pool the encoded vectors across positions. + flatten (bool): Whether to flatten the output array. + handle_aligned (bool): Whether to handle aligned sequences with gaps. + gap_fill_value (float): Value to use for gap positions. + """ + + _available = AVAILABLE + + def __init__( + self, + metadata_folder: str = None, + positions: Optional[List[int]] = None, + flatten: bool = False, + pool: bool = False, + handle_aligned: bool = True, + gap_fill_value: float = 0.0, + wt: Optional[Union[str, ProteinSequence]] = None, + **kwargs + ): + """ + Initialize the AAPropertiesEmbedding. + + Args: + metadata_folder (str): The folder where metadata is stored. + positions (Optional[List[int]]): Specific positions to encode. If None, all positions are encoded. + flatten (bool): Whether to flatten the output array. + pool (bool): Whether to pool the encoded vectors across positions. + handle_aligned (bool): Whether to handle aligned sequences with gaps. + gap_fill_value (float): Value to use for gap positions. + wt (Optional[Union[str, ProteinSequence]]): The wild type sequence, if any. + """ + super().__init__( + metadata_folder=metadata_folder, + wt=wt, + positions=positions, + pool=pool, + flatten=flatten, + handle_aligned=handle_aligned, + gap_fill_value=gap_fill_value, + **kwargs + ) + self.embedding_dim_ = 3 # 3 properties per amino acid + + def _fit(self, X: ProteinSequences, y: Optional[np.ndarray] = None) -> 'AAPropertiesEmbedding': + """ + Fit the embedder (no actual fitting needed as properties are predefined). + + Args: + X (ProteinSequences): The input protein sequences. + y (Optional[np.ndarray]): Ignored. Present for API consistency. + + Returns: + AAPropertiesEmbedding: The fitted embedder. + """ + self.fitted_ = True + return self + + def _transform(self, X: ProteinSequences) -> List[np.ndarray]: + """ + Transform the protein sequences into amino acid property embeddings. + + Args: + X (ProteinSequences): The input protein sequences. + + Returns: + List[np.ndarray]: The amino acid property embeddings for the sequences. + """ + all_embeddings = [] + + for seq in X: + seq_str = str(seq).upper() + seq_len = len(seq_str) + + # Create embedding matrix: (seq_len, 3) + embedding = np.zeros((1, seq_len, 3), dtype=np.float32) + + for i, aa in enumerate(seq_str): + if aa in AA_PROPERTIES: + embedding[0, i, :] = AA_PROPERTIES[aa] + else: + # Unknown amino acid - use zeros + logger.warning(f"Unknown amino acid '{aa}' in sequence {seq.id}, using zeros") + embedding[0, i, :] = [0.0, 0.0, 0.0] + + all_embeddings.append(embedding) + + # Return as list - PositionSpecificMixin will handle position selection, pooling, and alignment remapping + return all_embeddings + + def get_feature_names_out(self, input_features: Optional[List[str]] = None) -> List[str]: + """ + Get output feature names for transformation. + + Args: + input_features (Optional[List[str]]): Ignored. Present for API consistency. + + Returns: + List[str]: Output feature names. + """ + if not hasattr(self, 'fitted_'): + raise ValueError("Model has not been fitted yet. Call fit() before using this method.") + + positions = self.positions + property_names = ['hydrophobicity', 'charge', 'size'] + + if self.pool: + return [f"AAProps_{prop}" for prop in property_names] + elif self.flatten: + if positions is None: + raise ValueError("Cannot return feature names for flattened embeddings without specifying positions") + return [f"pos{p}_{prop}" for p in positions for prop in property_names] + else: + raise ValueError("Cannot return feature names for non-flattened non-pooled embeddings.") diff --git a/aide_predict/bespoke_models/embedders/esm2.py b/aide_predict/bespoke_models/embedders/esm2.py index 68b1f59..ffe15e0 100644 --- a/aide_predict/bespoke_models/embedders/esm2.py +++ b/aide_predict/bespoke_models/embedders/esm2.py @@ -149,20 +149,9 @@ def _transform(self, X: ProteinSequences) -> np.ndarray: raise ValueError("Cannot flatten variable length sequences without positions or pooling.") warnings.warn("Variable length sequences are being processed without positions or pooling, raw shapes will be output.") - mapping = None - if X.has_gaps: - # here we need to store a mapping such that if positions were specified we can map back to - # the aligned positions - mapping = X.get_alignment_mapping() - # convert mapping to have integer keys ascending from 0 - mapping = {i: m for i, m in enumerate(mapping.values())} - - X = X.with_no_gaps() - # raise if positions were not passed - here behavior is uncertain - if self.positions is None and not self.pool: - raise ValueError("Cannot return position-specific embeddings for sequences with gaps unless positions are specified or pooling is on.") - - base_index = 0 + # Note: gap handling is now managed by PositionSpecificMixin hooks + # X will arrive here without gaps if handle_aligned=True + bar = tqdm.tqdm(total=len(X), desc="Computing ESM2 embeddings") for batch in X.iter_batches(self.batch_size): batch_sequences = self._prepare_sequences(batch) @@ -181,49 +170,13 @@ def _transform(self, X: ProteinSequences) -> np.ndarray: # Remove special tokens (assuming first and last tokens are special) embeddings = [emb[1:-1] for emb in embeddings] - if self.positions is not None and mapping is None: - # here we have fixed length so we can just use positions - embeddings = [emb[self.positions] for emb in embeddings] - elif self.positions is not None and mapping is not None: - # here we have variable length sequences that were input as an aligned set, - # the user asked for positions in the alignment - aligned_embeddings = [] - for i, emb in enumerate(embeddings): - seq_mapping = mapping[base_index + i] - aligned_emb = np.zeros((len(self.positions), emb.shape[1])) - for j, pos in enumerate(self.positions): - if pos in seq_mapping: - aligned_pos = seq_mapping.index(pos) - aligned_emb[j] = emb[aligned_pos] - # If pos is not in seq_mapping, it remains a zero vector - aligned_embeddings.append(aligned_emb) - embeddings = aligned_embeddings - else: - # Here positions were not specified and either have fixed length or pooling - # is on - pass - - if self.pool: - if self.pool == 'mean' or self.pool is True: - embeddings = [emb.mean(axis=0) for emb in embeddings] - elif self.pool == 'max': - embeddings = [emb.max(axis=0) for emb in embeddings] - elif hasattr(np, self.pool): - # check if the pool is a numpy function - pool_func = getattr(np, self.pool) - embeddings = [pool_func(emb, axis=0) for emb in embeddings] - else: - raise ValueError(f"Invalid pooling method: {self.pool}") - - # add 0th dimension + # Add 0th dimension for stacking embeddings = [np.expand_dims(emb, 0) for emb in embeddings] all_embeddings.extend(embeddings) - base_index += len(batch) - bar.update(len(batch)) - # stack along 0 dimension + # Return as list - PositionSpecificMixin will handle position selection, pooling, and alignment remapping return all_embeddings def get_feature_names_out(self, input_features: Optional[List[str]] = None) -> List[str]: From 33b1b71db3fbd9f471a564ca55acbe4d99188be8 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Fri, 28 Nov 2025 11:58:58 -0700 Subject: [PATCH 04/11] PositionSpecificMixin will now auto handle stripping gaps from aligned sequences at runtime, giving them to the model, then mapping the outputs back to original alignment positions. Set handle_aligned=False for subclasses that actually use gaps during inference if inheritting from PositionSpecificMixin --- aide_predict/bespoke_models/base.py | 26 ++++++++++++------- .../embedders/msa_transformer.py | 2 +- aide_predict/bespoke_models/embedders/ohe.py | 9 ++++++- tests/test_bespoke_models/test_base.py | 2 +- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/aide_predict/bespoke_models/base.py b/aide_predict/bespoke_models/base.py index 09f785d..8e28344 100644 --- a/aide_predict/bespoke_models/base.py +++ b/aide_predict/bespoke_models/base.py @@ -924,6 +924,7 @@ def _init_handler(self, positions=None, pool=True, flatten=True, self.gap_fill_value = gap_fill_value # Temporary storage for alignment mapping during transform self._alignment_mapping = None + self._alignment_width = None def _is_ragged_array(self, arr): """Check if the input is a ragged array (list of arrays with different shapes).""" @@ -945,20 +946,16 @@ def _pre_transform_hook(self, X): # Handle aligned sequences if enabled if self.handle_aligned and X.has_gaps: + # Store original alignment width for later remapping + self._alignment_width = X.width # Store mapping in instance for post-transform hook self._alignment_mapping = X.get_alignment_mapping() # Convert mapping to have integer keys ascending from 0 self._alignment_mapping = {i: m for i, m in enumerate(self._alignment_mapping.values())} X = X.with_no_gaps() - - # Validate behavior is well-defined - if self.positions is None and not self.pool: - raise ValueError( - "Cannot return position-specific embeddings for sequences with gaps " - "unless positions are specified or pooling is enabled." - ) else: self._alignment_mapping = None + self._alignment_width = None return X @@ -1021,13 +1018,16 @@ def _post_transform_hook(self, result, X): if result is None or len(result) == 0: return result - # Remap to aligned positions if we have a mapping and positions were specified - if self._alignment_mapping is not None and self.positions is not None: + # Remap to aligned positions if we stripped gaps earlier + # If positions is None, remap to all positions in the alignment + if self._alignment_mapping is not None: + positions_to_map = self.positions if self.positions is not None else list(range(self._alignment_width)) result = self._remap_to_aligned_positions( - result, self._alignment_mapping, self.positions, self.gap_fill_value + result, self._alignment_mapping, positions_to_map, self.gap_fill_value ) # Clean up temporary storage self._alignment_mapping = None + self._alignment_width = None if self.pool: # get the pool function @@ -1040,6 +1040,12 @@ def _post_transform_hook(self, result, X): # Handle ragged arrays (different shapes) if self._is_ragged_array(result): + # If positions were specified, ragged output is not acceptable + if self.positions is not None: + raise ValueError( + f"Model returned ragged output (inconsistent shapes) but specific positions {self.positions} " + "were requested. Cannot extract consistent position-specific embeddings from ragged output." + ) if not self.pool: # return here ... the output is not useful for direct downtream use in eg sklearn but the user may want this logger.warning("Output is a ragged array and no pooling was specified. Returning list of arrays.") diff --git a/aide_predict/bespoke_models/embedders/msa_transformer.py b/aide_predict/bespoke_models/embedders/msa_transformer.py index 8ac2ff7..a50dc93 100644 --- a/aide_predict/bespoke_models/embedders/msa_transformer.py +++ b/aide_predict/bespoke_models/embedders/msa_transformer.py @@ -73,7 +73,7 @@ def __init__(self, metadata_folder: str=None, use_cache (bool): Whether to cache results to avoid redundant computations. wt (Optional[Union[str, ProteinSequence]]): The wild type sequence, if any. """ - super().__init__(metadata_folder=metadata_folder, wt=wt, positions=positions, pool=pool, flatten=flatten, use_cache=use_cache) + super().__init__(metadata_folder=metadata_folder, wt=wt, positions=positions, pool=pool, flatten=flatten, use_cache=use_cache, handle_aligned=False) self.layer = layer self.batch_size = batch_size self.device = device diff --git a/aide_predict/bespoke_models/embedders/ohe.py b/aide_predict/bespoke_models/embedders/ohe.py index 81eb3cc..1ecbfce 100644 --- a/aide_predict/bespoke_models/embedders/ohe.py +++ b/aide_predict/bespoke_models/embedders/ohe.py @@ -219,7 +219,14 @@ def __init__(self, metadata_folder: str, wt: Optional[Union[str, ProteinSequence Notes: WT is set to None to avoid normalization. For an embedder this is effectively a feature scaler which you should do manually if you want """ - super().__init__(metadata_folder=metadata_folder, wt=None, positions=positions, pool=False, flatten=flatten) + super().__init__( + metadata_folder=metadata_folder, + wt=None, + positions=positions, + pool=False, + flatten=flatten, + handle_aligned=False # DO NOT let position specific mixin handle alignments here, it will delete gaps and map to the called sequences ungapped + ) self._vocab = list(AA_SINGLE.union(GAP_CHARACTERS)) def _fit(self, X: ProteinSequences, y: Optional[np.ndarray] = None) -> 'OneHotAlignedEmbedding': diff --git a/tests/test_bespoke_models/test_base.py b/tests/test_bespoke_models/test_base.py index f53892a..154007e 100644 --- a/tests/test_bespoke_models/test_base.py +++ b/tests/test_bespoke_models/test_base.py @@ -258,7 +258,7 @@ def _transform(self, X): with pytest.raises(ValueError): full_dim_model.transform(seqs) - # Test with ragged array output without pooling + # Test with ragged array output without pooling and with pisitions asked for class RaggedModel(TestModel): def _transform(self, X): # Return a list of arrays with different shapes From 9bbbd646ee38fffcc9734eb8ab7d9b8363e81975 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Fri, 28 Nov 2025 13:05:34 -0700 Subject: [PATCH 05/11] aaproperties now uses aaindex to get properties --- .../bespoke_models/embedders/aa_properties.py | 175 +++++++++++++----- environment.yaml | 1 + 2 files changed, 128 insertions(+), 48 deletions(-) diff --git a/aide_predict/bespoke_models/embedders/aa_properties.py b/aide_predict/bespoke_models/embedders/aa_properties.py index 3fd37d3..df55585 100644 --- a/aide_predict/bespoke_models/embedders/aa_properties.py +++ b/aide_predict/bespoke_models/embedders/aa_properties.py @@ -5,10 +5,11 @@ * Company: National Renewable Energy Lab, Bioeneergy Science and Technology * License: MIT -Simple amino acid property embedder for testing position-specific functionality. +Amino acid property embedder using AAindex database for physicochemical properties. ''' import numpy as np -from typing import List, Union, Optional +from typing import List, Union, Optional, Dict, Tuple +from aaindex import aaindex1 from aide_predict.bespoke_models.base import ( ProteinModelWrapper, @@ -22,32 +23,81 @@ import logging logger = logging.getLogger(__name__) -AVAILABLE = MessageBool(True, "AAPropertiesEmbedding is always available") - - -# Simple physicochemical properties for the 20 standard amino acids -AA_PROPERTIES = { - 'A': [1.8, 0.0, 0.0], # Alanine: hydrophobicity, charge, size - 'C': [2.5, 0.0, 0.0], # Cysteine - 'D': [-3.5, -1.0, 0.0], # Aspartic acid - 'E': [-3.5, -1.0, 0.5], # Glutamic acid - 'F': [2.8, 0.0, 1.0], # Phenylalanine - 'G': [-0.4, 0.0, -1.0], # Glycine - 'H': [-3.2, 0.5, 0.5], # Histidine - 'I': [4.5, 0.0, 0.5], # Isoleucine - 'K': [-3.9, 1.0, 0.5], # Lysine - 'L': [3.8, 0.0, 0.5], # Leucine - 'M': [1.9, 0.0, 0.5], # Methionine - 'N': [-3.5, 0.0, 0.0], # Asparagine - 'P': [-1.6, 0.0, 0.0], # Proline - 'Q': [-3.5, 0.0, 0.5], # Glutamine - 'R': [-4.5, 1.0, 1.0], # Arginine - 'S': [-0.8, 0.0, -0.5], # Serine - 'T': [-0.7, 0.0, 0.0], # Threonine - 'V': [4.2, 0.0, 0.0], # Valine - 'W': [-0.9, 0.0, 1.5], # Tryptophan - 'Y': [-1.3, 0.0, 1.0], # Tyrosine -} +# Default AAindex indices to use for each property +DEFAULT_AAINDEX_PROPERTIES = [ + ('KYTJ820101', 'hydrophobicity'), # Kyte-Doolittle hydropathy index + ('KLEP840101', 'charge'), # Net charge + ('FASG760104', 'pKa_amino_terminus'), # pKa N-terminus + ('FASG760105', 'pKa_carboxyl_terminus'), # pKa C-terminus + ('GRAR740103', 'volume'), # Volume + ('GRAR740102', 'polarity'), # Polarity + ('FASG760101', 'molecular_weight'), # Molecular weight + ('ZIMJ680104', 'isoelectric_point'), # Isoelectric point + ('ZIMJ680102', 'bulkiness'), # Bulkiness (related to aromaticity/structure) +] + + +def _check_aaindex_available() -> MessageBool: + """Check if AAindex is available and can be accessed.""" + try: + from aaindex import aaindex1 + # Try to access a known index to verify it works + _ = aaindex1['KYTJ820101'] + return MessageBool(True, "AAindex is available") + except Exception as e: + return MessageBool(False, f"AAindex initialization failed: {str(e)}") + + +AVAILABLE = _check_aaindex_available() + + +def _build_aa_property_lookup( + aaindex_ids: List[Tuple[str, str]], + include_aromatic: bool = True +) -> Tuple[Dict[str, np.ndarray], List[str]]: + """ + Build a lookup table for amino acid properties from AAindex. + + Args: + aaindex_ids: List of tuples (aaindex_id, property_name) + include_aromatic: Whether to include a boolean aromatic property + + Returns: + Tuple of (lookup_dict, property_names) where: + - lookup_dict maps amino acid letter to property vector + - property_names is the list of property names in order + """ + # Standard amino acids + amino_acids = 'ACDEFGHIKLMNPQRSTVWY' + aromatic_aas = {'F', 'W', 'Y'} # Phenylalanine, Tryptophan, Tyrosine + + property_names = [name for _, name in aaindex_ids] + if include_aromatic: + property_names.append('aromatic') + + # Build lookup table + lookup = {} + for aa in amino_acids: + properties = [] + for aaindex_id, _ in aaindex_ids: + try: + record = aaindex1[aaindex_id] + value = record.values.get(aa) + if value is None: + logger.warning(f"No value for {aa} in {aaindex_id}, using 0.0") + value = 0.0 + properties.append(float(value)) + except Exception as e: + logger.warning(f"Error getting {aaindex_id} for {aa}: {e}, using 0.0") + properties.append(0.0) + + # Add aromatic boolean as 1.0 or 0.0 + if include_aromatic: + properties.append(1.0 if aa in aromatic_aas else 0.0) + + lookup[aa] = np.array(properties, dtype=np.float32) + + return lookup, property_names class AAPropertiesEmbedding( @@ -57,15 +107,22 @@ class AAPropertiesEmbedding( ProteinModelWrapper ): """ - A simple amino acid property embedder for testing position-specific functionality. + An amino acid property embedder using AAindex database. - This embedder converts each amino acid to a 3-dimensional vector based on: - - Hydrophobicity (Kyte-Doolittle scale approximation) - - Charge (at physiological pH) - - Size (relative volume) + This embedder converts each amino acid to a vector of physicochemical properties + from the AAindex database. By default, it uses: + - Hydrophobicity (Kyte-Doolittle scale) + - Net charge + - pKa of amino terminus + - pKa of carboxyl terminus + - Volume + - Polarity + - Molecular weight + - Isoelectric point + - Bulkiness + - Aromatic (boolean: 1.0 for F/W/Y, 0.0 otherwise) - This is a simple, fast embedder that can handle aligned sequences with gaps - and is useful for testing the PositionSpecificMixin functionality. + Custom properties can be specified by providing a list of (aaindex_id, property_name) tuples. Attributes: positions (Optional[List[int]]): Specific positions to encode. If None, all positions are encoded. @@ -73,6 +130,10 @@ class AAPropertiesEmbedding( flatten (bool): Whether to flatten the output array. handle_aligned (bool): Whether to handle aligned sequences with gaps. gap_fill_value (float): Value to use for gap positions. + aaindex_properties (List[Tuple[str, str]]): List of (aaindex_id, property_name) tuples. + include_aromatic (bool): Whether to include aromatic boolean property. + aa_property_lookup_ (Dict[str, np.ndarray]): Lookup table for amino acid properties. + property_names_ (List[str]): Names of the properties in order. """ _available = AVAILABLE @@ -86,6 +147,8 @@ def __init__( handle_aligned: bool = True, gap_fill_value: float = 0.0, wt: Optional[Union[str, ProteinSequence]] = None, + aaindex_properties: Optional[List[Tuple[str, str]]] = None, + include_aromatic: bool = True, **kwargs ): """ @@ -99,7 +162,15 @@ def __init__( handle_aligned (bool): Whether to handle aligned sequences with gaps. gap_fill_value (float): Value to use for gap positions. wt (Optional[Union[str, ProteinSequence]]): The wild type sequence, if any. + aaindex_properties (Optional[List[Tuple[str, str]]]): List of (aaindex_id, property_name) tuples. + If None, uses DEFAULT_AAINDEX_PROPERTIES. + include_aromatic (bool): Whether to include a boolean aromatic property (F, W, Y = 1.0, others = 0.0). """ + if aaindex_properties is None: + aaindex_properties = DEFAULT_AAINDEX_PROPERTIES + self.aaindex_properties = aaindex_properties + self.include_aromatic = include_aromatic + super().__init__( metadata_folder=metadata_folder, wt=wt, @@ -110,11 +181,10 @@ def __init__( gap_fill_value=gap_fill_value, **kwargs ) - self.embedding_dim_ = 3 # 3 properties per amino acid def _fit(self, X: ProteinSequences, y: Optional[np.ndarray] = None) -> 'AAPropertiesEmbedding': """ - Fit the embedder (no actual fitting needed as properties are predefined). + Fit the embedder by building the amino acid property lookup table. Args: X (ProteinSequences): The input protein sequences. @@ -123,6 +193,13 @@ def _fit(self, X: ProteinSequences, y: Optional[np.ndarray] = None) -> 'AAProper Returns: AAPropertiesEmbedding: The fitted embedder. """ + # Build lookup table from AAindex + self.aa_property_lookup_, self.property_names_ = _build_aa_property_lookup( + self.aaindex_properties, + include_aromatic=self.include_aromatic + ) + self.embedding_dim_ = len(self.property_names_) + self.fitted_ = True return self @@ -142,16 +219,19 @@ def _transform(self, X: ProteinSequences) -> List[np.ndarray]: seq_str = str(seq).upper() seq_len = len(seq_str) - # Create embedding matrix: (seq_len, 3) - embedding = np.zeros((1, seq_len, 3), dtype=np.float32) + # Create embedding matrix: (1, seq_len, n_properties) + embedding = np.zeros((1, seq_len, self.embedding_dim_), dtype=np.float32) for i, aa in enumerate(seq_str): - if aa in AA_PROPERTIES: - embedding[0, i, :] = AA_PROPERTIES[aa] + if aa in self.aa_property_lookup_: + embedding[0, i, :] = self.aa_property_lookup_[aa] else: - # Unknown amino acid - use zeros - logger.warning(f"Unknown amino acid '{aa}' in sequence {seq.id}, using zeros") - embedding[0, i, :] = [0.0, 0.0, 0.0] + # Unknown amino acid (including gaps) - use zeros or gap_fill_value + if aa == '-': + embedding[0, i, :] = self.gap_fill_value + else: + logger.warning(f"Unknown amino acid '{aa}' in sequence {seq.id}, using zeros") + embedding[0, i, :] = 0.0 all_embeddings.append(embedding) @@ -168,17 +248,16 @@ def get_feature_names_out(self, input_features: Optional[List[str]] = None) -> L Returns: List[str]: Output feature names. """ - if not hasattr(self, 'fitted_'): + if not hasattr(self, 'fitted_') or not self.fitted_: raise ValueError("Model has not been fitted yet. Call fit() before using this method.") positions = self.positions - property_names = ['hydrophobicity', 'charge', 'size'] if self.pool: - return [f"AAProps_{prop}" for prop in property_names] + return [f"AAProps_{prop}" for prop in self.property_names_] elif self.flatten: if positions is None: raise ValueError("Cannot return feature names for flattened embeddings without specifying positions") - return [f"pos{p}_{prop}" for p in positions for prop in property_names] + return [f"pos{p}_{prop}" for p in positions for prop in self.property_names_] else: raise ValueError("Cannot return feature names for non-flattened non-pooled embeddings.") diff --git a/environment.yaml b/environment.yaml index 3a36746..5613a67 100644 --- a/environment.yaml +++ b/environment.yaml @@ -19,5 +19,6 @@ dependencies: - h5py - tqdm - biopython==1.84 + - aaindex From cc0d5a556be481a530cf4ee1bea871281448edcc Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Tue, 26 May 2026 16:38:25 -0600 Subject: [PATCH 06/11] Onboarded ESMIF, new protein structure attribute for handleing multimers in structure conditioned models --- README.md | 7 + aide_predict/bespoke_models/__init__.py | 2 + .../bespoke_models/predictors/esm_if.py | 172 +++++++++ .../predictors/pretrained_transformers.py | 45 +-- .../utils/data_structures/sequences.py | 32 ++ .../utils/data_structures/structures.py | 162 ++++++++- requirements_files/requirements-esm-if.txt | 7 + .../test_predictors/test_esm_if_unit.py | 125 +++++++ .../test_not_base_models/test_esm_if_pred.py | 330 ++++++++++++++++++ .../test_data_structures/test_sequences.py | 59 +++- .../test_data_structures/test_structures.py | 191 +++++++++- 11 files changed, 1092 insertions(+), 40 deletions(-) create mode 100644 aide_predict/bespoke_models/predictors/esm_if.py create mode 100644 requirements_files/requirements-esm-if.txt create mode 100644 tests/test_bespoke_models/test_predictors/test_esm_if_unit.py create mode 100644 tests/test_not_base_models/test_esm_if_pred.py diff --git a/README.md b/README.md index 76e7cbd..bc295f1 100644 --- a/README.md +++ b/README.md @@ -90,6 +90,13 @@ You can always check which modules are installed/available to you by running `ge - Requires Structure - Requires independant SSEmb environment, see Installation. +9. [ESM-IF1*](https://proceedings.mlr.press/v162/hsu22a.html) + - Structure-conditioned autoregressive language model (GVP-Transformer) that scores variants given a backbone. Wildtype and mutant marginals are supported; `masked_marginal` is refused since the model has no bidirectional mask-scoring mode. + - Requires wild-type sequence and structure during inference + - Requires fixed-length sequences + - Supports multichain complexes via `ProteinStructure.context_chains` (set explicitly, or via `set_target_chain` / `with_target_chain`, or implicitly through `StructureMapper.get_protein_sequences(..., auto_context=True)`). + - Requires additional dependencies (see `requirements_files/requirements-esm-if.txt`, which adds `torch-geometric`, `torch-scatter`, `torch-cluster`, and pins `biotite<1.0` due to a fair-esm compatibility issue). + ### Embeddings for Downstream ML 1. One Hot Protein Embedding diff --git a/aide_predict/bespoke_models/__init__.py b/aide_predict/bespoke_models/__init__.py index 61b2128..fe3bcd2 100644 --- a/aide_predict/bespoke_models/__init__.py +++ b/aide_predict/bespoke_models/__init__.py @@ -13,6 +13,7 @@ from .predictors.vespa import VESPAWrapper from .predictors.eve import EVEWrapper from .predictors.ssemb import SSEmbWrapper +from .predictors.esm_if import ESMIFLikelihoodWrapper from .embedders.esm2 import ESM2Embedding from .embedders.ohe import OneHotAlignedEmbedding, OneHotProteinEmbedding @@ -32,6 +33,7 @@ VESPAWrapper, EVEWrapper, SSEmbWrapper, + ESMIFLikelihoodWrapper, # embedders ESM2Embedding, diff --git a/aide_predict/bespoke_models/predictors/esm_if.py b/aide_predict/bespoke_models/predictors/esm_if.py new file mode 100644 index 0000000..796fdd8 --- /dev/null +++ b/aide_predict/bespoke_models/predictors/esm_if.py @@ -0,0 +1,172 @@ +# aide_predict/bespoke_models/predictors/esm_if.py +''' +* Author: Evan Komp +* Created: 2026-05-26 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Wrapper around ESM-IF1 (ESM Inverse Folding), the GVP-Transformer from fair-esm. + +ESM-IF1 is autoregressive and conditioned on backbone structure. Wildtype and +mutant marginals are well-defined; masked_marginal is not (no bidirectional +mask scoring mode) and is refused at __init__. + +For multichain complexes, set ``context_chains`` on the attached +ProteinStructure (or call ``set_target_chain`` / ``with_target_chain``). The +wrapper concatenates context-chain coords with the standard 10-residue +NaN-padded gap used by fair-esm's ``score_sequence_in_complex``; the decoder +only emits predictions for the target chain. + +See: +Hsu, C. et al. Learning inverse folding from millions of predicted structures. +ICML 2022. https://doi.org/10.1101/2022.04.10.487779 +''' +import numpy as np +from typing import List, Optional, Union + +from aide_predict.bespoke_models.base import RequiresStructureMixin, RequiresWTToFunctionMixin +from aide_predict.bespoke_models.predictors.pretrained_transformers import LikelihoodTransformerBase, MarginalMethod +from aide_predict.utils.data_structures import ProteinSequences, ProteinSequence, ProteinStructure +from aide_predict.utils.common import MessageBool + +try: + import torch + import torch.nn.functional as F + import esm + from esm.inverse_folding.util import CoordBatchConverter + from esm.inverse_folding.multichain_util import _concatenate_coords + AVAILABLE = MessageBool(True, "ESM-IF is available.") +except ImportError: + AVAILABLE = MessageBool( + False, + "ESM-IF requires fair-esm and torch_geometric. Install via requirements_files/requirements-esm-if.txt", + ) + +import logging +logger = logging.getLogger(__name__) + + +class ESMIFLikelihoodWrapper(RequiresStructureMixin, RequiresWTToFunctionMixin, LikelihoodTransformerBase): + """ + ESM-IF1 zero-shot variant scorer. + + Inherits ExpectsNoFitMixin, RequiresFixedLengthMixin, CanRegressMixin, + RequiresWTDuringInferenceMixin, PositionSpecificMixin, CacheMixin from + LikelihoodTransformerBase. Adds RequiresStructureMixin and + RequiresWTToFunctionMixin so callers must supply a WT with attached + ProteinStructure. + """ + _available = AVAILABLE + + def __init__( + self, + metadata_folder: str = None, + model_checkpoint: str = 'esm_if1_gvp4_t16_142M_UR50', + marginal_method: Union[MarginalMethod, str] = MarginalMethod.WILDTYPE.value, + positions: Optional[List[int]] = None, + pool: bool = True, + flatten: bool = False, + wt: Optional[Union[str, ProteinSequence]] = None, + batch_size: int = 1, + use_cache: bool = False, + device: str = 'cpu', + ): + if isinstance(marginal_method, MarginalMethod): + marginal_method = marginal_method.value + if marginal_method == MarginalMethod.MASKED.value: + raise ValueError( + "ESM-IF is autoregressive; masked_marginal is not defined. " + "Use 'wildtype_marginal' or 'mutant_marginal'." + ) + super().__init__( + metadata_folder=metadata_folder, + marginal_method=marginal_method, + positions=positions, + pool=pool, + flatten=flatten, + use_cache=use_cache, + wt=wt, + batch_size=batch_size, + device=device, + ) + self.model_checkpoint = model_checkpoint + + def _fit(self, X: ProteinSequences, y: Optional[np.ndarray] = None) -> 'ESMIFLikelihoodWrapper': + self.fitted_ = True + return self + + def _load_model(self) -> None: + loader = getattr(esm.pretrained, self.model_checkpoint) + self.model_, self.alphabet_ = loader() + self.model_ = self.model_.eval().to(self.device) + self.batch_converter_ = CoordBatchConverter(self.alphabet_) + self._aa_to_index = {aa: self.alphabet_.get_idx(aa) for aa in "ACDEFGHIKLMNPQRSTVWY"} + self._vectorized_aa_to_index = np.vectorize(lambda x: self._aa_to_index.get(x, -1)) + + def _cleanup_model(self) -> None: + del self.model_ + del self.alphabet_ + del self.batch_converter_ + + def _resolve_structure(self, seq: ProteinSequence) -> ProteinStructure: + if seq.structure is not None: + return seq.structure + if self.wt is not None and self.wt.structure is not None: + return self.wt.structure + raise ValueError( + "ESM-IF requires structure; neither the input sequence nor the WT carries a ProteinStructure." + ) + + def _build_coords(self, structure: ProteinStructure) -> np.ndarray: + """Assemble coordinates with target chain first; multichain joins via fair-esm's NaN-gap convention.""" + target_coords = structure.get_chain_coords(structure.chain) + if structure.context_chains is None: + return target_coords + coords_dict = {structure.chain: target_coords} + for ctx in structure.context_chains: + coords_dict[ctx] = structure.get_chain_coords(ctx) + return _concatenate_coords(coords_dict, structure.chain) + + def _compute_log_likelihoods( + self, + X: ProteinSequences, + mask_positions: Optional[List[List[int]]] = None, + ) -> List[np.ndarray]: + if mask_positions: + for mp in mask_positions: + if mp: + raise ValueError( + "ESM-IF cannot honor mask_positions; choose wildtype_marginal or mutant_marginal." + ) + + results = [] + with torch.no_grad(): + for seq in X: + structure = self._resolve_structure(seq) + coords = self._build_coords(structure) + seq_str = str(seq).upper() + batch = [(coords, None, seq_str)] + coords_b, conf_b, _, tokens, padding_mask = self.batch_converter_(batch, device=self.device) + prev = tokens[:, :-1] + logits, _ = self.model_(coords_b, padding_mask, conf_b, prev) + # ESM-IF's alphabet sets append_eos=False, so logits length already + # equals len(seq_str). Slicing to len(seq_str) is a no-op in that + # case but stays correct if a future alphabet adds an EOS prediction. + logits_seq = logits[:, :, :len(seq_str)] + log_probs = F.log_softmax(logits_seq, dim=1) + # → [L, V] in float64 (validator uses atol=1e-5 on exp().sum()) + lp = log_probs[0].transpose(0, 1).cpu().numpy().astype(np.float64) + # Renormalize for float-precision safety so exp().sum(-1) hits 1.0 within atol. + lp = lp - np.log(np.exp(lp).sum(axis=-1, keepdims=True)) + results.append(lp) + return results + + def _index_log_probs(self, log_probs: np.ndarray, sequences: ProteinSequences) -> np.ndarray: + seq_array = np.array([list(str(seq).upper()) for seq in sequences]) + aa_indices = self._vectorized_aa_to_index(seq_array) + rows = np.arange(log_probs.shape[0]) + rows_expanded = np.expand_dims(rows, axis=0) + return log_probs[rows_expanded, aa_indices] + + def get_feature_names_out(self, input_features: Optional[List[str]] = None) -> List[str]: + return super().get_feature_names_out(input_features) diff --git a/aide_predict/bespoke_models/predictors/pretrained_transformers.py b/aide_predict/bespoke_models/predictors/pretrained_transformers.py index c3487ce..e7f20fa 100644 --- a/aide_predict/bespoke_models/predictors/pretrained_transformers.py +++ b/aide_predict/bespoke_models/predictors/pretrained_transformers.py @@ -299,33 +299,30 @@ def _compute_log_likelihoods_batches(self, X: ProteinSequences, mask_positions: def _compute_mutant_marginal(self, X: ProteinSequences) -> List[np.ndarray]: """ - Compute mutant marginal log likelihoods. + Score each variant by the model's per-position log-likelihood of its + own amino acid given its own context, with no WT subtraction and no + restriction to mutated positions. + + Matches ProteinGym's convention (compute_fitness_esm_if1.py) and is + the natural joint-likelihood interpretation for autoregressive + structure-conditioned models — from the variant's perspective there + is no privileged "mutated position"; the score is the full-sequence + likelihood under the model. Args: X (ProteinSequences): Input protein sequences. Returns: - List[np.ndarray]: List of mutant marginal log likelihoods for each sequence. + List[np.ndarray]: One ``(1, L)`` array per variant containing + per-position ``log p(variant_aa_i | variant_context_ List[np.ndarray]: @@ -431,8 +428,8 @@ def _post_process_likelihoods(self, log_likelihoods: List[np.ndarray], X: Protei log_likelihoods = np.vstack([ll[:, self.positions] for ll in log_likelihoods]) if self.pool: - # if all values are nana, that means it is equal to wildtype and nanmean will fail, return - # zero for that case + # If all values are NaN (variant identical to WT under wildtype_marginal), + # nanmean would warn — substitute 0.0 (a sensible "no-change" score). new_log_likelihoods = [] for ll in log_likelihoods: if np.all(np.isnan(ll)): @@ -440,16 +437,6 @@ def _post_process_likelihoods(self, log_likelihoods: List[np.ndarray], X: Protei else: new_log_likelihoods.append(np.nanmean(ll)) log_likelihoods = np.array(new_log_likelihoods) - - # Handle the edge case for mutant marginal with variable length sequences - if self.marginal_method == MarginalMethod.MUTANT.value and self.wt is not None and not X.aligned and X.width != len(self.wt): - # Compute WT log likelihood - wt_log_probs = self._compute_log_likelihoods_batches(ProteinSequences([self.wt]))[0] - wt_likelihood = self._index_log_probs(wt_log_probs, ProteinSequences([self.wt])) - wt_pooled = np.nanmean(wt_likelihood) - - # Subtract WT pooled likelihood from each pooled mutant likelihood - log_likelihoods = log_likelihoods - wt_pooled else: log_likelihoods = np.vstack(log_likelihoods) diff --git a/aide_predict/utils/data_structures/sequences.py b/aide_predict/utils/data_structures/sequences.py index b915288..abc618d 100644 --- a/aide_predict/utils/data_structures/sequences.py +++ b/aide_predict/utils/data_structures/sequences.py @@ -160,6 +160,38 @@ def with_no_gaps(self) -> 'ProteinSequence': """Return a new ProteinSequence with all gaps removed.""" return ProteinSequence("".join(c for c in self if c not in GAP_CHARACTERS), id=self._id, structure=self._structure, msa=self._msa) + + def with_target_chain(self, new_chain: str, auto_context: bool = True) -> 'ProteinSequence': + """ + Return a new ProteinSequence pointing at a different chain of the same structure. + + Clones the attached ProteinStructure (so the original is not mutated), + switches its primary chain via ``ProteinStructure.set_target_chain``, + and wraps the new chain's amino-acid string in a fresh ProteinSequence. + + Args: + new_chain: The chain ID to become the new primary chain. + auto_context: Forwarded to ``ProteinStructure.set_target_chain``; + when True, populates ``context_chains`` with the other chains + in the file. + + Returns: + ProteinSequence: A new instance with the new chain's sequence and + a cloned structure pointing at that chain. + + Raises: + ValueError: If no structure is attached. + """ + if self._structure is None: + raise ValueError("with_target_chain requires an attached ProteinStructure.") + cloned = ProteinStructure( + structure_file=self._structure.structure_file, + chain=self._structure.chain, + plddt_file=self._structure.plddt_file, + context_chains=self._structure.context_chains, + ) + cloned.set_target_chain(new_chain, auto_context=auto_context) + return ProteinSequence(cloned.get_sequence(), id=self._id, structure=cloned) @property def as_array(self) -> np.ndarray: diff --git a/aide_predict/utils/data_structures/structures.py b/aide_predict/utils/data_structures/structures.py index 9fc7ccb..9e8ca64 100644 --- a/aide_predict/utils/data_structures/structures.py +++ b/aide_predict/utils/data_structures/structures.py @@ -10,7 +10,7 @@ import glob import json from dataclasses import dataclass -from typing import Optional, Dict, List +from typing import Optional, Dict, List, Tuple import numpy as np from Bio.PDB import PDBParser, MMCIFParser, Structure, Chain from Bio.PDB.DSSP import dssp_dict_from_pdb_file @@ -25,11 +25,12 @@ class ProteinStructure: structure_file: str # Renamed from pdb_file to be more general chain: str = 'A' plddt_file: Optional[str] = None + context_chains: Optional[Tuple[str, ...]] = None def __post_init__(self): if not os.path.exists(self.structure_file): raise FileNotFoundError(f"Structure file not found: {self.structure_file}") - + if self.plddt_file and not os.path.exists(self.plddt_file): raise FileNotFoundError(f"pLDDT file not found: {self.plddt_file}") @@ -38,6 +39,22 @@ def __post_init__(self): self._dssp: Optional[Dict[str, str]] = None self._file_format: Optional[str] = None + if self.context_chains is not None: + self.context_chains = tuple(self.context_chains) + if self.chain in self.context_chains: + raise ValueError( + f"context_chains must not include the primary chain '{self.chain}'." + ) + # Validate against the full chain list so users can explicitly pass + # ligand/cofactor chains as context if they want. + available = set(self.get_all_chain_ids(protein_only=False)) + missing = [c for c in self.context_chains if c not in available] + if missing: + raise ValueError( + f"context_chains {missing} not found in {self.structure_file}. " + f"Available chains: {sorted(available)}" + ) + @property def file_format(self) -> str: """ @@ -170,7 +187,99 @@ def get_residue_positions(self) -> List[int]: """ chain = self.get_chain() return [residue.id[1] for residue in chain if residue.id[0] == " "] - + + def get_all_chain_ids(self, protein_only: bool = True) -> List[str]: + """ + Return chain IDs present in the structure file (first BioPython model). + + Most PDB files include non-protein "chains" — ligands, waters, ions, + glycans — that aren't useful as ESM-IF context. With the default + ``protein_only=True`` only chains with at least one canonical amino-acid + residue (id[0] == ' ' and resname in THREE_TO_ONE_AA) are returned. Set + ``protein_only=False`` to get the unfiltered list (used internally for + validation so users can still pass non-protein chain IDs explicitly to + ``context_chains`` if they want). + + Args: + protein_only: When True, filter out chains with no canonical + amino-acid residues. Default True. + + Returns: + List[str]: Chain IDs in BioPython iteration order. + """ + structure = self._get_parser().get_structure("protein", self.structure_file) + chain_ids = [] + for chain in structure[0]: + if protein_only and not any( + residue.id[0] == " " and residue.resname in THREE_TO_ONE_AA + for residue in chain + ): + continue + chain_ids.append(chain.id) + return chain_ids + + def get_chain_coords(self, chain_id: str) -> np.ndarray: + """ + Extract N / CA / C backbone coordinates for a chain. + + Missing atoms are filled with NaN, matching the convention used by + ``esm.inverse_folding.util.extract_coords_from_structure`` — the GVP + encoder auto-masks NaN coords. + + Args: + chain_id: The chain to extract. May be ``self.chain`` or any + entry of ``self.context_chains``. + + Returns: + np.ndarray of shape ``[L, 3, 3]`` and dtype ``float32`` where the + second axis is ordered (N, CA, C). + """ + structure = self._get_parser().get_structure("protein", self.structure_file) + chain = structure[0][chain_id] + coords = [] + for residue in chain: + if residue.id[0] != " " or residue.resname not in THREE_TO_ONE_AA: + continue + residue_coords = np.full((3, 3), np.nan, dtype=np.float32) + for i, atom_name in enumerate(("N", "CA", "C")): + if atom_name in residue: + residue_coords[i] = residue[atom_name].coord + coords.append(residue_coords) + return np.stack(coords) if coords else np.zeros((0, 3, 3), dtype=np.float32) + + def set_target_chain(self, new_chain: str, auto_context: bool = True) -> None: + """ + Switch the primary chain in-place, optionally repopulating ``context_chains``. + + Args: + new_chain: The chain ID to become the new primary chain. Must + exist in the structure file. + auto_context: When True (default), set ``self.context_chains`` to + a tuple of all *other* chains present in the file. When False, + ``self.context_chains`` is set to None. + + Raises: + ValueError: If ``new_chain`` is not present in the structure file. + """ + # Validate against the full chain list (users may target any chain in the file). + available_all = self.get_all_chain_ids(protein_only=False) + if new_chain not in available_all: + raise ValueError( + f"Chain '{new_chain}' not found in {self.structure_file}. " + f"Available chains: {available_all}" + ) + self.chain = new_chain + if auto_context: + # Auto-populate only protein chains as context — ligands/waters would + # contribute no meaningful structural signal and just bloat the input. + available_protein = self.get_all_chain_ids(protein_only=True) + others = tuple(c for c in available_protein if c != new_chain) + self.context_chains = others if others else None + else: + self.context_chains = None + self._sequence = None + self._dssp = None + @classmethod def from_af2_folder(cls, folder_path: str, chain: str = 'A') -> 'ProteinStructure': """ @@ -249,10 +358,13 @@ def pdb_file(self, value: str) -> None: self.structure_file = value def __hash__(self) -> int: - return hash((self.structure_file, self.chain, self.plddt_file)) + return hash((self.structure_file, self.chain, self.plddt_file, self.context_chains)) def __repr__(self) -> str: - return f"ProteinStructure(structure_file='{self.structure_file}', chain='{self.chain}', format='{self.file_format}')" + base = f"ProteinStructure(structure_file='{self.structure_file}', chain='{self.chain}', format='{self.file_format}'" + if self.context_chains is not None: + base += f", context_chains={self.context_chains}" + return base + ")" class StructureMapper: @@ -357,15 +469,47 @@ def get_available_structures(self) -> List[str]: """ return list(self.structure_map.keys()) - def get_protein_sequences(self) -> 'ProteinSequences': + def get_protein_sequences( + self, + target_chain: Union[str, os.PathLike] = 'A', + auto_context: bool = True, + ) -> 'ProteinSequences': """ - Get a ProteinSequences object containing all available protein sequences. + Build a ProteinSequences from every structure discovered in the folder. + + Args: + target_chain: Either a single chain ID applied uniformly to every + structure (e.g. 'A', 'B'), or a filesystem path to a JSON file + mapping ``{structure_id: chain_id}`` for per-file overrides. + The path form is detected via ``os.path.isfile``; otherwise the + value is treated as a chain-ID literal. Files not listed in the + JSON fall back to chain 'A'. The structure_id keys match those + used by ``_scan_folder`` (filename without extension, or AF2 + folder name). + auto_context: When True (default), every produced ProteinStructure + has ``context_chains`` populated with the other chains present + in its file. Set to False for strict single-chain mode. Returns: - ProteinSequences: A ProteinSequences object containing all available protein sequences. + ProteinSequences: One ProteinSequence per discovered structure, + with sequence drawn from the resolved primary chain. """ from aide_predict.utils.data_structures import ProteinSequences, ProteinSequence - return ProteinSequences([ProteinSequence(struct.get_sequence(), id=struct_id, structure=struct) for struct_id, struct in self.structure_map.items()]) + + chain_map: Dict[str, str] = {} + default_chain: str = 'A' + if isinstance(target_chain, (str, os.PathLike)) and os.path.isfile(str(target_chain)): + with open(target_chain) as f: + chain_map = json.load(f) + else: + default_chain = str(target_chain) + + result = [] + for struct_id, struct in self.structure_map.items(): + desired = chain_map.get(struct_id, default_chain) + struct.set_target_chain(desired, auto_context=auto_context) + result.append(ProteinSequence(struct.get_sequence(), id=struct_id, structure=struct)) + return ProteinSequences(result) def __repr__(self): """ diff --git a/requirements_files/requirements-esm-if.txt b/requirements_files/requirements-esm-if.txt new file mode 100644 index 0000000..6d05a84 --- /dev/null +++ b/requirements_files/requirements-esm-if.txt @@ -0,0 +1,7 @@ +fair-esm +transformers[torch] +torch-geometric +torch-scatter +torch-cluster +# fair-esm imports filter_backbone from biotite.structure; that name was removed in biotite 1.0+. +biotite<1.0 diff --git a/tests/test_bespoke_models/test_predictors/test_esm_if_unit.py b/tests/test_bespoke_models/test_predictors/test_esm_if_unit.py new file mode 100644 index 0000000..cfffccb --- /dev/null +++ b/tests/test_bespoke_models/test_predictors/test_esm_if_unit.py @@ -0,0 +1,125 @@ +# tests/test_bespoke_models/test_predictors/test_esm_if_unit.py +''' +* Author: Evan Komp +* Created: 2026-05-26 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Static unit tests for ESMIFLikelihoodWrapper that do not load the model. +Tests that exercise actual model forward passes live in +tests/test_not_base_models/test_esm_if_pred.py (excluded from default CI). +''' +import pytest +import tempfile +import os + +from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper +from aide_predict.bespoke_models.predictors.pretrained_transformers import MarginalMethod +from aide_predict.utils.data_structures import ProteinSequence, ProteinSequences, ProteinStructure +from aide_predict.utils.checks import check_model_compatibility + + +SAMPLE_PDB = """ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N +ATOM 2 CA ALA A 1 1.458 0.000 0.000 1.00 0.00 C +ATOM 3 C ALA A 1 2.009 1.362 0.000 1.00 0.00 C +ATOM 4 O ALA A 1 1.702 2.144 0.907 1.00 0.00 O +ATOM 5 N GLY A 2 2.831 1.687 -0.987 1.00 0.00 N +ATOM 6 CA GLY A 2 3.396 3.037 -1.009 1.00 0.00 C +ATOM 7 C GLY A 2 2.362 4.089 -1.408 1.00 0.00 C +ATOM 8 O GLY A 2 2.730 5.261 -1.509 1.00 0.00 O +END +""" + + +@pytest.fixture +def temp_pdb(): + with tempfile.NamedTemporaryFile(mode='w', suffix='.pdb', delete=False) as f: + f.write(SAMPLE_PDB) + yield f.name + os.unlink(f.name) + + +class TestESMIFMixinFlags: + """Class-level mixin flags don't require fair-esm to be installed.""" + + def test_requires_structure(self): + assert ESMIFLikelihoodWrapper._requires_structure is True + + def test_requires_wt_to_function(self): + assert ESMIFLikelihoodWrapper._requires_wt_to_function is True + + def test_requires_wt_during_inference(self): + assert ESMIFLikelihoodWrapper._requires_wt_during_inference is True + + def test_per_position_capable(self): + assert ESMIFLikelihoodWrapper._per_position_capable is True + + def test_requires_fixed_length(self): + assert ESMIFLikelihoodWrapper._requires_fixed_length is True + + def test_can_regress(self): + assert ESMIFLikelihoodWrapper._can_regress is True + + def test_expects_no_fit(self): + assert ESMIFLikelihoodWrapper._expects_no_fit is True + + +@pytest.mark.optional +class TestESMIFInit: + """Initialization tests require fair-esm to be importable so AVAILABLE is True.""" + + def test_masked_marginal_string_refused(self, temp_pdb): + wt = ProteinSequence("AG", structure=ProteinStructure(temp_pdb)) + with pytest.raises(ValueError, match="masked_marginal is not defined"): + ESMIFLikelihoodWrapper(marginal_method='masked_marginal', wt=wt) + + def test_masked_marginal_enum_refused(self, temp_pdb): + wt = ProteinSequence("AG", structure=ProteinStructure(temp_pdb)) + with pytest.raises(ValueError, match="masked_marginal is not defined"): + ESMIFLikelihoodWrapper(marginal_method=MarginalMethod.MASKED, wt=wt) + + def test_wildtype_marginal_accepted(self, temp_pdb): + wt = ProteinSequence("AG", structure=ProteinStructure(temp_pdb)) + wrapper = ESMIFLikelihoodWrapper(marginal_method='wildtype_marginal', wt=wt) + assert wrapper.marginal_method == 'wildtype_marginal' + + def test_mutant_marginal_accepted(self, temp_pdb): + wt = ProteinSequence("AG", structure=ProteinStructure(temp_pdb)) + wrapper = ESMIFLikelihoodWrapper(marginal_method='mutant_marginal', wt=wt) + assert wrapper.marginal_method == 'mutant_marginal' + + def test_enum_normalized_to_string(self, temp_pdb): + wt = ProteinSequence("AG", structure=ProteinStructure(temp_pdb)) + wrapper = ESMIFLikelihoodWrapper(marginal_method=MarginalMethod.WILDTYPE, wt=wt) + # Base class compares against MarginalMethod.*.value (strings), so + # the enum must be normalized to its string value. + assert wrapper.marginal_method == MarginalMethod.WILDTYPE.value + + +class TestESMIFCompatibility: + """check_model_compatibility uses class attributes only — does not load fair-esm.""" + + def test_incompatible_without_structure(self): + wt = ProteinSequence("ACDEFGHIKLMNPQRSTVWY") + result = check_model_compatibility( + training_sequences=None, + testing_sequences=None, + wt=wt, + ) + assert "ESMIFLikelihoodWrapper" in result["incompatible"] + + def test_compatible_with_structure(self, temp_pdb): + # Tiny PDB has 'AG'; need a wt whose length matches the structure or compatibility + # logic flags fixed-length mismatches separately. For this test we only verify + # the structure-presence path doesn't immediately disqualify the model when + # fair-esm is unavailable. If fair-esm is missing, _available is False which + # also makes the model incompatible; skip in that case. + if not ESMIFLikelihoodWrapper._available: + pytest.skip("ESM-IF deps not installed; cannot exercise the compatible path.") + wt = ProteinSequence("AG", structure=ProteinStructure(temp_pdb)) + result = check_model_compatibility( + training_sequences=None, + testing_sequences=None, + wt=wt, + ) + assert "ESMIFLikelihoodWrapper" in result["compatible"] diff --git a/tests/test_not_base_models/test_esm_if_pred.py b/tests/test_not_base_models/test_esm_if_pred.py new file mode 100644 index 0000000..64c59c6 --- /dev/null +++ b/tests/test_not_base_models/test_esm_if_pred.py @@ -0,0 +1,330 @@ +# tests/test_not_base_models/test_esm_if_pred.py +''' +* Author: Evan Komp +* Created: 2026-05-26 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Integration tests for ESMIFLikelihoodWrapper against the ENVZ_ECOLI Ghose +DMS benchmark. Requires fair-esm + torch_geometric + torch_scatter + +torch_cluster + biotite<1.0; excluded from default CI by pytest.ini. +''' +import os +import tempfile + +import pytest +import numpy as np +import pandas as pd +from scipy.stats import spearmanr + +from aide_predict.utils.data_structures import ProteinSequences, ProteinSequence, ProteinStructure, StructureMapper + +import torch +if torch.cuda.is_available(): + DEVICE = "cuda" +elif torch.backends.mps.is_available(): + DEVICE = "mps" +else: + DEVICE = "cpu" + + +def _build_two_chain_pdb(src_path: str, dst_path: str, offset: float = 100.0) -> None: + """ + Append a copy of chain A as chain B, translated by `offset` Å along x. + Used for multichain smoke tests; chain B is far enough that it should have + negligible structural influence on chain A's per-residue predictions. + """ + with open(src_path) as f: + lines = f.readlines() + + a_atoms = [line for line in lines if line.startswith("ATOM") and line[21] == "A"] + if not a_atoms: + raise RuntimeError(f"No chain A atoms found in {src_path}") + next_serial = max(int(line[6:11]) for line in a_atoms) + 1 + + b_atoms = [] + for line in a_atoms: + x = float(line[30:38]) + offset + new_line = ( + line[:6] + + f"{next_serial:>5}" + + line[11:21] + + "B" + + line[22:30] + + f"{x:8.3f}" + + line[38:] + ) + b_atoms.append(new_line) + next_serial += 1 + + with open(dst_path, "w") as f: + for line in lines: + if line.startswith("END"): + f.writelines(b_atoms) + f.write(f"TER {next_serial:>5}\n") + f.write(line) + + +def _run_proteingym_benchmark(pdb_path: str, csv_path: str, wt_sequence: str, marginal_method: str = 'mutant_marginal') -> float: + """Run an ESM-IF marginal-method scoring against a ProteinGym-style CSV and return Spearman.""" + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + structure = ProteinStructure(pdb_path) + wt = ProteinSequence(wt_sequence, structure=structure) + + assay_data = pd.read_csv(csv_path) + sequences = ProteinSequences.from_list(assay_data['mutated_sequence'].tolist()) + scores = assay_data['DMS_score'].tolist() + + model = ESMIFLikelihoodWrapper( + marginal_method=marginal_method, + device=DEVICE, + pool=True, + wt=wt, + metadata_folder='./tmp/esm_if', + ) + model.fit(sequences) + predictions = model.predict(sequences) + assert np.isfinite(predictions).all() + return float(spearmanr(scores, predictions)[0]) + + +@pytest.mark.optional +def test_esm_if_envz_ecoli_spearman(): + """ENVZ_ECOLI Ghose 2023, mutant_marginal scoring. ProteinGym reference for + ESM-IF1 on this assay is ~0.115; the wrapper's mutant_marginal now uses + ProteinGym's exact formula (mean per-position log-likelihood of the variant + under its own context, no WT subtraction) so the numbers should match.""" + wt_sequence = "LADDRTLLMAGVSHDLRTPLTRIRLATEMMSEQDGYLAESINKDIEECNAIIEQFIDYLR" + spearman = _run_proteingym_benchmark( + pdb_path='tests/data/ENVZ_ECOLI.pdb', + csv_path=os.path.join('tests', 'data', 'ENVZ_ECOLI_Ghose_2023.csv'), + wt_sequence=wt_sequence, + ) + print(f"\n>>> ESM-IF mutant_marginal Spearman on ENVZ_ECOLI (Ghose 2023): {spearman:.4f} (ProteinGym ref: 0.115)") + assert abs(spearman - 0.115) < 0.03 + + +def _proteingym_score_sequences(pdb_path: str, chain: str, sequences: list[str], device: str) -> "np.ndarray": + """ + Replicate ProteinGym's compute_fitness_esm_if1.py scoring formula on a list + of full variant sequences, using fair-esm directly (no aide wrapper). + + Returns one scalar per variant: ``ll_fullseq = -mean(cross_entropy)`` averaged + over all non-pad target positions of the variant. + """ + import esm + from esm.inverse_folding.util import CoordBatchConverter, load_coords + + model, alphabet = esm.pretrained.esm_if1_gvp4_t16_142M_UR50() + model = model.eval().to(device) + coords, _ = load_coords(pdb_path, chain) + batch_converter = CoordBatchConverter(alphabet) + + out = np.empty(len(sequences), dtype=np.float64) + with torch.no_grad(): + for i, seq in enumerate(sequences): + batch = [(coords, None, seq)] + coords_b, conf_b, _, tokens, padding_mask = batch_converter(batch, device=device) + prev = tokens[:, :-1] + target = tokens[:, 1:] + target_padding_mask = (target == alphabet.padding_idx) + logits, _ = model(coords_b, padding_mask, conf_b, prev) + loss = torch.nn.functional.cross_entropy(logits, target, reduction='none') + loss_np = loss[0].cpu().numpy() + mask_np = target_padding_mask[0].cpu().numpy() + ll = -np.sum(loss_np * ~mask_np) / np.sum(~mask_np) + out[i] = ll + del model, alphabet + if device == 'cuda': + torch.cuda.empty_cache() + return out + + +@pytest.mark.optional +def test_proteingym_reference_envz_ecoli(): + """Bypass the wrapper and replicate ProteinGym's exact scoring formula on + our ENVZ_ECOLI test data. If this reproduces ~0.115, the wrapper's lower + Spearman is a scoring-formula difference (Meier-style WT marginal vs + ProteinGym's full-sequence mean log-likelihood), not a wrapper bug.""" + assay = pd.read_csv(os.path.join('tests', 'data', 'ENVZ_ECOLI_Ghose_2023.csv')) + seqs = assay['mutated_sequence'].tolist() + scores = assay['DMS_score'].tolist() + preds = _proteingym_score_sequences( + pdb_path='tests/data/ENVZ_ECOLI.pdb', chain='A', sequences=seqs, device=DEVICE, + ) + spearman = float(spearmanr(scores, preds)[0]) + print(f"\n>>> ProteinGym-formula Spearman on ENVZ_ECOLI: {spearman:.4f} (reference: 0.115)") + + +@pytest.mark.optional +def test_proteingym_reference_gfp_aequvi(): + """Same as above for GFP_AEQVI. If this reproduces ~0.713, the wrapper's + lower Spearman is purely a formula choice — not a wrapper bug.""" + assay = pd.read_csv(os.path.join('tests', 'data', 'GFP_AEQVI_Sarkisyan_2016.csv')) + seqs = assay['mutated_sequence'].tolist() + scores = assay['DMS_score'].tolist() + preds = _proteingym_score_sequences( + pdb_path='tests/data/GFP_AEQVI.pdb', chain='A', sequences=seqs, device=DEVICE, + ) + spearman = float(spearmanr(scores, preds)[0]) + print(f"\n>>> ProteinGym-formula Spearman on GFP_AEQVI: {spearman:.4f} (reference: 0.713)") + + +@pytest.mark.optional +def test_esm_if_gfp_aequvi_spearman(): + """GFP_AEQVI Sarkisyan 2016, mutant_marginal scoring. ProteinGym reference + for ESM-IF1 on this assay is ~0.713; with the redefined mutant_marginal we + expect to be within float-noise of that.""" + structure_for_wt = ProteinStructure('tests/data/GFP_AEQVI.pdb') + wt_sequence = structure_for_wt.get_sequence() + spearman = _run_proteingym_benchmark( + pdb_path='tests/data/GFP_AEQVI.pdb', + csv_path=os.path.join('tests', 'data', 'GFP_AEQVI_Sarkisyan_2016.csv'), + wt_sequence=wt_sequence, + ) + print(f"\n>>> ESM-IF mutant_marginal Spearman on GFP_AEQVI (Sarkisyan 2016): {spearman:.4f} (ProteinGym ref: 0.713)") + assert abs(spearman - 0.713) < 0.03 + + +@pytest.mark.optional +def test_esm_if_mutant_marginal_subset(): + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + structure = ProteinStructure('tests/data/ENVZ_ECOLI.pdb') + wt = ProteinSequence( + "LADDRTLLMAGVSHDLRTPLTRIRLATEMMSEQDGYLAESINKDIEECNAIIEQFIDYLR", + structure=structure, + ) + assay_data = pd.read_csv(os.path.join('tests', 'data', 'ENVZ_ECOLI_Ghose_2023.csv')) + # 50 variants is enough to verify the path; one forward pass per variant is ~100ms on CPU. + subset = assay_data.head(50) + sequences = ProteinSequences.from_list(subset['mutated_sequence'].tolist()) + + model = ESMIFLikelihoodWrapper( + marginal_method='mutant_marginal', + device=DEVICE, + pool=True, + wt=wt, + metadata_folder='./tmp/esm_if', + ) + model.fit(sequences) + predictions = model.predict(sequences) + assert predictions.shape == (50, 1) + assert np.all(np.isfinite(predictions)) + + +@pytest.mark.optional +def test_esm_if_position_specific(): + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + structure = ProteinStructure('tests/data/ENVZ_ECOLI.pdb') + wt = ProteinSequence( + "LADDRTLLMAGVSHDLRTPLTRIRLATEMMSEQDGYLAESINKDIEECNAIIEQFIDYLR", + structure=structure, + ) + assay_data = pd.read_csv(os.path.join('tests', 'data', 'ENVZ_ECOLI_Ghose_2023.csv')) + sequences = ProteinSequences.from_list(assay_data['mutated_sequence'].head(20).tolist()) + + model = ESMIFLikelihoodWrapper( + marginal_method='wildtype_marginal', + positions=[8, 9, 10], + pool=False, + device=DEVICE, + wt=wt, + metadata_folder='./tmp/esm_if', + ) + model.fit(sequences) + predictions = model.predict(sequences) + assert predictions.shape == (20, 3) + + +@pytest.mark.optional +def test_esm_if_multichain_smoke(tmp_path): + """Adding a far-translated copy of chain A as chain B should leave the per-position + deltas approximately unchanged (no structural interaction at 100 Å).""" + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + single_chain_pdb = 'tests/data/ENVZ_ECOLI.pdb' + multichain_pdb = str(tmp_path / "envz_2chain.pdb") + _build_two_chain_pdb(single_chain_pdb, multichain_pdb, offset=100.0) + + wt_sequence = "LADDRTLLMAGVSHDLRTPLTRIRLATEMMSEQDGYLAESINKDIEECNAIIEQFIDYLR" + + # Single-chain reference. + single_struct = ProteinStructure(single_chain_pdb) + single_wt = ProteinSequence(wt_sequence, structure=single_struct) + + # Multichain version with the duplicated chain B as structural context. + multi_struct = ProteinStructure(multichain_pdb, chain='A', context_chains=('B',)) + multi_wt = ProteinSequence(wt_sequence, structure=multi_struct) + + # Score a tiny subset of variants under both setups. + assay_data = pd.read_csv(os.path.join('tests', 'data', 'ENVZ_ECOLI_Ghose_2023.csv')) + sequences = ProteinSequences.from_list(assay_data['mutated_sequence'].head(10).tolist()) + + common_kwargs = dict( + marginal_method='wildtype_marginal', + device=DEVICE, + pool=True, + metadata_folder='./tmp/esm_if', + ) + single_model = ESMIFLikelihoodWrapper(wt=single_wt, **common_kwargs) + single_model.fit(sequences) + single_preds = single_model.predict(sequences).ravel() + + multi_model = ESMIFLikelihoodWrapper(wt=multi_wt, **common_kwargs) + multi_model.fit(sequences) + multi_preds = multi_model.predict(sequences).ravel() + + # Both should be finite and similarly ranked. + assert np.all(np.isfinite(single_preds)) + assert np.all(np.isfinite(multi_preds)) + # Correlation between single- and multi-chain scores should be high since + # chain B is far away and shouldn't change autoregressive context-conditional + # log probs at chain A meaningfully. + if len(single_preds) > 1 and np.std(single_preds) > 0 and np.std(multi_preds) > 0: + rho = spearmanr(single_preds, multi_preds).correlation + print(f"Single-vs-multichain Spearman: {rho}") + assert rho > 0.7 + + +@pytest.mark.optional +def test_esm_if_structure_mapper_smoke(tmp_path): + """End-to-end: discover a structure folder, build ProteinSequences with multichain + context, score variants with ESM-IF.""" + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + multichain_pdb = str(tmp_path / "envz_2chain.pdb") + _build_two_chain_pdb('tests/data/ENVZ_ECOLI.pdb', multichain_pdb, offset=100.0) + + folder = tmp_path / "structures" + folder.mkdir() + target_path = folder / "envz.pdb" + target_path.write_text(open(multichain_pdb).read()) + + mapper = StructureMapper(str(folder)) + wts = mapper.get_protein_sequences(target_chain='A') + assert len(wts) == 1 + wt = wts[0] + assert wt.structure.chain == 'A' + assert wt.structure.context_chains == ('B',) + + assay_data = pd.read_csv(os.path.join('tests', 'data', 'ENVZ_ECOLI_Ghose_2023.csv')) + sequences = ProteinSequences.from_list(assay_data['mutated_sequence'].head(5).tolist()) + + model = ESMIFLikelihoodWrapper( + marginal_method='wildtype_marginal', + device=DEVICE, + pool=True, + wt=wt, + metadata_folder='./tmp/esm_if', + ) + model.fit(sequences) + predictions = model.predict(sequences).ravel() + assert predictions.shape == (5,) + assert np.all(np.isfinite(predictions)) + + +if __name__ == "__main__": + test_esm_if_wildtype_marginal() diff --git a/tests/test_utils/test_data_structures/test_sequences.py b/tests/test_utils/test_data_structures/test_sequences.py index 9c9819b..731e595 100644 --- a/tests/test_utils/test_data_structures/test_sequences.py +++ b/tests/test_utils/test_data_structures/test_sequences.py @@ -15,7 +15,7 @@ import pandas as pd from unittest.mock import patch, MagicMock -from aide_predict.utils.data_structures import ProteinCharacter, ProteinSequence, ProteinSequences, ProteinSequencesOnFile +from aide_predict.utils.data_structures import ProteinCharacter, ProteinSequence, ProteinSequences, ProteinSequencesOnFile, ProteinStructure from aide_predict.utils.constants import AA_SINGLE, GAP_CHARACTERS, NON_CONONICAL_AA_SINGLE from Bio.PDB import PDBIO, Structure, Model, Chain, Residue, Atom @@ -1225,6 +1225,63 @@ def test_msa_with_protein_sequences_operations(self, sample_msa_file): # The new sequence should not have +class TestProteinSequenceMultichain: + """Tests for ProteinSequence.with_target_chain, which requires a multichain structure.""" + + @pytest.fixture + def multichain_pdb(self, tmp_path): + pdb_content = """ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N +ATOM 2 CA ALA A 1 1.458 0.000 0.000 1.00 0.00 C +ATOM 3 C ALA A 1 2.009 1.362 0.000 1.00 0.00 C +ATOM 4 O ALA A 1 1.702 2.144 0.907 1.00 0.00 O +ATOM 5 N GLY A 2 2.831 1.687 -0.987 1.00 0.00 N +ATOM 6 CA GLY A 2 3.396 3.037 -1.009 1.00 0.00 C +ATOM 7 C GLY A 2 2.362 4.089 -1.408 1.00 0.00 C +ATOM 8 O GLY A 2 2.730 5.261 -1.509 1.00 0.00 O +TER 9 GLY A 2 +ATOM 10 N CYS B 1 20.000 20.000 20.000 1.00 0.00 N +ATOM 11 CA CYS B 1 21.458 20.000 20.000 1.00 0.00 C +ATOM 12 C CYS B 1 22.009 21.362 20.000 1.00 0.00 C +ATOM 13 O CYS B 1 21.702 22.144 20.907 1.00 0.00 O +ATOM 14 N ASP B 2 22.831 21.687 19.013 1.00 0.00 N +ATOM 15 CA ASP B 2 23.396 23.037 18.991 1.00 0.00 C +ATOM 16 C ASP B 2 22.362 24.089 18.592 1.00 0.00 C +ATOM 17 O ASP B 2 22.730 25.261 18.491 1.00 0.00 O +TER 18 ASP B 2 +END +""" + pdb_path = tmp_path / "complex.pdb" + pdb_path.write_text(pdb_content) + return str(pdb_path) + + def test_with_target_chain_switches_sequence_and_structure(self, multichain_pdb): + struct = ProteinStructure(multichain_pdb, chain='A') + seq = ProteinSequence("AG", id="prot1", structure=struct) + new_seq = seq.with_target_chain('B') + + # New sequence reflects chain B. + assert str(new_seq) == "CD" + assert new_seq.id == "prot1" + assert new_seq.structure.chain == 'B' + assert new_seq.structure.context_chains == ('A',) + + # Original is untouched (immutability + structure clone). + assert str(seq) == "AG" + assert seq.structure.chain == 'A' + assert seq.structure is not new_seq.structure + + def test_with_target_chain_no_auto_context(self, multichain_pdb): + struct = ProteinStructure(multichain_pdb, chain='A') + seq = ProteinSequence("AG", id="prot1", structure=struct) + new_seq = seq.with_target_chain('B', auto_context=False) + assert new_seq.structure.context_chains is None + + def test_with_target_chain_requires_structure(self): + seq = ProteinSequence("ACDE", id="no_struct") + with pytest.raises(ValueError, match="requires an attached"): + seq.with_target_chain('B') + + # Run the tests if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file diff --git a/tests/test_utils/test_data_structures/test_structures.py b/tests/test_utils/test_data_structures/test_structures.py index a627bc1..bd39a00 100644 --- a/tests/test_utils/test_data_structures/test_structures.py +++ b/tests/test_utils/test_data_structures/test_structures.py @@ -101,6 +101,140 @@ def test_from_af2_folder_no_suitable_files(self, tmp_path): with pytest.raises(FileNotFoundError): ProteinStructure.from_af2_folder(str(tmp_path)) + @pytest.fixture + def temp_multichain_pdb(self): + with tempfile.NamedTemporaryFile(mode='w', suffix='.pdb', delete=False) as f: + f.write("""ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N +ATOM 2 CA ALA A 1 1.458 0.000 0.000 1.00 0.00 C +ATOM 3 C ALA A 1 2.009 1.362 0.000 1.00 0.00 C +ATOM 4 O ALA A 1 1.702 2.144 0.907 1.00 0.00 O +ATOM 5 N GLY A 2 2.831 1.687 -0.987 1.00 0.00 N +ATOM 6 CA GLY A 2 3.396 3.037 -1.009 1.00 0.00 C +ATOM 7 C GLY A 2 2.362 4.089 -1.408 1.00 0.00 C +ATOM 8 O GLY A 2 2.730 5.261 -1.509 1.00 0.00 O +TER 9 GLY A 2 +ATOM 10 N CYS B 1 20.000 20.000 20.000 1.00 0.00 N +ATOM 11 CA CYS B 1 21.458 20.000 20.000 1.00 0.00 C +ATOM 12 C CYS B 1 22.009 21.362 20.000 1.00 0.00 C +ATOM 13 O CYS B 1 21.702 22.144 20.907 1.00 0.00 O +ATOM 14 N ASP B 2 22.831 21.687 19.013 1.00 0.00 N +ATOM 15 CA ASP B 2 23.396 23.037 18.991 1.00 0.00 C +ATOM 16 C ASP B 2 22.362 24.089 18.592 1.00 0.00 C +ATOM 17 O ASP B 2 22.730 25.261 18.491 1.00 0.00 O +TER 18 ASP B 2 +END +""") + yield f.name + os.unlink(f.name) + + def test_context_chains_tuple_normalization(self, temp_multichain_pdb): + # Lists are normalized to tuples for hashability. + s = ProteinStructure(temp_multichain_pdb, chain='A', context_chains=['B']) + assert s.context_chains == ('B',) + + def test_context_chains_default_none(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A') + assert s.context_chains is None + + def test_context_chains_in_hash(self, temp_multichain_pdb): + s1 = ProteinStructure(temp_multichain_pdb, chain='A') + s2 = ProteinStructure(temp_multichain_pdb, chain='A', context_chains=('B',)) + assert hash(s1) != hash(s2) + + def test_context_chains_in_repr(self, temp_multichain_pdb): + s1 = ProteinStructure(temp_multichain_pdb, chain='A') + s2 = ProteinStructure(temp_multichain_pdb, chain='A', context_chains=('B',)) + assert 'context_chains' not in repr(s1) + assert "context_chains=('B',)" in repr(s2) + + def test_context_chains_cannot_include_primary(self, temp_multichain_pdb): + with pytest.raises(ValueError, match="primary chain"): + ProteinStructure(temp_multichain_pdb, chain='A', context_chains=('A',)) + + def test_context_chains_missing_raises(self, temp_multichain_pdb): + with pytest.raises(ValueError, match="not found"): + ProteinStructure(temp_multichain_pdb, chain='A', context_chains=('Z',)) + + def test_get_all_chain_ids(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A') + assert set(s.get_all_chain_ids()) == {'A', 'B'} + + @pytest.fixture + def temp_pdb_with_ligand_chain(self): + # Chain A has protein residues (ALA, GLY); chain Z has only HETATM (a heme + # cofactor). Default get_all_chain_ids should hide Z. + with tempfile.NamedTemporaryFile(mode='w', suffix='.pdb', delete=False) as f: + f.write("""ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N +ATOM 2 CA ALA A 1 1.458 0.000 0.000 1.00 0.00 C +ATOM 3 C ALA A 1 2.009 1.362 0.000 1.00 0.00 C +ATOM 4 O ALA A 1 1.702 2.144 0.907 1.00 0.00 O +ATOM 5 N GLY A 2 2.831 1.687 -0.987 1.00 0.00 N +ATOM 6 CA GLY A 2 3.396 3.037 -1.009 1.00 0.00 C +ATOM 7 C GLY A 2 2.362 4.089 -1.408 1.00 0.00 C +ATOM 8 O GLY A 2 2.730 5.261 -1.509 1.00 0.00 O +TER 9 GLY A 2 +HETATM 10 FE HEM Z 1 30.000 30.000 30.000 1.00 0.00 FE +HETATM 11 NA HEM Z 1 30.500 30.000 30.000 1.00 0.00 N +END +""") + yield f.name + os.unlink(f.name) + + def test_get_all_chain_ids_filters_non_protein(self, temp_pdb_with_ligand_chain): + s = ProteinStructure(temp_pdb_with_ligand_chain, chain='A') + # Default: protein chains only. + assert s.get_all_chain_ids() == ['A'] + # Opt-in to the full list. + assert set(s.get_all_chain_ids(protein_only=False)) == {'A', 'Z'} + + def test_context_chains_accept_non_protein_explicit(self, temp_pdb_with_ligand_chain): + # Power users can still attach a ligand chain as context explicitly. + s = ProteinStructure( + temp_pdb_with_ligand_chain, chain='A', context_chains=('Z',), + ) + assert s.context_chains == ('Z',) + + def test_set_target_chain_auto_context_skips_non_protein(self, temp_pdb_with_ligand_chain): + s = ProteinStructure(temp_pdb_with_ligand_chain, chain='A') + # Re-target chain A (no-op for chain but exercises the auto_context path). + s.set_target_chain('A') + # The HEM cofactor chain Z must NOT be pulled in as auto context. + assert s.context_chains is None + + def test_get_chain_coords_shape_and_dtype(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A') + coords = s.get_chain_coords('A') + assert coords.shape == (2, 3, 3) + assert coords.dtype == np.float32 + assert not np.any(np.isnan(coords)) + np.testing.assert_allclose(coords[0, 0], [0.0, 0.0, 0.0], rtol=1e-5) + + def test_get_chain_coords_context_chain(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A', context_chains=('B',)) + coords_b = s.get_chain_coords('B') + assert coords_b.shape == (2, 3, 3) + np.testing.assert_allclose(coords_b[0, 0], [20.0, 20.0, 20.0], rtol=1e-5) + + def test_set_target_chain(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A') + assert s.get_sequence() == 'AG' + s.set_target_chain('B') + assert s.chain == 'B' + assert s.context_chains == ('A',) + # Cache must be invalidated so the new chain's sequence is returned. + assert s.get_sequence() == 'CD' + + def test_set_target_chain_no_auto_context(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A', context_chains=('B',)) + s.set_target_chain('B', auto_context=False) + assert s.chain == 'B' + assert s.context_chains is None + + def test_set_target_chain_invalid(self, temp_multichain_pdb): + s = ProteinStructure(temp_multichain_pdb, chain='A') + with pytest.raises(ValueError, match="not found"): + s.set_target_chain('Z') + class TestStructureMapper: @pytest.fixture @@ -169,4 +303,59 @@ def test_repr(self, temp_structure_folder): def test_assign_structures_invalid_input(self, temp_structure_folder): mapper = StructureMapper(temp_structure_folder) with pytest.raises(ValueError): - mapper.assign_structures("invalid input") \ No newline at end of file + mapper.assign_structures("invalid input") + + @pytest.fixture + def temp_multichain_folder(self): + with tempfile.TemporaryDirectory() as tmpdirname: + with open(os.path.join(tmpdirname, "complex.pdb"), "w") as f: + f.write("""ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N +ATOM 2 CA ALA A 1 1.458 0.000 0.000 1.00 0.00 C +ATOM 3 C ALA A 1 2.009 1.362 0.000 1.00 0.00 C +ATOM 4 O ALA A 1 1.702 2.144 0.907 1.00 0.00 O +ATOM 5 N GLY A 2 2.831 1.687 -0.987 1.00 0.00 N +ATOM 6 CA GLY A 2 3.396 3.037 -1.009 1.00 0.00 C +ATOM 7 C GLY A 2 2.362 4.089 -1.408 1.00 0.00 C +ATOM 8 O GLY A 2 2.730 5.261 -1.509 1.00 0.00 O +TER 9 GLY A 2 +ATOM 10 N CYS B 1 20.000 20.000 20.000 1.00 0.00 N +ATOM 11 CA CYS B 1 21.458 20.000 20.000 1.00 0.00 C +ATOM 12 C CYS B 1 22.009 21.362 20.000 1.00 0.00 C +ATOM 13 O CYS B 1 21.702 22.144 20.907 1.00 0.00 O +ATOM 14 N ASP B 2 22.831 21.687 19.013 1.00 0.00 N +ATOM 15 CA ASP B 2 23.396 23.037 18.991 1.00 0.00 C +ATOM 16 C ASP B 2 22.362 24.089 18.592 1.00 0.00 C +ATOM 17 O ASP B 2 22.730 25.261 18.491 1.00 0.00 O +TER 18 ASP B 2 +END +""") + yield tmpdirname + + def test_get_protein_sequences_uniform_chain(self, temp_multichain_folder): + mapper = StructureMapper(temp_multichain_folder) + seqs = mapper.get_protein_sequences(target_chain='A') + assert len(seqs) == 1 + assert str(seqs[0]) == 'AG' + assert seqs[0].structure.chain == 'A' + assert seqs[0].structure.context_chains == ('B',) + + def test_get_protein_sequences_other_chain(self, temp_multichain_folder): + mapper = StructureMapper(temp_multichain_folder) + seqs = mapper.get_protein_sequences(target_chain='B') + assert str(seqs[0]) == 'CD' + assert seqs[0].structure.chain == 'B' + assert seqs[0].structure.context_chains == ('A',) + + def test_get_protein_sequences_json_map(self, temp_multichain_folder, tmp_path): + mapper = StructureMapper(temp_multichain_folder) + chain_map = {"complex": "B"} + json_path = tmp_path / "chain_map.json" + json_path.write_text(json.dumps(chain_map)) + seqs = mapper.get_protein_sequences(target_chain=str(json_path)) + assert seqs[0].structure.chain == 'B' + assert str(seqs[0]) == 'CD' + + def test_get_protein_sequences_no_auto_context(self, temp_multichain_folder): + mapper = StructureMapper(temp_multichain_folder) + seqs = mapper.get_protein_sequences(target_chain='A', auto_context=False) + assert seqs[0].structure.context_chains is None \ No newline at end of file From 5ca0894c32862ab7649ce6377792e4385ff56b80 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Wed, 27 May 2026 14:27:52 -0600 Subject: [PATCH 07/11] base class for composites which do some sort of transformation arounf an inner predictor. Not as a seperate transformer because they might need metadata that the inner predictor does not directly output so cannot be used in standard sklearn transformer. Also added z scoring shift a la multi evolve --- aide_predict/bespoke_models/__init__.py | 4 + .../bespoke_models/composites/__init__.py | 14 + .../bespoke_models/composites/zscore.py | 180 +++++++++++++ aide_predict/utils/scoring.py | 193 ++++++++++++++ .../test_composites/__init__.py | 0 .../test_composites/test_zscore.py | 239 ++++++++++++++++++ .../test_zscore_composite.py | 103 ++++++++ tests/test_utils/test_scoring.py | 159 ++++++++++++ 8 files changed, 892 insertions(+) create mode 100644 aide_predict/bespoke_models/composites/__init__.py create mode 100644 aide_predict/bespoke_models/composites/zscore.py create mode 100644 aide_predict/utils/scoring.py create mode 100644 tests/test_bespoke_models/test_composites/__init__.py create mode 100644 tests/test_bespoke_models/test_composites/test_zscore.py create mode 100644 tests/test_not_base_models/test_zscore_composite.py create mode 100644 tests/test_utils/test_scoring.py diff --git a/aide_predict/bespoke_models/__init__.py b/aide_predict/bespoke_models/__init__.py index fe3bcd2..cf98037 100644 --- a/aide_predict/bespoke_models/__init__.py +++ b/aide_predict/bespoke_models/__init__.py @@ -14,6 +14,7 @@ from .predictors.eve import EVEWrapper from .predictors.ssemb import SSEmbWrapper from .predictors.esm_if import ESMIFLikelihoodWrapper +from .composites.zscore import ZScoreRescaledScorer from .embedders.esm2 import ESM2Embedding from .embedders.ohe import OneHotAlignedEmbedding, OneHotProteinEmbedding @@ -35,6 +36,9 @@ SSEmbWrapper, ESMIFLikelihoodWrapper, + # composites + ZScoreRescaledScorer, + # embedders ESM2Embedding, OneHotAlignedEmbedding, diff --git a/aide_predict/bespoke_models/composites/__init__.py b/aide_predict/bespoke_models/composites/__init__.py new file mode 100644 index 0000000..a6392c6 --- /dev/null +++ b/aide_predict/bespoke_models/composites/__init__.py @@ -0,0 +1,14 @@ +# aide_predict/bespoke_models/composites/__init__.py +''' +* Author: Evan Komp +* Created: 2026-05-26 + +Composite ProteinModelWrappers — wrappers that take another wrapper as input +and add additional behavior. The composition pattern lets us keep +ProteinSequences (and the AA-identity information they carry) in scope +across multi-stage scoring pipelines without resorting to sklearn metadata +routing or DataFrame contracts. +''' +from .zscore import ZScoreRescaledScorer + +__all__ = ["ZScoreRescaledScorer"] diff --git a/aide_predict/bespoke_models/composites/zscore.py b/aide_predict/bespoke_models/composites/zscore.py new file mode 100644 index 0000000..44e8904 --- /dev/null +++ b/aide_predict/bespoke_models/composites/zscore.py @@ -0,0 +1,180 @@ +# aide_predict/bespoke_models/composites/zscore.py +''' +* Author: Evan Komp +* Created: 2026-05-26 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +ZScoreRescaledScorer — a composition wrapper that runs any per-variant-scalar +scorer (e.g. an ESM-IF / ESM2 / SaProt wrapper in pool=True mode) on a set of +variants and rescales the per-variant log-ratios within an amino-acid group +(destination residue or substitution type), matching MULTI-evolve's z-score +rescaling step. + +Designed for SSM workflows: fit on ``wt.saturation_mutagenesis()``, then +transform any subset (or the SSM itself) to get both the raw log-ratio and +its AA-group z-score per variant. +''' +from typing import Dict, List, Optional, Tuple, Union +from typing_extensions import Literal + +import numpy as np +import pandas as pd + +from aide_predict.bespoke_models.base import ( + ProteinModelWrapper, + RequiresWTToFunctionMixin, + RequiresWTDuringInferenceMixin, + RequiresFixedLengthMixin, + CanRegressMixin, + ShouldRefitOnSequencesMixin, +) +from aide_predict.utils.data_structures import ProteinSequence, ProteinSequences +from aide_predict.utils.scoring import per_variant_mutation_info, zscore_by_aa_group + + +class ZScoreRescaledScorer( + RequiresWTToFunctionMixin, + RequiresWTDuringInferenceMixin, + RequiresFixedLengthMixin, + CanRegressMixin, + ShouldRefitOnSequencesMixin, + ProteinModelWrapper, +): + """ + Composition wrapper: runs an inner per-variant-scalar scorer and adds + AA-group z-score rescaling. + + Mechanics: + - ``inner_scorer`` is any ProteinModelWrapper whose ``transform(X)`` + returns one scalar per variant (e.g. any wildtype_marginal-mode + LikelihoodTransformerBase subclass with ``pool=True``). + - At ``fit``, the inner scorer is fit on X, then its scores are paired + with each variant's (wt_aa, mut_aa) identity to compute per-group + (mean, std). Group keys are stored on the wrapper. + - At ``transform``, the inner is rerun on the new X; per-variant scalars + are paired with mutation identity; the stored group stats produce a + z-score per variant. Output shape ``(N, 2)`` with columns + ``[logratio, z_logratio]``. + - ``score_table(X)`` returns the full per-variant DataFrame for the + human-readable view (variant_idx, variant_id, mutation, position, + wt_aa, mut_aa, logratio, z_logratio, n_mutations). + + Multi-mutation variants are handled per ``multi_mutant_strategy`` — + ``"skip"`` (default) gives them NaN z-scores while leaving their + log-ratios untouched; ``"raise"`` errors if any are present. + + Note: groups with fewer than ``min_group_size`` samples at fit time + contribute no stats; variants in those groups get NaN z-scores at + transform. This matches MULTI-evolve's threshold of 5. + """ + + def __init__( + self, + inner_scorer: ProteinModelWrapper, + grouping: Literal["destination_residue", "substitution_type"] = "destination_residue", + min_group_size: int = 5, + multi_mutant_strategy: Literal["skip", "raise"] = "skip", + metadata_folder: Optional[str] = None, + wt: Optional[Union[str, ProteinSequence]] = None, + **kwargs, + ): + if hasattr(inner_scorer, "pool") and inner_scorer.pool is False: + raise ValueError( + "ZScoreRescaledScorer expects the inner scorer to produce one " + "scalar per variant; got inner_scorer.pool=False (per-position " + "output). Set inner_scorer.pool=True." + ) + if grouping not in ("destination_residue", "substitution_type"): + raise ValueError( + f"grouping must be 'destination_residue' or 'substitution_type', got {grouping!r}" + ) + if multi_mutant_strategy not in ("skip", "raise"): + raise ValueError( + f"multi_mutant_strategy must be 'skip' or 'raise', got {multi_mutant_strategy!r}" + ) + + super().__init__(metadata_folder=metadata_folder, wt=wt, **kwargs) + + self.inner_scorer = inner_scorer + self.grouping = grouping + self.min_group_size = min_group_size + self.multi_mutant_strategy = multi_mutant_strategy + + # Sync WT into the inner scorer if it doesn't have one yet. + if self.wt is not None and getattr(inner_scorer, "wt", None) is None: + inner_scorer.wt = self.wt + + def _fit(self, X: ProteinSequences, y: Optional[np.ndarray] = None) -> "ZScoreRescaledScorer": + self.inner_scorer.fit(X, y) + scalar_scores = self._call_inner(X) + df = self._build_table(X, scalar_scores) + self._check_multi_mutants(df) + _, self.group_stats_ = zscore_by_aa_group( + df, + grouping=self.grouping, + score_col="logratio", + out_col="z_logratio", + min_group_size=self.min_group_size, + ) + return self + + def _transform(self, X: ProteinSequences) -> np.ndarray: + df = self.score_table(X) + return df[["logratio", "z_logratio"]].to_numpy() + + def score_table(self, X: ProteinSequences) -> pd.DataFrame: + """ + Return the full per-variant DataFrame with columns + ``[variant_idx, variant_id, n_mutations, mutation, position, wt_aa, + mut_aa, logratio, z_logratio]``. + + Uses the group stats stored at ``fit`` time to compute z-scores; + variants belonging to a group that wasn't seen during fit get NaN + for ``z_logratio``. + """ + if not hasattr(self, "group_stats_"): + raise RuntimeError("ZScoreRescaledScorer must be fit before score_table is called.") + scalar_scores = self._call_inner(X) + df = self._build_table(X, scalar_scores) + self._check_multi_mutants(df) + df, _ = zscore_by_aa_group( + df, + grouping=self.grouping, + score_col="logratio", + out_col="z_logratio", + min_group_size=self.min_group_size, + fit_stats=self.group_stats_, + ) + return df + + def _call_inner(self, X: ProteinSequences) -> np.ndarray: + out = np.asarray(self.inner_scorer.transform(X)) + if out.ndim == 2 and out.shape[1] == 1: + out = out.ravel() + if out.ndim != 1: + raise ValueError( + "Inner scorer returned shape " + f"{out.shape}; ZScoreRescaledScorer expects one scalar per variant. " + "Configure the inner scorer with pool=True." + ) + if out.shape[0] != len(X): + raise ValueError( + f"Inner scorer returned {out.shape[0]} scores for {len(X)} variants." + ) + return out + + def _build_table(self, X: ProteinSequences, scores: np.ndarray) -> pd.DataFrame: + info = per_variant_mutation_info(X, self.wt) + info["logratio"] = scores + return info + + def _check_multi_mutants(self, df: pd.DataFrame) -> None: + if self.multi_mutant_strategy == "raise": + multi = df[df["n_mutations"] > 1] + if len(multi) > 0: + raise ValueError( + f"{len(multi)} multi-mutation variant(s) present with " + "multi_mutant_strategy='raise'. Set multi_mutant_strategy='skip' " + "to tolerate them (their z-scores will be NaN)." + ) diff --git a/aide_predict/utils/scoring.py b/aide_predict/utils/scoring.py new file mode 100644 index 0000000..692b588 --- /dev/null +++ b/aide_predict/utils/scoring.py @@ -0,0 +1,193 @@ +# aide_predict/utils/scoring.py +''' +* Author: Evan Komp +* Created: 2026-05-26 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Scorer-agnostic primitives for per-variant mutation accounting and +amino-acid-grouped z-score rescaling. + +These two functions together let any per-variant scalar score (e.g. the +output of a wildtype_marginal scorer in pool=True mode) be rescaled against +the distribution of scores within a destination-AA or substitution-type +group — the rescaling the MULTI-evolve ensemble paper uses on top of +ESM-1v/ESM2/ESM-IF outputs. +''' +from typing import Dict, List, Optional, Tuple +from typing_extensions import Literal + +import numpy as np +import pandas as pd + +from aide_predict.utils.data_structures import ProteinSequence, ProteinSequences + + +def per_variant_mutation_info( + sequences: ProteinSequences, + wt: ProteinSequence, +) -> pd.DataFrame: + """ + Extract single-point mutation identity for each variant in ``sequences``. + + Returns a DataFrame with one row per input variant containing: + + - ``variant_idx`` (int): row index in ``sequences``. + - ``variant_id`` (Optional[str]): ``sequences[v].id`` (None if unset). + - ``n_mutations`` (int): number of positions differing from WT. + - ``mutation`` (Optional[str]): e.g. ``"A123L"`` (1-indexed), only set when + ``n_mutations == 1``; ``None`` otherwise. + - ``position`` (Optional[int]): 1-indexed mutated position, only set when + ``n_mutations == 1``. + - ``wt_aa`` (Optional[str]): WT amino acid at the mutated position, only + set when ``n_mutations == 1``. + - ``mut_aa`` (Optional[str]): variant amino acid at the mutated position, + only set when ``n_mutations == 1``. + + Multi-mutation variants still get a row (with the per-mutation columns + None) so callers can choose to drop, skip, or error on them — the + decision is left to the consumer (typically ``ZScoreRescaledScorer``). + + Args: + sequences: ProteinSequences of variants to inspect. + wt: Wild-type sequence to compare against. Must match the length of + each variant. + + Returns: + pd.DataFrame with the columns described above, one row per variant. + """ + rows = [] + for v_idx, seq in enumerate(sequences): + if len(seq) != len(wt): + raise ValueError( + f"Variant {v_idx} has length {len(seq)} but WT has length {len(wt)}. " + "per_variant_mutation_info requires same-length sequences." + ) + diffs = seq.mutated_positions(wt) # 0-indexed + row = { + "variant_idx": v_idx, + "variant_id": seq.id, + "n_mutations": len(diffs), + "mutation": None, + "position": None, + "wt_aa": None, + "mut_aa": None, + } + if len(diffs) == 1: + p = diffs[0] + row["position"] = p + 1 + row["wt_aa"] = str(wt[p]) + row["mut_aa"] = str(seq[p]) + row["mutation"] = f"{row['wt_aa']}{row['position']}{row['mut_aa']}" + rows.append(row) + return pd.DataFrame(rows) + + +def zscore_by_aa_group( + df: pd.DataFrame, + grouping: Literal["destination_residue", "substitution_type"] = "destination_residue", + score_col: str = "score", + out_col: Optional[str] = None, + min_group_size: int = 5, + fit_stats: Optional[Dict[Tuple[str, ...], Tuple[float, float]]] = None, +) -> Tuple[pd.DataFrame, Dict[Tuple[str, ...], Tuple[float, float]]]: + """ + Add an AA-group z-score column to ``df`` (matches MULTI-evolve's rescaling). + + Groups are keyed by: + - ``"destination_residue"`` (default) → ``(mut_aa,)`` — bins all '→P' + mutations together. + - ``"substitution_type"`` → ``(wt_aa, mut_aa)`` — bins all 'A→P' + mutations together. + + Per-group statistics are ``(mean, std)`` with pandas' default sample + standard deviation (N-1 denominator). Groups with fewer than + ``min_group_size`` samples are not z-scored — output values stay NaN. + Rows missing the grouping columns (e.g. multi-mutation variants without + a defined ``mut_aa``) are also NaN'd. + + Args: + df: Input DataFrame. Must contain ``score_col``, ``mut_aa``, and + (if ``grouping=="substitution_type"``) ``wt_aa``. + grouping: Group definition. + score_col: Column to z-score. + out_col: Output column name. Defaults to ``f"z_{score_col}"``. + min_group_size: Groups with fewer samples produce NaN z-scores + (matches MULTI-evolve's threshold of 5). + fit_stats: If provided, skip stat-computation and apply these + ``(mean, std)`` per group instead. Enables the sklearn-style + fit/transform split where stats are learned on a calibration + set (typically the full SSM) and applied to new variants. + + Returns: + Tuple of (DataFrame with the new column added, dict of group_key → + (mean, std)). The returned stats dict is whatever ``fit_stats`` was + passed in, or — if ``fit_stats is None`` — the freshly computed + stats so the caller can persist them for later transform calls. + """ + if score_col not in df.columns: + raise KeyError(f"score_col '{score_col}' not in df.columns ({list(df.columns)})") + if "mut_aa" not in df.columns: + raise KeyError("zscore_by_aa_group requires a 'mut_aa' column on df") + if grouping == "substitution_type" and "wt_aa" not in df.columns: + raise KeyError("zscore_by_aa_group with grouping='substitution_type' requires a 'wt_aa' column") + + if out_col is None: + out_col = f"z_{score_col}" + + # Build per-row hashable group keys. Using a plain list avoids + # pandas.DataFrame.apply / Series-of-tuple equality quirks that surface + # when comparing tuple keys against a stored fit_stats dict at transform + # time. + mut_col = df["mut_aa"].tolist() + if grouping == "destination_residue": + keys: List[Optional[Tuple[str, ...]]] = [ + (m,) if (m is not None and not (isinstance(m, float) and np.isnan(m))) else None + for m in mut_col + ] + else: + wt_col = df["wt_aa"].tolist() + keys = [ + (w, m) if ( + w is not None and m is not None + and not (isinstance(w, float) and np.isnan(w)) + and not (isinstance(m, float) and np.isnan(m)) + ) else None + for w, m in zip(wt_col, mut_col) + ] + + if fit_stats is None: + # Compute group stats by aggregating row indices per key. + key_to_rows: Dict[Tuple[str, ...], List[int]] = {} + for i, k in enumerate(keys): + if k is None: + continue + key_to_rows.setdefault(k, []).append(i) + + stats: Dict[Tuple[str, ...], Tuple[float, float]] = {} + scores = df[score_col].to_numpy() + for k, idxs in key_to_rows.items(): + vals = scores[idxs] + vals = vals[~np.isnan(vals)] if vals.dtype.kind == "f" else vals + if len(vals) < min_group_size: + continue + std = float(np.std(vals, ddof=1)) + if not np.isfinite(std) or std == 0: + continue + stats[k] = (float(np.mean(vals)), std) + else: + stats = fit_stats + + df = df.copy() + z_out = np.full(len(df), np.nan) + scores = df[score_col].to_numpy() + for i, k in enumerate(keys): + if k is None or k not in stats: + continue + s = scores[i] + if isinstance(s, float) and np.isnan(s): + continue + mean, std = stats[k] + z_out[i] = (s - mean) / std + df[out_col] = z_out + return df, stats diff --git a/tests/test_bespoke_models/test_composites/__init__.py b/tests/test_bespoke_models/test_composites/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_bespoke_models/test_composites/test_zscore.py b/tests/test_bespoke_models/test_composites/test_zscore.py new file mode 100644 index 0000000..e33b323 --- /dev/null +++ b/tests/test_bespoke_models/test_composites/test_zscore.py @@ -0,0 +1,239 @@ +# tests/test_bespoke_models/test_composites/test_zscore.py +''' +* Author: Evan Komp +* Created: 2026-05-26 + +Unit tests for ZScoreRescaledScorer using a mock inner scorer. These do not +depend on fair-esm / transformers and run under default CI. +''' +from typing import List, Optional + +import numpy as np +import pandas as pd +import pytest + +from aide_predict.bespoke_models.base import ( + ProteinModelWrapper, + CanRegressMixin, + RequiresFixedLengthMixin, + ExpectsNoFitMixin, + RequiresWTToFunctionMixin, +) +from aide_predict.bespoke_models.composites.zscore import ZScoreRescaledScorer +from aide_predict.utils.data_structures import ProteinSequence, ProteinSequences + + +WT_STR = "ACDEFGHIKLMNPQRSTVWY" # 20-mer one of each canonical AA + + +class _MockScorer( + RequiresFixedLengthMixin, + CanRegressMixin, + ExpectsNoFitMixin, + ProteinModelWrapper, +): + """ + Deterministic per-variant scalar scorer. For each variant, returns a + score equal to ``ord(mut_aa) * 0.1 + pos * 0.01`` so different (mut_aa) + groups have systematically different score distributions — useful to + verify the z-score rescaling logic. + """ + + def __init__(self, metadata_folder=None, wt=None, pool=True, **kwargs): + super().__init__(metadata_folder=metadata_folder, wt=wt, **kwargs) + self.pool = pool + + def _fit(self, X, y=None): + self.fitted_ = True + return self + + def _transform(self, X): + out = [] + for seq in X: + diffs = seq.mutated_positions(self.wt) + if len(diffs) == 0: + out.append(0.0) + else: + # base on first mutation; for multi-mutants we sum + score = 0.0 + for p in diffs: + score += ord(str(seq[p])) * 0.1 + p * 0.01 + out.append(score) + return np.array(out).reshape(-1, 1) + + +@pytest.fixture +def wt(): + return ProteinSequence(WT_STR, id="wt") + + +@pytest.fixture +def ssm(wt): + """All 19*L single-point variants of WT.""" + return wt.saturation_mutagenesis() + + +@pytest.fixture +def composite(wt): + inner = _MockScorer(wt=wt, pool=True) + return ZScoreRescaledScorer(inner_scorer=inner, wt=wt) + + +class TestMixinFlags: + def test_requires_wt(self): + assert ZScoreRescaledScorer._requires_wt_to_function is True + + def test_requires_fixed_length(self): + assert ZScoreRescaledScorer._requires_fixed_length is True + + def test_can_regress(self): + assert ZScoreRescaledScorer._can_regress is True + + def test_fit_required(self): + # Composite needs X at fit time to compute group stats, so it does NOT + # set _expects_no_fit. (Removed ExpectsNoFitMixin from the inheritance.) + assert ZScoreRescaledScorer._expects_no_fit is False + + +class TestInit: + def test_pool_false_rejected(self, wt): + inner = _MockScorer(wt=wt, pool=False) + with pytest.raises(ValueError, match="pool=False"): + ZScoreRescaledScorer(inner_scorer=inner, wt=wt) + + def test_invalid_grouping(self, wt): + inner = _MockScorer(wt=wt, pool=True) + with pytest.raises(ValueError, match="grouping must be"): + ZScoreRescaledScorer(inner_scorer=inner, wt=wt, grouping="bogus") + + def test_invalid_multi_strategy(self, wt): + inner = _MockScorer(wt=wt, pool=True) + with pytest.raises(ValueError, match="multi_mutant_strategy"): + ZScoreRescaledScorer(inner_scorer=inner, wt=wt, multi_mutant_strategy="explode") + + def test_inner_wt_synced(self, wt): + inner = _MockScorer(pool=True) # no wt + composite = ZScoreRescaledScorer(inner_scorer=inner, wt=wt) + # inner.wt should now be the composite's wt + assert composite.inner_scorer.wt == wt + + +class TestFitTransform: + def test_fit_transform_shape(self, composite, ssm): + composite.fit(ssm) + arr = composite.transform(ssm) + # 20-residue WT × 19 substitutions = 380 variants + assert arr.shape == (len(ssm), 2) + + def test_score_table_columns(self, composite, ssm): + composite.fit(ssm) + df = composite.score_table(ssm) + expected = { + "variant_idx", "variant_id", "n_mutations", + "mutation", "position", "wt_aa", "mut_aa", + "logratio", "z_logratio", + } + assert expected.issubset(set(df.columns)) + assert len(df) == len(ssm) + + def test_z_logratio_centered_per_group(self, composite, ssm): + composite.fit(ssm) + df = composite.score_table(ssm) + # Group by destination AA → each group's z-scores should have mean ≈ 0 + # and std ≈ 1 (sample std, N-1). + for dest_aa, sub in df.groupby("mut_aa"): + if len(sub) >= 5: + z = sub["z_logratio"].dropna() + if len(z) > 1: + np.testing.assert_allclose(z.mean(), 0.0, atol=1e-6) + np.testing.assert_allclose(z.std(ddof=1), 1.0, atol=1e-6) + + def test_score_table_before_fit_raises(self, composite, ssm): + with pytest.raises(RuntimeError, match="must be fit"): + composite.score_table(ssm) + + def test_substitution_type_grouping(self, wt): + inner = _MockScorer(wt=wt, pool=True) + comp = ZScoreRescaledScorer( + inner_scorer=inner, wt=wt, grouping="substitution_type" + ) + ssm_local = wt.saturation_mutagenesis() + comp.fit(ssm_local) + df = comp.score_table(ssm_local) + # Substitution-type groups: (wt_aa, mut_aa). WT has 20 distinct AAs and + # 19 substitutions each → 20×19 = 380 unique pairs, each with 1 sample, + # so each group is below min_group_size=5 → all z-scores NaN. + assert df["z_logratio"].isna().all() + + def test_fit_one_transform_another(self, wt): + """Stats learned on the SSM should be reused on a smaller variant set.""" + inner = _MockScorer(wt=wt, pool=True) + comp = ZScoreRescaledScorer(inner_scorer=inner, wt=wt) + ssm = wt.saturation_mutagenesis() + comp.fit(ssm) + stats_after_fit = dict(comp.group_stats_) + + small = ProteinSequences(ssm.data[:10]) + arr = comp.transform(small) + assert arr.shape == (10, 2) + # Stats should not have changed. + assert comp.group_stats_ == stats_after_fit + + +class TestMultiMutantStrategy: + def test_skip_default(self, wt): + inner = _MockScorer(wt=wt, pool=True) + comp = ZScoreRescaledScorer(inner_scorer=inner, wt=wt, multi_mutant_strategy="skip") + # mix single + double mutants + variants = ProteinSequences([ + ProteinSequence("ACDEFGHIKLMNPQRSTVWA"), # Y20A → ord('A')=65 + ProteinSequence("ACDEFGHIKLMNPQRSTVWP"), # Y20P + ProteinSequence("ACDEFGHIKLMNPQRSTVWG"), # Y20G + ProteinSequence("ACDEFGHIKLMNPQRSTVWS"), # Y20S + ProteinSequence("ACDEFGHIKLMNPQRSTVWT"), # Y20T + ProteinSequence("AADEFGHIKLMNPQRSTVWA"), # double mutant C2A,Y20A + ]) + comp.fit(variants) + df = comp.score_table(variants) + # Double mutant row's z_logratio should be NaN. + assert pd.isna(df.iloc[5]["z_logratio"]) + # But its logratio is finite (inner scorer scored it). + assert np.isfinite(df.iloc[5]["logratio"]) + + def test_raise(self, wt): + inner = _MockScorer(wt=wt, pool=True) + comp = ZScoreRescaledScorer(inner_scorer=inner, wt=wt, multi_mutant_strategy="raise") + variants = ProteinSequences([ + ProteinSequence("AADEFGHIKLMNPQRSTVWA"), # double mutant + ]) + with pytest.raises(ValueError, match="multi-mutation"): + comp.fit(variants) + + +class TestEdgeCases: + def test_group_below_min_size(self, wt): + inner = _MockScorer(wt=wt, pool=True) + comp = ZScoreRescaledScorer(inner_scorer=inner, wt=wt, min_group_size=5) + # Only 3 destination-residue groups, each with 3 samples → all below min size. + variants = ProteinSequences([ + ProteinSequence("CCDEFGHIKLMNPQRSTVWY"), # A1C + ProteinSequence("ACCEFGHIKLMNPQRSTVWY"), # D3C + ProteinSequence("ACDEFCHIKLMNPQRSTVWY"), # G6C + ]) + comp.fit(variants) + df = comp.score_table(variants) + assert df["z_logratio"].isna().all() + + def test_wt_identical_variant(self, wt): + inner = _MockScorer(wt=wt, pool=True) + comp = ZScoreRescaledScorer(inner_scorer=inner, wt=wt) + variants = ProteinSequences([ + ProteinSequence(WT_STR), # identical to WT + ]) + # No fit; just transform after a separate fit on real SSM: + ssm = wt.saturation_mutagenesis() + comp.fit(ssm) + df = comp.score_table(variants) + assert df.iloc[0]["n_mutations"] == 0 + # z_logratio NaN, logratio = inner's prediction (0.0 for our mock) + assert pd.isna(df.iloc[0]["z_logratio"]) diff --git a/tests/test_not_base_models/test_zscore_composite.py b/tests/test_not_base_models/test_zscore_composite.py new file mode 100644 index 0000000..662d86e --- /dev/null +++ b/tests/test_not_base_models/test_zscore_composite.py @@ -0,0 +1,103 @@ +# tests/test_not_base_models/test_zscore_composite.py +''' +* Author: Evan Komp +* Created: 2026-05-26 + +Integration tests for ZScoreRescaledScorer wrapping the real ESM-IF wrapper. +Excluded from default CI (@pytest.mark.optional) because it loads fair-esm. +''' +import os + +import numpy as np +import pandas as pd +import pytest +import torch + +from aide_predict.utils.data_structures import ( + ProteinSequence, + ProteinSequences, + ProteinStructure, +) +from aide_predict.bespoke_models.composites.zscore import ZScoreRescaledScorer + + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +ENVZ_WT = "LADDRTLLMAGVSHDLRTPLTRIRLATEMMSEQDGYLAESINKDIEECNAIIEQFIDYLR" + + +@pytest.mark.optional +def test_zscore_wrapping_esm_if_envz_ssm(): + """End-to-end: wrap ESM-IF wildtype_marginal with z-score rescaling on the + ENVZ_ECOLI SSM. Check shape, columns, and that within-group z-scores have + the expected mean ≈ 0 / std ≈ 1 behavior.""" + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + structure = ProteinStructure(os.path.join("tests", "data", "ENVZ_ECOLI.pdb")) + wt = ProteinSequence(ENVZ_WT, structure=structure, id="ENVZ_WT") + + inner = ESMIFLikelihoodWrapper( + marginal_method="wildtype_marginal", + wt=wt, + pool=True, + device=DEVICE, + ) + composite = ZScoreRescaledScorer( + inner_scorer=inner, + grouping="destination_residue", + wt=wt, + ) + + ssm = wt.saturation_mutagenesis() + assert len(ssm) == len(ENVZ_WT) * 19 # 60 * 19 = 1140 + + composite.fit(ssm) + arr = composite.transform(ssm) + assert arr.shape == (len(ssm), 2) + assert np.isfinite(arr[:, 0]).all() # raw log-ratios should all be finite + + df = composite.score_table(ssm) + expected_cols = { + "variant_idx", "variant_id", "n_mutations", + "mutation", "position", "wt_aa", "mut_aa", + "logratio", "z_logratio", + } + assert expected_cols.issubset(df.columns) + assert (df["n_mutations"] == 1).all() + # Every group in an SSM has 60 samples (one per position) so all groups + # exceed min_group_size=5 and every variant gets a z_logratio. + assert df["z_logratio"].notna().all() + + # Within each destination-AA group, z-scores should be centered and unit-std. + for dest_aa, sub in df.groupby("mut_aa"): + z = sub["z_logratio"] + np.testing.assert_allclose(z.mean(), 0.0, atol=1e-6) + np.testing.assert_allclose(z.std(ddof=1), 1.0, atol=1e-6) + + +@pytest.mark.optional +def test_zscore_fit_one_transform_subset(): + """Fit on the full SSM, transform a subset — stats must persist and be applied.""" + from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper + + structure = ProteinStructure(os.path.join("tests", "data", "ENVZ_ECOLI.pdb")) + wt = ProteinSequence(ENVZ_WT, structure=structure, id="ENVZ_WT") + + composite = ZScoreRescaledScorer( + inner_scorer=ESMIFLikelihoodWrapper( + marginal_method="wildtype_marginal", wt=wt, pool=True, device=DEVICE, + ), + wt=wt, + ) + + ssm = wt.saturation_mutagenesis() + composite.fit(ssm) + learned_stats = dict(composite.group_stats_) + + subset = ProteinSequences(ssm.data[:20]) + arr = composite.transform(subset) + assert arr.shape == (20, 2) + assert np.isfinite(arr).all() + + # Stats must not have been recomputed during transform. + assert composite.group_stats_ == learned_stats diff --git a/tests/test_utils/test_scoring.py b/tests/test_utils/test_scoring.py new file mode 100644 index 0000000..220a57e --- /dev/null +++ b/tests/test_utils/test_scoring.py @@ -0,0 +1,159 @@ +# tests/test_utils/test_scoring.py +''' +* Author: Evan Komp +* Created: 2026-05-26 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Unit tests for the scorer-agnostic primitives in aide_predict.utils.scoring. +''' +import numpy as np +import pandas as pd +import pytest + +from aide_predict.utils.data_structures import ProteinSequence, ProteinSequences +from aide_predict.utils.scoring import per_variant_mutation_info, zscore_by_aa_group + + +WT_STR = "ACDEFGHIKL" + + +@pytest.fixture +def wt(): + return ProteinSequence(WT_STR, id="wt") + + +class TestPerVariantMutationInfo: + def test_single_point_variants(self, wt): + variants = ProteinSequences([ + ProteinSequence("ACDEYGHIKL", id="F5Y"), # single mutation at pos 5 + ProteinSequence("ACDEFGHIKP", id="L10P"), # single mutation at pos 10 + ]) + df = per_variant_mutation_info(variants, wt) + assert len(df) == 2 + assert df.iloc[0]["mutation"] == "F5Y" + assert df.iloc[0]["position"] == 5 + assert df.iloc[0]["wt_aa"] == "F" + assert df.iloc[0]["mut_aa"] == "Y" + assert df.iloc[0]["n_mutations"] == 1 + assert df.iloc[1]["mutation"] == "L10P" + assert df.iloc[1]["position"] == 10 + + def test_multi_mutation_variant(self, wt): + variants = ProteinSequences([ + ProteinSequence("ACDEYGHIKP", id="F5Y_L10P"), # two mutations + ]) + df = per_variant_mutation_info(variants, wt) + assert df.iloc[0]["n_mutations"] == 2 + assert df.iloc[0]["mutation"] is None + assert df.iloc[0]["position"] is None + assert df.iloc[0]["wt_aa"] is None + assert df.iloc[0]["mut_aa"] is None + + def test_identical_to_wt(self, wt): + variants = ProteinSequences([ProteinSequence(WT_STR, id="wt_copy")]) + df = per_variant_mutation_info(variants, wt) + assert df.iloc[0]["n_mutations"] == 0 + assert df.iloc[0]["mutation"] is None + + def test_variant_id_preserved(self, wt): + variants = ProteinSequences([ + ProteinSequence("ACDEYGHIKL", id="some_id"), + ProteinSequence("ACDEFGHIKP"), # no id + ]) + df = per_variant_mutation_info(variants, wt) + assert df.iloc[0]["variant_id"] == "some_id" + assert df.iloc[1]["variant_id"] is None + + def test_length_mismatch_raises(self, wt): + variants = ProteinSequences([ProteinSequence("ACDEF", id="short")]) + with pytest.raises(ValueError, match="length"): + per_variant_mutation_info(variants, wt) + + +class TestZScoreByAAGroup: + @pytest.fixture + def synthetic_df(self): + # Build a df with two destination groups, each with 5 samples. + # Group "→P": scores = [1, 2, 3, 4, 5], mean=3, std=sqrt(2.5) ≈ 1.5811 + # Group "→Q": scores = [10, 20, 30, 40, 50], mean=30, std=sqrt(250) ≈ 15.811 + rows = [] + for s in [1, 2, 3, 4, 5]: + rows.append({"score": float(s), "wt_aa": "A", "mut_aa": "P"}) + for s in [10, 20, 30, 40, 50]: + rows.append({"score": float(s), "wt_aa": "A", "mut_aa": "Q"}) + return pd.DataFrame(rows) + + def test_destination_residue_grouping(self, synthetic_df): + df_out, stats = zscore_by_aa_group( + synthetic_df, grouping="destination_residue", score_col="score" + ) + # Two groups: ('P',) and ('Q',) + assert set(stats.keys()) == {("P",), ("Q",)} + # P group: mean=3, std≈1.5811 + np.testing.assert_allclose(stats[("P",)][0], 3.0) + np.testing.assert_allclose(stats[("P",)][1], np.std([1, 2, 3, 4, 5], ddof=1)) + # Z scores: row 0 (score=1) → (1-3)/1.5811 ≈ -1.2649 + z_p_first = (1.0 - 3.0) / np.std([1, 2, 3, 4, 5], ddof=1) + np.testing.assert_allclose(df_out.iloc[0]["z_score"], z_p_first) + # Row 5 (score=10, group Q) → (10-30)/std_q + z_q_first = (10.0 - 30.0) / np.std([10, 20, 30, 40, 50], ddof=1) + np.testing.assert_allclose(df_out.iloc[5]["z_score"], z_q_first) + + def test_substitution_type_grouping(self, synthetic_df): + df_out, stats = zscore_by_aa_group( + synthetic_df, grouping="substitution_type", score_col="score" + ) + # Two groups: ('A','P') and ('A','Q') + assert set(stats.keys()) == {("A", "P"), ("A", "Q")} + + def test_min_group_size_filter(self): + # Group with 4 samples → below threshold of 5, z-scores stay NaN. + df = pd.DataFrame([ + {"score": 1.0, "wt_aa": "A", "mut_aa": "P"}, + {"score": 2.0, "wt_aa": "A", "mut_aa": "P"}, + {"score": 3.0, "wt_aa": "A", "mut_aa": "P"}, + {"score": 4.0, "wt_aa": "A", "mut_aa": "P"}, + ]) + df_out, stats = zscore_by_aa_group(df, score_col="score", min_group_size=5) + assert ("P",) not in stats + assert df_out["z_score"].isna().all() + + def test_custom_out_col(self, synthetic_df): + df_out, _ = zscore_by_aa_group( + synthetic_df, score_col="score", out_col="custom_z" + ) + assert "custom_z" in df_out.columns + assert "z_score" not in df_out.columns + + def test_fit_stats_override(self, synthetic_df): + # Provide pre-computed stats and confirm they're used as-is. + fit_stats = {("P",): (100.0, 1.0), ("Q",): (200.0, 1.0)} + df_out, returned_stats = zscore_by_aa_group( + synthetic_df, score_col="score", fit_stats=fit_stats + ) + # P group score=1 → (1-100)/1 = -99 + np.testing.assert_allclose(df_out.iloc[0]["z_score"], -99.0) + # Q group score=10 → (10-200)/1 = -190 + np.testing.assert_allclose(df_out.iloc[5]["z_score"], -190.0) + assert returned_stats is fit_stats + + def test_missing_mut_aa_column_raises(self): + df = pd.DataFrame({"score": [1.0, 2.0]}) + with pytest.raises(KeyError, match="mut_aa"): + zscore_by_aa_group(df, score_col="score") + + def test_missing_wt_aa_for_substitution_type_raises(self): + df = pd.DataFrame({"score": [1.0], "mut_aa": ["P"]}) + with pytest.raises(KeyError, match="wt_aa"): + zscore_by_aa_group(df, grouping="substitution_type", score_col="score") + + def test_nan_mut_aa_rows_left_nan(self): + # Multi-mutation rows (mut_aa = None) shouldn't get z-scored. + rows = [{"score": float(s), "wt_aa": "A", "mut_aa": "P"} for s in range(1, 6)] + rows.append({"score": 99.0, "wt_aa": None, "mut_aa": None}) + df = pd.DataFrame(rows) + df_out, stats = zscore_by_aa_group(df, score_col="score") + assert pd.isna(df_out.iloc[5]["z_score"]) + # The first 5 rows ARE z-scored (P group has 5 samples). + assert df_out.iloc[:5]["z_score"].notna().all() From 3d4f1326e65c4b95eb97df976bdf859bd7952618 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Wed, 10 Jun 2026 11:55:05 -0600 Subject: [PATCH 08/11] esm1f, zscore, combine into multi evolve style, update docs to show it as an example --- README.md | 7 + component_specification.md | 12 +- docs/index.rst | 1 + docs/user_guide/data_structures.md | 31 ++++ docs/user_guide/multi_evolve_ensemble.md | 155 ++++++++++++++++ docs/user_guide/zero_shot.md | 27 +++ scripts/zscore_ssm_demo.py | 219 +++++++++++++++++++++++ 7 files changed, 451 insertions(+), 1 deletion(-) create mode 100644 docs/user_guide/multi_evolve_ensemble.md create mode 100644 scripts/zscore_ssm_demo.py diff --git a/README.md b/README.md index bc295f1..2ce6b2f 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ You can always check which modules are installed/available to you by running `ge ### Data Structures and Utilities - Protein Sequence and Structure data structures + - `ProteinStructure` supports multichain complexes: a primary chain plus `context_chains` supplied to structure-aware models (e.g. ESM-IF1) as 3D context. Use `set_target_chain` / `with_target_chain` to switch the scored chain and auto-populate context. - `StructureMapper` - A utility for mapping a folder of PDB structures to sequences ### Prediction Models @@ -97,6 +98,12 @@ You can always check which modules are installed/available to you by running `ge - Supports multichain complexes via `ProteinStructure.context_chains` (set explicitly, or via `set_target_chain` / `with_target_chain`, or implicitly through `StructureMapper.get_protein_sequences(..., auto_context=True)`). - Requires additional dependencies (see `requirements_files/requirements-esm-if.txt`, which adds `torch-geometric`, `torch-scatter`, `torch-cluster`, and pins `biotite<1.0` due to a fair-esm compatibility issue). +### Composite / Ensemble Utilities + +1. `ZScoreRescaledScorer` + - Wraps any per-variant scorer and rescales its log-ratios within an amino-acid group (destination residue or substitution type) — the z-score ("Z") ranking from the [MULTI-evolve](https://www.science.org/doi/10.1126/science.adr8628) ensemble. Behaves as a single scikit-learn transformer (`fit`/`transform`), and `score_table` returns the per-mutation DataFrame. + - See the "Ensemble Variant Nomination" user guide for a full structure + sequence first-round scan. Lower-level primitives (`per_variant_mutation_info`, `zscore_by_aa_group`) live in `aide_predict.utils.scoring`. + ### Embeddings for Downstream ML 1. One Hot Protein Embedding diff --git a/component_specification.md b/component_specification.md index 63bca21..ae23843 100644 --- a/component_specification.md +++ b/component_specification.md @@ -167,11 +167,15 @@ def as_array(self) -> np.ndarray **Key Methods**: ```python -def __init__(self, pdb_file: str, chain: str = 'A', plddt_file: Optional[str] = None) +def __init__(self, structure_file: str, chain: str = 'A', plddt_file: Optional[str] = None, + context_chains: Optional[Tuple[str, ...]] = None) # context_chains: other chains supplied as 3D context to structure-aware models (multichain complexes) def get_sequence(self) -> str def get_plddt(self) -> Optional[np.ndarray] def get_dssp(self) -> Dict[str, str] def validate_sequence(self, protein_sequence: str) -> bool +def get_all_chain_ids(self, protein_only: bool = True) -> List[str] +def get_chain_coords(self, chain_id: str) -> np.ndarray # [L, 3, 3] backbone (N, CA, C) +def set_target_chain(self, new_chain: str, auto_context: bool = True) -> None @classmethod def from_af2_folder(cls, folder_path: str, chain: str = 'A') -> 'ProteinStructure' ``` @@ -215,6 +219,12 @@ The package includes several bespoke models implemented using the `ProteinModelW - MSA Transformer (Embeddings and Log Likelihood Predictor) - SaProt (Embeddings and Log Likelihood Predictor) - EVMutation (Predicts hamiltonian changes using pairwise potentials) +- ESM-IF1 (Structure-conditioned autoregressive Log Likelihood Predictor; supports multichain complexes) +- HMM, VESPA, EVE, SSEmb (additional zero-shot predictors) + +The package also provides composite wrappers built on the same API: + +- ZScoreRescaledScorer (wraps a per-variant scorer and adds amino-acid-group z-score rescaling; see `aide_predict.utils.scoring` for the underlying primitives) These models follow the structure and API defined by the base classes and mixins. diff --git a/docs/index.rst b/docs/index.rst index 924da99..790e299 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -24,6 +24,7 @@ User Guide user_guide/zero_shot.md user_guide/supervised.md user_guide/saturation_mutagenesis.md + user_guide/multi_evolve_ensemble.md user_guide/pipelines.md user_guide/caching.md user_guide/position_specific.md diff --git a/docs/user_guide/data_structures.md b/docs/user_guide/data_structures.md index f999df2..6617341 100644 --- a/docs/user_guide/data_structures.md +++ b/docs/user_guide/data_structures.md @@ -78,6 +78,10 @@ seq2 = seq2.align(seq) # align to an existing alignment seq2 = seq2.align(msa) # Align to existing MSA +# For a sequence backed by a multichain structure, point it at a different +# chain (returns a NEW ProteinSequence with a cloned structure; the original +# is unchanged). auto_context populates the other chains as structural context. +seqB = seq.with_target_chain("B", auto_context=True) ``` ## ProteinSequences @@ -187,6 +191,28 @@ chain_obj = structure.get_chain() # Get specific chain positions = structure.get_residue_positions() # Get residue numbers ``` +### Multichain complexes + +For complexes (e.g. AlphaFold3 multimers), a `ProteinStructure` carries one +**primary chain** (the one whose sequence is scored) plus optional +`context_chains` — additional chains supplied to structure-aware models like +ESM-IF1 as 3D context but not themselves scored. + +```python +# List the chains present in the file +structure.get_all_chain_ids() # e.g. ['A', 'B', 'C'] + +# Switch the primary chain in place; auto_context makes every OTHER chain context +structure.set_target_chain("A", auto_context=True) +structure.context_chains # -> ('B', 'C') + +# Set context chains explicitly at construction time instead +structure = ProteinStructure("complex.pdb", chain="A", context_chains=("B", "C")) + +# Backbone coords for any chain as an [L, 3, 3] (N, CA, C) array +coords = structure.get_chain_coords("A") +``` + ## StructureMapper `StructureMapper` helps manage multiple structures and map them to sequences, particularly useful when working with structure-aware models. @@ -214,6 +240,11 @@ available_ids = mapper.get_available_structures() # Get ProteinSequences with structures sequences = mapper.get_protein_sequences() + +# Choose a target chain for every structure, or pass a JSON path mapping +# {structure_id: chain_id} for per-file overrides. auto_context (default True) +# wires the other chains of each file in as structural context. +sequences = mapper.get_protein_sequences(target_chain="A", auto_context=True) ``` The StructureMapper is particularly useful when working with structure-aware models like SaProt, which can use structure information to improve predictions: diff --git a/docs/user_guide/multi_evolve_ensemble.md b/docs/user_guide/multi_evolve_ensemble.md new file mode 100644 index 0000000..c8beff9 --- /dev/null +++ b/docs/user_guide/multi_evolve_ensemble.md @@ -0,0 +1,155 @@ +--- +title: Ensemble Variant Nomination +--- + +# Ensemble Variant Nomination (MULTI-evolve first round) + +## Overview + +This guide reproduces the first round of the [MULTI-evolve](https://www.science.org/doi/10.1126/science.adr8628) zero-shot nomination workflow using AIDE primitives. The idea is to score a full saturation-mutagenesis (SSM) library along two complementary tracks and combine them: + +- a **structure track** — `ESMIFLikelihoodWrapper` (ESM-IF1), conditioned on the backbone; +- a **sequence track** — an ensemble of sequence PLMs (`ESM2LikelihoodWrapper` loading the ESM-1v / ESM-2 checkpoints), averaged. + +Each track is scored as a per-mutation log-ratio, then **rescaled within an amino-acid group** (`zscore_by_aa_group`) so that mutations to, say, proline are compared against other mutations to proline rather than against the whole library. Finally we take a **position-exclusive top-N** under four rankings (fold-change and z-score for each track) and **union** them — at most one mutation per residue per method. + +All of the pieces below are public API. A ready-to-run driver that implements this exact workflow (including multimer handling and plots) lives at `scripts/multi_evolve_ssm_scan.py`. + +## 1. Build the WT and its SSM library + +```python +from aide_predict import ProteinSequence, ProteinSequences + +# from_pdb extracts the chain sequence and attaches the ProteinStructure. +wt = ProteinSequence.from_pdb("my_enzyme.pdb", chain="A", id="my_enzyme") + +# All single-point variants (19 * L). MULTI-evolve drops Met1 by convention +# (mutated_positions is 0-indexed, so position 0 is residue 1). +ssm = wt.saturation_mutagenesis() +ssm = ProteinSequences([s for s in ssm if s.mutated_positions(wt)[0] != 0]) +``` + +## 2. Structure track — ESM-IF1 + +```python +from aide_predict import ESMIFLikelihoodWrapper + +esm_if = ESMIFLikelihoodWrapper( + wt=wt, + marginal_method="wildtype_marginal", # masked_marginal is refused (autoregressive) + pool=True, # one scalar (log-ratio) per variant + device="cuda", +) +esm_if.fit() +esm_if_logratio = esm_if.predict(ssm).ravel() +``` + +## 3. Sequence track — PLM ensemble + +The full MULTI-evolve sequence track is five ESM-1v checkpoints plus `esm2_t36_3B_UR50D`. Load each one, score, then free it before the next so peak GPU memory stays bounded. Two checkpoints are shown here for brevity: + +```python +import numpy as np +from aide_predict import ESM2LikelihoodWrapper + +PLM_CHECKPOINTS = ["esm1v_t33_650M_UR90S_1", "esm2_t36_3B_UR50D"] +plm_logratios = [] +for ckpt in PLM_CHECKPOINTS: + plm = ESM2LikelihoodWrapper( + wt=wt, model_checkpoint=ckpt, + marginal_method="wildtype_marginal", pool=True, device="cuda", + ) + plm.fit() + plm_logratios.append(plm.predict(ssm).ravel()) + del plm # free before loading the next checkpoint + +plm_logratios = np.vstack(plm_logratios) # (n_plms, n_variants) +mean_plm_logratio = plm_logratios.mean(axis=0) # Track A score +plm_pass_count = (plm_logratios > 0).sum(axis=0) # consensus tier for Seq-FC +``` + +## 4. Assemble the table and z-score each track + +`per_variant_mutation_info` gives the `(wt_aa, mut_aa, position, mutation)` scaffold; `zscore_by_aa_group` adds the AA-grouped z-score column and returns the per-group `(mean, std)` stats it learned. + +```python +from aide_predict.utils.scoring import per_variant_mutation_info, zscore_by_aa_group + +df = per_variant_mutation_info(ssm, wt) +df["esm_if_logratio"] = esm_if_logratio +df["mean_plm_logratio"] = mean_plm_logratio +df["plm_pass_count"] = plm_pass_count + +# Rescale within destination-residue groups (the MULTI-evolve default). +df, _ = zscore_by_aa_group(df, grouping="destination_residue", + score_col="esm_if_logratio", out_col="z_esm_if_logratio", + min_group_size=5) +df, _ = zscore_by_aa_group(df, grouping="destination_residue", + score_col="mean_plm_logratio", out_col="z_mean_plm_logratio", + min_group_size=5) +``` + +## 5. Four-method position-exclusive top-N, then union + +```python +import pandas as pd + +def position_exclusive_topn(df, sort_cols, ascendings, n): + """Top-n rows, at most one per residue position, skipping NaN sort keys.""" + used, picked = set(), [] + for _, row in df.sort_values(sort_cols, ascending=ascendings).iterrows(): + if any(pd.isna(row[c]) for c in sort_cols) or row["position"] in used: + continue + picked.append(row); used.add(row["position"]) + if len(picked) >= n: + break + return pd.DataFrame(picked) + +N = 24 +methods = { + "seq_fc": (["plm_pass_count", "mean_plm_logratio"], [False, False]), # consensus + magnitude + "struct_fc": (["esm_if_logratio"], [False]), + "seq_z": (["z_mean_plm_logratio"], [False]), + "struct_z": (["z_esm_if_logratio"], [False]), +} +picks = {m: position_exclusive_topn(df, cols, asc, N) for m, (cols, asc) in methods.items()} + +# Union: one row per nominated mutation, with a per-method indicator column. +nominations = pd.concat(picks.values()).drop_duplicates("mutation")[["mutation"]].copy() +for m, sub in picks.items(): + nominations[f"{m}_picked"] = nominations["mutation"].isin(sub["mutation"]).astype(int) +nominations["n_methods_picked"] = nominations.filter(like="_picked").sum(axis=1) +nominations = nominations.sort_values("n_methods_picked", ascending=False) +``` + +`nominations` now holds the unioned variant short-list; `n_methods_picked` (1–4) tells you how many of the four rankings agreed on each mutation. + +## Reusable composite: `ZScoreRescaledScorer` + +The z-score step (steps 2–4 for a single scorer) is also packaged as a scikit-learn-style transformer you can wrap around *any* per-variant scorer. It fits the AA-group stats on a calibration set and applies them at transform time: + +```python +from aide_predict import ESMIFLikelihoodWrapper, ZScoreRescaledScorer + +inner = ESMIFLikelihoodWrapper(wt=wt, marginal_method="wildtype_marginal", pool=True) +scorer = ZScoreRescaledScorer(inner_scorer=inner, grouping="destination_residue", + min_group_size=5, wt=wt) + +scorer.fit(ssm) +arr = scorer.transform(ssm) # (N, 2): columns [logratio, z_logratio] +table = scorer.score_table(ssm) # rich DataFrame: mutation, position, wt_aa, + # mut_aa, logratio, z_logratio, ... +``` + +`grouping` switches between `"destination_residue"` (bin by `mut_aa`) and `"substitution_type"` (bin by the `(wt_aa, mut_aa)` pair); `multi_mutant_strategy` controls whether multi-mutation variants are tolerated (`"skip"`, giving them a NaN z-score) or rejected (`"raise"`). + +## Multichain complexes + +ESM-IF1 can condition the target chain on the rest of a complex. If `my_enzyme.pdb` were a multimer, point the WT at the chain you want to score and let the others become structural context: + +```python +wt = ProteinSequence.from_pdb("complex.pdb", chain="A") +wt = wt.with_target_chain("A", auto_context=True) # other chains -> ESM-IF context +``` + +For a homo-/hetero-multimer you typically loop the target chain over each protein chain and run the scan once per chain. See [Data Structures](data_structures.md) for `set_target_chain`, `with_target_chain`, and `context_chains`. diff --git a/docs/user_guide/zero_shot.md b/docs/user_guide/zero_shot.md index 629e252..4e145d7 100644 --- a/docs/user_guide/zero_shot.md +++ b/docs/user_guide/zero_shot.md @@ -160,6 +160,33 @@ scores = model.predict(mutants) SSEmb is especially effective for scoring mutations in proteins with known structures and rich evolutionary information. +### ESM-IF1 + +[ESM-IF1](https://proceedings.mlr.press/v162/hsu22a.html) is a structure-conditioned, autoregressive language model (the GVP-Transformer) that scores variants given a protein backbone: + +```python +from aide_predict import ESMIFLikelihoodWrapper, ProteinSequence + +# Build the wild type from a structure (sequence + ProteinStructure attached) +wt = ProteinSequence.from_pdb("structures/structure.pdb", chain="A") + +# wildtype_marginal or mutant_marginal; masked_marginal is refused +# (ESM-IF has no bidirectional mask-scoring mode) +model = ESMIFLikelihoodWrapper( + wt=wt, + marginal_method="wildtype_marginal", +) + +# No training needed +model.fit() + +# Score mutations +mutants = wt.saturation_mutagenesis() +scores = model.predict(mutants) +``` + +ESM-IF1 also conditions on the rest of a complex when the attached structure carries `context_chains`. See [Ensemble Variant Nomination](multi_evolve_ensemble.md) for a worked structure + sequence ensemble and [Data Structures](data_structures.md) for the multichain helpers. + ## Evolutionary Models ### HMM diff --git a/scripts/zscore_ssm_demo.py b/scripts/zscore_ssm_demo.py new file mode 100644 index 0000000..d8c9f83 --- /dev/null +++ b/scripts/zscore_ssm_demo.py @@ -0,0 +1,219 @@ +""" +Walkthrough: ESM-IF SSM scores with AA-grouped z-score rescaling. + +Wraps an ESM-IF wildtype_marginal scorer with ZScoreRescaledScorer, runs the +SSM of a small protein (ENVZ_ECOLI from tests/data, 60 residues), and prints +the per-mutation table with both raw log-ratios and destination-AA z-scores. + +Designed to be stepped through in a debugger — single linear function. + +Run:: + + python scripts/zscore_ssm_demo.py + +Output:: + + scripts/zscore_ssm_demo.csv (per-mutation table) + scripts/zscore_ssm_demo_plots/*.png (distribution + SSM plots) +""" +import os +import sys +from pathlib import Path + +HERE = Path(__file__).resolve().parent +REPO_ROOT = HERE.parent +sys.path.insert(0, str(REPO_ROOT)) + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import torch + +os.environ.setdefault("KEEP_MODEL_ON_DEVICE", "1") + +from aide_predict.utils.data_structures import ( + ProteinSequence, + ProteinSequences, + ProteinStructure, +) +from aide_predict.bespoke_models.predictors.esm_if import ESMIFLikelihoodWrapper +from aide_predict.bespoke_models.composites.zscore import ZScoreRescaledScorer + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +PDB = REPO_ROOT / "tests" / "data" / "ENVZ_ECOLI.pdb" +WT_SEQ = "LADDRTLLMAGVSHDLRTPLTRIRLATEMMSEQDGYLAESINKDIEECNAIIEQFIDYLR" +OUT_CSV = HERE / "zscore_ssm_demo.csv" +PLOT_DIR = HERE / "zscore_ssm_demo_plots" + +AA_ORDER = list("ACDEFGHIKLMNPQRSTVWY") + + +def _save(fig, name): + PLOT_DIR.mkdir(exist_ok=True) + path = PLOT_DIR / name + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f" wrote {path}") + + +def make_plots(df, wt_seq): + """Render distribution + SSM-style plots of the rescaled SSM.""" + sns.set_theme(context="notebook", style="whitegrid") + + # 1. Overlaid histograms of raw vs z-scored values. + fig, axes = plt.subplots(1, 2, figsize=(11, 4)) + axes[0].hist(df["logratio"], bins=40, color="steelblue", edgecolor="white") + axes[0].axvline(0, color="black", lw=0.8, ls="--") + axes[0].set(title="Raw log-ratio", xlabel="logratio", ylabel="count") + axes[1].hist(df["z_logratio"], bins=40, color="darkorange", edgecolor="white") + axes[1].axvline(0, color="black", lw=0.8, ls="--") + axes[1].set(title="Z-rescaled log-ratio (by destination AA)", + xlabel="z_logratio", ylabel="count") + fig.suptitle("SSM score distributions: before vs after z-score rescaling") + fig.tight_layout() + _save(fig, "01_distribution_raw_vs_zscore.png") + + # 2. Per-destination-AA violins — raw scale shows why rescaling is needed + # (AA-specific offsets/spreads), z scale collapses them to a common frame. + fig, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True) + sns.violinplot(data=df, x="mut_aa", y="logratio", order=AA_ORDER, + inner="quartile", ax=axes[0], color="steelblue") + axes[0].axhline(0, color="black", lw=0.6, ls="--") + axes[0].set(title="Raw log-ratio by destination AA", xlabel="", ylabel="logratio") + sns.violinplot(data=df, x="mut_aa", y="z_logratio", order=AA_ORDER, + inner="quartile", ax=axes[1], color="darkorange") + axes[1].axhline(0, color="black", lw=0.6, ls="--") + axes[1].set(title="Z-rescaled log-ratio by destination AA", + xlabel="destination AA", ylabel="z_logratio") + fig.tight_layout() + _save(fig, "02_per_destination_aa_violins.png") + + # 3. Scatter of raw vs z, colored by destination AA — visualizes how each + # AA group's offset+scale gets removed. + fig, ax = plt.subplots(figsize=(7, 6)) + palette = sns.color_palette("husl", n_colors=len(AA_ORDER)) + for aa, color in zip(AA_ORDER, palette): + sub = df[df["mut_aa"] == aa] + ax.scatter(sub["logratio"], sub["z_logratio"], s=14, alpha=0.7, + color=color, label=aa) + ax.axhline(0, color="black", lw=0.5, ls="--") + ax.axvline(0, color="black", lw=0.5, ls="--") + ax.set(title="Raw vs z-rescaled log-ratio (color = destination AA)", + xlabel="logratio", ylabel="z_logratio") + ax.legend(ncol=2, fontsize=8, loc="best", frameon=True) + fig.tight_layout() + _save(fig, "03_raw_vs_zscore_scatter.png") + + # 4. Classic SSM heatmaps (position x destination AA). Two panels — raw + # then z — share the same layout so they're directly comparable. + pivot_raw = df.pivot(index="mut_aa", columns="position", values="logratio") \ + .reindex(AA_ORDER) + pivot_z = df.pivot(index="mut_aa", columns="position", values="z_logratio") \ + .reindex(AA_ORDER) + + fig, axes = plt.subplots(2, 1, figsize=(max(12, 0.18 * len(wt_seq)), 9), + sharex=True) + vmax_raw = np.nanmax(np.abs(pivot_raw.values)) + sns.heatmap(pivot_raw, ax=axes[0], cmap="RdBu_r", center=0, + vmin=-vmax_raw, vmax=vmax_raw, yticklabels=AA_ORDER, + cbar_kws={"label": "logratio"}) + axes[0].set(title="Raw log-ratio SSM", ylabel="destination AA", xlabel="") + vmax_z = np.nanmax(np.abs(pivot_z.values)) + sns.heatmap(pivot_z, ax=axes[1], cmap="RdBu_r", center=0, + vmin=-vmax_z, vmax=vmax_z, yticklabels=AA_ORDER, + cbar_kws={"label": "z_logratio"}) + axes[1].set(title="Z-rescaled SSM (by destination AA)", + ylabel="destination AA", xlabel="position") + for ax in axes: + ax.set_yticklabels(AA_ORDER, rotation=0, fontsize=8) + + # Tick every 5 residues with the WT identity, to keep the x-axis readable. + positions = sorted(df["position"].unique()) + tick_idx = [i for i, p in enumerate(positions) if (p - 1) % 5 == 0] + tick_labels = [f"{wt_seq[positions[i] - 1]}{positions[i]}" for i in tick_idx] + for ax in axes: + ax.set_xticks([i + 0.5 for i in tick_idx]) + ax.set_xticklabels(tick_labels, rotation=90, fontsize=8) + fig.tight_layout() + _save(fig, "04_ssm_heatmaps.png") + + # 5. Per-position summary: best (max z) and mean z. Highlights "hotspots" + # — positions that tolerate or prefer substitutions on the rescaled scale. + per_pos = df.groupby("position")["z_logratio"].agg(["max", "mean", "min"]) + fig, ax = plt.subplots(figsize=(max(12, 0.18 * len(wt_seq)), 4)) + ax.plot(per_pos.index, per_pos["max"], label="max z", color="darkorange") + ax.plot(per_pos.index, per_pos["mean"], label="mean z", color="steelblue") + ax.plot(per_pos.index, per_pos["min"], label="min z", color="firebrick") + ax.axhline(0, color="black", lw=0.5, ls="--") + ax.set(title="Per-position z_logratio summary", + xlabel="position", ylabel="z_logratio") + ax.legend() + fig.tight_layout() + _save(fig, "05_per_position_summary.png") + + +def main(): + # 1. Build WT with structure attached. + structure = ProteinStructure(str(PDB)) + wt = ProteinSequence(WT_SEQ, structure=structure, id="ENVZ_WT") + print(f"WT length: {len(wt)} (chains: {structure.get_all_chain_ids()})") + + # 2. Build the inner scorer (ESM-IF wildtype_marginal, pool=True → one + # scalar per variant). Any per-variant-scalar scorer works here; ESM-IF + # was chosen because it's tested. + inner = ESMIFLikelihoodWrapper( + marginal_method="wildtype_marginal", + wt=wt, + pool=True, + device=DEVICE, + metadata_folder=str(HERE / "tmp_esm_if"), + ) + + # 3. Wrap with the z-score rescaler. Default grouping is by destination + # residue (→P bins). Switch via grouping='substitution_type' for A→P bins. + composite = ZScoreRescaledScorer( + inner_scorer=inner, + grouping="destination_residue", + wt=wt, + ) + + # 4. Generate the full SSM and score it in one fit_transform-like call. + # composite.fit(ssm) → runs inner.transform(ssm) once, learns + # group stats from the SSM distribution. + # composite.score_table(ssm) → returns the rich DataFrame. + ssm = wt.saturation_mutagenesis() + print(f"Generated {len(ssm)} SSM variants ({len(WT_SEQ)} positions x 19 AAs)") + + composite.fit(ssm) + df = composite.score_table(ssm) + print(f"\nGroup stats keys (destination AAs): " + f"{sorted(g[0] for g in composite.group_stats_.keys())}") + + # 5. Save and summarize. + df.to_csv(OUT_CSV, index=False) + print(f"\nWrote {len(df)} rows to {OUT_CSV}") + + print("\nTop 5 by z_logratio (most favored under destination-AA rescaling):") + print(df.sort_values("z_logratio", ascending=False) + .head()[["mutation", "logratio", "z_logratio"]] + .to_string(index=False)) + + print("\nTop 5 by raw logratio (most favored under absolute log-ratio):") + print(df.sort_values("logratio", ascending=False) + .head()[["mutation", "logratio", "z_logratio"]] + .to_string(index=False)) + + # Sanity: within each destination group the z-scores have mean ≈ 0, std ≈ 1. + print("\nPer-destination-AA z-score sanity (should be ~0 / ~1):") + summary = df.groupby("mut_aa")["z_logratio"].agg(["mean", "std", "count"]) + print(summary.round(4).to_string()) + + # 6. Plots. + print(f"\nWriting plots to {PLOT_DIR}/ ...") + make_plots(df, WT_SEQ) + + +if __name__ == "__main__": + main() From 678ec3bdc046c5a437dab2853576d1a4a740f283 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Wed, 10 Jun 2026 12:11:58 -0600 Subject: [PATCH 09/11] pandas version fix --- aide_predict/bespoke_models/predictors/hmm.py | 2 +- aide_predict/utils/scoring.py | 8 +++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/aide_predict/bespoke_models/predictors/hmm.py b/aide_predict/bespoke_models/predictors/hmm.py index 9de2c9d..20c7e23 100644 --- a/aide_predict/bespoke_models/predictors/hmm.py +++ b/aide_predict/bespoke_models/predictors/hmm.py @@ -174,7 +174,7 @@ def _transform(self, X: ProteinSequences) -> np.ndarray: if not os.path.exists(out_tbl) or os.path.getsize(out_tbl) == 0: raise RuntimeError("HMMsearch failed to produce output") - data = pd.read_csv(out_tbl, delim_whitespace=True, comment='#', header=None) + data = pd.read_csv(out_tbl, sep=r'\s+', comment='#', header=None) scores = np.zeros((len(X), 1)) for i, seq in enumerate(X): diff --git a/aide_predict/utils/scoring.py b/aide_predict/utils/scoring.py index 692b588..5ba83cf 100644 --- a/aide_predict/utils/scoring.py +++ b/aide_predict/utils/scoring.py @@ -80,7 +80,13 @@ def per_variant_mutation_info( row["mut_aa"] = str(seq[p]) row["mutation"] = f"{row['wt_aa']}{row['position']}{row['mut_aa']}" rows.append(row) - return pd.DataFrame(rows) + df = pd.DataFrame(rows) + # pandas can coerce None -> NaN when inferring object columns (notably for a + # mixed str/None column under pandas >= 3.0). Keep the documented contract + # that unset string fields are None, not NaN. + for col in ("variant_id", "mutation", "wt_aa", "mut_aa"): + df[col] = df[col].astype(object).where(df[col].notna(), None) + return df def zscore_by_aa_group( From 2e844b36d6b4c463bf7c5f6304ce476d48538b87 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Thu, 11 Jun 2026 14:53:32 -0600 Subject: [PATCH 10/11] test files --- .coveragerc | 1 + tests/test_bespoke_models/test_base.py | 59 ++++++- .../test_composites/test_zscore.py | 24 +++ .../test_embedders/test_aa_properties.py | 154 ++++++++++++++++++ .../test_data_structures/test_structures.py | 47 +++++- tests/test_utils/test_scoring.py | 28 ++++ 6 files changed, 311 insertions(+), 2 deletions(-) create mode 100644 tests/test_bespoke_models/test_embedders/test_aa_properties.py diff --git a/.coveragerc b/.coveragerc index 9645051..19d9c09 100644 --- a/.coveragerc +++ b/.coveragerc @@ -20,6 +20,7 @@ omit = aide_predict/bespoke_models/predictors/vespa.py aide_predict/bespoke_models/predictors/eve.py aide_predict/bespoke_models/predictors/ssemb.py + aide_predict/bespoke_models/predictors/esm_if.py aide_predict/bespoke_models/embedders/esm2.py aide_predict/bespoke_models/embedders/msa_transformer.py diff --git a/tests/test_bespoke_models/test_base.py b/tests/test_bespoke_models/test_base.py index 154007e..ce7289d 100644 --- a/tests/test_bespoke_models/test_base.py +++ b/tests/test_bespoke_models/test_base.py @@ -972,4 +972,61 @@ def _transform(self, X): result_no_mixin = no_mixin_model.transform(sequences) # Check results - should automatically normalize by WT length (4) - np.testing.assert_array_equal(result_no_mixin, np.array([[0], [1]])) \ No newline at end of file + np.testing.assert_array_equal(result_no_mixin, np.array([[0], [1]])) + +def test_position_specific_gap_handling(tmp_path): + """PositionSpecificMixin auto-strips gaps from aligned input, runs the model + on ungapped sequences, and remaps per-position output back onto alignment + columns (filling gap columns with gap_fill_value).""" + + class GapModel(PositionSpecificMixin, ProteinModelWrapper): + def _fit(self, X, y=None): + self.fitted_ = True + return self + + def _transform(self, X): + # X is ungapped here (the mixin stripped gaps in the pre-hook). + # Return per-sequence (1, L, 2) arrays whose row i is [i, i] so the + # remap can be checked exactly against the original ungapped index. + out = [] + for seq in X: + emb = np.zeros((1, len(seq), 2), dtype=np.float32) + for i in range(len(seq)): + emb[0, i, :] = i + out.append(emb) + return out + + seqs = ProteinSequences([ + ProteinSequence("AC-DE"), # ungapped ACDE; aligned cols 0,1,3,4 + ProteinSequence("A-CDE"), # ungapped ACDE; aligned cols 0,2,3,4 + ]) + + # --- Case A: full alignment width, no pooling, distinct gap fill (-1) --- + model = GapModel(metadata_folder=str(tmp_path), pool=False, flatten=False, + gap_fill_value=-1.0) + model.fit([]) + result = model.transform(seqs) + assert result.shape == (2, 5, 2) # (n_seqs, alignment_width, dim) + np.testing.assert_array_equal( + result[0], [[0, 0], [1, 1], [-1, -1], [2, 2], [3, 3]]) + np.testing.assert_array_equal( + result[1], [[0, 0], [-1, -1], [1, 1], [2, 2], [3, 3]]) + + # --- Case B: pooling across the (remapped) alignment positions --- + model_pool = GapModel(metadata_folder=str(tmp_path), pool=True, + gap_fill_value=-1.0) + model_pool.fit([]) + pooled = model_pool.transform(seqs) + assert pooled.shape == (2, 2) + # mean of [0,1,-1,2,3] == 1.0 (same for both sequences) + np.testing.assert_allclose(pooled, np.ones((2, 2))) + + # --- Case C: specific alignment positions subset --- + model_pos = GapModel(metadata_folder=str(tmp_path), positions=[0, 3], + pool=False, gap_fill_value=-1.0) + model_pos.fit([]) + sub = model_pos.transform(seqs) + assert sub.shape == (2, 2, 2) # (n_seqs, n_positions, dim) + # aligned col 0 -> ungapped idx 0; aligned col 3 -> ungapped idx 2 (both seqs) + np.testing.assert_array_equal(sub[0], [[0, 0], [2, 2]]) + np.testing.assert_array_equal(sub[1], [[0, 0], [2, 2]]) diff --git a/tests/test_bespoke_models/test_composites/test_zscore.py b/tests/test_bespoke_models/test_composites/test_zscore.py index e33b323..04d9ab0 100644 --- a/tests/test_bespoke_models/test_composites/test_zscore.py +++ b/tests/test_bespoke_models/test_composites/test_zscore.py @@ -237,3 +237,27 @@ def test_wt_identical_variant(self, wt): assert df.iloc[0]["n_mutations"] == 0 # z_logratio NaN, logratio = inner's prediction (0.0 for our mock) assert pd.isna(df.iloc[0]["z_logratio"]) + + +class TestInnerOutputValidation: + """_call_inner rejects inner scorers that don't produce one scalar per variant.""" + + def test_non_1d_output_raises(self, wt): + class _ThreeDimScorer(_MockScorer): + def transform(self, X): + return np.ones((len(X), 2, 2)) # not one scalar per variant + + comp = ZScoreRescaledScorer(inner_scorer=_ThreeDimScorer(wt=wt, pool=True), wt=wt) + ssm = wt.saturation_mutagenesis() + with pytest.raises(ValueError, match="one scalar per variant"): + comp.fit(ssm) + + def test_wrong_count_output_raises(self, wt): + class _WrongCountScorer(_MockScorer): + def transform(self, X): + return np.ones(len(X) + 1) # wrong number of scores + + comp = ZScoreRescaledScorer(inner_scorer=_WrongCountScorer(wt=wt, pool=True), wt=wt) + ssm = wt.saturation_mutagenesis() + with pytest.raises(ValueError, match="scores for"): + comp.fit(ssm) diff --git a/tests/test_bespoke_models/test_embedders/test_aa_properties.py b/tests/test_bespoke_models/test_embedders/test_aa_properties.py new file mode 100644 index 0000000..6d2a702 --- /dev/null +++ b/tests/test_bespoke_models/test_embedders/test_aa_properties.py @@ -0,0 +1,154 @@ +# tests/test_bespoke_models/test_embedders/test_aa_properties.py +''' +* Author: Evan Komp +* Created: 2026-06-10 +* Company: National Renewable Energy Lab, Bioeneergy Science and Technology +* License: MIT + +Unit tests for AAPropertiesEmbedding. Depends only on `aaindex` (a base +dependency in environment.yaml), so these run under default CI. +''' +import numpy as np +import pytest + +from aide_predict.utils.data_structures import ProteinSequences, ProteinSequence +from aide_predict.bespoke_models.embedders.aa_properties import ( + AAPropertiesEmbedding, + _build_aa_property_lookup, + DEFAULT_AAINDEX_PROPERTIES, +) + +CANONICAL = "ACDEFGHIKLMNPQRSTVWY" + + +class TestBuildAAPropertyLookup: + def test_lookup_contents(self): + lookup, names = _build_aa_property_lookup(DEFAULT_AAINDEX_PROPERTIES) + # All 20 canonical amino acids present. + assert set(lookup.keys()) == set(CANONICAL) + # 9 default properties + the aromatic boolean. + assert len(names) == len(DEFAULT_AAINDEX_PROPERTIES) + 1 + assert names[-1] == "aromatic" + # Each vector matches the number of property names. + for aa in CANONICAL: + assert lookup[aa].shape == (len(names),) + assert lookup[aa].dtype == np.float32 + # Aromatic flag: F/W/Y -> 1.0, others -> 0.0. + assert lookup["F"][-1] == 1.0 + assert lookup["W"][-1] == 1.0 + assert lookup["Y"][-1] == 1.0 + assert lookup["A"][-1] == 0.0 + + def test_lookup_without_aromatic(self): + lookup, names = _build_aa_property_lookup( + DEFAULT_AAINDEX_PROPERTIES, include_aromatic=False + ) + assert len(names) == len(DEFAULT_AAINDEX_PROPERTIES) + assert "aromatic" not in names + assert lookup["A"].shape == (len(DEFAULT_AAINDEX_PROPERTIES),) + + +class TestAAPropertiesEmbedding: + @pytest.fixture + def sample_sequences(self): + return ProteinSequences([ + ProteinSequence("ACDEF"), + ProteinSequence("GHIKL"), + ProteinSequence("MNPQR"), + ]) + + def test_initialization_defaults(self, tmp_path): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path)) + assert emb.include_aromatic is True + assert emb.aaindex_properties == DEFAULT_AAINDEX_PROPERTIES + assert emb.pool is False + assert emb.flatten is False + assert emb.positions is None + + def test_initialization_custom_properties(self, tmp_path): + custom = [("KYTJ820101", "hydrophobicity")] + emb = AAPropertiesEmbedding( + metadata_folder=str(tmp_path), aaindex_properties=custom, + include_aromatic=False, pool=True, + ) + assert emb.aaindex_properties == custom + assert emb.include_aromatic is False + assert emb.pool is True + + def test_fit_builds_lookup(self, tmp_path, sample_sequences): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path)) + emb.fit(sample_sequences) + assert emb.embedding_dim_ == len(DEFAULT_AAINDEX_PROPERTIES) + 1 + assert len(emb.property_names_) == emb.embedding_dim_ + assert set(emb.aa_property_lookup_.keys()) == set(CANONICAL) + + def test_transform_per_position(self, tmp_path, sample_sequences): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), pool=False) + emb.fit(sample_sequences) + out = emb.transform(sample_sequences) + assert out.shape == (3, 5, len(DEFAULT_AAINDEX_PROPERTIES) + 1) + + def test_transform_pooled(self, tmp_path, sample_sequences): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), pool=True) + emb.fit(sample_sequences) + out = emb.transform(sample_sequences) + assert out.shape == (3, len(DEFAULT_AAINDEX_PROPERTIES) + 1) + + def test_transform_flatten_with_positions(self, tmp_path, sample_sequences): + dim = len(DEFAULT_AAINDEX_PROPERTIES) + 1 + emb = AAPropertiesEmbedding( + metadata_folder=str(tmp_path), positions=[0, 2], flatten=True, pool=False, + ) + emb.fit(sample_sequences) + out = emb.transform(sample_sequences) + assert out.shape == (3, 2 * dim) + + def test_include_aromatic_false_dim(self, tmp_path, sample_sequences): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), include_aromatic=False) + emb.fit(sample_sequences) + assert emb.embedding_dim_ == len(DEFAULT_AAINDEX_PROPERTIES) + + def test_get_feature_names_pooled(self, tmp_path, sample_sequences): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), pool=True) + emb.fit(sample_sequences) + names = emb.get_feature_names_out() + assert len(names) == len(DEFAULT_AAINDEX_PROPERTIES) + 1 + assert all(n.startswith("AAProps_") for n in names) + + def test_get_feature_names_flatten(self, tmp_path, sample_sequences): + dim = len(DEFAULT_AAINDEX_PROPERTIES) + 1 + emb = AAPropertiesEmbedding( + metadata_folder=str(tmp_path), positions=[0, 2], flatten=True, + ) + emb.fit(sample_sequences) + names = emb.get_feature_names_out() + assert len(names) == 2 * dim + assert names[0].startswith("pos0_") + + def test_get_feature_names_before_fit_raises(self, tmp_path): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), pool=True) + with pytest.raises(ValueError, match="fitted"): + emb.get_feature_names_out() + + def test_get_feature_names_non_pool_non_flatten_raises(self, tmp_path, sample_sequences): + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), pool=False, flatten=False) + emb.fit(sample_sequences) + with pytest.raises(ValueError): + emb.get_feature_names_out() + + def test_aligned_gapped_input(self, tmp_path): + # handle_aligned=True (default): gaps are stripped, scored, then remapped + # back to alignment columns with gap_fill_value. + dim = len(DEFAULT_AAINDEX_PROPERTIES) + 1 + emb = AAPropertiesEmbedding(metadata_folder=str(tmp_path), pool=False, gap_fill_value=0.0) + seqs = ProteinSequences([ + ProteinSequence("AC-DE"), + ProteinSequence("A-CDE"), + ]) + emb.fit(seqs) + out = emb.transform(seqs) + # Output spans the full alignment width (5). + assert out.shape == (2, 5, dim) + # The gap column for each sequence is filled with gap_fill_value (0.0). + np.testing.assert_array_equal(out[0, 2], np.zeros(dim, dtype=np.float32)) + np.testing.assert_array_equal(out[1, 1], np.zeros(dim, dtype=np.float32)) diff --git a/tests/test_utils/test_data_structures/test_structures.py b/tests/test_utils/test_data_structures/test_structures.py index bd39a00..dd1ff41 100644 --- a/tests/test_utils/test_data_structures/test_structures.py +++ b/tests/test_utils/test_data_structures/test_structures.py @@ -235,6 +235,42 @@ def test_set_target_chain_invalid(self, temp_multichain_pdb): with pytest.raises(ValueError, match="not found"): s.set_target_chain('Z') + @pytest.fixture + def temp_pdb_missing_atom(self): + # Chain A residue 1 (ALA) is missing its CA atom -> the CA row of the + # backbone tensor must be NaN. + with tempfile.NamedTemporaryFile(mode='w', suffix='.pdb', delete=False) as f: + f.write("""ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N +ATOM 2 C ALA A 1 2.009 1.362 0.000 1.00 0.00 C +ATOM 3 O ALA A 1 1.702 2.144 0.907 1.00 0.00 O +ATOM 4 N GLY A 2 2.831 1.687 -0.987 1.00 0.00 N +ATOM 5 CA GLY A 2 3.396 3.037 -1.009 1.00 0.00 C +ATOM 6 C GLY A 2 2.362 4.089 -1.408 1.00 0.00 C +ATOM 7 O GLY A 2 2.730 5.261 -1.509 1.00 0.00 O +TER 8 GLY A 2 +END +""") + yield f.name + os.unlink(f.name) + + def test_get_chain_coords_nan_for_missing_atom(self, temp_pdb_missing_atom): + s = ProteinStructure(temp_pdb_missing_atom, chain='A') + coords = s.get_chain_coords('A') + assert coords.shape == (2, 3, 3) + # Residue 1 missing CA -> CA row (index 1) is all NaN; N and C present. + assert np.all(np.isnan(coords[0, 1])) + assert not np.any(np.isnan(coords[0, 0])) # N present + assert not np.any(np.isnan(coords[0, 2])) # C present + # Residue 2 (GLY) is complete. + assert not np.any(np.isnan(coords[1])) + + def test_get_chain_coords_empty_for_non_protein_chain(self, temp_pdb_with_ligand_chain): + # Chain Z holds only HETATM records -> no canonical residues -> empty tensor. + s = ProteinStructure(temp_pdb_with_ligand_chain, chain='A', context_chains=('Z',)) + coords = s.get_chain_coords('Z') + assert coords.shape == (0, 3, 3) + assert coords.dtype == np.float32 + class TestStructureMapper: @pytest.fixture @@ -358,4 +394,13 @@ def test_get_protein_sequences_json_map(self, temp_multichain_folder, tmp_path): def test_get_protein_sequences_no_auto_context(self, temp_multichain_folder): mapper = StructureMapper(temp_multichain_folder) seqs = mapper.get_protein_sequences(target_chain='A', auto_context=False) - assert seqs[0].structure.context_chains is None \ No newline at end of file + assert seqs[0].structure.context_chains is None + + def test_get_protein_sequences_json_map_fallback(self, temp_multichain_folder, tmp_path): + # The JSON map doesn't list "complex", so it falls back to chain 'A'. + mapper = StructureMapper(temp_multichain_folder) + json_path = tmp_path / "chain_map.json" + json_path.write_text(json.dumps({"some_other_structure": "B"})) + seqs = mapper.get_protein_sequences(target_chain=str(json_path)) + assert seqs[0].structure.chain == 'A' + assert str(seqs[0]) == 'AG' \ No newline at end of file diff --git a/tests/test_utils/test_scoring.py b/tests/test_utils/test_scoring.py index 220a57e..399d4dd 100644 --- a/tests/test_utils/test_scoring.py +++ b/tests/test_utils/test_scoring.py @@ -157,3 +157,31 @@ def test_nan_mut_aa_rows_left_nan(self): assert pd.isna(df_out.iloc[5]["z_score"]) # The first 5 rows ARE z-scored (P group has 5 samples). assert df_out.iloc[:5]["z_score"].notna().all() + + def test_integer_score_column(self): + # An integer-dtype score column exercises the non-float branch of the + # NaN filter (vals.dtype.kind != "f"). + df = pd.DataFrame([{"score": s, "wt_aa": "A", "mut_aa": "P"} for s in [1, 2, 3, 4, 5]]) + assert df["score"].dtype.kind == "i" + df_out, stats = zscore_by_aa_group(df, score_col="score") + assert ("P",) in stats + # Group centered: mean of z-scores ~ 0. + np.testing.assert_allclose(df_out["z_score"].mean(), 0.0, atol=1e-9) + + def test_zero_variance_group_skipped(self): + # All-identical scores -> std == 0 -> group is skipped, z stays NaN. + df = pd.DataFrame([{"score": 7.0, "wt_aa": "A", "mut_aa": "P"} for _ in range(5)]) + df_out, stats = zscore_by_aa_group(df, score_col="score") + assert ("P",) not in stats + assert df_out["z_score"].isna().all() + + def test_nan_score_with_valid_key_skipped(self): + # A row with a valid group key but a NaN score is skipped at apply time + # (its z stays NaN) while the rest of the group is z-scored normally. + rows = [{"score": float(s), "wt_aa": "A", "mut_aa": "P"} for s in [1, 2, 3, 4, 5]] + rows.append({"score": np.nan, "wt_aa": "A", "mut_aa": "P"}) + df = pd.DataFrame(rows) + df_out, stats = zscore_by_aa_group(df, score_col="score") + assert ("P",) in stats + assert pd.isna(df_out.iloc[5]["z_score"]) + assert df_out.iloc[:5]["z_score"].notna().all() From 5190cad422c054a7ad3b9133d99085c794996bc1 Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Tue, 16 Jun 2026 14:03:21 -0600 Subject: [PATCH 11/11] tests --- tests/test_bespoke_models/test_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_bespoke_models/test_base.py b/tests/test_bespoke_models/test_base.py index ce7289d..61e7b96 100644 --- a/tests/test_bespoke_models/test_base.py +++ b/tests/test_bespoke_models/test_base.py @@ -1023,7 +1023,7 @@ def _transform(self, X): # --- Case C: specific alignment positions subset --- model_pos = GapModel(metadata_folder=str(tmp_path), positions=[0, 3], - pool=False, gap_fill_value=-1.0) + pool=False, flatten=False, gap_fill_value=-1.0) model_pos.fit([]) sub = model_pos.transform(seqs) assert sub.shape == (2, 2, 2) # (n_seqs, n_positions, dim)