starling.inference.generation.generate_backend
- generate_backend(sequence_dict, conformations, device, steps, sampler, return_structures, batch_size, num_cpus_mds, num_mds_init, output_directory, return_data, verbose, show_progress_bar, show_per_step_progress_bar, pdb_trajectory, model_manager=<starling.inference.model_loading.ModelManager object>, ionic_strength=150, constraint=None, encoder_path=None, ddpm_path=None)[source]
Backend function for generating the distance maps using STARLING.
NOTE - this function does VERY littel sanity checking; to actually perform predictions use starling.frontend.ensemble_generation. This is NOT a user facing function!
- Parameters:
sequence_dict (dict) – A dictionary with the sequence names as the key and the sequences as the values. These names will be used to write any output files (if writing is requested).
ddpm (str) – The path to the DDPM model
device (str) – The device to use for predictions.
steps (int) – The number of steps to run the DDPM model.
ddim (bool) – Whether to use DDIM for sampling.
return_structures (bool) – Whether to return the 3D structure.
batch_size (int) – The batch size to use for sampling.
num_cpus_mds (int) – The number of CPUs to use for MDS. There is no point specifying more than the default number of MDS runs performed (defined in configs)
output_directory (str or None) – If None, no output is saved. If not None, will save the output to the specified path. This includes the distance maps and if return_structures=True, the 3D structures. The distance maps are saved as .npy files with the names <sequence_name>_STARLING_DM.npy and the structures are save with the file names <sequence_name>_STARLING.xtc and <sequence_name>_STARLING.pdb.
return_data (bool) – If True, will return the distance maps and structures (if generated) as a dictionary regardless of the output_directory. If False, will return None. Note the reason to set this to None is if you’re predicting a large set of sequences this will save memory.
verbose (bool) – Whether to print verbose output. Default is False.
show_progress_bar (bool) – Whether to show a progress bar. Default is True.
show_per_step_progress_bar (bool, optional) – whether to show progress bar per step.
pdb_trajectory (bool) – Whether to save the trajectory as a PDB file. Default is False.
model_manager (ModelManager) – A ModelManager object to manage loaded models. This lets us avoid loading the model iteratively when calling generate multiple times in a single session. Default is model_manager, which is initialized outside of this function code block. To update the path to the models, update the paths in config.py, which are read into the ModelManager object located the model_loading.py
encoder_path (str, optional) – Path to a custom encoder model checkpoint file to use instead of the default. Default is None, which uses the default model path from configs.py.
ddpm_path (str, optional) – Path to a custom diffusion model checkpoint file to use instead of the default. Default is None, which uses the default model path from configs.py.
- Returns:
A dict with the sequence names as the key and a starling.ensembl.Ensemble objects for each sequence as values.
If output_directory is not none, the output will save to the specified path.
- Return type:
dict or None