sabr-kit 0.4.2__tar.gz → 0.4.3.dev1__tar.gz

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: sabr-kit
-Version: 0.4.2
+Version: 0.4.3.dev1
 Summary: Structure-based Antibody Renumbering
 Project-URL: Homepage, https://github.com/delalamo/sabr
 Project-URL: Issues, https://github.com/delalamo/sabr/issues
@@ -174,18 +174,22 @@ the CLI with mmCIF output.
 - In `sabr` mode, gap extension is `-0.175027` and gap opening is `-2.525591`.
   In `softalign` mode, they are `0.1942468136548996` and
   `-2.5441808700561523`, respectively, as stored in the repository asset.
-- CDR gap distribution is always applied.
-- No deterministic light-chain DE-loop or C-terminal correction is applied.
+- Deterministic CDR gap distribution and DE-loop correction are always
+  applied. Between IMGT anchors 79 and 85, DE-loop residues fill 80 first,
+  then 84 back through 81; additional residues are inserted after 82.
+- No deterministic C-terminal correction is applied.
 - Automatic chain selection aligns against H, K, and L references and uses
   the highest score, with deterministic H/K/L tie order.
 - scFv mode appends H:K, H:L, K:H, and L:H reference candidates in that order.
+- Composite candidates do not apply gap-open or gap-extension costs to query
+  linker residues aligned at the boundary between their two references.
 - Composite candidates receive normal affine gap-open and gap-extension costs
   for unaligned query and reference termini when their selection scores are
   compared; the underlying alignments and raw alignment scores are unchanged.
 
 A structural gap is detected when the C–N distance between consecutive
-residues exceeds 2.66 Å. A gap skips only the affected CDR correction and
-emits a warning; other regions continue normally.
+residues exceeds 2.66 Å. A gap skips only the affected CDR or DE-loop
+correction and emits a warning; other regions continue normally.
 
 T-cell receptors are not an officially supported SAbR target. For
 experimental low-level use, align a TCR against the K reference because that

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/README.md

@@ -130,18 +130,22 @@ the CLI with mmCIF output.
 - In `sabr` mode, gap extension is `-0.175027` and gap opening is `-2.525591`.
   In `softalign` mode, they are `0.1942468136548996` and
   `-2.5441808700561523`, respectively, as stored in the repository asset.
-- CDR gap distribution is always applied.
-- No deterministic light-chain DE-loop or C-terminal correction is applied.
+- Deterministic CDR gap distribution and DE-loop correction are always
+  applied. Between IMGT anchors 79 and 85, DE-loop residues fill 80 first,
+  then 84 back through 81; additional residues are inserted after 82.
+- No deterministic C-terminal correction is applied.
 - Automatic chain selection aligns against H, K, and L references and uses
   the highest score, with deterministic H/K/L tie order.
 - scFv mode appends H:K, H:L, K:H, and L:H reference candidates in that order.
+- Composite candidates do not apply gap-open or gap-extension costs to query
+  linker residues aligned at the boundary between their two references.
 - Composite candidates receive normal affine gap-open and gap-extension costs
   for unaligned query and reference termini when their selection scores are
   compared; the underlying alignments and raw alignment scores are unchanged.
 
 A structural gap is detected when the C–N distance between consecutive
-residues exceeds 2.66 Å. A gap skips only the affected CDR correction and
-emits a warning; other regions continue normally.
+residues exceeds 2.66 Å. A gap skips only the affected CDR or DE-loop
+correction and emits a warning; other regions continue normally.
 
 T-cell receptors are not an officially supported SAbR target. For
 experimental low-level use, align a TCR against the K reference because that

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/src/sabr/alignment.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 """Reference alignment using the validated affine Smith-Waterman method."""
 
 import functools
@@ -15,16 +14,18 @@ from sabr.corrections import apply_corrections
 LOGGER = logging.getLogger(__name__)
 
 
