Samples using given checkpoint on a datamodule predictset. This method is wrapped in optional @task_wrapper decorator, that controls the behavior during failure. Useful for multiruns, saving info about the crash, etc. :param cfg: DictConfig configuration composed by Hydra. :return:
(cfg: DictConfig)
| 136 | |
| 137 | @task_wrapper |
| 138 | def sample(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: |
| 139 | """Samples using given checkpoint on a datamodule predictset. |
| 140 | |
| 141 | This method is wrapped in optional @task_wrapper decorator, that controls the behavior during |
| 142 | failure. Useful for multiruns, saving info about the crash, etc. |
| 143 | |
| 144 | :param cfg: DictConfig configuration composed by Hydra. |
| 145 | :return: Tuple[dict, dict] with metrics and dict with all instantiated objects. |
| 146 | """ |
| 147 | assert cfg.ckpt_path, "Please provide a checkpoint path with which to sample!" |
| 148 | assert os.path.exists(cfg.ckpt_path), f"Checkpoint path {cfg.ckpt_path} does not exist!" |
| 149 | assert ( |
| 150 | cfg.sampling_task in AVAILABLE_SAMPLING_TASKS |
| 151 | ), f"Sampling task {cfg.sampling_task} is not one of the following available tasks: {AVAILABLE_SAMPLING_TASKS}." |
| 152 | assert (cfg.input_receptor is not None and cfg.input_ligand is not None) or ( |
| 153 | cfg.csv_path is not None and os.path.exists(cfg.csv_path) |
| 154 | ), "Please provide either an input receptor and ligand or a CSV file with receptor and ligand sequences/filepaths." |
| 155 | |
| 156 | # set seed for random number generators in pytorch, numpy and python.random |
| 157 | if cfg.get("seed"): |
| 158 | L.seed_everything(cfg.seed, workers=True) |
| 159 | |
| 160 | log.info( |
| 161 | f"Setting `float32_matmul_precision` to {cfg.model.cfg.task.float32_matmul_precision}." |
| 162 | ) |
| 163 | torch.set_float32_matmul_precision(precision=cfg.model.cfg.task.float32_matmul_precision) |
| 164 | |
| 165 | # Establish model input arguments |
| 166 | with open_dict(cfg): |
| 167 | # NOTE: Structure trajectories will not be visualized when performing auxiliary estimation only |
| 168 | cfg.model.cfg.prior_type = cfg.prior_type |
| 169 | cfg.model.cfg.task.detect_covalent = cfg.detect_covalent |
| 170 | cfg.model.cfg.task.use_template = cfg.use_template |
| 171 | cfg.model.cfg.task.csv_path = cfg.csv_path |
| 172 | cfg.model.cfg.task.input_receptor = cfg.input_receptor |
| 173 | cfg.model.cfg.task.input_ligand = cfg.input_ligand |
| 174 | cfg.model.cfg.task.input_template = cfg.input_template |
| 175 | cfg.model.cfg.task.visualize_generated_samples = ( |
| 176 | cfg.visualize_sample_trajectories and not cfg.auxiliary_estimation_only |
| 177 | ) |
| 178 | cfg.model.cfg.task.auxiliary_estimation_only = cfg.auxiliary_estimation_only |
| 179 | if cfg.latent_model is not None: |
| 180 | with open_dict(cfg): |
| 181 | cfg.model.cfg.latent_model = cfg.latent_model |
| 182 | with open_dict(cfg): |
| 183 | if cfg.start_time == "auto": |
| 184 | cfg.start_time = 1.0 |
| 185 | else: |
| 186 | cfg.start_time = float(cfg.start_time) |
| 187 | |
| 188 | log.info("Converting sampling inputs into a <SamplingDataset>") |
| 189 | dataloaders: List[DataLoader] = [ |
| 190 | DataLoader( |
| 191 | SamplingDataset(cfg), |
| 192 | batch_size=1, |
| 193 | shuffle=False, |
| 194 | num_workers=0, |
| 195 | pin_memory=False, |
no test coverage detected