import os import sys from collections import defaultdict from pathlib import Path from typing import Dict, List, Tuple, Union import numpy as np import argparse import json import hydra import lightning as L import torch from dotenv import load_dotenv from loguru import logger from torch.utils.data import DataLoader from .datasets.gen_dataset import GenDataset from .proteina import Proteina from .partial_autoencoder.autoencoder import AutoEncoder from .utils.pdb_utils import write_prot_to_pdb from lightning.pytorch.utilities import rank_zero_only rank = int(os.environ.get('LOCAL_RANK', 0)) ## 2026 UPDATE CONFIG_ROOT = Path(__file__).resolve().parents[1] / "configs" def parse_args_and_cfg() -> Tuple[Dict, Dict, str]: """ Parses command line arguments and loads the corresponding config file. Returns: Command line arguments (dict) Config file (dict) config_name (string) """ parser = argparse.ArgumentParser(description="Job info") parser.add_argument( "--config_name", type=str, default="inference_base", help="Name of the config yaml file.", ) parser.add_argument( "--config_number", type=int, default=-1, help="Number of the config yaml file." ) parser.add_argument( "--job_id", type=int, default=0, help="Job id for this config to determine which split to use.", ) parser.add_argument( "--config_subdir", type=str, help="(Optional) Name of directory with config files, if not included uses base inference config.\ Likely only used when submitting to the cluster with script.", ) parser.add_argument( "--data_path", type=str, help="Name of the data path", ) args = parser.parse_args() if args.data_path is not None: os.environ["DATA_PATH"] = args.data_path # Inference config # If config_subdir is None then use base inference config # Otherwise use config_subdir/some_config if args.config_subdir is None: config_path = str(CONFIG_ROOT) else: config_path = str(CONFIG_ROOT / args.config_subdir) with hydra.initialize_config_dir(config_dir=config_path, version_base=hydra.__version__): # If number provided use it, otherwise name if args.config_number != -1: config_name = f"inf_{args.config_number}" else: config_name = args.config_name cfg = hydra.compose(config_name=config_name) logger.info(f"Inference config {cfg}") return args, cfg, config_name def setup( cfg: Dict, create_root: bool = True, config_name: str = ".", job_id: int = 0 ) -> str: """ Checks if metrics being computed are compatible, sets the right seed, and creates the root directory where the run will store things. Returns: Path of the root directory (string) """ logger.info(" ".join(sys.argv)) assert ( torch.cuda.is_available() ), "CUDA not available" # Needed for ESMfold and designability logger.add( sys.stdout, format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}", ) # Send to stdout assert ( not ( cfg.generation.metric.compute_designability or cfg.generation.metric.compute_novelty_pdb or cfg.generation.metric.compute_novelty_afdb ) or not cfg.generation.metric.compute_fid ), "Designability/Novelty cannot be computed together with FID" # Set root path for this inference run if "motif_task_name" in cfg.generation.dataset: root_path = ( f"./inference/{config_name}_{cfg.generation.dataset.motif_task_name}" ) else: root_path = f"./inference/{config_name}" if create_root: os.makedirs(root_path, exist_ok=True) else: if not os.path.exists(root_path): raise ValueError("Results path %s does not exist" % root_path) # Set seed cfg.seed = cfg.seed + job_id # Different seeds for different splits ids logger.info(f"Seeding everything to seed {cfg.seed}") L.seed_everything(cfg.seed) return root_path def check_cfg_validity(cfg_data: Dict, cfg_sample_args: Dict) -> None: """ Checks if guidance arguments (CFG and AG) are valid. """ # Logging CFG if cfg_sample_args.guidance_w != 1.0: logger.info( f"Guidance is turned on with guidance weight {cfg_sample_args.guidance_w} and autoguidance ratio {cfg_sample_args.ag_ratio}." ) assert ( cfg_sample_args.ag_ratio >= 0.0 and cfg_sample_args.ag_ratio <= 1.0 ), f"Autoguidance ratio should be between 0 and 1, but now is {cfg_sample_args.ag_ratio}." assert (cfg_sample_args.ag_ratio == 0.0) or ( cfg_sample_args.ag_ckpt_path is not None ), f"Autoguidance checkpoint path should be provided" else: logger.info(f"Guidance is turned off.") # Logging conditional generation if cfg_sample_args.fold_cond: logger.info("Conditional generation is turned on.") assert ( cfg_data.empirical_distribution_cfg.len_cath_code_path is not None ), "Empirical (len, cath_code) distribution file should be provided when using conditional generation." else: logger.info("Conditional generation is turned off.") assert ( cfg_data.empirical_distribution_cfg.len_cath_code_path is None ), "Empirical (len, cath_code) distribution file shouldn't be provided when using unconditional generation." def load_ag_ckpt(cfg: Dict) -> Union[None, torch.nn.Module]: """ Loads the neural network for the "bad" checkpoint in autoguidance, if requested. Returns: A nn module, if autogudance enabled. """ nn_ag = None if cfg.ag_ratio > 0 and cfg.guidance_w != 1.0: logger.info( f"Using autoguidance with guidance weight {cfg.guidance_w} and autoguidance ratio {cfg.ag_ratio} based on the checkpoint {cfg.ag_ckpt_path}" ) ckpt_ag_file = cfg.ag_ckpt_path assert os.path.exists(ckpt_ag_file), f"Not a valid checkpoint {ckpt_ag_file}" model_ag = Proteina.load_from_checkpoint(ckpt_ag_file, strict=False) # OPTIMIZATION: Remove encoder from autoguidance model autoencoder during generation (only decoder needed) if model_ag.autoencoder is not None: logger.info( "Removing autoencoder encoder from autoguidance model during generation to save memory" ) del model_ag.autoencoder.encoder model_ag.autoencoder.encoder = None nn_ag = model_ag.nn return nn_ag def load_ckpt_n_configure_inference(cfg: Dict) -> Proteina: """ Loads the model, potentially the autoguidance checkpoint as well, if requested. Returns: Model (Proteina) """ # Load model from checkpoint ckpt_path = cfg.ckpt_path ckpt_file = os.path.join(ckpt_path, cfg.ckpt_name) logger.info(f"Using checkpoint {ckpt_file}") assert os.path.exists(ckpt_file), f"Not a valid checkpoint {ckpt_file}" model = Proteina.load_from_checkpoint(ckpt_file, strict=False, autoencoder_ckpt_path=cfg.get("autoencoder_ckpt_path", None)) # Set inference variables and potentially load autoguidance nn_ag = load_ag_ckpt(cfg.generation.args) model.configure_inference(cfg.generation, nn_ag=nn_ag) return model def split_by_job(cfg: Dict, job_id: int, njobs: int) -> Dict: """ Since generation may be split across multiple jobs, this function determines how many samples are produced per job. Then, it sets the right value in the config dict, and returns the updated config. Returns: Config updated with the correct number of samples to generate. """ nsamples = cfg.dataset.nsamples nsamples_per_split = (nsamples - 1) // njobs + 1 if nsamples_per_split * job_id >= nsamples: logger.info(f"Job id {job_id} get 0 samples. Finishing job...") exit(0) else: cfg.dataset.nsamples = min( nsamples_per_split, nsamples - nsamples_per_split * job_id ) return cfg def binder_split_by_job(cfg: Dict, job_id: int, njobs: int) -> Dict: """ Since generation may be split across multiple jobs, this function determines how many samples are produced per job. Then, it sets the right value in the config dict, and returns the updated config. Returns: Config updated with the correct number of samples to generate. """ nsamples = cfg.dataset.nlens_cfg.random_lens[2] nsamples_per_split = (nsamples - 1) // njobs + 1 if nsamples_per_split * job_id >= nsamples: logger.info(f"Job id {job_id} get 0 samples. Finishing job...") exit(0) else: cfg.dataset.nlens_cfg.random_lens[2] = min( nsamples_per_split, nsamples - nsamples_per_split * job_id ) return cfg def save_predictions( root_path: str, predictions: List[List[Tuple[torch.tensor]]], job_id: int = 0, chain_indexes: np.ndarray = None, cath_codes: List[List[List[str]]] = None, ) -> None: """ Saves generated samples. Args: root_path: root directory where samples will be stored (within subdirectories)/ predictions: List of lists of tuples. Each tuple represents a sample, has to components (coors [n, 37, 3], aatype [n]) job_id: job number, used to store files. chain_indexes: chain indexes for each sample, used to store files. cath_codes: conditional sampling... """ predictions = [sample for sublist in predictions for sample in sublist] # List[tuple] where each tuple is (coors [n, 37, 3], aatype [n]) samples_per_length = defaultdict(int) for j, pred in enumerate(predictions): coors_atom37, residue_type = pred # [n, 37, 3] and [n] n = coors_atom37.shape[-3] if chain_indexes: chain_index = chain_indexes[j].numpy() else: chain_index = None # Create directory where everything related to this sample will be stored suffix = "" # LL 2026 UPDATE #dir_name = f"job_{job_id}_n_{n}_id_{samples_per_length[n]}{suffix}" dir_name = f"job_{job_id}_n_{n}_id_{samples_per_length[n]}{suffix}_rank{rank}" samples_per_length[n] += 1 sample_root_path = os.path.join( root_path, dir_name ) os.makedirs(sample_root_path, exist_ok=False) # Save generated structure as pdb fname = dir_name + ".pdb" pdb_path = os.path.join(sample_root_path, fname) write_prot_to_pdb( prot_pos=coors_atom37.float().detach().cpu().numpy(), aatype=residue_type.detach().cpu().numpy(), file_path=pdb_path, chain_index=chain_index, overwrite=True, no_indexing=True, ) def save_motif_predictions( root_path: str, predictions: List[List[Tuple[torch.tensor]]], job_id: int = 0, motif_pdb_name: str = None, ) -> None: predictions = [sample for sublist in predictions for sample in sublist] print([(p[0].shape, p[1].shape) for p in predictions]) samples_per_length = defaultdict(int) for j, pred in enumerate(predictions): coors_atom37, residue_type = pred # [n, 37, 3] and [n] n = coors_atom37.shape[-3] dir_name = f"job_{job_id}_id_{j}_motif_{motif_pdb_name}" samples_per_length[n] += 1 sample_root_path = os.path.join(root_path, dir_name) os.makedirs(sample_root_path, exist_ok=False) fname = dir_name + ".pdb" pdb_path = os.path.join(sample_root_path, fname) write_prot_to_pdb( prot_pos=coors_atom37.float().detach().cpu().numpy(), aatype=residue_type.detach().cpu().numpy(), file_path=pdb_path, overwrite=True, no_indexing=True, ) def main(): load_dotenv() # Parse arguments, load appropriate config, and set up root path args, cfg, config_name = parse_args_and_cfg() # cfg.run_name_ motif_cond = cfg.generation.args.get("motif_cond", False) target_cond = cfg.generation.args.get("target_cond", False) cfg.generation.args.get("multi_cond", False) cfg.generation.args.get("fold_cond", False) njobs = cfg.get("gen_njobs", 1) root_path = setup( cfg, create_root=True, config_name=config_name, job_id=args.job_id ) # Exit if results from analysis already exist (assumes samples already there) # File to store analysis (next step, this is generate) results csv_filename = f"results_{config_name}_{args.job_id}.csv" csv_path = os.path.join(root_path, "..", csv_filename) # Exit if results from analysis already exist if os.path.exists(csv_path): logger.info(f"Results already exist at {csv_path}. Exiting generate.py.") sys.exit(0) cfg_gen = cfg.generation check_cfg_validity(cfg_gen.dataset, cfg_gen.args) # Load model model = load_ckpt_n_configure_inference(cfg) # Create generation dataset cfg_gen = split_by_job(cfg_gen, args.job_id, njobs) # Motif-specific dataset creation if motif_cond or ("motif_task_name" in cfg.generation.dataset): motif_csv_path = os.path.join( root_path, f"{cfg_gen.dataset.get('motif_task_name', 'motif')}_{args.job_id}_motif_info.csv", ) """ Motif Configuration Examples: The motif dataset supports two modes for specifying which atoms to include: 1. **Atom-level specification** (precise control): motif_dict_cfg: my_motif: motif_pdb_path: "path/to/motif.pdb" motif_atom_spec: "A64: [O, CG]; A65: [N, CA]; A66: [CB, CD]" # atom_selection_mode is ignored when motif_atom_spec is provided 2. **Residue/range-based specification** (automatic atom selection): motif_dict_cfg: my_motif: motif_pdb_path: "path/to/motif.pdb" contig_string: "A1-7/A28-79" atom_selection_mode: "tip_atoms" # NEW: Choose atom selection mode Available atom_selection_mode options: - "ca_only": Only CA atoms (default, fastest) - "all": All available atoms (most complete motif) - "backbone": Backbone atoms only (N, CA, C, O) - "sidechain": Sidechain atoms only - "tip_atoms": Tip atoms of sidechains (e.g., OH for Ser, NH2 for Arg) - "random": Random subset of available atoms If atom_selection_mode is not specified, defaults to "ca_only" for backward compatibility. """ dataset = GenDataset(motif_csv_path=motif_csv_path, **cfg_gen.dataset) else: dataset = GenDataset(**cfg_gen.dataset) dataloader = DataLoader(dataset, batch_size=1, shuffle=False) # Sample model #trainer = L.Trainer(accelerator="gpu", devices=1) ##LL 2026 UPDATE trainer = L.Trainer(accelerator="gpu", devices=-1) predictions = trainer.predict(model, dataloader) chain_indexes = None if motif_cond or ("motif_task_name" in cfg.generation.dataset): save_motif_predictions( root_path, predictions, job_id=args.job_id, motif_pdb_name=cfg_gen.dataset.get("motif_task_name", None), ) import shutil motif_csv = f"./{cfg_gen.dataset.get('motif_task_name', '')}_motif_info.csv" if os.path.exists(motif_csv): shutil.copy(motif_csv, root_path) else: save_predictions( root_path, predictions, job_id=args.job_id, chain_indexes=chain_indexes, cath_codes=dataset.cath_codes, ) if __name__ == "__main__": main()