# 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)