Method run_inference

src/simplefold/wrapper.py:321–367

(self, batch, model, plddt_model, device)

        return batch, structure, record

    def run_inference(self, batch, model, plddt_model, device):
        # run inference for target protein
        if self.backend == "torch":
            noise = torch.randn_like(batch["coords"]).to(device)
        elif self.backend == "mlx":
            noise = mx.random.normal(batch["coords"].shape)
        out_dict = self.sampler.sample(model, self.flow, noise, batch)

        plddt_out_module = plddt_model["plddt_out_module"]
        plddt_latent_module = plddt_model["plddt_latent_module"]

        if plddt_latent_module is None or plddt_out_module is None:
            plddts = None
        else:
            if self.backend == "torch":
                t = torch.ones(batch["coords"].shape[0], device=device)
                # use unscaled coords to extract latent for pLDDT prediction
                out_feat = plddt_latent_module(
                    out_dict["denoised_coords"].detach(), t, batch
                )
                plddt_out_dict = plddt_out_module(
                    out_feat["latent"].detach(),
                    batch,
                )
            elif self.backend == "mlx":
                t = mx.ones(batch["coords"].shape[0])
                # use unscaled coords to extract latent for pLDDT prediction
                out_feat = plddt_latent_module(out_dict["denoised_coords"], t, batch)
                plddt_out_dict = plddt_out_module(
                    out_feat["latent"],
                    batch,
                )
            # scale pLDDT to [0, 100]
            plddts = plddt_out_dict["plddt"] * 100.0

        out_dict = self.processor.postprocess(out_dict, batch)
        # sampled_coord = out_dict['denoised_coords'].detach()
        if self.backend == "torch":
            sampled_coord = out_dict["denoised_coords"].detach()
        else:
            sampled_coord = out_dict["denoised_coords"]

        return {
            "sampled_coord": sampled_coord,
            "pad_mask": batch["atom_pad_mask"],
            "plddts": plddts,
        }

    def save_result(self, structure, record, results, out_name):
        sampled_coord = results["sampled_coord"]