time_gen_scorer.py

class TimeGenScorer:
    def __init__(self, n_times, train=None, cuda=False):
        self.test_times = range(n_times)
        if train is None:
            train = slice(None, None, 1)
        self.train_times = self.test_times[train]
        self.cuda = cuda

    def __call__(self, Y_true, Y_pred):

        R = np.zeros((len(self.train_times), len(self.test_times), Y_true.shape[1]))
        
        if self.cuda:
            Y_true = torch.from_numpy(Y_true).cuda()
            Y_pred = torch.from_numpy(Y_pred).cuda()

        Y_true = Y_true - Y_true.mean(0)
        Y_pred = Y_pred - Y_pred.mean(0)

        SX2 = (Y_true ** 2).sum(0) ** 0.5
        SY2 = (Y_pred ** 2).sum(0) ** 0.5

        for k, t1 in enumerate(tqdm(self.train_times)):
            for t2 in self.test_times:
                SXY = (Y_true[:, :, t1] * Y_pred[:, :, t2]).sum(0)
                r = SXY / (SX2[:, t1] * SY2[:, t2])
                if self.cuda:
                    r = r.cpu().numpy()
                R[k, t2] = r

        if Y_true.shape[1] == 1:
            R = R[:, :, 0]
        return R