c49bd59c7d43af4a3baadbfe17d6d77a6a6040fd,torch_geometric/datasets/particle.py,TrackMLParticleTrackingDataset,get,#TrackMLParticleTrackingDataset#Any#,58

Before Change

        y = pandas.read_csv(truth_path, usecols=["particle_id"],
                            dtype=np.int64)
        y = torch.from_numpy(y.values).squeeze()
        _, y = y.unique(return_inverse=True)
        y = y - 1  // Zero entries are marked as invalid via `-1`.

        return TrackingData(x=x, pos=pos, y=y)

After Change

        y = pandas.read_csv(truth_path, usecols=["particle_id"],
                            dtype=np.int64)
        y = torch.from_numpy(y.values).squeeze()
        y_row = y.unique(return_inverse=True)[1].sub_(1)
        mask = y_row != -1
        y_row = y_row[mask]
        y_col = torch.arange(y.size(0))[mask]
        perm = (y_row * y.size(0) + y_col).argsort()
        y_index = torch.stack([y_row[perm], y_col[perm]], dim=0)

        return TrackingData(x=x, pos=pos, y_index=y_index)

Instances