Dev

.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

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
@@ -90,6 +91,19 @@ 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).
+
+### 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

aide_predict/bespoke_models/__init__.py

@@ -13,13 +13,16 @@
 from .predictors.vespa import VESPAWrapper
 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
 from .embedders.msa_transformer import MSATransformerEmbedding
 from .embedders.saprot import SaProtEmbedding
 from .embedders.kmer import KmerEmbedding
 from .embedders.ssemb import SSEmbEmbedding
+from .embedders.aa_properties import AAPropertiesEmbedding
 
 
 TOOLS = [
@@ -31,6 +34,10 @@
     VESPAWrapper,
     EVEWrapper,
     SSEmbWrapper,
+    ESMIFLikelihoodWrapper,
+
+    # composites
+    ZScoreRescaledScorer,
 
     # embedders
     ESM2Embedding,
@@ -40,4 +47,5 @@
     SaProtEmbedding,
     KmerEmbedding,
     SSEmbEmbedding,
+    AAPropertiesEmbedding,
 ]

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,22 @@ 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
+        self._alignment_width = None
 
     def _is_ragged_array(self, arr):
         """Check if the input is a ragged array (list of arrays with different shapes)."""
@@ -934,9 +943,67 @@ 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 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()
+        else:
+            self._alignment_mapping = None
+            self._alignment_width = 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 +1017,18 @@ def _post_transform_hook(self, result, X):
         """
         if result is None or len(result) == 0:
             return result
+
+        # 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, 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
             if isinstance(self.pool, str):
@@ -961,8 +1040,16 @@ 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:
-                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]

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"]