Class ConfidencePrediction

diffpack/task.py:224–283

@R.register("tasks.ConfidencePrediction")
class ConfidencePrediction(TorsionalDiffusion):
    eps = 1e-10
    _option_members = {"task", "criterion", "metric"}

    def __init__(self, sigma_embedding: nn.Module,
                 model: nn.Module,
                 confidence_model: nn.Module,
                 torsion_mlp_hidden_dims: list,
                 schedule_1pi_periodic: SO2VESchedule,
                 schedule_2pi_periodic: SO2VESchedule,
                 num_sample: int = 5,
                 num_mlp_layer: int = 1,
                 graph_construction_model: Optional[Any] = None,
                 verbose: int = 0,
                 train_chi_id: Optional[Any] = None):
        super().__init__(sigma_embedding,
                         model,
                         torsion_mlp_hidden_dims,
                         schedule_1pi_periodic,
                         schedule_2pi_periodic,
                         graph_construction_model,
                         verbose,
                         train_chi_id)
        self.confidence_model = confidence_model
        self.num_sample = num_sample
        self.mlp = layers.MLP(self.confidence_model.output_dim,
                              [self.confidence_model.output_dim] * num_mlp_layer + [1])

    def predict_rmsd(self, batch, all_loss=None, metric=None):
        protein = batch['graph']
        if self.graph_construction_model:
            protein = self.graph_construction_model(protein)
        atom_feature = self.confidence_model(protein, protein.node_feature.float())["node_feature"]
        residue_feature = scatter_mean(atom_feature, protein.atom2residue, dim=0,
                                       dim_size=protein.num_residue)  # [num_residue, feature_dim]
        pred = self.mlp(residue_feature).squeeze(-1)  # [num_residue]
        return pred

    @torch.no_grad()
    def generate(self, batch, randomize=True):
        protein = batch['graph']
        if randomize:
            protein = rotamer.randomize(protein)

        best_protein = protein.clone()
        best_rmsd = torch.zeros(protein.num_residue, device=self.device) + 1e6
        for _ in tqdm(range(self.num_sample), desc="Confidence sampling"):
            batch = super().generate(batch, randomize=True)  # TODO: do we need to randomize?
            protein = batch['graph']
            rmsd = self.predict_rmsd(batch)
            residue_update_mask = rmsd < best_rmsd  # [num_residue]
            atom_update_mask = residue_update_mask[protein.atom2residue]  # [num_atom]
            best_protein.node_position[atom_update_mask] = protein.node_position[atom_update_mask]
            best_rmsd[residue_update_mask] = rmsd[residue_update_mask]

        best_batch = {
            "graph": best_protein,
            "rmsd": best_rmsd