-def _rotate_for_dp(x, NINF):
+def _rotate_for_dp(x, NINF, free_gap_boundary=-1):
     """Rotate a matrix for striped dynamic programming."""
     a, b = x.shape
     ar = jnp.arange(a)[::-1, None]
     br = jnp.arange(b)[None, :]
     i, j = (br - ar) + (a - 1), (ar + br) // 2
     n, m = (a + b - 1), (a + b) // 2
+    free_gap = jnp.broadcast_to(br == free_gap_boundary, (a, b))
     output = {
         "x": jnp.full([n, m], NINF).at[i, j].set(x),
         "o": (jnp.arange(n) + a % 2) % 2,
+        "free_gap": jnp.full([n, m], False).at[i, j].set(free_gap),
     }
     prev = (jnp.full((m, 3), NINF), jnp.full((m, 3), NINF))
     return output, prev, (i, j)
@@ -71,9 +72,10 @@ def _affine_score(
     temperature,
     gap_extend,
     gap_open,
+    free_gap_boundary=-1,
     NINF=-1e30,
 ):
-    """Return the smooth affine Smith-Waterman score for one matrix."""
+    """Return a score, optionally making one reference gap boundary free."""
     right_penalties = jnp.asarray(
         [
             gap_open,
@@ -105,7 +107,11 @@ def _affine_score(
             _pad(h1[1:], ([0, 1], [0, 0]), NINF),
         )
         right += right_penalties
-        down += down_penalties
+        down += jnp.where(
+            stripe["free_gap"][:, None],
+            0,
+            down_penalties,
+        )
         right = right[:, :2]
 
         h0 = jnp.stack(
@@ -119,7 +125,9 @@ def _affine_score(
         return (h1, h0), h0
 
     similarities, mask = _apply_length_mask(similarities, lengths, NINF)
-    stripes, previous, indices = _rotate_for_dp(similarities[:-1, :-1], NINF)
+    stripes, previous, indices = _rotate_for_dp(
+        similarities[:-1, :-1], NINF, free_gap_boundary
+    )
     scores = jax.lax.scan(step, previous, stripes, unroll=2)[-1][indices]
     return _soft_maximum(
         scores + similarities[1:, 1:, None],
@@ -130,7 +138,10 @@ def _affine_score(
 
 
 _AFFINE_ALIGNMENT = jax.jit(
-    jax.vmap(jax.value_and_grad(_affine_score), (0, 0, None, None, None))
+    jax.vmap(
+        jax.value_and_grad(_affine_score),
+        (0, 0, None, None, None, None),
+    )
 )
 
 
@@ -218,48 +229,6 @@ def _alignment_path(alignment: np.ndarray) -> np.ndarray:
     return path
 
 
-def _validate_alignment(alignment: np.ndarray, chain_type: str) -> None:
-    """Reject ambiguous paths before they are converted to numbering states."""
-    path = _alignment_path(alignment)
-    rows = path[:, 0]
-    columns = path[:, 1]
-
-    first_row = int(rows[0])
-    first_position = int(columns[0]) + 1
-    if first_row and first_position - first_row < 1:
-        raise ValueError(
-            "N-terminal residues would require non-positive numbering; "
-            "use residue_range to select the antibody domain."
-        )
-
-    regions = list(constants.IMGT_LOOPS.values())
-    if chain_type in ("K", "L"):
-        regions.append((79, 84))
-    for index, (left_row, right_row) in enumerate(zip(rows, rows[1:])):
-        if right_row == left_row + 1:
-            continue
-        left_position = int(columns[index]) + 1
-        right_position = int(columns[index + 1]) + 1
-        if not any(
-            start <= left_position <= end and start <= right_position <= end
-            for start, end in regions
-        ):
-            raise ValueError(
-                f"Unassigned query rows {left_row + 1}-{right_row - 1} "
-                f"are bracketed by IMGT {left_position} and "
-                f"{right_position}; use residue_range to select one "
-                "antibody domain."
-            )
-
-    last_row = int(rows[-1])
-    last_position = int(columns[-1]) + 1
-    if last_row < alignment.shape[0] - 1 and last_position < 125:
-        raise ValueError(
-            f"Unassigned trailing query rows follow IMGT {last_position}; "
-            "use residue_range to select the antibody domain."
-        )
-
-
 def _validate_scfv_alignment(
     alignment: np.ndarray, representation: str
 ) -> None:
@@ -271,7 +240,6 @@ def _validate_scfv_alignment(
         )
     _alignment_path(alignment)
 
-    domain_rows = []
     for domain_index, chain_type in enumerate(representation.split(":")):
         start = domain_index * constants.IMGT_MAX_POSITION
         end = start + constants.IMGT_MAX_POSITION
@@ -281,21 +249,15 @@ def _validate_scfv_alignment(
             raise ValueError(
                 f"scFv {chain_type} domain contains no assigned residues."
             )
-        domain_rows.append((int(assigned_rows[0]), int(assigned_rows[-1])))
-        if domain_index == 0:
-            _validate_alignment(domain[: assigned_rows[-1] + 1], chain_type)
-        else:
-            _validate_alignment(domain[assigned_rows[0] :], chain_type)
-
-    if domain_rows[0][1] >= domain_rows[1][0]:
-        raise ValueError("scFv domain assignments overlap or are out of order.")
 
 
 def _align_reference(
     query: np.ndarray,
     reference: np.ndarray,
     mode: str = "sabr",
+    free_gap_boundary: int = -1,
 ):
+    """Align one reference, optionally allowing a free query insertion."""
     anchor = np.zeros((1, reference.shape[1]), dtype=reference.dtype)
     augmented_reference = np.concatenate((anchor, reference, anchor), axis=0)
     query_batch = jnp.asarray(query[None, :])
@@ -309,6 +271,7 @@ def _align_reference(
         constants.DEFAULT_TEMPERATURE,
         gap_extend,
         gap_open,
+        free_gap_boundary,
     )
     return (
         np.asarray(soft_alignment[0])[:, 1:-1],
@@ -367,7 +330,21 @@ def align(
     best = None
     for candidate in candidates:
         reference, positions = references[candidate]
-        reduced, similarity, score = _align_reference(query, reference, mode)
+        if ":" in candidate:
+            free_gap_boundary = sum(
+                position <= constants.IMGT_MAX_POSITION
+                for position in positions
+            )
+            reduced, similarity, score = _align_reference(
+                query,
+                reference,
+                mode,
+                free_gap_boundary=free_gap_boundary,
+            )
+        else:
+            reduced, similarity, score = _align_reference(
+                query, reference, mode
+            )
         if (
             not np.isfinite(score)
             or not np.isfinite(reduced).all()
@@ -419,5 +396,5 @@ def align(
         _validate_scfv_alignment(corrected, selected_type)
     else:
         corrected = apply_corrections(full_alignment, gap_indices=gap_indices)
-        _validate_alignment(corrected, selected_type)
+        _alignment_path(corrected)
     return corrected, selected_type, score

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/src/sabr/constants.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 """Constants and configuration values for SAbR.
 
 This module defines constants used throughout the SAbR package including:
@@ -25,12 +24,6 @@ PEPTIDE_BOND_LENGTH = 1.33
 PEPTIDE_BOND_MAX_DISTANCE = 2 * PEPTIDE_BOND_LENGTH
 MAX_SELECTED_RESIDUES = 1024
 
-# Backbone indices for MPNN
-BACKBONE_N_IDX = 0
-BACKBONE_CA_IDX = 1
-BACKBONE_C_IDX = 2
-BACKBONE_CB_IDX = 3
-
 # Gap scores (from NPZ)
 SW_GAP_EXTEND = -0.175027
 SW_GAP_OPEN = -2.525591

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/src/sabr/corrections.py

@@ -1,4 +1,4 @@
-"""Deterministic CDR corrections."""
+"""Deterministic CDR and DE-loop corrections."""
 
 import logging
 import warnings
@@ -10,12 +10,6 @@ from sabr import constants
 LOGGER = logging.getLogger(__name__)
 
 
-def _has_gap_in_region(
-    gap_indices: frozenset[int], start_row: int, end_row: int
-) -> bool:
-    return any(index in gap_indices for index in range(start_row, end_row))
-
-
 def _aligned_row_near(aln: np.ndarray, target_col: int) -> int | None:
     """Return the row aligned at or within two columns of a target."""
     for offset in (0, -1, 1, -2, 2):
@@ -49,7 +43,9 @@ def _skip_for_structural_gap(
     region_name: str,
 ) -> bool:
     """Warn and return true when a regional correction crosses a gap."""
-    if gap_indices and _has_gap_in_region(gap_indices, start_row, end_row):
+    if gap_indices and any(
+        index in gap_indices for index in range(start_row, end_row)
+    ):
         message = (
             f"Skipping {region_name} deterministic correction: structural "
             f"gap detected between rows {start_row} and {end_row}; using "
@@ -143,14 +139,78 @@ def correct_cdr_loop(
     return aln
 
 
+def de_loop_positions(n_residues: int) -> list[int]:
+    """Return IMGT positions for residues between anchors 79 and 85.
+
+    Position 80 is filled first, followed by positions 84 through 81 from
+    right to left. Once positions 80 through 84 are occupied, additional
+    residues become insertions on position 82.
+    """
+    if n_residues <= 0:
+        return []
+    if n_residues <= 5:
+        return [80, *range(86 - n_residues, 85)]
+    return [80, 81, 82, *([82] * (n_residues - 5)), 83, 84]
+
+
+def correct_de_loop(
+    aln: np.ndarray,
+    gap_indices: frozenset[int] | None = None,
+) -> np.ndarray:
+    """Assign the residues between IMGT 79 and 85 by loop length."""
+    anchor_79_row = _aligned_row_near(aln, 78)
+    anchor_85_row = _aligned_row_near(aln, 84)
+
+    if anchor_79_row is None or anchor_85_row is None:
+        LOGGER.warning(
+            "Skipping DE loop correction; missing anchor near IMGT 79 or 85."
+        )
+        return aln
+    if anchor_79_row >= anchor_85_row:
+        LOGGER.warning(
+            "Skipping DE loop correction; anchor 79 row (%d) is not before "
+            "anchor 85 row (%d).",
+            anchor_79_row,
+            anchor_85_row,
+        )
+        return aln
+    if _skip_for_structural_gap(
+        gap_indices, anchor_79_row, anchor_85_row, "DE loop (79-85)"
+    ):
+        return aln
+
+    intermediate_rows = list(range(anchor_79_row + 1, anchor_85_row))
+    positions = de_loop_positions(len(intermediate_rows))
+
+    # Clear the learned assignments for this region before rebuilding it.
+    # Additional occurrences of 82 remain unassigned here; alignment_to_states
+    # converts those orphan rows to 82A, 82B, and later insertion states.
+    aln[intermediate_rows, :] = 0
+    assigned_positions = set()
+    for row, position in zip(intermediate_rows, positions):
+        if position in assigned_positions:
+            continue
+        aln[row, position - 1] = 1
+        assigned_positions.add(position)
+
+    if intermediate_rows:
+        LOGGER.info(
+            "DE loop correction: assigned %d residue(s) to %s.",
+            len(intermediate_rows),
+            positions,
+        )
+    return aln
+
+
 def apply_corrections(
     aln: np.ndarray,
     gap_indices: frozenset[int] | None = None,
 ) -> np.ndarray:
-    """Apply deterministic CDR corrections."""
+    """Apply deterministic CDR and DE-loop corrections."""
     for loop_name, (cdr_start, cdr_end) in constants.IMGT_LOOPS.items():
-        correct_cdr_loop(
+        aln = correct_cdr_loop(
             aln, loop_name, cdr_start, cdr_end, gap_indices=gap_indices
         )
 
+    correct_de_loop(aln, gap_indices=gap_indices)
     return aln

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/src/sabr/model.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 """MPNN (Message Passing Neural Network) encoder for protein structures.
 
 This module provides the ENC class which encodes protein backbone structures
@@ -130,20 +129,15 @@ class ProteinFeatures(hk.Module):
     def __init__(
         self,
         edge_features: int,
-        node_features: int,
         num_positional_embeddings: int = 16,
         num_rbf: int = 16,
         top_k: int = 30,
         augment_eps: float = 0.0,
-        num_chain_embeddings: int = 16,
     ):
         super(ProteinFeatures, self).__init__()
-        self.edge_features = edge_features
-        self.node_features = node_features
         self.top_k = top_k
         self.augment_eps = augment_eps
         self.num_rbf = num_rbf
-        self.num_positional_embeddings = num_positional_embeddings
 
         self.embeddings = PositionalEncodings(num_positional_embeddings)
         # edge_in = num_positional_embeddings + num_rbf * 25 (for reference)
@@ -270,7 +264,7 @@ class PositionWiseFeedForward(hk.Module):
         self.act = Gelu
 
     def __call__(self, h_V):
-        h = self.act(self.W_in(h_V), approximate=False)
+        h = self.act(self.W_in(h_V))
         h = self.W_out(h)
         return h
 
@@ -294,19 +288,14 @@ class EncLayer(hk.Module):
     def __init__(
         self,
         num_hidden: int,
-        num_in: int,
         dropout: float = 0.1,
-        num_heads: int = None,
         scale: int = 30,
         name: str = None,
     ):
         super(EncLayer, self).__init__()
         self.num_hidden = num_hidden
-        self.num_in = num_in
         self.scale = scale
 
-        self.safe_key = SafeKey(hk.next_rng_key())
-
         self.dropout1 = DropoutCust(dropout)
         self.dropout2 = DropoutCust(dropout)
         self.dropout3 = DropoutCust(dropout)
@@ -387,7 +376,6 @@ class ENC:
 
     def __init__(
         self,
-        node_features: int,
         edge_features: int,
         hidden_dim: int,
         num_encoder_layers: int = 1,
@@ -398,7 +386,6 @@ class ENC:
         """Initialize the MPNN encoder.
 
         Args:
-            node_features: Dimension of node features.
             edge_features: Dimension of edge features.
             hidden_dim: Hidden dimension for layers.
             num_encoder_layers: Number of encoder layers.
@@ -407,12 +394,8 @@ class ENC:
             dropout: Dropout rate.
         """
         super(ENC, self).__init__()
-        self.node_features = node_features
-        self.edge_features = edge_features
-        self.hidden_dim = hidden_dim
 
         self.features = ProteinFeatures(
-            node_features,
             edge_features,
             top_k=k_neighbors,
             augment_eps=augment_eps,
@@ -420,9 +403,7 @@ class ENC:
 
         self.W_e = hk.Linear(hidden_dim, with_bias=True, name="W_e")
         self.encoder_layers = [
-            EncLayer(
-                hidden_dim, hidden_dim * 2, dropout=dropout, name="enc" + str(i)
-            )
+            EncLayer(hidden_dim, dropout=dropout, name="enc" + str(i))
             for i in range(num_encoder_layers)
         ]
 
@@ -481,7 +462,6 @@ def load_parameters(mode: str = "sabr") -> dict:
 
 def _encode(coords, mask, chain_ids, residue_indices):
     encoder = ENC(
-        constants.EMBED_DIM,
         constants.EMBED_DIM,
         constants.EMBED_DIM,
         constants.N_MPNN_LAYERS,

{sabr_kit-0.4.2 → sabr_kit-0.4.3.dev1}/src/sabr/structure.py

@@ -352,11 +352,6 @@ def _new_residue_ids(data: _ChainData, numbered: list) -> dict:
     mapping = {}
     first_row, first_number = numbered[0][:2]
     first_assigned_number = first_number - first_row
-    if first_assigned_number < 1:
-        raise ValueError(
-            "N-terminal residues would require non-positive numbering; "
-            "use residue_range to select the antibody domain."
-        )
     for query_index in range(first_row):
         mapping[data.residue_indices[query_index]] = (
             first_assigned_number + query_index,