import torch
import torchvision
# Download and load training and test datasets
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True)
# Scale data and subtract global mean
def scale_and_center(x_train):
 std = x_train.std()
 x_train = x_train / (std * n_row)
 global_mean = x_train.mean(axis=0, keepdims=True)
 x_train = x_train - global_mean
 return x_train
n_samples, n_row, n_col = trainset.data.shape
n_dim = n_row * n_col
x_train = trainset.data.reshape(-1, n_dim).float()
y_train = trainset.targets
from sklearn.decomposition import PCA
pca = PCA(n_components=6)
pca.fit(x_train)
pca_filters = pca.components_
from ot.dr import wda
Pwda, projwda = wda(x_train.numpy(), y_train.numpy(), p=6, reg=0.01,
 P0=pca_filters.T)