INNER CODE UNIT · Python

get

QizhiPei/FABind · FABind/fabind/data.py:49

    def get(self, idx):
        line = self.data.iloc[idx]
        pocket_com = line['pocket_com']
        use_compound_com = line['use_compound_com']
        use_whole_protein = line['use_whole_protein'] if "use_whole_protein" in line.index else self.use_whole_protein
        group = line['group'] if "group" in line.index else 'train'
        if group == 'train' and use_compound_com:
            add_noise_to_com = self.add_noise_to_com
        elif group == 'train' and not use_compound_com:
            add_noise_to_com = self.noise_for_predicted_pocket
        else:
            add_noise_to_com = None

        if group == 'train':
            random_rotation = True
        elif group == 'test' and self.test_random_rotation:
            random_rotation = True
        else:

View source record →