AF3 OpenFoldMultimerDataset类解读

AlphaFold3 data_modules 模块的 OpenFoldMultimerDataset 类是 OpenFoldDataset 类的子类，专门用于 多链蛋白质（Multimer） 数据集的训练。它通过引入 AlphaFold Multimer 论文 中描述的过滤步骤，来实现多链蛋白质的训练。这个类扩展了父类的功能，特别是为了处理多链蛋白质的数据预处理和样本过滤。

源代码：

class OpenFoldMultimerDataset(OpenFoldDataset):
    """
    Create a torch Dataset object for multimer training and 
    add filtering steps described in AlphaFold Multimer's paper:
    https://www.biorxiv.org/content/10.1101/2021.10.04.463034v2.full.pdf Supplementary section 7.1 
    """

    def __init__(self,
                 datasets: Sequence[OpenFoldSingleMultimerDataset],
                 probabilities: Sequence[float],
                 epoch_len: int,
                 generator: torch.Generator = None,
                 _roll_at_init: bool = True
                 ):
        super(OpenFoldMultimerDataset, self).__init__(datasets=datasets,
                                                      probabilities=probabilities,
                                                      epoch_len=epoch_len,
                                                      generator=generator,
                                                      _roll_at_init=_roll_at_init)

    @staticmethod
    def deterministic_train_filter(
        cache_entry: Any,
        is_distillation: bool,
        max_resolution: float = 9.,
        max_single_aa_prop: float = 0.8,
        minimum_number_of_residues: int = 200,
        *args, **kwargs
    ) -> bool:
        """
        Implement multimer training filtering criteria described in
        https://www.biorxiv.org/content/10.1101/2021.10.04.463034v2.full.pdf Supplementary section 7.1
        """
        resolution = cache_entry.get("resolution", None)
        seqs = cache_entry["seqs"]

        return all([resolution_filter(resolution=resolution,
                                      max_resolution=max_resolution),
                    aa_count_filter(seqs=seqs,
                                    max_single_aa_prop=max_single_aa_prop),
                    (not is_distillation or all_seq_len_filter(seqs=seqs,