AF3 FeaturePipeline类解读

AlphaFold3 feature_pipeline 模块 FeaturePipeline 类是一个封装类，通过调用函数np_example_to_features 实现整个数据处理流程。

源代码：

def np_to_tensor_dict(np_example: Mapping[str, np.ndarray],features: Sequence[str],
) -> TensorDict:"""Creates dict of tensors from a dict of NumPy arrays.Args:np_example: A dict of NumPy feature arrays.features: A list of strings of feature names to be returned in the dataset.Returns:A dictionary of features mapping feature names to features. Only the givenfeatures are returned, all other ones are filtered out."""# torch generates warnings if feature is already a torch Tensorto_tensor = lambda t: torch.tensor(t) if type(t) != torch.Tensor else t.clone().detach()tensor_dict = {k: to_tensor(v) for k, v in np_example.items() if k in features}return tensor_dictdef make_data_config(config: ml_collections.ConfigDict,mode: str,num_res: int,
) -> Tuple[ml_collections.ConfigDict, List[str]]:cfg = copy.deepcopy(config)mode_cfg = cfg[mode]# with cfg.unlocked():if mode_cfg.crop_size is None:mode_cfg.crop_size = num_resfeature_names = cfg.common.unsupervised_features# Add seqemb related features if using seqemb mode.if cfg.seqemb_mode.enabled:feature_names += cfg.common.seqemb_featuresif cfg.common.use_templates:feature_names += cfg.common.template_featuresif cfg[mode].supervised:feature_names += cfg.supervised.supervised_featuresreturn cfg, feature_namesdef np_example_to_features(np_example: FeatureDict,config: ml_collections.ConfigDict,mode: str,is_multimer: bool = False
):np_example = dict(np_example)seq_length = np_example["seq_length"]num_res = int(seq_length[0]) if seq_length.ndim != 0 else int(seq_length)cfg, feature_names = make_data_config(config, mode=mode, num_res=num_res)if "deletion_matrix_int" in np_example:np_example["deletion_matrix"] = np_example.pop("deletion_matrix_int").astype(np.float32)tensor_dict = np_to_tensor_dict(np_example=np_example, features=feature_names)with torch.no_grad():if is_multimer:features = input_pipeline_multimer.process_tensors_from_config(tensor_dict,cfg.common,cfg[mode],)else:features = input_pipeline.process_tensors_from_config(tensor_dict,cfg.common,cfg[mode],)if mode == "train":p = torch.rand(1).item()use_clamped_fape_value