from typing import Any import torch import time import os import subprocess import shutil import random import wandb import numpy as np import pandas as pd import logging from pytorch_lightning import LightningModule from Bio import PDB # import esm from biotite.sequence.io import fasta from utils.experiments import write_prot_to_pdb, save_traj, create_full_prot from utils import metrics from models.proteinflow import ProteinFlow from models.classifier import ProtClassifier from models.classifier_wrapper import ClasfModule from utils import all_atom from utils import so3Utils as su from utils import residue_constants as rc from utils import pdbUtils as du from utils.flows import Interpolant from utils.modelUtils import t_stratified_loss, to_numpy from pytorch_lightning.loggers.wandb import WandbLogger from dataset import protein class ProteinFlowModulev2(LightningModule): def __init__(self, cfg, classifier_cfg=None): super().__init__() self._print_logger = logging.getLogger(__name__) self._exp_cfg = cfg.experiment self._model_cfg = cfg.model self._data_cfg = cfg.data self._interpolant_cfg = cfg.interpolant # self._classf_cfg = classifier_cfg # Set-up vector field prediction model self.model = ProteinFlow(cfg.model) # Set-up interpolant self.interpolant = Interpolant(cfg.interpolant) # Classifier self.loaded_classifier = False # self.load_classifiers(self._classf_cfg) self._sample_write_dir = self._exp_cfg.checkpointer.dirpath os.makedirs(self._sample_write_dir, exist_ok=True) self.validation_epoch_metrics = [] self.validation_epoch_samples = [] self.save_hyperparameters() print(f"Model is initiated on GPU: {torch.cuda.current_device()}") def on_train_start(self): self._epoch_start_time = time.time() def on_train_epoch_end(self): epoch_time = (time.time() - self._epoch_start_time) / 60.0 self.log( 'train/epoch_time_minutes', epoch_time, on_step=False, on_epoch=True, prog_bar=False ) self._epoch_start_time = time.time() def model_step(self, noisy_batch: Any): training_cfg = self._exp_cfg.training loss_mask = noisy_batch['res_mask'] if training_cfg.min_plddt_mask is not None: plddt_mask = noisy_batch['res_plddt'] > training_cfg.min_plddt_mask loss_mask *= plddt_mask num_batch, num_res = loss_mask.shape # Ground truth labels gt_trans_1 = noisy_batch['trans_1'] gt_rotmats_1 = noisy_batch['rotmats_1'] rotmats_t = noisy_batch['rotmats_t'] gt_rot_vf = su.calc_rot_vf( rotmats_t, gt_rotmats_1.type(torch.float32)) gt_bb_atoms = all_atom.to_atom37(gt_trans_1, gt_rotmats_1)[:, :, :3] # Timestep used for normalization. t = noisy_batch['t'] norm_scale = 1 - torch.min( t[..., None], torch.tensor(training_cfg.t_normalize_clip)) # Model output predictions. model_output = self.model(noisy_batch) pred_trans_1 = model_output['pred_trans'] pred_rotmats_1 = model_output['pred_rotmats'] pred_rots_vf = su.calc_rot_vf(rotmats_t, pred_rotmats_1) # Backbone atom loss pred_bb_atoms = all_atom.to_atom37(pred_trans_1, pred_rotmats_1)[:, :, :3] gt_bb_atoms *= training_cfg.bb_atom_scale / norm_scale[..., None] pred_bb_atoms *= training_cfg.bb_atom_scale / norm_scale[..., None] loss_denom = torch.sum(loss_mask, dim=-1) * 3 bb_atom_loss = torch.sum( (gt_bb_atoms - pred_bb_atoms) ** 2 * loss_mask[..., None, None], dim=(-1, -2, -3) ) / loss_denom # Translation VF loss trans_error = (gt_trans_1 - pred_trans_1) / norm_scale * training_cfg.trans_scale trans_loss = training_cfg.translation_loss_weight * torch.sum( trans_error ** 2 * loss_mask[..., None], dim=(-1, -2) ) / loss_denom # Rotation VF loss rots_vf_error = (gt_rot_vf - pred_rots_vf) / norm_scale rots_vf_loss = training_cfg.rotation_loss_weights * torch.sum( rots_vf_error ** 2 * loss_mask[..., None], dim=(-1, -2) ) / loss_denom # Pairwise distance loss gt_flat_atoms = gt_bb_atoms.reshape([num_batch, num_res * 3, 3]) gt_pair_dists = torch.linalg.norm( gt_flat_atoms[:, :, None, :] - gt_flat_atoms[:, None, :, :], dim=-1) pred_flat_atoms = pred_bb_atoms.reshape([num_batch, num_res * 3, 3]) pred_pair_dists = torch.linalg.norm( pred_flat_atoms[:, :, None, :] - pred_flat_atoms[:, None, :, :], dim=-1) flat_loss_mask = torch.tile(loss_mask[:, :, None], (1, 1, 3)) flat_loss_mask = flat_loss_mask.reshape([num_batch, num_res * 3]) flat_res_mask = torch.tile(loss_mask[:, :, None], (1, 1, 3)) flat_res_mask = flat_res_mask.reshape([num_batch, num_res * 3]) gt_pair_dists = gt_pair_dists * flat_loss_mask[..., None] pred_pair_dists = pred_pair_dists * flat_loss_mask[..., None] pair_dist_mask = flat_loss_mask[..., None] * flat_res_mask[:, None, :] dist_mat_loss = torch.sum( (gt_pair_dists - pred_pair_dists) ** 2 * pair_dist_mask, dim=(1, 2)) dist_mat_loss /= (torch.sum(pair_dist_mask, dim=(1, 2)) - num_res) se3_vf_loss = trans_loss + rots_vf_loss auxiliary_loss = (bb_atom_loss + dist_mat_loss) * ( t[:, 0] > training_cfg.aux_loss_t_pass ) auxiliary_loss *= self._exp_cfg.training.aux_loss_weight se3_vf_loss += auxiliary_loss if torch.isnan(se3_vf_loss).any(): raise ValueError('NaN loss encountered') return { "bb_atom_loss": bb_atom_loss, "trans_loss": trans_loss, "dist_mat_loss": dist_mat_loss, "auxiliary_loss": auxiliary_loss, "rots_vf_loss": rots_vf_loss, "se3_vf_loss": se3_vf_loss } def validation_step(self, batch: Any, batch_idx: int): res_mask = batch['res_mask'] self.interpolant.set_device(res_mask.device) num_batch, num_res = res_mask.shape samples = self.interpolant.sample( num_batch, num_res, self.model, )[0][-1].numpy() batch_metrics = [] for i in range(num_batch): # Write out sample to PDB file final_pos = samples[i] saved_path = write_prot_to_pdb( final_pos, os.path.join( self._sample_write_dir, f'sample_{i}_idx_{batch_idx}_len_{num_res}.pdb'), no_indexing=True ) if isinstance(self.logger, WandbLogger): self.validation_epoch_samples.append( [saved_path, self.global_step, wandb.Molecule(saved_path)] ) mdtraj_metrics = metrics.calc_mdtraj_metrics(saved_path) ca_idx = rc.atom_order['CA'] ca_ca_metrics = metrics.calc_ca_ca_metrics(final_pos[:, ca_idx]) batch_metrics.append((mdtraj_metrics | ca_ca_metrics)) batch_metrics = pd.DataFrame(batch_metrics) self.validation_epoch_metrics.append(batch_metrics) def on_validation_epoch_end(self): if len(self.validation_epoch_samples) > 0: self.logger.log_table( key='valid/samples', columns=["sample_path", "global_step", "Protein"], data=self.validation_epoch_samples) self.validation_epoch_samples.clear() val_epoch_metrics = pd.concat(self.validation_epoch_metrics) for metric_name, metric_val in val_epoch_metrics.mean().to_dict().items(): self._log_scalar( f'valid/{metric_name}', metric_val, on_step=False, on_epoch=True, prog_bar=False, batch_size=len(val_epoch_metrics), ) self.validation_epoch_metrics.clear() def _log_scalar( self, key, value, on_step=True, on_epoch=False, prog_bar=True, batch_size=None, sync_dist=False, rank_zero_only=True ): if sync_dist and rank_zero_only: raise ValueError('Unable to sync dist when rank_zero_only=True') self.log( key, value, on_step=on_step, on_epoch=on_epoch, prog_bar=prog_bar, batch_size=batch_size, sync_dist=sync_dist, rank_zero_only=rank_zero_only ) def training_step(self, batch: Any, stage: int): self.stage = 'train' step_start_time = time.time() self.interpolant.set_device(batch['res_mask'].device) noisy_batch = self.interpolant.corrupt_batch(batch) if self._interpolant_cfg.self_condition and random.random() > 0.5: with torch.no_grad(): model_sc = self.model(noisy_batch) noisy_batch['trans_sc'] = model_sc['pred_trans'] batch_losses = self.model_step(noisy_batch) num_batch = batch_losses['bb_atom_loss'].shape[0] total_losses = { k: torch.mean(v) for k, v in batch_losses.items() } for k, v in total_losses.items(): self._log_scalar( f"train/{k}", v, prog_bar=False, batch_size=num_batch) # Losses to track. Stratified across t. t = torch.squeeze(noisy_batch['t']) self._log_scalar( "train/t", np.mean(to_numpy(t)), prog_bar=False, batch_size=num_batch) for loss_name, loss_dict in batch_losses.items(): stratified_losses = t_stratified_loss( t, loss_dict, loss_name=loss_name) for k, v in stratified_losses.items(): self._log_scalar( f"train/{k}", v, prog_bar=False, batch_size=num_batch) # Training throughput self._log_scalar( "train/length", batch['res_mask'].shape[1], prog_bar=False, batch_size=num_batch) self._log_scalar( "train/batch_size", num_batch, prog_bar=False) step_time = time.time() - step_start_time self._log_scalar( "train/examples_per_second", num_batch / step_time) train_loss = ( total_losses[self._exp_cfg.training.loss] + total_losses['auxiliary_loss'] ) self._log_scalar( "train/loss", train_loss, batch_size=num_batch) return train_loss def configure_optimizers(self): return torch.optim.AdamW( params=self.model.parameters(), **self._exp_cfg.optimizer ) def load_classifiers(self, cfg, requires_grad = True): self._classf_cfg = cfg self.cls_model = ClasfModule.load_from_checkpoint( checkpoint_path=self._classf_cfg.ckpt_path, map_location=f'cuda:{torch.cuda.current_device()}' ) self._pmpnn_dir = self._infer_cfg.pmpnn_dir #self.cls_model = ProtClassifier(self._classifier_cfg) #self.cls_model.load_state_dict(torch.load(self._classifier_cfg.ckpt_path)) #self.cls_model.eval() #self.cls_model.to(self.device) for param in self.cls_model.parameters(): param.requires_grad = requires_grad def load_folding_model(self): print(f"Current GPU of folding model is {torch.cuda.current_device()}") from transformers import AutoTokenizer,EsmForProteinFolding self._tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1") self._folding_model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True) self._folding_model = self._folding_model.to(f'cuda:{torch.cuda.current_device()}') self._folding_model.esm = self._folding_model.esm.half() """ self._folding_model = esm.pretrained.esmfold_v1() self._folding_model = self._folding_model.eval() self._folding_model = self._folding_model.to(f'cuda:{torch.cuda.current_device()}') """ def run_self_consistency( self, decoy_pdb_dir: str, reference_pdb_path: str, motif_mask = None, run_folding=True, ): device = f'cuda:{torch.cuda.current_device()}' # Run ProteinMPNN output_path = os.path.join(decoy_pdb_dir, "parsed_pdbs.jsonl") process = subprocess.Popen( [ "python", f"{self._pmpnn_dir}/helper_scripts/parse_multiple_chains.py", f"--input_path={decoy_pdb_dir}", f"--output_path={output_path}", ] ) _ = process.wait() num_tries = 0 ret = -1 pmpnn_args = [ "python", f"{self._pmpnn_dir}/protein_mpnn_run.py", "--out_folder", decoy_pdb_dir, "--jsonl_path", output_path, "--num_seq_per_target", str(self._samples_cfg.seq_per_sample), "--sampling_temp", "0.1", "--seed", str(self._infer_cfg.seed), "--batch_size", "1", ] pmpnn_args.append("--device") pmpnn_args.append(str(torch.cuda.current_device())) while ret < 0: try: process = subprocess.Popen( pmpnn_args, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT ) ret = process.wait() except Exception as e: num_tries += 1 self._log.info(f"Failed ProteinMPNN. Attempt {num_tries}/5 {e}") torch.cuda.empty_cache() if num_tries < 4: raise e mpnn_fasta_path = os.path.join( decoy_pdb_dir, "seqs", os.path.basename(reference_pdb_path).replace(".pdb", ".fa") ) if not run_folding: return mpnn_fasta_path # --- ESMFold and metrics part below (unchanged) --- mpnn_results = { "tm_score": [], "sample_path": [], "header": [], "sequence": [], "rmsd": [], } if motif_mask is not None: mpnn_results["motif_rmsd"] = [] esmf_dir = os.path.join(decoy_pdb_dir, "esmf") os.makedirs(esmf_dir, exist_ok=True) fasta_seqs = fasta.FastaFile.read(mpnn_fasta_path) sample_feats = du.parse_pdb_feats("sample", reference_pdb_path) for i, (header, string) in enumerate(fasta_seqs.items()): # Run ESMFold esmf_sample_path = os.path.join(esmf_dir, f"sample_{i}.pdb") _ = self.run_folding(string, esmf_sample_path) esmf_feats = du.parse_pdb_feats("folded_sample", esmf_sample_path) sample_seq = du.aatype_to_seq(sample_feats["aatype"]) # Calculate scTM and ESMFold outputs with reference _, tm_score = metrics.calc_tm_score( sample_feats["bb_positions"], esmf_feats["bb_positions"], sample_seq, sample_seq, ) rmsd = metrics.calc_aligned_rmsd( sample_feats["bb_positions"], esmf_feats["bb_positions"] ) if motif_mask is not None: sample_motif = sample_feats["bb_positions"][motif_mask] of_motif = esmf_feats["bb_positions"][motif_mask] motif_rmsd = metrics.calc_aligned_rmsd(sample_motif, of_motif) mpnn_results["motif_rmsd"].append(motif_rmsd) mpnn_results["rmsd"].append(rmsd) mpnn_results["tm_score"].append(tm_score) mpnn_results["sample_path"].append(esmf_sample_path) mpnn_results["header"].append(header) mpnn_results["sequence"].append(string) # Save results to CSV csv_path = os.path.join(decoy_pdb_dir, "sc_results.csv") mpnn_results = pd.DataFrame(mpnn_results) mpnn_results.to_csv(csv_path) def run_folding(self, sequence, save_path): def convert_outputs_to_pdb(outputs): from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein from transformers.models.esm.openfold_utils.feats import atom14_to_atom37 final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs) outputs = {k: v.to("cpu").numpy() for k, v in outputs.items()} final_atom_positions = final_atom_positions.cpu().numpy() final_atom_mask = outputs["atom37_atom_exists"] pdbs = [] for i in range(outputs["aatype"].shape[0]): aa = outputs["aatype"][i] pred_pos = final_atom_positions[i] mask = final_atom_mask[i] resid = outputs["residue_index"][i] + 1 pred = OFProtein( aatype=aa, atom_positions=pred_pos, atom_mask=mask, residue_index=resid, b_factors=outputs["plddt"][i], chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None, ) pdbs.append(to_pdb(pred)) return pdbs # Tokenize input tokenized_input = self._tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'] tokenized_input = tokenized_input.to(f'cuda:{torch.cuda.current_device()}') with torch.no_grad(): output = self._folding_model(tokenized_input) """ with torch.no_grad(): # print(sequence) output = self._folding_model.infer_pdb(sequence) """ output = convert_outputs_to_pdb(output) with open(save_path, "w") as f: f.write("".join(output)) return output def predict_step(self, batch, batch_idx): device = f'cuda:{torch.cuda.current_device()}' interpolant = Interpolant(self._infer_cfg.interpolant) interpolant.set_device(device) sample_length = batch['num_res'].item() diffuse_mask = torch.ones(1, sample_length) sample_id = batch['sample_id'].item() sample_dir = os.path.join( self._output_dir, f'length_{sample_length}', f'sample_{sample_id}') top_sample_csv_path = os.path.join(sample_dir, 'top_sample.csv') if os.path.exists(top_sample_csv_path): self._print_logger.info( f'Skipping instance {sample_id} length {sample_length}') return atom37_traj, model_traj, _ = interpolant.sample_clf( 1, sample_length, self.model, self.cls_model ) os.makedirs(sample_dir, exist_ok=True) bb_traj = to_numpy(torch.concat(atom37_traj, dim=0)) traj_paths = save_traj( bb_traj[-1], bb_traj, np.flip(to_numpy(torch.concat(model_traj, dim=0)), axis=0), to_numpy(diffuse_mask)[0], output_dir=sample_dir, ) # Run ProteinMPNN pdb_path = traj_paths["sample_path"] sc_output_dir = os.path.join(sample_dir, "self_consistency") os.makedirs(sc_output_dir, exist_ok=True) shutil.copy( pdb_path, os.path.join(sc_output_dir, os.path.basename(pdb_path)) ) # Run self consistency _ = self.run_self_consistency(sc_output_dir, pdb_path, motif_mask=None) def evaluate_structure_quality(self, pdb_path, reference_pdb_path=None, fixed_residues=None): """Evaluate the quality of a generated protein structure. Args: pdb_path: Path to the PDB file to evaluate reference_pdb_path: Optional path to a reference structure for comparison fixed_residues: List of residue indices that were fixed (PDB numbering) Returns: Dictionary containing various quality metrics: - Basic geometric validation (bond lengths, angles) - Secondary structure content - Ramachandran plot statistics - RMSD and TM-score to reference if provided - RMSD of fixed residues if provided """ quality_metrics = {} # Calculate basic MDTraj metrics mdtraj_metrics = metrics.calc_mdtraj_metrics(pdb_path) quality_metrics.update(mdtraj_metrics) # Calculate CA-CA metrics sample_feats = du.parse_pdb_feats("sample", pdb_path) ca_idx = rc.atom_order['CA'] # Debug logging logger = logging.getLogger(__name__) logger.info(f"PDB features keys: {sample_feats.keys()}") logger.info(f"bb_positions shape: {sample_feats['bb_positions'].shape}") logger.info(f"bb_positions type: {type(sample_feats['bb_positions'])}") # Ensure we have the correct shape for CA positions [N, 3] ca_positions = sample_feats["bb_positions"] if len(ca_positions.shape) == 1: # If we got a flattened array, reshape it ca_positions = ca_positions.reshape(-1, 3) elif len(ca_positions.shape) > 2: # If we got extra dimensions, take just the CA positions ca_positions = ca_positions[:, ca_idx] # Convert to numpy if needed if isinstance(ca_positions, torch.Tensor): ca_positions = ca_positions.detach().cpu().numpy() logger.info(f"Final ca_positions shape: {ca_positions.shape}") # Calculate CA-CA metrics ca_ca_metrics = metrics.calc_ca_ca_metrics(ca_positions) quality_metrics.update(ca_ca_metrics) # Compare to reference structure if provided if reference_pdb_path is not None: ref_feats = du.parse_pdb_feats("reference", reference_pdb_path) ref_seq = du.aatype_to_seq(ref_feats["aatype"]) sample_seq = du.aatype_to_seq(sample_feats["aatype"]) # Calculate TM-score _, tm_score = metrics.calc_tm_score( sample_feats["bb_positions"], ref_feats["bb_positions"], sample_seq, ref_seq ) quality_metrics["tm_score"] = tm_score # Calculate RMSD rmsd = metrics.calc_aligned_rmsd( sample_feats["bb_positions"], ref_feats["bb_positions"] ) quality_metrics["rmsd_to_ref"] = rmsd # Calculate RMSD for fixed residues if provided if fixed_residues is not None: # Map PDB residue numbers to indices in residue_index array sample_res_indices = sample_feats['residue_index'] ref_res_indices = ref_feats['residue_index'] sample_fixed_indices = [] ref_fixed_indices = [] for resnum in fixed_residues: sample_matches = np.where(sample_res_indices == resnum)[0] ref_matches = np.where(ref_res_indices == resnum)[0] if len(sample_matches) == 0 or len(ref_matches) == 0: logger.warning(f"Residue number {resnum} not found in sample or reference residue_indices.") continue sample_fixed_indices.append(sample_matches[0]) ref_fixed_indices.append(ref_matches[0]) if sample_fixed_indices and ref_fixed_indices: sample_fixed = sample_feats["bb_positions"][sample_fixed_indices] ref_fixed = ref_feats["bb_positions"][ref_fixed_indices] # Calculate RMSD for fixed residues fixed_rmsd = metrics.calc_aligned_rmsd(sample_fixed, ref_fixed) quality_metrics["fixed_residues_rmsd"] = fixed_rmsd logger.info(f"RMSD for fixed residues {fixed_residues}: {fixed_rmsd:.3f} Å") else: logger.warning(f"No valid fixed residue indices found for RMSD calculation.") return quality_metrics def analyze_sample_diversity(self, sample_pdbs, reference_pdb=None): """Analyze the diversity of a set of generated protein samples. Args: sample_pdbs: List of paths to PDB files of generated samples reference_pdb: Optional path to reference structure Returns: Dictionary containing diversity metrics: - Pairwise RMSD statistics between samples - RMSD to reference if provided - Structure clustering analysis - Secondary structure diversity """ diversity_metrics = {} # Calculate all pairwise RMSDs between samples num_samples = len(sample_pdbs) pairwise_rmsds = [] for i in range(num_samples): sample_i_feats = du.parse_pdb_feats("sample_i", sample_pdbs[i]) for j in range(i+1, num_samples): sample_j_feats = du.parse_pdb_feats("sample_j", sample_pdbs[j]) rmsd = metrics.calc_aligned_rmsd( sample_i_feats["bb_positions"], sample_j_feats["bb_positions"] ) pairwise_rmsds.append(rmsd) diversity_metrics["mean_pairwise_rmsd"] = np.mean(pairwise_rmsds) diversity_metrics["std_pairwise_rmsd"] = np.std(pairwise_rmsds) diversity_metrics["min_pairwise_rmsd"] = np.min(pairwise_rmsds) diversity_metrics["max_pairwise_rmsd"] = np.max(pairwise_rmsds) # Compare all samples to reference if provided if reference_pdb is not None: ref_feats = du.parse_pdb_feats("reference", reference_pdb) ref_seq = du.aatype_to_seq(ref_feats["aatype"]) ref_rmsds = [] ref_tm_scores = [] for sample_pdb in sample_pdbs: sample_feats = du.parse_pdb_feats("sample", sample_pdb) sample_seq = du.aatype_to_seq(sample_feats["aatype"]) # Calculate RMSD to reference rmsd = metrics.calc_aligned_rmsd( sample_feats["bb_positions"], ref_feats["bb_positions"] ) ref_rmsds.append(rmsd) # Calculate TM-score to reference _, tm_score = metrics.calc_tm_score( sample_feats["bb_positions"], ref_feats["bb_positions"], sample_seq, ref_seq ) ref_tm_scores.append(tm_score) diversity_metrics["mean_rmsd_to_ref"] = np.mean(ref_rmsds) diversity_metrics["std_rmsd_to_ref"] = np.std(ref_rmsds) diversity_metrics["mean_tm_score_to_ref"] = np.mean(ref_tm_scores) diversity_metrics["std_tm_score_to_ref"] = np.std(ref_tm_scores) return diversity_metrics def prepare_conditional_inputs(self, pdb_path, fixed_residues=None, chain_id='A'): """Prepare inputs for conditional sampling from a PDB file. Args: pdb_path: Path to the PDB file containing the partial/reference structure fixed_residues: List of residue indices to fix (PDB numbering, e.g. 628, 629, ...) chain_id: Chain ID to use from the PDB file (default='A') Returns: Dictionary containing: - fixed_positions: [N, 3] tensor of fixed atom positions - fixed_mask: [N] boolean mask indicating which positions are fixed - num_res: Total number of residues - residue_indices: Original residue indices from PDB file """ # Parse PDB features, excluding HETATM entries pdb_feats = du.parse_pdb_feats("reference", pdb_path, exclude_hetatm=True) # Get number of residues num_res = pdb_feats["aatype"].shape[0] device = next(self.parameters()).device # Debug logging logger = logging.getLogger(__name__) logger.info(f"PDB features keys: {pdb_feats.keys()}") # Get backbone positions - bb_positions should already be CA positions from parse_chain_feats if 'bb_positions' not in pdb_feats: raise ValueError("No backbone positions found in PDB features") fixed_positions = pdb_feats['bb_positions'] # Already [N, 3] for CA atoms # Get original residue indices from PDB residue_indices = pdb_feats['residue_index'] # Create fixed mask if fixed_residues is None: # Fix all residues fixed_mask = torch.ones(num_res, dtype=torch.bool, device=device) else: # Map PDB residue numbers to indices in residue_indices array fixed_indices = [] for resnum in fixed_residues: matches = np.where(residue_indices == resnum)[0] if len(matches) == 0: raise ValueError(f"Residue number {resnum} not found in PDB residue_indices: {residue_indices}") fixed_indices.append(matches[0]) fixed_mask = torch.zeros(num_res, dtype=torch.bool, device=device) fixed_mask[fixed_indices] = True # Convert positions to tensor and ensure correct shape [N, 3] fixed_positions = torch.tensor(fixed_positions, dtype=torch.float32, device=device) # Add debug information logger.info(f"fixed_positions shape: {fixed_positions.shape}") logger.info(f"fixed_mask shape: {fixed_mask.shape}") if len(fixed_positions.shape) != 2 or fixed_positions.shape[1] != 3: raise ValueError(f"Expected fixed_positions shape [N, 3], got {fixed_positions.shape}") return { 'fixed_positions': fixed_positions, # Shape: [N, 3] 'fixed_mask': fixed_mask, # Shape: [N] 'num_res': num_res, 'residue_indices': residue_indices } def sample_with_fixed_residues( self, pdb_path, fixed_residues=None, num_samples=1, temperature=1.0, chain_id='A', output_dir=None, clf_model=None, guidance_scale=0.2, target_class=1 ): """Generate protein samples while keeping specified residues fixed. Args: pdb_path: Path to PDB file with reference structure fixed_residues: List of residue indices to fix (1-indexed) num_samples: Number of samples to generate temperature: Temperature for sampling diversity chain_id: Chain ID to use from PDB output_dir: Directory to save samples (if None, uses self._sample_write_dir) clf_model: Optional classifier model for guidance guidance_scale: Scale factor for classifier guidance (default=0.2) target_class: Target class for classifier guidance (default=1) Returns: List of paths to generated PDB files """ # Prepare inputs inputs = self.prepare_conditional_inputs( pdb_path, fixed_residues=fixed_residues, chain_id=chain_id ) # Initialize interpolant with current device device = next(self.parameters()).device self.interpolant.set_device(device) # Prepare inputs with correct dimensions # Add batch dimension to fixed_positions: [N, 3] -> [num_samples, N, 3] fixed_positions = inputs['fixed_positions'].unsqueeze(0).expand(num_samples, -1, -1) # Convert fixed_mask to float and add necessary dimensions fixed_mask = inputs['fixed_mask'].to(device) # Create additional model inputs batch = { 'res_mask': torch.ones(num_samples, inputs['num_res'], device=device), 'flow_mask': ~fixed_mask, # Only flow non-fixed positions 'fixed_positions': fixed_positions, 'fixed_mask': fixed_mask, } # Run conditional sampling with optional classifier guidance atom37_traj, clean_atom37_traj, clean_traj = self.interpolant.sample_conditional( num_batch=num_samples, num_res=inputs['num_res'], model=self.model, fixed_positions=fixed_positions, # Shape: [num_samples, N, 3] fixed_mask=fixed_mask, # Shape: [N] temperature=temperature, clf_model=clf_model, guidance_scale=guidance_scale, target_class=target_class ) # Save samples sample_paths = [] save_dir = output_dir if output_dir is not None else self._sample_write_dir os.makedirs(save_dir, exist_ok=True) # Get original residue indices from input PDB parser = PDB.PDBParser(QUIET=True) structure = parser.get_structure('ref', pdb_path) residue_indices = np.array([residue.id[1] for residue in structure[0][chain_id]]) for i in range(num_samples): sample_path = os.path.join( save_dir, f'conditional_sample_{i}.pdb' ) # Set b-factors to indicate fixed residues b_factors = torch.zeros((inputs['num_res'], 37), device=device) b_factors[fixed_mask] = 100.0 # Convert tensors to numpy arrays before writing sample_coords = clean_atom37_traj[i].detach().cpu().numpy() b_factors = b_factors.detach().cpu().numpy() # Ensure sample_coords has shape [N, 37, 3] for a single model if len(sample_coords.shape) == 4: # If shape is [1, N, 37, 3] sample_coords = sample_coords[0] # Take first (and only) model # Create atom mask atom37_mask = np.sum(np.abs(sample_coords), axis=-1) > 1e-7 # Create protein object with original residue indices prot = create_full_prot( sample_coords, atom37_mask, b_factors=b_factors, residue_indices=residue_indices ) # Write protein to PDB pdb_str = protein.to_pdb(prot, model=1, add_end=False) with open(sample_path, 'w') as f: f.write(pdb_str) f.write('END\n') sample_paths.append(sample_path) return sample_paths