pCoMole / cas9 /constraints.py
Maximilian Holsman
Claude Opus 5
Add Cas9 task
12fea4a
Raw History Blame Contribute Delete
27.3 kB
# constraints.py
import os
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple
import torch
from rdkit import Chem
from rdkit.Chem import Descriptors, rdMolDescriptors
# Make selfies optional (only needed for SMILES constraints, not protein constraints)
try:
import selfies as sf
except ImportError:
sf = None
# from .is_peptidomimetic import is_peptidomimetic_not_natural
import pdb
# ---- Optional import: Cas9 domain detector (HNH + RuvC) ----
# Expects domain_detector.py to live in the same package as this constraints.py
try:
from .domain_detector import DomainDetector, DomainDetectorConfig, HitCriteria
except Exception as e:
import traceback
print(f"Warning: Failed to import DomainDetector: {e}")
traceback.print_exc()
DomainDetector = None
DomainDetectorConfig = None
HitCriteria = None
def detokenize_output(x, cfg, tokenizer, bos_id, eos_id, pad_id):
"""
Convert a single generated sequence (1, L) back to string.
"""
seq = x[0].tolist()
# strip padding
seq = [tok for tok in seq if tok != pad_id]
# strip BOS/EOS
if len(seq) > 0 and seq[0] == bos_id:
seq = seq[1:]
if len(seq) > 0 and seq[-1] == eos_id:
seq = seq[:-1]
if cfg.task == 'protein':
# esm tokenizer has batch_decode
return tokenizer.batch_decode([seq], skip_special_tokens=True)[0]
elif cfg.task in ('smiles', 'selfies'):
return tokenizer.decode(seq)
else:
return " ".join(map(str, seq))
class Peptidomimetic:
def __init__(self):
pass
def __call__(self, smiles_tokens, smiles_seq):
flag, audit = is_peptidomimetic_not_natural(smiles_seq)
return 1 if flag else 0
# SMARTS
_AMIDE = Chem.MolFromSmarts("[CX3](=[OX1])[NX3]") # amide C(=O)-N (includes carbamates/ureas as centers)
_SULFONAMIDE = Chem.MolFromSmarts("[SX4](=[OX1])(=[OX1])[NX3]") # sulfonamide S(=O)2-N (optional isostere)
def _centers_and_N(mol, include_sulfonamide=True):
centers, n2c = set(), {}
for c_idx, _, n_idx in mol.GetSubstructMatches(_AMIDE):
centers.add(c_idx)
n2c.setdefault(n_idx, set()).add(c_idx)
if include_sulfonamide:
for s_idx, _, _, n_idx in mol.GetSubstructMatches(_SULFONAMIDE):
centers.add(s_idx)
n2c.setdefault(n_idx, set()).add(s_idx)
return centers, n2c
def _is_sp3_carbon(atom):
return atom.GetAtomicNum() == 6 and atom.GetHybridization() == Chem.HybridizationType.SP3
def _backbone_edges(mol, centers, n2c):
"""
Edge between centers i--j exists if there is a single-bond path:
center_i — N — C(sp3) — center_j (length 3)
or center_i — N — (O|S|N) — C(sp3) — center_j (length 4)
"""
edges = set()
for n_idx, left_centers in n2c.items():
n = mol.GetAtomWithIdx(n_idx)
for b1 in n.GetBonds():
if b1.GetBondType() != Chem.BondType.SINGLE:
continue
a1 = b1.GetOtherAtom(n)
# case A: N—C(sp3)
if _is_sp3_carbon(a1):
for b2 in a1.GetBonds():
if b2.GetBondType() != Chem.BondType.SINGLE:
continue
other = b2.GetOtherAtom(a1)
j = other.GetIdx()
if j in centers:
for i in left_centers:
if i != j:
edges.add(tuple(sorted((i, j))))
# case B: N—(O|S|N)—C(sp3)
if a1.GetAtomicNum() in (8, 16, 7): # O, S, N
for b2 in a1.GetBonds():
if b2.GetBondType() != Chem.BondType.SINGLE:
continue
a2 = b2.GetOtherAtom(a1)
if _is_sp3_carbon(a2):
for b3 in a2.GetBonds():
if b3.GetBondType() != Chem.BondType.SINGLE:
continue
other = b3.GetOtherAtom(a2)
j = other.GetIdx()
if j in centers:
for i in left_centers:
if i != j:
edges.add(tuple(sorted((i, j))))
return edges
def _largest_component(centers, edges):
if not centers:
return set()
adj = {i: set() for i in centers}
for i, j in edges:
if i in adj and j in adj:
adj[i].add(j)
adj[j].add(i)
seen, best = set(), set()
for s in adj:
if s in seen:
continue
stack, comp = [s], set()
while stack:
u = stack.pop()
if u in comp:
continue
comp.add(u)
seen.add(u)
stack.extend(v for v in adj[u] if v not in comp)
if len(comp) > len(best):
best = comp
return best
def backbone_metrics(smiles: str, include_sulfonamide=True):
mol = Chem.MolFromSmiles(smiles)
if mol is None:
raise ValueError(f"Invalid SMILES: {smiles}")
centers, n2c = _centers_and_N(mol, include_sulfonamide)
edges = _backbone_edges(mol, centers, n2c)
main_nodes = _largest_component(centers, edges)
main_edges = {e for e in edges if e[0] in main_nodes and e[1] in main_nodes}
return {
"backbone_centers_main": len(main_nodes),
"backbone_links_main": len(main_edges),
"MW": Descriptors.MolWt(mol),
"HAC": mol.GetNumAtoms(),
"TPSA": rdMolDescriptors.CalcTPSA(mol),
"rotatable_bonds": rdMolDescriptors.CalcNumRotatableBonds(mol),
}
def compare_shorter_smaller(orig_smiles: str, cand_smiles: str, include_sulfonamide=True,
mw_margin: float = 0.0, hac_margin: int = 0):
o = backbone_metrics(orig_smiles, include_sulfonamide)
c = backbone_metrics(cand_smiles, include_sulfonamide)
shorter_backbone = (c["backbone_centers_main"] < o["backbone_centers_main"]) or \
(c["backbone_links_main"] < o["backbone_links_main"])
smaller_bulk = (c["MW"] <= o["MW"] - mw_margin) and (c["HAC"] <= o["HAC"] - hac_margin)
return {"original": o, "candidate": c, "decisions": {
"shorter_backbone": shorter_backbone,
"smaller_bulk": smaller_bulk,
"overall_shorter_or_smaller": (shorter_backbone or smaller_bulk),
"require_both_shorter_and_smaller": (shorter_backbone and smaller_bulk),
}}
class Length:
def __init__(self, orig_smiles, cfg, smiles_tokenizer):
self.cfg = cfg
self.smiles_tokenizer = smiles_tokenizer
self.orig_smiles = sf.decoder(sf.encoder(orig_smiles))
self.L0 = smiles_tokenizer(orig_smiles, return_tensors='pt')['input_ids'].shape[1]
print("Initial Length: ", self.L0)
def __call__(self, smiles_tokens, smiles_seq):
res = compare_shorter_smaller(self.orig_smiles, smiles_seq)
return int(res["decisions"]["overall_shorter_or_smaller"] and smiles_tokens.shape[1] >= self.L0 // 2)
# return int(smiles_tokens.shape[1] <= self.L0 and smiles_tokens.shape[1] >= self.L0 // 2)
class ProteinLength:
"""
Length constraint for protein sequences.
Identical logic to the SMILES Length constraint, but using protein sequence length.
By default requires candidate to be strictly shorter than original AND at least half
the original length. If allow_equal_length is True, candidates equal to the original
length also pass the upper bound (still must be at least L0 // 2).
"""
def __init__(self, orig_protein_seq, cfg, tokenizer, allow_equal_length: bool = False):
self.cfg = cfg
self.tokenizer = tokenizer
self.allow_equal_length = bool(allow_equal_length)
self.orig_protein_seq = orig_protein_seq.strip().replace(" ", "").replace("\n", "")
# Store length in amino acids (equivalent to tokenized length for proteins)
self.L0 = len(self.orig_protein_seq)
print("Initial Length: ", self.L0)
if self.allow_equal_length:
print("ProteinLength: allowing equal length (candidates may have length == L0)")
def __call__(self, protein_tokens, protein_seq):
"""
Check protein sequence length against original.
- Candidate must be at least half the original length (L0 // 2).
- Upper bound: strictly shorter than L0 by default; if allow_equal_length, L_cand <= L0.
Args:
protein_tokens: Unused (kept for interface compatibility, could use for tokenized length)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if constraints pass, 0 otherwise
"""
seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
L_cand = len(seq_clean)
if self.allow_equal_length:
is_within_upper = L_cand <= self.L0
else:
is_within_upper = L_cand < self.L0
is_at_least_half = L_cand >= self.L0 // 2
return int(is_within_upper and is_at_least_half)
class TargetLength:
"""
Target length constraint for protein sequences.
Requires candidate sequence length to exactly match the target length.
"""
def __init__(self, target_length: int):
"""
Args:
target_length: The exact target length in amino acids
"""
if target_length <= 0:
raise ValueError(f"target_length must be positive, got {target_length}")
self.target_length = int(target_length)
print(f"Target Length: {self.target_length}")
def __call__(self, protein_tokens, protein_seq):
"""
Check if protein sequence length exactly matches the target length.
Args:
protein_tokens: Unused (kept for interface compatibility)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if length matches target, 0 otherwise
"""
seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
L_cand = len(seq_clean)
return int(L_cand == self.target_length)
def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
Args:
seqs: List of protein sequence strings
Returns:
List of bool: True if length matches target, False otherwise
"""
clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
return [len(s) == self.target_length for s in clean]
class MinTargetLength:
"""
Minimum target length constraint for protein sequences.
Requires candidate sequence length to be greater than or equal to the minimum target length.
"""
def __init__(self, min_target_length: int):
"""
Args:
min_target_length: The minimum allowed length in amino acids (inclusive)
"""
if min_target_length <= 0:
raise ValueError(f"min_target_length must be positive, got {min_target_length}")
self.min_target_length = int(min_target_length)
print(f"Min Target Length: {self.min_target_length}")
def __call__(self, protein_tokens, protein_seq):
"""
Check if protein sequence length is greater than or equal to the minimum target length.
Args:
protein_tokens: Unused (kept for interface compatibility)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if length >= min_target_length, 0 otherwise
"""
seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
L_cand = len(seq_clean)
return int(L_cand >= self.min_target_length)
def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
Args:
seqs: List of protein sequence strings
Returns:
List of bool: True if length >= min_target_length, False otherwise
"""
clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
return [len(s) >= self.min_target_length for s in clean]
class MaxTargetLength:
"""
Maximum target length constraint for protein sequences.
Requires candidate sequence length to be less than or equal to the maximum target length.
"""
def __init__(self, max_target_length: int):
"""
Args:
max_target_length: The maximum allowed length in amino acids (inclusive)
"""
if max_target_length <= 0:
raise ValueError(f"max_target_length must be positive, got {max_target_length}")
self.max_target_length = int(max_target_length)
print(f"Max Target Length: {self.max_target_length}")
def __call__(self, protein_tokens, protein_seq):
"""
Check if protein sequence length is less than or equal to the maximum target length.
Args:
protein_tokens: Unused (kept for interface compatibility)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if length <= max_target_length, 0 otherwise
"""
seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
L_cand = len(seq_clean)
return int(L_cand <= self.max_target_length)
def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
Args:
seqs: List of protein sequence strings
Returns:
List of bool: True if length <= max_target_length, False otherwise
"""
clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
return [len(s) <= self.max_target_length for s in clean]
class PAMMatchingConstraint:
"""
Terminal constraint for protein sequences: predicted PAM must match target PAM.
Requires that the predicted PAM of the generated sequence exactly matches the target PAM.
For positions where target PAM is 'N', any nucleotide is acceptable.
"""
def __init__(self, pam_matching_obj):
"""
Args:
pam_matching_obj: PAMMatching objective object that has predict_pam method and target_pam attribute
"""
if not hasattr(pam_matching_obj, 'predict_pam'):
raise ValueError("pam_matching_obj must have a predict_pam method")
if not hasattr(pam_matching_obj, 'target_pam'):
raise ValueError("pam_matching_obj must have a target_pam attribute")
self.pam_matching_obj = pam_matching_obj
self.target_pam = pam_matching_obj.target_pam.upper()
print(f"PAM Matching Constraint: sequences must have predicted PAM matching target: {self.target_pam}")
def _pam_matches(self, predicted_pam: str, target_pam: str) -> bool:
"""
Check if predicted PAM matches target PAM.
For positions where target PAM is 'N', any nucleotide is acceptable.
Args:
predicted_pam: Predicted PAM string (10 nucleotides)
target_pam: Target PAM string (10 nucleotides, may contain 'N')
Returns:
bool: True if predicted PAM matches target PAM, False otherwise
"""
# if len(predicted_pam) != len(target_pam):
# return False
# for pred_nuc, target_nuc in zip(predicted_pam, target_pam):
# if target_nuc == 'N':
# # N in target means any nucleotide is acceptable
# continue
# if pred_nuc != target_nuc:
# return False
# return True
return predicted_pam.upper() == target_pam.upper()
def __call__(self, protein_tokens, protein_seq):
"""
Check if protein sequence's predicted PAM matches the target PAM.
Args:
protein_tokens: Unused (kept for interface compatibility)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if predicted PAM matches target PAM, 0 otherwise
"""
try:
# Check if the model is available
if not hasattr(self.pam_matching_obj, 'model') or self.pam_matching_obj.model is None:
print(f"[WARNING] PAMMatchingConstraint: model is not initialized. Cannot predict PAM.")
return 0
if not hasattr(self.pam_matching_obj, 'tokenizer') or self.pam_matching_obj.tokenizer is None:
print(f"[WARNING] PAMMatchingConstraint: tokenizer is not initialized. Cannot predict PAM.")
return 0
predicted_pams = self.pam_matching_obj.predict_pam([protein_seq])
if not predicted_pams:
print(f"[WARNING] PAMMatchingConstraint: predict_pam returned empty list for sequence.")
return 0
predicted_pam = predicted_pams[0]
matches = self._pam_matches(predicted_pam, self.target_pam)
if not matches:
# Debug output: show what was predicted vs expected
print(f"[DEBUG] PAMMatchingConstraint: predicted '{predicted_pam}' != target '{self.target_pam}'")
return int(matches)
except Exception as e:
import traceback
print(f"[ERROR] PAMMatchingConstraint failed: {e}")
traceback.print_exc()
return 0
def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
Args:
seqs: List of protein sequence strings
Returns:
List of bool: True if predicted PAM matches target PAM, False otherwise
"""
try:
# Check if the model is available
if not hasattr(self.pam_matching_obj, 'model') or self.pam_matching_obj.model is None:
print(f"[WARNING] PAMMatchingConstraint.predict_batch: model is not initialized. Cannot predict PAM.")
return [False] * len(seqs)
if not hasattr(self.pam_matching_obj, 'tokenizer') or self.pam_matching_obj.tokenizer is None:
print(f"[WARNING] PAMMatchingConstraint.predict_batch: tokenizer is not initialized. Cannot predict PAM.")
return [False] * len(seqs)
predicted_pams = self.pam_matching_obj.predict_pam(seqs)
results = []
for i, predicted_pam in enumerate(predicted_pams):
if not predicted_pam:
print(f"[WARNING] PAMMatchingConstraint.predict_batch: empty prediction for sequence {i}")
results.append(False)
else:
results.append(self._pam_matches(predicted_pam, self.target_pam))
return results
except Exception as e:
import traceback
print(f"[ERROR] PAMMatchingConstraint.predict_batch failed: {e}")
traceback.print_exc()
return [False] * len(seqs)
class Cas9ScoreThreshold:
"""
Terminal constraint for protein sequences: Cas9 classifier score must be above a threshold.
Requires that the Cas9 classifier score (probability) is strictly greater than the threshold.
"""
def __init__(self, cas9_classifier_obj, threshold: float):
"""
Args:
cas9_classifier_obj: Cas9Classification objective object that has get_scores method
threshold: Minimum Cas9 score threshold (0-1). Sequences must have score > threshold.
"""
if not hasattr(cas9_classifier_obj, 'get_scores'):
raise ValueError("cas9_classifier_obj must have a get_scores method")
if not (0.0 <= threshold <= 1.0):
raise ValueError(f"threshold must be in [0, 1], got {threshold}")
self.cas9_classifier_obj = cas9_classifier_obj
self.threshold = float(threshold)
print(f"Cas9 Score Threshold Constraint: sequences must have Cas9 score > {self.threshold}")
def __call__(self, protein_tokens, protein_seq):
"""
Check if protein sequence's Cas9 score is above the threshold.
Args:
protein_tokens: Unused (kept for interface compatibility)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if Cas9 score > threshold, 0 otherwise
"""
scores = self.cas9_classifier_obj.get_scores([protein_seq])
if not scores:
return 0
score = scores[0]
return int(score > self.threshold)
def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
Args:
seqs: List of protein sequence strings
Returns:
List of bool: True if Cas9 score > threshold, False otherwise
"""
scores = self.cas9_classifier_obj.get_scores(seqs)
return [score > self.threshold for score in scores]
class PAMMatchingProbabilityThreshold:
"""
Terminal constraint for protein sequences: PAM matching probability must be above a threshold.
Uses the PAMMatching objective model to compute probability for its configured target PAM and
requires probability > threshold.
"""
def __init__(self, pam_matching_obj, threshold: float):
"""
Args:
pam_matching_obj: PAMMatching-like object with target_pam and get_score_for_pam methods
threshold: Minimum PAM matching probability threshold (0-1). Sequences must have score > threshold.
"""
if not hasattr(pam_matching_obj, "get_score_for_pam"):
raise ValueError("pam_matching_obj must have a get_score_for_pam method")
if not hasattr(pam_matching_obj, "target_pam"):
raise ValueError("pam_matching_obj must have a target_pam attribute")
if not (0.0 <= threshold <= 1.0):
raise ValueError(f"threshold must be in [0, 1], got {threshold}")
self.pam_matching_obj = pam_matching_obj
self.threshold = float(threshold)
self.target_pam = pam_matching_obj.target_pam
print(
"PAM Matching Probability Threshold Constraint: sequences must have "
f"PAM matching probability > {self.threshold} for target PAM {self.target_pam}"
)
def __call__(self, protein_tokens, protein_seq):
"""
Check if protein sequence's PAM matching probability is above the threshold.
Args:
protein_tokens: Unused (kept for interface compatibility)
protein_seq: Protein sequence string (single sequence)
Returns:
int: 1 if PAM matching probability > threshold, 0 otherwise
"""
scores = self.pam_matching_obj.get_score_for_pam(
[protein_seq], self.target_pam, use_temperature_scaling=False
)
if not scores:
return 0
score = scores[0]
return int(score > self.threshold)
def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
Args:
seqs: List of protein sequence strings
Returns:
List of bool: True if PAM matching probability > threshold, False otherwise
"""
scores = self.pam_matching_obj.get_score_for_pam(
seqs, self.target_pam, use_temperature_scaling=False
)
return [score > self.threshold for score in scores]
# -------------------------------------------------------------------------
# Cas9 nuclease domain completeness detector (HNH + RuvC), recall-first
# -------------------------------------------------------------------------
def _sha1(s: str) -> str:
return hashlib.sha1(s.encode("utf-8")).hexdigest()
class Cas9DomainCompleteness:
"""
Hard constraint for protein sequences:
pass iff sequence contains BOTH HNH and RuvC domains (domain presence),
using the Option-B detector (Stage1 fast scan + Stage2 --max rescue).
Intended use:
- As a *domain presence* gate (high recall)
- Combine with an MLP classifier to enforce "is Cas9" / other plausibility checks.
"""
def __init__(
self,
hmm_path: str,
cpu: int = 16,
# Stage1 acceptance
ev1: float = 1e-2,
minlen1: int = 35,
# Stage2 acceptance (recall-first)
ev2: float = 10.0,
minlen2: int = 20,
# Optional: relax Stage1 HMMER filters (usually leave False)
stage1_relax_filters: bool = False,
# Cache results by sequence content to avoid repeated hmmscan calls
enable_cache: bool = True,
):
if DomainDetector is None:
raise ImportError(
"Could not import DomainDetector. Ensure domain_detector.py is in the same package as constraints.py "
"and that its imports succeed."
)
self.hmm_path = hmm_path
self.enable_cache = enable_cache
self._cache: Dict[str, int] = {}
cfg = DomainDetectorConfig(
hmm_path=hmm_path,
cpu=cpu,
stage1=HitCriteria(ievalue_max=ev1, min_dom_len=minlen1),
stage2=HitCriteria(ievalue_max=ev2, min_dom_len=minlen2),
)
if stage1_relax_filters:
cfg = DomainDetectorConfig(
hmm_path=hmm_path,
cpu=cpu,
stage1=HitCriteria(ievalue_max=ev1, min_dom_len=minlen1),
stage2=HitCriteria(ievalue_max=ev2, min_dom_len=minlen2),
stage1_F1=1.0, stage1_F2=1.0, stage1_F3=1.0,
stage1_report_E=1e6, stage1_report_domE=1e6,
)
self.detector = DomainDetector(cfg)
def __call__(self, protein_tokens, protein_seq: str) -> int:
# protein_tokens is unused; kept to match existing constraint signature
seq = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
if not seq:
return 0
if self.enable_cache:
key = _sha1(seq)
if key in self._cache:
return self._cache[key]
passed = int(self.detector.predict_one(seq, id_="query"))
if self.enable_cache:
self._cache[_sha1(seq)] = passed
return passed
def predict_batch(self, seqs: Sequence[str], ids: Optional[Sequence[str]] = None) -> List[bool]:
"""
Batch prediction helper for evaluation / faster integration.
"""
clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
return self.detector.predict_batch(clean, ids=ids)