def fit(self, train_loader, dim=None, test_loader=None):
'''Train the model on the provided data loader.
Arguments:
train_loader (DataLoader): the training data
dim (int): the target dimensionality
test_loader (DataLoader): the data for validation
'''
self.x_mean, self.y_mean = _get_mean(train_loader)
self.cxx, self.cxy, self.cyy = _get_covariance(train_loader,
self.x_mean,
self.y_mean)
if self.symmetrize:
self.cxx = 0.5 * (self.cxx + self.cyy)
self.cyy.copy_(self.cxx)
self.cxy = 0.5 * (self.cxy + self.cxy.t())
self.transformer = _Transform(x_mean=self.x_mean,
x_covariance=self.cxx,
y_mean=self.y_mean,
y_covariance=self.cyy)
self.ixx = self.transformer.x.mul
self.iyy = self.transformer.y.mul
u, s, v = _svd(self.ixx.mm(self.cxy.mm(self.iyy)))
if dim is None:
dim = s.size()[0]
self.decoder_matrix = v[:, :dim]
self.encoder_matrix = u.t()[:dim, :]
if self.kinetic_map:
self.encoder_matrix = _diag(s[:dim]).mm(self.encoder_matrix)
else:
self.decoder_matrix = self.decoder_matrix.mm(_diag(s[:dim]))
self.koopman_matrix = self.decoder_matrix.mm(self.encoder_matrix)
return self.get_loss(train_loader), self.get_loss(test_loader)
def sqrtinv(matrix):
eigval, eigvec = DecomposeRSPDMatrix.apply(matrix)
diag = _diag(1.0 / _sqrt(eigval))
return _mm(eigvec, _mm(diag, eigvec.t()))