diff --git a/hbdesigner/data/hbnet.py b/hbdesigner/data/hbnet.py index 336cec5..7e40bae 100755 --- a/hbdesigner/data/hbnet.py +++ b/hbdesigner/data/hbnet.py @@ -3,6 +3,7 @@ from multiprocessing import Pool from typing import Dict, List, Sequence, Tuple, Optional import os +import shutil import numpy as np import pyrosetta import time @@ -32,8 +33,12 @@ from hbdesigner.data.features import impute_CB from hbdesigner.data.protein import Protein -REDUCE_EXE=os.path.join(os.path.dirname(os.path.dirname(__file__)), "reduce/reduce_src/reduce") -HET_DICT = os.path.join(os.path.dirname(os.path.dirname(__file__)), "reduce/reduce_wwPDB_het_dict.txt") +# Reduce is installed via conda/pixi (see pyproject.toml); override with $REDUCE_EXE if it isn't on PATH +REDUCE_EXE = os.environ.get("REDUCE_EXE") or shutil.which("reduce") or "reduce" +HET_DICT = os.environ.get("REDUCE_HET_DICT") or os.path.join( + os.path.dirname(os.path.dirname(os.path.realpath(shutil.which(REDUCE_EXE) or REDUCE_EXE))), + "share/reduce/reduce_wwPDB_het_dict.txt", +) def batch_to_proteins(batch: gd.Batch) -> Tuple[List[Protein], List[gd.Data]]: @@ -200,7 +205,7 @@ def get_satisfaction( def rosetta_hbond_detect( p: Protein, max_energy: float = 0.0, optH: bool = True, optH_MCA: bool = True, -) -> List[Tuple[int, int]]: +) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]], Dict[str, float]]: """ Detect Rosetta HBonds based on energy threshold. @@ -209,6 +214,12 @@ def rosetta_hbond_detect( max_energy (float): Max energy to consider a 'valid' HBond. Larger is more permissive. Default is 0.0. optH (bool): Whether to let Rosetta re-optimize the hydrogens in the input Protein. Default is True. optH_MCA (bool): Whether to use more rigorous but slower optH MCA protocol. Default is True. + + Returns: + pair_list (List[Tuple[int, int]]): (acc_res, don_res) pairs for sidechain-sidechain HBonds. + bb_sc_pair_list (List[Tuple[int, int]]): (acc_res, don_res) pairs for backbone-sidechain HBonds + (either the acceptor or the donor hydrogen is a backbone atom). + stats (Dict[str, float]): Energy/satisfaction stats for the pose. """ t0 = time.time() # Configure scoring settings, if not already set @@ -221,6 +232,7 @@ def rosetta_hbond_detect( # Set up polyG backbone for energy calc pose, polyG_pose = Pose(), Pose() pair_list = [] + bb_sc_pair_list = [] try: pose_from_pdbstring(pose, p.to_pdb(unk_to_gly=True)) # Make PolyG pose for comparison @@ -233,20 +245,24 @@ def rosetta_hbond_detect( except RuntimeError: print("Rosetta failed to accept PDB!") - return pair_list + return pair_list, bb_sc_pair_list, {} # Need to apply FA scorefxn to populate HBondSet scorefxn = get_fa_scorefxn() scorefxn(pose) - hbondset = pose.get_hbonds(exclude_bb=True, exclude_bsc=True, exclude_scb=True) + # Exclude bb-bb only here; sc-sc and bb-sc bonds are split below by atom identity + hbondset = pose.get_hbonds(exclude_bb=True) hbonds = hbondset.hbonds() - # Collect HBonds + # Collect HBonds, splitting sidechain-sidechain from backbone-sidechain for bond in hbonds: # Convert from Pose to Protein numbering acc_res = int(bond.acc_res()) - 1 don_res = int(bond.don_res()) - 1 - pair_list.append((acc_res, don_res)) + if bond.acc_atm_is_backbone() or bond.don_hatm_is_backbone(): + bb_sc_pair_list.append((acc_res, don_res)) + else: + pair_list.append((acc_res, don_res)) seq, seq_polyG = pose.sequence(), polyG_pose.sequence() net_idx = [i for i, (s, sg) in enumerate(zip(seq, seq_polyG)) if s != sg] @@ -275,7 +291,7 @@ def rosetta_hbond_detect( stats.update(get_satisfaction(pose)) t1 = time.time() stats["score_time"] = (t1 - t0) - return pair_list, stats + return pair_list, bb_sc_pair_list, stats def biotite_hbond_detect( @@ -364,8 +380,14 @@ def score_protein(p: Protein) -> Dict[str, float]: Dict[str, float]: Dict of single protein metrics and their labels. """ + # Guide atom (if present) is stored in hetatm_dict, which masking below drops + guide_atom_xyz = None + if p.hetatm_dict != {} and "V1" in p.hetatm_dict["atom_name"]: + guide_idx = np.where(p.hetatm_dict["atom_name"] == "V1")[0][0] + guide_atom_xyz = p.hetatm_dict["atom_xyz"][guide_idx] + # Get bonds and energy metrics from Rosetta - bond_list, rosetta_stats = rosetta_hbond_detect( + bond_list, bb_sc_bond_list, rosetta_stats = rosetta_hbond_detect( p, max_energy=0.0, optH=True, @@ -387,7 +409,7 @@ def score_protein(p: Protein) -> Dict[str, float]: chains_used = np.unique(p.chain_index).size used_all_chains = float(chains_used == chains_total) - # Count motif nodes and edges + # Count motif nodes and edges (sidechain-sidechain HBonds) hbond_res = [] for b in bond_list: hbond_res.extend(list(b)) @@ -396,12 +418,26 @@ def score_protein(p: Protein) -> Dict[str, float]: total_res = p.n_res hbond_res_frac = hbond_res / total_res + # Same counts for backbone-sidechain HBonds + hbond_bb_sc_res = [] + for b in bb_sc_bond_list: + hbond_bb_sc_res.extend(list(b)) + hbond_bb_sc_res = np.unique(hbond_bb_sc_res).size + hbond_bb_sc_bonds = len(bb_sc_bond_list) + hbond_bb_sc_res_frac = hbond_bb_sc_res / total_res + # Calculate max pairwise distance b/w motif Cb atoms bb_xyz = p.atom27_xyz[:, :4, :] cb_xyz = impute_CB(bb_xyz[:, 0], bb_xyz[:, 1], bb_xyz[:, 2]) cd = cdist(cb_xyz, cb_xyz) max_Cb_dist = np.max(np.triu(cd)) + # Distance from network Cb centroid to guide atom, if one was provided + displacement = None + if guide_atom_xyz is not None: + cb_centroid = np.mean(cb_xyz, axis=0) + displacement = float(np.linalg.norm(cb_centroid - guide_atom_xyz)) + # Check if residue aatypes are valid valid_res = np.array([rc.restype_order[r] for r in rc.restype_hb_sc]) vmask = np.sum(p.aatype == valid_res[:, None], axis=0) > 0 @@ -421,11 +457,15 @@ def score_protein(p: Protein) -> Dict[str, float]: "hbond_res_frac": hbond_res_frac, "hbond_res": hbond_res, "hbond_bonds": hbond_bonds, + "hbond_bb_sc_res_frac": hbond_bb_sc_res_frac, + "hbond_bb_sc_res": hbond_bb_sc_res, + "hbond_bb_sc_bonds": hbond_bb_sc_bonds, "chains_used": float(chains_used), "used_all_chains": used_all_chains, "total_res": total_res, "max_Cb_dist": max_Cb_dist, "valid_res_frac": valid_res_frac, + "displacement": displacement, } score_info.update(rosetta_stats) @@ -708,15 +748,19 @@ def safe_run(p: Protein, mode: str = "pack"): elif mode == "minimize-cart": return minimize_task(p, True) elif mode == "reduce": - return run_reduce(p, his=True, flip=False) + out_pdb = run_reduce(p.to_pdb(unk_to_gly=True), his=True, flip=False) + red_p = Protein.from_pdb_string(out_pdb, discard_Hs=False).pad(n=p.n_res) + if hasattr(p, "pack_time"): + red_p.pack_time = p.pack_time + return red_p elif mode == "hydride": return run_hydride(p) elif mode == "pdb2pqr": return run_pdb2pqr(p) else: return minimize_task(p, False) - except Exception: - print(f"Job distributor failed to run {mode} on {p}, returning unpacked protein.") + except Exception as e: + print(f"Job distributor failed to run {mode} on {p} ({type(e).__name__}: {e}), returning unpacked protein.") return p @@ -885,21 +929,85 @@ def calc_seq_rec_batched( assert (pred_seq != rc.restype_num).sum() == (true_seq != rc.restype_num).sum() # Calculate pos recovery per sample + # NOTE: scatter with reduce="sum" on a bool tensor stays bool (saturates at True), + # so all summed quantities below must be cast to float first. mask = true_pos.clone() pos_rec_per_samp = scatter( - pred_pos[mask] == true_pos[mask], aatype_batch[mask], dim=0, reduce="sum" + (pred_pos[mask] == true_pos[mask]).to(torch.float32), aatype_batch[mask], dim=0, reduce="sum" ) # [B] - pos_per_samp = scatter(pred_pos, aatype_batch, dim=0, reduce="sum") # [B] + pos_per_samp = scatter(pred_pos.to(torch.float32), aatype_batch, dim=0, reduce="sum") # [B] pos_rec = pos_rec_per_samp / (pos_per_samp + 1e-8) # [B] # Calculate seq recovery per sample rec_mask = (pred_seq[mask] == true_seq[mask]) * (pred_pos[mask] == true_pos[mask]) - seq_rec_per_samp = scatter(rec_mask, aatype_batch[mask], dim=0, reduce="sum") # [B] + seq_rec_per_samp = scatter(rec_mask.to(torch.float32), aatype_batch[mask], dim=0, reduce="sum") # [B] seq_rec = seq_rec_per_samp / (pos_per_samp + 1e-8) # [B] return pos_rec, seq_rec +def calc_unordered_seq_rec_batched( + pred_pos: torch.tensor, + true_pos: torch.tensor, + pred_seq: torch.tensor, + true_seq: torch.tensor, + aatype_batch: torch.tensor, +) -> torch.tensor: + """ + Calculates unordered (position-agnostic) seq recovery on batched predictions. + + Unlike calc_seq_rec_batched, this ignores which specific position each residue + was assigned to and instead checks whether the predicted network's multiset of + residue types matches the native network's multiset of residue types. + + Arguments: + pred_pos (torch.tensor): Predicted position vector of shape [L] for the full batch. + true_pos (torch.tensor): True position vector of shape [L] for the full batch. + pred_seq (torch.tensor): Predicted seq vector of shape [L] for the full batch. + true_seq (torch.tensor): True seq vector of shape [L] for the full batch. + aatype_batch (torch.tensor): Batch indexing vector of shape [L] for scatter ops. + + Returns: + torch.tensor: unordered seq recovery of each sample, shaped [B]. + """ + # Check that vectors are comparable + assert ( + pred_pos.numel() + == true_pos.numel() + == pred_seq.numel() + == true_seq.numel() + == aatype_batch.numel() + ) + num_graphs = int(aatype_batch.max().item()) + 1 + + # Count occurrences of each aatype per sample by flattening (sample, aatype) into + # a single scatter index, for both the true and predicted networks + true_idx = aatype_batch[true_pos] * rc.restype_num + true_seq[true_pos] + true_counts = scatter( + torch.ones_like(true_idx, dtype=torch.float32), + true_idx, + dim=0, + reduce="sum", + dim_size=num_graphs * rc.restype_num, + ).view(num_graphs, rc.restype_num) + + pred_idx = aatype_batch[pred_pos] * rc.restype_num + pred_seq[pred_pos] + pred_counts = scatter( + torch.ones_like(pred_idx, dtype=torch.float32), + pred_idx, + dim=0, + reduce="sum", + dim_size=num_graphs * rc.restype_num, + ).view(num_graphs, rc.restype_num) + + # Multiset intersection: for each sample, how many residues of each aatype + # are shared between the predicted and native networks, regardless of position + matched_per_samp = torch.minimum(true_counts, pred_counts).sum(dim=1) # [B] + pos_per_samp = scatter(pred_pos.float(), aatype_batch, dim=0, reduce="sum", dim_size=num_graphs) # [B] + unordered_seq_rec = matched_per_samp / (pos_per_samp + 1e-8) # [B] + return unordered_seq_rec + + def initialize_rosetta(mode: str = "fast", seed: int = None) -> None: """ Initialize Rosetta for packing/scoring use. diff --git a/hbdesigner/model/hbdesign_model.py b/hbdesigner/model/hbdesign_model.py index afd47ce..1daa6c2 100755 --- a/hbdesigner/model/hbdesign_model.py +++ b/hbdesigner/model/hbdesign_model.py @@ -579,25 +579,29 @@ def sample( b: gd.Batch, res_sample_temp: float = 0.1, seq_sample_temp: float = 0.1, + ignore_seq_cond_mask: bool = False, ) -> Dict[str, torch.Tensor]: - """Currently used by Trainer.sample_batch in seq design mode""" + """Currently used by Trainer.sample_batch in seq design mode + + Args: + ignore_seq_cond_mask (bool): If True, don't hard-constrain sampling to the + specific aatype budget from seq cond. Sampling is still restricted to + the 11 HB-capable residues, but seq cond only informs the model via its + soft conditioning embedding rather than masking out disallowed aatypes. + Defaults to False (existing hard-constraint behavior). + """ - # Clear all seq and sidechain info - b.aatype[:] = rc.restype_num + # Clear sidechain info. Note: b.aatype and b.done_mask are left as-is, since any + # anchor residues (done_mask == 1) already carry their true aatype and must not + # be re-decoded below. b.atom14_xyz[:, 4:, :] = 0.0 b.atom14_mask[:, 4:] = 0.0 - # Revert masks back to start of decoding - net_sizes = ( - scatter(b.nll_mask, b.aatype_batch, dim=0) - + scatter(b.done_mask, b.aatype_batch, dim=0) - ).long() - b.nll_mask[:] = 0 - b.done_mask[:] = 0 - b.net_res_num = net_sizes - - # Create initial sequence - seq = (rc.restype_num * torch.ones_like(b.aatype)).to(torch.long) # [B] + # Only the not-yet-decided residues remain to be sampled; anchors are already done + b.net_res_num = scatter(b.nll_mask, b.aatype_batch, dim=0).long() + + # Create initial sequence: anchor residues keep their true aatype, rest start as UNK + seq = b.aatype.clone().to(torch.long) # [B] net_res_probs = torch.zeros_like(b.aatype).to(torch.float32) # [B] seq_probs = torch.zeros_like(b.aatype).to(torch.float32) # [B] @@ -716,7 +720,12 @@ def get_processed_node_embeddings(b, seq, cond_info): seq_logits = self.seq_layer(protein_nodes[net_res]) # For any samples w/seq cond, mask out all other logits to ensure correct aatype chosen - aatypes_not_allowed = (aatype_cond_counts <= 0.0).to(bool) + if ignore_seq_cond_mask: + # Still restrict to HB-capable residues, but don't hard-enforce the seq cond budget + aatypes_not_allowed = torch.ones_like(aatype_cond_counts, dtype=torch.bool) + aatypes_not_allowed[:, polars] = False + else: + aatypes_not_allowed = (aatype_cond_counts <= 0.0).to(bool) seq_logits[aatypes_not_allowed[net_res_mask]] = -torch.inf seq_pred, seq_p = self.sample_seq( diff --git a/hbdesigner/scripts/evaluate_hbdesigner.py b/hbdesigner/scripts/evaluate_hbdesigner.py index 5bb4c22..c15a2de 100644 --- a/hbdesigner/scripts/evaluate_hbdesigner.py +++ b/hbdesigner/scripts/evaluate_hbdesigner.py @@ -101,6 +101,37 @@ default=1, help="Number of repeats for error bars. Defaults to 1.", ) + parser.add_argument( + "--guide_atom_cond", + action="store_true", + help="Whether to enable guide atom conditioning for all eval samples. Defaults to False (off).", + ) + parser.add_argument( + "--seq_cond", + action="store_true", + help="Whether to enable sequence conditioning for all eval samples. Defaults to False (off).", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Whether to save per-network stats to a CSV (hbdes_eval_data.csv). Defaults to False.", + ) + parser.add_argument( + "--ignore_seq_cond_mask", + action="store_true", + help="Whether to disable the hard seq cond aatype constraint during sampling, still " + "restricting to the 11 HB-capable residues but without forcing the conditioned " + "identities/counts. Defaults to False (off).", + ) + parser.add_argument( + "--n_anchor_res", + type=int, + default=0, + help="Number of network residues to anchor (treat as already decided) before " + "sampling completes the rest of the network. Must be between 0 and " + "hbdesigner.max_res - 1. Networks with fewer than n_anchor_res + 1 residues are " + "skipped. Defaults to 0 (design the whole network from scratch).", + ) args = parser.parse_args() print("Args:", args) @@ -122,10 +153,17 @@ design_config.model.hbdesigner.rescore = False design_config.model.hbdesigner.rescore_filter = False + if not (0 <= args.n_anchor_res <= design_config.model.hbdesigner.max_res - 1): + parser.error( + "--n_anchor_res must be between 0 and " + f"{design_config.model.hbdesigner.max_res - 1} (hbdesigner.max_res - 1), " + f"got {args.n_anchor_res}." + ) + # Conditioning info params - design_config.model.hbdesigner.guide_atom_pct = 0.0 + design_config.model.hbdesigner.guide_atom_pct = 1.0 if args.guide_atom_cond else 0.0 design_config.model.hbdesigner.guide_atom_sigma = 4.0 - design_config.model.hbdesigner.seq_cond_pct = 0.0 + design_config.model.hbdesigner.seq_cond_pct = 1.0 if args.seq_cond else 0.0 design_config.model.hbdesigner.seq_cond_unk_pct = 0.0 # Eval params @@ -171,7 +209,7 @@ for r in range(args.repeats): t0 = time.time() seed_everything(seeds[r]) - test_dl = design_trainer.build_test_data_loader() + test_dl = design_trainer.build_test_data_loader(n_anchor_res=args.n_anchor_res) test_dl_pack = pack_trainer.build_test_data_loader() print(f"Starting test loop # {r + 1} / {args.repeats}") @@ -183,8 +221,9 @@ first_n=args.first_n, seq_sample_temp=args.seq_temp, res_sample_temp=args.pos_temp, - verbose=False, + verbose=args.verbose, pack_trainer=pack_trainer, + ignore_seq_cond_mask=args.ignore_seq_cond_mask, ) t1 = time.time() elapsed = round(t1 - t0) diff --git a/hbdesigner/scripts/evaluate_hbpacker.py b/hbdesigner/scripts/evaluate_hbpacker.py index 5ac6a81..f6eada3 100644 --- a/hbdesigner/scripts/evaluate_hbpacker.py +++ b/hbdesigner/scripts/evaluate_hbpacker.py @@ -81,6 +81,11 @@ default=1, help="Number of repeats for error bars. Defaults to 1.", ) + parser.add_argument( + "--verbose", + action="store_true", + help="Whether to save per-network stats to a CSV (hbpacker_eval_data.csv). Defaults to False.", + ) args = parser.parse_args() print("Args:", args) @@ -134,7 +139,7 @@ n_workers=args.num_workers, dump=args.dump, first_n=args.first_n, - verbose=False, + verbose=args.verbose, ) t1 = time.time() elapsed = round(t1 - t0) diff --git a/hbdesigner/scripts/train_hbdesigner.py b/hbdesigner/scripts/train_hbdesigner.py index 80cd307..ce56621 100755 --- a/hbdesigner/scripts/train_hbdesigner.py +++ b/hbdesigner/scripts/train_hbdesigner.py @@ -29,6 +29,7 @@ get_guide_atom, get_seq_cond, calc_seq_rec_batched, + calc_unordered_seq_rec_batched, crop_by_distance, score_protein, ) @@ -55,10 +56,14 @@ class HBDesignerDataset(torch.utils.data.IterableDataset): """ - def __init__(self, cfg: TrainConfig, split: str = "train"): + def __init__(self, cfg: TrainConfig, split: str = "train", n_anchor_res: int = 0): # Load metadata self.cfg = cfg self.split = split + # Number of network residues to treat as already-decided on the test split, so + # that featurize() leaves them out of nll_mask and sample() completes the rest. + # Not part of TrainConfig since it's only meaningful for eval scripts. + self.n_anchor_res = n_anchor_res if self.cfg.model.model_name == "HBDesigner": self.model_cfg = self.cfg.model.hbdesigner @@ -155,7 +160,7 @@ def get(self, idx: int) -> Union[gd.Data, None]: if p.n_res > self.MAX_LENGTH: return None - return self.featurize(p, hbnet_arr) + return self.featurize(p, hbnet_arr, pdb_ch=pdb_ch) def _sample_native_net( self, p: Protein, g: nx.DiGraph @@ -352,13 +357,15 @@ def _load_assembly(self, pdb_ch: str) -> Union[Tuple[Protein, np.ndarray], None] hbnet_arr = hbnet_arr[chain_mask] return p, hbnet_arr - def featurize(self, p: Protein, hbnet_arr: np.ndarray) -> gd.Data: + def featurize(self, p: Protein, hbnet_arr: np.ndarray, pdb_ch: str = "") -> gd.Data: """ Featurize Protein and Graph info into HBDesigner model inputs. Arguments: p (Protein): Protein object. hbnet_arr (np.ndarray): Copy of p.aatype with all non network residues set to GLY. + pdb_ch (str): PDB chain identifier (e.g. "1ABC_A") this sample was drawn from, + carried through purely for bookkeeping/logging. Defaults to "". Returns: gd.Data: torch_geometric Data object with featurized protein. @@ -379,9 +386,12 @@ def featurize(self, p: Protein, hbnet_arr: np.ndarray) -> gd.Data: # Do step and residue sampling for training n_res = hbnet_pos.size - # Always run test from tstep 0 to get accurate cond info + # On test, anchor n_anchor_res residues (already decided) and complete the rest. + # Skip networks too small to have any non-anchor residues left to design. if self.split == "test": - tstep = 0 + tstep = self.n_anchor_res + if n_res <= tstep: + return else: tstep = random.randint(0, n_res - 1) @@ -486,6 +496,7 @@ def featurize(self, p: Protein, hbnet_arr: np.ndarray) -> gd.Data: ) # [1, 21] protein_data["c_idx"] = protein_data["chain_index"] + protein_data["pdb_ch"] = pdb_ch return protein_data @staticmethod @@ -681,8 +692,8 @@ def _make_data_loader( worker_init_fn=worker_init, ) - def build_test_data_loader(self) -> DataLoader: - self.test_data = HBDesignerDataset(self.cfg, split="test") + def build_test_data_loader(self, n_anchor_res: int = 0) -> DataLoader: + self.test_data = HBDesignerDataset(self.cfg, split="test", n_anchor_res=n_anchor_res) return self._make_data_loader(self.test_data) @torch.no_grad() @@ -691,6 +702,7 @@ def sample_batch( batch: gd.Batch, seq_sample_temp: float = 0.1, res_sample_temp: float = 0.1, + ignore_seq_cond_mask: bool = False, ) -> Tuple[gd.Batch, Dict[int, any]]: """ Sample batch with self.model and update the batch object w/network information. @@ -699,6 +711,8 @@ def sample_batch( batch (gd.Batch): Full batch of data for sampling. seq_sample_temp (float): Sampling temperature for seq (aatype) decoding. res_sample_temp (float): Sampling temperature for res (position/stop action) decoding. + ignore_seq_cond_mask (bool): If True, don't hard-constrain sampling to the seq cond + aatype budget (still restricted to the 11 HB-capable residues). Defaults to False. Returns: gd.Batch: Updated batch ready for packing/scoring. @@ -709,6 +723,7 @@ def sample_batch( batch.to(self.cfg.device), seq_sample_temp=seq_sample_temp, res_sample_temp=res_sample_temp, + ignore_seq_cond_mask=ignore_seq_cond_mask, ) # Apply results to each scaffold in the batch b_list = batch.to_data_list() @@ -728,6 +743,7 @@ def test_batch( seq_sample_temp: float = 0.1, res_sample_temp: float = 0.1, n_workers: int = 1, + ignore_seq_cond_mask: bool = False, ) -> Dict[str, Any]: """ Sample, pack, and score a Batch of Proteins using HBNet metrics. @@ -739,6 +755,8 @@ def test_batch( seq_sample_temp (float): Sampling temperature for restype decoding. Defaults to 0.1. res_sample_temp (float): Sampling temperature for network position decoding. Defaults to 0.1. n_workers (int): Number of workers for Rosetta parallel packing. Defaults to 1. + ignore_seq_cond_mask (bool): If True, don't hard-constrain sampling to the seq cond + aatype budget (still restricted to the 11 HB-capable residues). Defaults to False. Returns: info (Dict[str, any]): Set of metrics for the batch. @@ -749,11 +767,15 @@ def test_batch( native_seq = deepcopy(b.aatype_gt) native_pos = native_seq != rc.restype_num aatype_batch = b.aatype_batch.clone() + # Anchor residues (already decided pre-sampling) aren't actually "predicted", + # so recovery metrics restricted to newly-designed residues exclude them below + anchor_mask = b.done_mask.clone().bool() b, results = self.sample_batch( b, seq_sample_temp=seq_sample_temp, res_sample_temp=res_sample_temp, + ignore_seq_cond_mask=ignore_seq_cond_mask, ) pred_seq = b.aatype.clone().to(native_seq.device) pred_seq[pred_seq == rc.restype_order["G"]] = rc.restype_num @@ -766,14 +788,44 @@ def test_batch( test_info["pos_rec"] = pos_rec.tolist() test_info["seq_rec"] = seq_rec.tolist() - # Collect avg prob/log-prob of positions + # Unordered seq recovery: does the predicted network contain the same aatypes + # as the native network, regardless of position/order within the network + unordered_seq_rec = calc_unordered_seq_rec_batched( + pred_pos, native_pos, pred_seq, native_seq, aatype_batch + ) + test_info["unordered_seq_rec"] = unordered_seq_rec.tolist() + + # Same recovery metrics restricted to newly-designed residues, excluding anchors. + # With n_anchor_res=0 these are identical to the metrics above. + anchor_mask = anchor_mask.to(native_pos.device) + new_native_pos = native_pos & ~anchor_mask + new_pred_pos = pred_pos & ~anchor_mask + pos_rec_new, seq_rec_new = calc_seq_rec_batched( + new_pred_pos, new_native_pos, pred_seq, native_seq, aatype_batch + ) + test_info["pos_rec_new"] = pos_rec_new.tolist() + test_info["seq_rec_new"] = seq_rec_new.tolist() + + unordered_seq_rec_new = calc_unordered_seq_rec_batched( + new_pred_pos, new_native_pos, pred_seq, native_seq, aatype_batch + ) + test_info["unordered_seq_rec_new"] = unordered_seq_rec_new.tolist() + + # Collect avg prob/log-prob of positions, and predicted network size with/without anchors net_res_probs = [] seq_probs = [] + n_res_pred = [] + n_res_pred_new = [] for i in range(b.num_graphs): net_res_probs.append(np.mean(results[i]["net_res_probs"].tolist())) seq_probs.append(np.mean(results[i]["seq_probs"].tolist())) + samp_mask = aatype_batch == i + n_res_pred.append(int(pred_pos[samp_mask].sum())) + n_res_pred_new.append(int(new_pred_pos[samp_mask].sum())) test_info["net_res_probs"] = net_res_probs test_info["seq_probs"] = seq_probs + test_info["n_res_pred"] = n_res_pred + test_info["n_res_pred_new"] = n_res_pred_new # Need to adjust gd.Batch to be compatible with HBPacker pack_ds = pack_trainer.test_data @@ -810,6 +862,9 @@ def test_batch( else: test_info[key].append(value) + # Track which PDB/chain each row came from + test_info.setdefault("pdb_ch", []).append(b_c.pdb_ch) + return test_info, proteins @torch.no_grad() @@ -824,6 +879,7 @@ def test_loop( dump: bool = False, first_n: int = None, verbose: bool = False, + ignore_seq_cond_mask: bool = False, ) -> Dict[str, any]: """ Run test loop on the specified DataLoader. @@ -839,6 +895,8 @@ def test_loop( n_workers (int): Number of workers available for Rosetta packing job. Defaults to 1. dump (bool): Whether to dump the packed PDBs after scoring. Defaults to False. verbose (bool): Whether to dump per-network stats to CSV. Defaults to False. + ignore_seq_cond_mask (bool): If True, don't hard-constrain sampling to the seq cond + aatype budget (still restricted to the 11 HB-capable residues). Defaults to False. Returns: Dict[str, any]: Dictionary of test metrics. @@ -853,6 +911,7 @@ def test_loop( res_sample_temp=res_sample_temp, n_workers=n_workers, pack_trainer=pack_trainer, + ignore_seq_cond_mask=ignore_seq_cond_mask, ) # Save PDBs to disk, if requested if dump: @@ -886,11 +945,16 @@ def test_loop( df = pd.DataFrame.from_dict(valid_info) df.to_csv("hbdes_eval_data.csv") + # Identifier column, not a metric to average + valid_info.pop("pdb_ch", None) + for k, v in valid_info.items(): - v = np.array(v) + # astype(float) turns None entries (e.g. displacement w/o a guide atom) into NaN + v = np.array(v).astype(float) if first_n is not None: v = v[:first_n] - valid_info[k] = np.nanmean(v) + with np.errstate(invalid="ignore"): + valid_info[k] = np.nanmean(v) valid_info["n_samples"] = first_n return valid_info @@ -901,7 +965,7 @@ def load_model_state(self, ckpt_path: str) -> None: map_location = {"cuda:0": f"cuda:{self.rank}"} else: map_location = "cpu" - state = torch.load(ckpt_path, map_location=map_location) + state = torch.load(ckpt_path, map_location=map_location, weights_only=False) self.model.load_state_dict(state["model_state_dict"]) self.model.to(self.device) diff --git a/hbdesigner/scripts/train_hbpacker.py b/hbdesigner/scripts/train_hbpacker.py index d69f8cd..e55ff56 100644 --- a/hbdesigner/scripts/train_hbpacker.py +++ b/hbdesigner/scripts/train_hbpacker.py @@ -59,13 +59,15 @@ class HBPackerDataset(HBDesignerDataset): def __init__(self, cfg: TrainConfig, split: str = "train"): super().__init__(cfg, split) - def featurize(self, p: Protein, hbnet_arr: np.ndarray) -> gd.Data: + def featurize(self, p: Protein, hbnet_arr: np.ndarray, pdb_ch: str = "") -> gd.Data: """ Featurize Protein and Graph info into HBPacker model inputs. Arguments: p (Protein): Protein object. hbnet_arr (np.ndarray): Copy of p.aatype with all non network residues set to GLY. + pdb_ch (str): PDB chain identifier (e.g. "1ABC_A") this sample was drawn from, + carried through purely for bookkeeping/logging. Defaults to "". Returns: gd.Data: torch_geometric Data object with featurized protein. @@ -171,9 +173,12 @@ def featurize(self, p: Protein, hbnet_arr: np.ndarray) -> gd.Data: ), # [L, 4, 2] # Masks and cond info chi_mask=torch.from_numpy(chi_mask).to(torch.float32), # [L] + # No symmetry info available here; treat each packable residue as its own group + symmetry_idx=torch.arange(chi_mask_bool.sum(), dtype=torch.long), # [n_packable] ) protein_data["c_idx"] = protein_data["chain_index"] + protein_data["pdb_ch"] = pdb_ch return protein_data @staticmethod @@ -284,10 +289,13 @@ def convert_batch_design_to_pack(b: gd.Batch, pack_crop: float = 10.0) -> Sequen torch.float32 ), # [L, 4] chi_mask=torch.from_numpy(chi_mask).to(torch.float32), # [L] + # No symmetry info available here; treat each packable residue as its own group + symmetry_idx=torch.arange(chi_mask_bool.sum(), dtype=torch.long), # [n_packable] ) protein_data["c_idx"] = protein_data["chain_index"] protein_data["guide_atom_xyz"] = b_i.guide_atom_xyz protein_data["aatype_cond"] = b_i.aatype_cond + protein_data["pdb_ch"] = b_i.pdb_ch protein_data_list.append(protein_data) return protein_data_list @@ -472,6 +480,9 @@ def test_batch( else: test_info[key].append(value) + # Track which PDB/chain each row came from + test_info.setdefault("pdb_ch", []).append(b_c.pdb_ch) + return test_info, proteins @torch.no_grad() @@ -537,13 +548,18 @@ def test_loop( # Take mean for each metric if verbose: df = pd.DataFrame.from_dict(valid_info) - df.to_csv("hbdes_eval_data.csv") + df.to_csv("hbpacker_eval_data.csv") + + # Identifier column, not a metric to average + valid_info.pop("pdb_ch", None) for k, v in valid_info.items(): - v = np.array(v) + # astype(float) turns None entries (e.g. displacement w/o a guide atom) into NaN + v = np.array(v).astype(float) if first_n is not None: v = v[:first_n] - valid_info[k] = np.nanmean(v) + with np.errstate(invalid="ignore"): + valid_info[k] = np.nanmean(v) valid_info["n_samples"] = first_n return valid_info @@ -667,14 +683,18 @@ def compute_packing_metrics(self, p: Protein, b: gd.Data) -> Dict[str, Any]: pack_info = {} # Collect true and pred chi values true_chi_rad = sincos_to_angle(b.chi_sincos_gt) - p_mask = p.atom27_mask[:, 5] != 0 - pred_coords = torch.from_numpy(p.atom27_xyz[p_mask, :14]).to(torch.float32) + # Network positions. Packing never reorders/drops residues, so this mask + # (from the pre-packing batch) lines up 1:1 with p's residues. We rely on + # this instead of p.atom27_mask[:, 5] != 0, which under-selects whenever a + # native network residue's sidechain past CB is unresolved in the source PDB. + mask = b.chi_mask.bool() + mask_np = mask.numpy() + pred_coords = torch.from_numpy(p.atom27_xyz[mask_np, :14]).to(torch.float32) pred_chi_rad = calc_sc_dihedrals( - pred_coords, torch.from_numpy(p.aatype[p_mask]), return_mask=False + pred_coords, torch.from_numpy(p.aatype[mask_np]), return_mask=False ) # Collect relevant chi mask - mask = b.chi_mask.bool() chi_mask = b.sc_dihedral_mask[mask] dev = next(self.model.parameters()).device diff --git a/model_weights/design_020.yaml b/model_weights/design_020.yaml index cbc0c79..126c012 100644 --- a/model_weights/design_020.yaml +++ b/model_weights/design_020.yaml @@ -42,7 +42,7 @@ model: nll_smoothing: 0.0 loss_type: focal focal_gamma: 2.0 - data_location: /work/users/d/i/dieckhau/pdb_2021aug02_hbdesigner4/ + data_location: /work/users/d/i/dieckhau/datasets/pdb_2021aug02_hbdesigner4/ inter_weight: 1.0 hbnet_pct: 0.0 batch_size: 10000 diff --git a/model_weights/design_020_cpu.yaml b/model_weights/design_020_cpu.yaml index 97f646a..ed99b94 100644 --- a/model_weights/design_020_cpu.yaml +++ b/model_weights/design_020_cpu.yaml @@ -42,7 +42,7 @@ model: nll_smoothing: 0.0 loss_type: focal focal_gamma: 2.0 - data_location: /work/users/d/i/dieckhau/pdb_2021aug02_hbdesigner4/ + data_location: /work/users/d/i/dieckhau/datasets/pdb_2021aug02_hbdesigner4/ inter_weight: 1.0 hbnet_pct: 0.0 batch_size: 10000 diff --git a/model_weights/pack.yaml b/model_weights/pack.yaml index 3264bd2..7ae0c62 100644 --- a/model_weights/pack.yaml +++ b/model_weights/pack.yaml @@ -71,7 +71,7 @@ model: orient_msd_weight: 0.2 reweight_chi_mse: true num_recycles: 3 - data_location: /work/users/d/i/dieckhau/pdb_2021aug02_hbdesigner4/ + data_location: /work/users/d/i/dieckhau/datasets/pdb_2021aug02_hbdesigner4/ inter_weight: 1.0 hbnet_pct: 0.0 batch_size: 10000