How does the AlphaFlow sample the conformation?

[smallfishabc]

Hi everyone,
I have noticed that AlphaFlow will generate the same conformation all the time if I run the predict.py without using --resample option.

This resample option is actually resampling the MSA instead of using random process in the diffusion module.

The resample option is only used in the line 122 of the predict.py file.

if args.subsample or args.resample:
      # Get item function                
      item = valset[i] # resample MSA

The diffusion process is in the new_model.py starting form line 368. These diffusion code in the original AlphaFlow model will be found in the wrapper.py with the inference funcions.

        if schedule is None:
            schedule = np.array([1.0, 0.75, 0.5, 0.25, 0.1, 0]) 
        outputs = []
        prev_outputs = None
        for t, s in zip(schedule[:-1], schedule[1:]):
            output = self.model(batch, prev_outputs=prev_outputs)
            pseudo_beta = pseudo_beta_fn(batch['aatype'], output['final_atom_positions'], None)
            outputs.append({**output, **batch})
            noisy = rmsdalign(pseudo_beta, noisy)
            noisy = (s / t) * noisy + (1 - s / t) * pseudo_beta
            batch['noised_pseudo_beta_dists'] = torch.sum((noisy.unsqueeze(-2) - noisy.unsqueeze(-3)) ** 2, dim=-1)**0.5
            batch['t'] = torch.ones(1, device=noisy.device) * s # first one doesn't get the time embedding, last one is ignored :)
            if self_cond:
                prev_outputs = output

My hypothesis is that AlphaFlow is generating structure with MSA resample instead of diffusion. Could you have a look at it?
Thank you.
Best
Feng Yu