Dataset for sampling.
| 53 | |
| 54 | |
| 55 | class SamplingDataset(torch.utils.data.Dataset): |
| 56 | """Dataset for sampling.""" |
| 57 | |
| 58 | def __init__(self, cfg: DictConfig): |
| 59 | """Initializes the SamplingDataset.""" |
| 60 | if cfg.sampling_task == "batched_structure_sampling": |
| 61 | if cfg.csv_path is not None: |
| 62 | # handle variable CSV inputs |
| 63 | df_rows = [] |
| 64 | self.df = pd.read_csv(cfg.csv_path) |
| 65 | for _, row in self.df.iterrows(): |
| 66 | sample_id = row.id |
| 67 | input_receptor = row.input_receptor |
| 68 | input_ligand = row.input_ligand |
| 69 | input_template = row.input_template |
| 70 | assert input_receptor is not None, "Receptor path is required for sampling." |
| 71 | if input_ligand is not None: |
| 72 | if input_ligand.endswith(".sdf"): |
| 73 | ligand_paths = create_temp_ligand_frag_files(input_ligand) |
| 74 | else: |
| 75 | ligand_paths = list(input_ligand.split("|")) |
| 76 | else: |
| 77 | ligand_paths = None # handle `null` ligand input |
| 78 | if not input_receptor.endswith(".pdb"): |
| 79 | log.warning( |
| 80 | "Assuming the provided receptor input is a protein sequence. Creating a dummy PDB file." |
| 81 | ) |
| 82 | create_full_pdb_with_zero_coordinates( |
| 83 | input_receptor, os.path.join(cfg.out_path, f"input_{sample_id}.pdb") |
| 84 | ) |
| 85 | input_receptor = os.path.join(cfg.out_path, f"input_{sample_id}.pdb") |
| 86 | df_row = { |
| 87 | "sample_id": sample_id, |
| 88 | "rec_path": input_receptor, |
| 89 | "lig_paths": ligand_paths, |
| 90 | } |
| 91 | if input_template is not None: |
| 92 | df_row["input_template"] = input_template |
| 93 | df_rows.append(df_row) |
| 94 | self.df = pd.DataFrame(df_rows) |
| 95 | else: |
| 96 | sample_id = cfg.sample_id |
| 97 | input_receptor = cfg.input_receptor |
| 98 | input_ligand = cfg.input_ligand |
| 99 | if input_ligand is not None: |
| 100 | if input_ligand.endswith(".sdf"): |
| 101 | ligand_paths = create_temp_ligand_frag_files(input_ligand) |
| 102 | else: |
| 103 | ligand_paths = list(input_ligand.split("|")) |
| 104 | else: |
| 105 | ligand_paths = None # handle `null` ligand input |
| 106 | if not input_receptor.endswith(".pdb"): |
| 107 | log.warning( |
| 108 | "Assuming the provided receptor input is a protein sequence. Creating a dummy PDB file." |
| 109 | ) |
| 110 | create_full_pdb_with_zero_coordinates( |
| 111 | input_receptor, os.path.join(cfg.out_path, "input.pdb") |
| 112 | ) |