Initialize a `FlowDock` module. :param cfg: A model configuration dictionary.
(
self,
cfg: DictConfig,
)
| 79 | """A geometric conditional flow matching model for protein-ligand docking.""" |
| 80 | |
| 81 | def __init__( |
| 82 | self, |
| 83 | cfg: DictConfig, |
| 84 | ) -> None: |
| 85 | """Initialize a `FlowDock` module. |
| 86 | |
| 87 | :param cfg: A model configuration dictionary. |
| 88 | """ |
| 89 | super().__init__() |
| 90 | self.cfg = cfg |
| 91 | self.ligand_cfg = cfg.mol_encoder |
| 92 | self.protein_cfg = cfg.protein_encoder |
| 93 | self.relational_reasoning_cfg = cfg.relational_reasoning |
| 94 | self.contact_cfg = cfg.contact_predictor |
| 95 | self.score_cfg = cfg.score_head |
| 96 | self.confidence_cfg = cfg.confidence |
| 97 | self.affinity_cfg = cfg.affinity |
| 98 | self.global_cfg = cfg.task |
| 99 | self.protatm_padding_dim = self.protein_cfg.atom_padding_dim # := 37 |
| 100 | self.max_n_edges = self.global_cfg.edge_crop_size |
| 101 | self.latent_model = cfg.latent_model |
| 102 | self.prior_type = cfg.prior_type |
| 103 | |
| 104 | # VDW radius mapping, in Angstrom |
| 105 | self.atnum2vdw = torch.nn.Parameter( |
| 106 | torch.tensor(get_vdw_radii_array() / 100.0), |
| 107 | requires_grad=False, |
| 108 | ) |
| 109 | self.atnum2vdw_uff = torch.nn.Parameter( |
| 110 | torch.tensor(get_vdw_radii_array_uff() / 100.0), |
| 111 | requires_grad=False, |
| 112 | ) |
| 113 | |
| 114 | # graph hyperparameters |
| 115 | self.BINDING_SITE_CUTOFF = 6.0 |
| 116 | self.INTERNAL_TARGET_VARIANCE_SCALE = self.global_cfg.internal_max_sigma |
| 117 | self.GLOBAL_TARGET_VARIANCE_SCALE = self.global_cfg.global_max_sigma |
| 118 | self.CONTACT_SCALE = 5.0 # fixed hyperparameter |
| 119 | |
| 120 | ( |
| 121 | standard_aa_template_featset, |
| 122 | standard_aa_graph_featset, |
| 123 | ) = get_standard_aa_features() |
| 124 | self.standard_aa_template_featset = inplace_to_torch(standard_aa_template_featset) |
| 125 | self.standard_aa_molgraph_featset = inplace_to_torch( |
| 126 | collate_numpy_samples(standard_aa_graph_featset) |
| 127 | ) |
| 128 | |
| 129 | # load pretrained weights as desired |
| 130 | from_pretrained = ( |
| 131 | self.ligand_cfg.from_pretrained |
| 132 | or self.protein_cfg.from_pretrained |
| 133 | or self.relational_reasoning_cfg.from_pretrained |
| 134 | or self.contact_cfg.from_pretrained |
| 135 | or self.score_cfg.from_pretrained |
| 136 | or self.confidence_cfg.from_pretrained |
| 137 | or self.affinity_cfg.from_pretrained |
| 138 | ) |
nothing calls this directly
no test coverage detected