File size: 2,454 Bytes
94391f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | # Copyright (c) DP Technology.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import numpy as np
from functools import lru_cache
from unicore.data import BaseWrapperDataset
class NormalizeDataset(BaseWrapperDataset):
def __init__(self, dataset, coordinates, normalize_coord=True):
self.dataset = dataset
self.coordinates = coordinates
self.normalize_coord = normalize_coord # normalize the coordinates.
self.set_epoch(None)
def set_epoch(self, epoch, **unused):
super().set_epoch(epoch)
self.epoch = epoch
@lru_cache(maxsize=16)
def __cached_item__(self, index: int, epoch: int):
dd = self.dataset[index].copy()
coordinates = dd[self.coordinates]
# normalize
if self.normalize_coord:
coordinates = coordinates - coordinates.mean(axis=0)
dd[self.coordinates] = coordinates.astype(np.float32)
return dd
def __getitem__(self, index: int):
return self.__cached_item__(index, self.epoch)
class NormalizeDockingPoseDataset(BaseWrapperDataset):
def __init__(
self,
dataset,
coordinates,
pocket_coordinates,
center_coordinates="center_coordinates",
):
self.dataset = dataset
self.coordinates = coordinates
self.pocket_coordinates = pocket_coordinates
self.center_coordinates = center_coordinates
self.set_epoch(None)
def set_epoch(self, epoch, **unused):
super().set_epoch(epoch)
self.epoch = epoch
@lru_cache(maxsize=16)
def __cached_item__(self, index: int, epoch: int):
dd = self.dataset[index].copy()
coordinates = dd[self.coordinates]
pocket_coordinates = dd[self.pocket_coordinates]
# normalize coordinates and pocket coordinates ,align with pocket center coordinates
center_coordinates = pocket_coordinates.mean(axis=0)
coordinates = coordinates - center_coordinates
pocket_coordinates = pocket_coordinates - center_coordinates
dd[self.coordinates] = coordinates.astype(np.float32)
dd[self.pocket_coordinates] = pocket_coordinates.astype(np.float32)
dd[self.center_coordinates] = center_coordinates.astype(np.float32)
return dd
def __getitem__(self, index: int):
return self.__cached_item__(index, self.epoch)
|