AlphaFold3 protein_datamodule 模块 ProteinDataModule 类继承自 PyTorch Lightning 数据模块（LightningDataModule），负责 ProteinFlow 数据的准备、加载、拆分、变换等逻辑封装在一起，便于训练过程中的统一管理和复现。

这个类承担了 AlphaFold3 训练和评估过程中的 数据准备、划分、转换、加载 四个核心任务：

源代码：

class ProteinDataModule(LightningDataModule):
    """`LightningDataModule` for the Protein Data Bank.

    A `LightningDataModule` implements 7 key methods:

    ```python
        def prepare_data(self):
        # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP).
        # Download data, pre-process, split, save to disk, etc...

        def setup(self, stage):
        # Things to do on every process in DDP.
        # Load data, set variables, etc...

        def train_dataloader(self):
        # return train dataloader

        def val_dataloader(self):
        # return validation dataloader

        def test_dataloader(self):
        # return test dataloader

        def predict_dataloader(self):
        # return predict dataloader

        def teardown(self, stage):
        # Called on every process in DDP.
        # Clean up after fit or test.
    ```

    This allows you to share a full dataset without explaining how to download,
    split, transform and process the data.

    Read the docs:
        https://lightning.ai/docs/pytorch/latest/data/datamodule.html
    """

    def __init__(
            self,
            data_dir: str = "./data/",
            resolution_thr: float = 3.5,
            min_seq_id: float = 0.3,
            crop_size: int = 384,
            max_length: int = 10_000,
            use_fraction: float = 1.0,
            entry_type: str = "chain",
            classes_to_exclude: Optional[List[str]] = None,
            mask_residues: bool = False,
            lower_limit: int = 15,
            upper_limit: int = 100,
            mask_frac: Optional[float] = None,
            mask_sequential: bool = False,
            mask_whole_chains: bool = False,
            force_binding_sites_frac: float = 0.15,
            batch_size: int = 64,
            num_workers: int = 0,
            pin_memory: bool = False,
            debug: bool = False
    ) -> None:
        """Initialize a `ProteinDataModule`.

        :param resolution_thr: Resolution threshold for PDB structures
        :param min_seq_id: Minimum sequence identity for MMSeq2 clustering
        :param crop_size: The number of residues to crop the proteins to.
        :param max_length: Entries with total length of chains larger than max_length will be disregarded.
        :param use_fraction: the fraction of the clusters to use (first N in alphabetic order)
        :param entry_type: {"biounit", "chain", "pair"} the type of entries to generate ("biounit" for biounit-level
                            complexes, "chain" for chain-level, "pair" for chain-chain pairs (all pairs that are seen
                            in the same biounit and have intersecting coordinate clouds))
        :param classes_to_exclude: a list of classes to exclude from the dataset (select from "single_chains",
                                   "heteromers", "homomers")
        :param mask_residues: if True, the masked residues will be added to the output
        :param lower_limit: the lower limit of the number of residues to mask
        :param upper_limit: the upper limit of the number of residues to mask