def test_axpby(self): ttOp1 = pitts_py.TensorTrainOperator_double([2, 3, 4], [3, 2, 2]) ttOp1.setTTranks(2) ttOp2 = pitts_py.TensorTrainOperator_double([2, 3, 4], [3, 2, 2]) pitts_py.randomize(ttOp1) pitts_py.randomize(ttOp2) ttOp12 = pitts_py.TensorTrainOperator_double([2, 3, 4], [3, 2, 2]) pitts_py.copy(ttOp2, ttOp12) pitts_py.axpby(0.33, ttOp1, -0.97, ttOp12) ttX = pitts_py.TensorTrain_double([3, 2, 2]) ttX.setTTranks(2) pitts_py.randomize(ttX) ttY = pitts_py.TensorTrain_double([2, 3, 4]) pitts_py.apply(ttOp12, ttX, ttY) ttY1 = pitts_py.TensorTrain_double([2, 3, 4]) pitts_py.apply(ttOp1, ttX, ttY1) ttY2 = pitts_py.TensorTrain_double([2, 3, 4]) pitts_py.apply(ttOp2, ttX, ttY2) ttY12 = pitts_py.TensorTrain_double([2, 3, 4]) pitts_py.copy(ttY2, ttY12) nrm = pitts_py.axpby(0.33, ttY1, -0.97, ttY12) err = pitts_py.axpby(-nrm, ttY12, 1., ttY) self.assertLess(err, 1.e-8)
def test_apply_zero(self): ttOp = pitts_py.TensorTrainOperator_double([2, 3, 2], [3, 4, 1]) ttX = pitts_py.TensorTrain_double([3, 4, 1]) ttY = pitts_py.TensorTrain_double([2, 3, 2]) ttOp.setZero() pitts_py.randomize(ttX) pitts_py.apply(ttOp, ttX, ttY) self.assertEqual(0, pitts_py.norm2(ttY))
def test_dot(self): tt1 = pitts_py.TensorTrain_double([2, 5, 3]) tt2 = pitts_py.TensorTrain_double([2, 5, 3]) tt1.setUnit([0, 1, 2]) tt2.setUnit([0, 2, 2]) np.testing.assert_almost_equal(1., pitts_py.dot(tt1, tt1)) np.testing.assert_almost_equal(0., pitts_py.dot(tt1, tt2)) np.testing.assert_almost_equal(1., pitts_py.dot(tt2, tt2))
def test_copy(self): tt = pitts_py.TensorTrain_double([2, 4, 3]) tt2 = pitts_py.TensorTrain_double([2, 4, 3]) tt.setUnit([0, 1, 2]) pitts_py.copy(tt, tt2) tt.setUnit([1, 0, 0]) fullTensor = pitts_py.toDense(tt2) fullTensor_ref = np.zeros([2, 4, 3]) fullTensor_ref[0, 1, 2] = 1 np.testing.assert_array_almost_equal(fullTensor_ref, fullTensor)
def test_apply_identity(self): ttOp = pitts_py.TensorTrainOperator_double([5, 3, 3], [5, 3, 3]) ttX = pitts_py.TensorTrain_double([5, 3, 3]) ttY = pitts_py.TensorTrain_double([5, 3, 3]) ttOp.setEye() pitts_py.randomize(ttX) pitts_py.apply(ttOp, ttX, ttY) err = pitts_py.axpby(-1, ttX, 1, ttY) self.assertLess(err, 1.e-8)
def test_setOnes(self): tt = pitts_py.TensorTrain_double([2, 3, 4]) tt.setTTranks([2, 4]) tt.setOnes() self.assertEqual([1, 1], tt.getTTranks()) fullTensor = pitts_py.toDense(tt) np.testing.assert_array_almost_equal(np.ones([2, 3, 4]), fullTensor)
def test_norm(self): tt = pitts_py.TensorTrain_double([2, 5, 3]) tt.setTTranks([2, 2]) pitts_py.randomize(tt) np.testing.assert_almost_equal(np.sqrt(pitts_py.dot(tt, tt)), pitts_py.norm2(tt))
def tt_gmres(AOp, b, nrm_b, eps=1.e-6, maxIter=20, verbose=True): """ Tensor-train GMRES algorithm without restart """ # assumes b is normalized and nrm_b is the desired rhs norm # define initial subspace beta = nrm_b curr_beta = beta V = [b] m = maxIter H = np.mat(np.zeros((m + 1, m), order='F')) if verbose: print("TT-GMRES: initial residual norm: %g, max. rank: %d" % (beta, np.max(b.getTTranks()))) for j in range(m): delta = eps / (curr_beta / beta) w = pitts_py.TensorTrain_double(b.dimensions()) w_nrm = AOp(V[j], w, delta / m) for i in range(j + 1): H[i, j] = w_nrm * pitts_py.dot(w, V[i]) w_nrm = pitts_py.axpby(-H[i, j], V[i], w_nrm, w, delta / m) if verbose: print("TT-GMRES: iteration %d, new Krylov vector max. rank: %d" % (j, np.max(w.getTTranks()))) H[j + 1, j] = w_nrm V = V + [w] Hj = H[:j + 2, :j + 1] betae = np.zeros(j + 2) betae[0] = beta # solving Hj * y = beta e_1 y, curr_beta, rank, s = np.linalg.lstsq(Hj, betae, rcond=None) curr_beta = np.sqrt(curr_beta[0]) if curr_beta.size > 0 else 0 if verbose: print("TT-GMRES: LSTSQ resirual norm: %g " % (curr_beta / beta)) if curr_beta / beta <= eps: break x = pitts_py.TensorTrain_double(b.dimensions()) x.setZero() nrm_x = 0 for i in range(len(y)): nrm_x = pitts_py.axpby(y[i], V[i], nrm_x, x, eps / m) if verbose: print("TT-GMRES: solution max rank %d" % np.max(x.getTTranks())) return x, nrm_x
def test_example(self): pitts_py.initialize() # do something interesting... tt = pitts_py.TensorTrain_double([5,5,5]) pitts_py.randomize(tt) pitts_py.finalize()
def test_setSubTensor_invalidShape(self): tt = pitts_py.TensorTrain_double([3, 2, 5]) with self.assertRaises(IndexError): tt.setSubTensor(10, [[[1], [2], [3]]]) with self.assertRaises(ValueError): tt.setSubTensor(0, [1, 2, 3]) tt.setSubTensor(0, [[[1], [2], [3]]]) with self.assertRaises(ValueError): tt.setSubTensor(1, [[[1], [2], [3]]])
def test_setUnit(self): tt = pitts_py.TensorTrain_double([2, 3, 4]) tt.setTTranks([2, 4]) tt.setUnit([1, 0, 3]) self.assertEqual([1, 1], tt.getTTranks()) fullTensor = pitts_py.toDense(tt) fullTensor_ref = np.zeros([2, 3, 4]) fullTensor_ref[1, 0, 3] = 1 np.testing.assert_array_almost_equal(fullTensor_ref, fullTensor)
def test_getSubTensor_unit(self): tt = pitts_py.TensorTrain_double([3, 2, 5]) tt.setUnit([1, 0, 2]) t1 = tt.getSubTensor(0) t2 = tt.getSubTensor(1) t3 = tt.getSubTensor(2) np.testing.assert_array_almost_equal([[[0], [1], [0]]], t1) np.testing.assert_array_almost_equal([[[1], [0]]], t2) np.testing.assert_array_almost_equal([[[0], [0], [1], [0], [0]]], t3)
def test_getSubTensor_zeros(self): tt = pitts_py.TensorTrain_double([3, 2, 5]) tt.setZero() t1 = tt.getSubTensor(0) t2 = tt.getSubTensor(1) t3 = tt.getSubTensor(2) np.testing.assert_array_almost_equal(np.zeros([1, 3, 1]), t1) np.testing.assert_array_almost_equal(np.zeros([1, 2, 1]), t2) np.testing.assert_array_almost_equal(np.zeros([1, 5, 1]), t3)
def test_normalize(self): tt = pitts_py.TensorTrain_double([2, 5, 3]) tt.setTTranks([2, 2]) pitts_py.randomize(tt) norm_ref = pitts_py.norm2(tt) norm = pitts_py.normalize(tt) np.testing.assert_almost_equal(norm_ref, norm) np.testing.assert_almost_equal(1., pitts_py.norm2(tt))
def test_setGetSubTensor_large(self): tt = pitts_py.TensorTrain_double([50, 100, 20]) tt.setTTranks([2, 3]) pitts_py.randomize(tt) t1_ref = np.random.rand(1, 50, 2) t2_ref = np.random.rand(2, 100, 3) t3_ref = np.random.rand(3, 20, 1) tt.setSubTensor(0, t1_ref) tt.setSubTensor(1, t2_ref) tt.setSubTensor(2, t3_ref) np.testing.assert_array_almost_equal(t1_ref, tt.getSubTensor(0)) np.testing.assert_array_almost_equal(t2_ref, tt.getSubTensor(1)) np.testing.assert_array_almost_equal(t3_ref, tt.getSubTensor(2))
def cdist2_TTapprox(X, dims_samples, dims_features, rank_tolerance, max_rank, debug=False): """approximate the pair-wise squared distances of all rows using a TT decomposition of the data""" # convert to TT format dims = dims_samples + dims_features n_samples = np.prod(dims_samples) n_features = np.prod(dims_features) Xm = pitts_py.MultiVector_double(n_samples*n_features//dims[-1], dims[-1]) work = pitts_py.MultiVector_double() Xm_view = np.array(Xm, copy=False) # scale by norm -> length 1 Xnorms = np.linalg.norm(X,axis=1) Xm_view[...] = (np.diag(1 / Xnorms) @ X).reshape(Xm_view.shape, order='F') Xtt = pitts_py.fromDense(Xm, work, dims, rankTolerance=0.01, maxRank=200) if debug: X_approx = pitts_py.toDense(Xtt).reshape((n_samples, n_features), order='F') * Xnorms[:, np.newaxis] plt.scatter(X[:,0], X[:,1]) plt.scatter(X_approx[:,0], X_approx[:,1]) plt.savefig('scatter_01_approx.png') plt.close() plt.scatter(X[:,1], X[:,2]) plt.scatter(X_approx[:,1], X_approx[:,2]) plt.savefig('scatter_12_approx.png') plt.close() plt.scatter(X[:,3], X[:,4]) plt.scatter(X_approx[:,3], X_approx[:,4]) plt.savefig('scatter_34_approx.png') plt.close() # setup approximated (US) from SVD (U S V^T) d_samples = len(dims_samples) Xtt_ranks = Xtt.getTTranks() k = Xtt_ranks[d_samples-1] US_tt = pitts_py.TensorTrain_double(dims_samples + [k,]) US_tt.setTTranks(Xtt_ranks[:d_samples-1] + [k,]) for i in range(d_samples): US_tt.setSubTensor(i, Xtt.getSubTensor(i)) # last sub-tensor is just the identity to get an additional direction US_tt.setSubTensor(d_samples, np.eye(k,k).reshape((k,k,1))) US_approx = pitts_py.toDense(US_tt).reshape((n_samples, k), order='F') if debug: plt.matshow(US_approx) plt.savefig('US_approx.png') plt.close() US_approx = np.diag(Xnorms) @ US_approx return cdist2(US_approx)
def test_setGetSubTensor(self): tt = pitts_py.TensorTrain_double([3, 2, 5]) tt.setTTranks([2, 3]) pitts_py.randomize(tt) t1_ref = np.array([1, 2, 3, 4, 5, 6]).reshape([1, 3, 2]) t2_ref = np.array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]).reshape([2, 2, 3]) t3_ref = np.array([ 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115 ]).reshape([3, 5, 1]) tt.setSubTensor(0, t1_ref) tt.setSubTensor(1, t2_ref) tt.setSubTensor(2, t3_ref) np.testing.assert_array_almost_equal(t1_ref, tt.getSubTensor(0)) np.testing.assert_array_almost_equal(t2_ref, tt.getSubTensor(1)) np.testing.assert_array_almost_equal(t3_ref, tt.getSubTensor(2))
def test_randomize(self): tt = pitts_py.TensorTrain_double([2, 5, 3]) self.assertEqual([1, 1], tt.getTTranks()) fullTensor1 = pitts_py.toDense(tt) pitts_py.randomize(tt) self.assertEqual([1, 1], tt.getTTranks()) fullTensor2 = pitts_py.toDense(tt) pitts_py.randomize(tt) fullTensor3 = pitts_py.toDense(tt) tt.setTTranks([2, 3]) pitts_py.randomize(tt) self.assertEqual([2, 3], tt.getTTranks()) fullTensor4 = pitts_py.toDense(tt) # check for big enough differences... self.assertTrue(np.linalg.norm(fullTensor1 - fullTensor2) > 1.e-4) self.assertTrue(np.linalg.norm(fullTensor2 - fullTensor3) > 1.e-4) self.assertTrue(np.linalg.norm(fullTensor3 - fullTensor4) > 1.e-4)
def test_rightNormalize(self): tt = pitts_py.TensorTrain_double([2, 5, 3]) tt.setTTranks([2, 2]) pitts_py.randomize(tt) norm_ref = pitts_py.norm2(tt) norm = pitts_py.rightNormalize(tt) np.testing.assert_almost_equal(norm_ref, norm) np.testing.assert_almost_equal(1., pitts_py.norm2(tt)) t1 = tt.getSubTensor(0) t2 = tt.getSubTensor(1) t3 = tt.getSubTensor(2) r1 = tt.getTTranks()[0] r2 = tt.getTTranks()[1] t1_mat = t1.reshape([1, t1.size // 1]) t2_mat = t2.reshape([r1, t2.size // r1]) t3_mat = t3.reshape([r2, t3.size // r2]) np.testing.assert_array_almost_equal( np.eye(1, 1), np.dot(t1_mat, t1_mat.transpose())) np.testing.assert_array_almost_equal( np.eye(r1, r1), np.dot(t2_mat, t2_mat.transpose())) np.testing.assert_array_almost_equal( np.eye(2, 2), np.dot(t3_mat, t3_mat.transpose()))
eye_i = TTOp_dummy.getSubTensor(iDim) tridi_i = np.zeros((n_i,n_i)) for i in range(n_i): for j in range(n_i): if i == j: tridi_i[i,j] = 2. / (n_i+1) elif i+1 == j or i-1 == j: tridi_i[i,j] = -1. / (n_i+1) else: tridi_i[i,j] = 0 TTOp_dummy.setSubTensor(iDim, tridi_i.reshape(1,n_i,n_i,1)) pitts_py.axpby(1, TTOp_dummy, 1, TTOp) TTOp_dummy.setSubTensor(iDim, eye_i) return TTOp TTOp = LaplaceOperator([10,]*5) x0 = pitts_py.TensorTrain_double(TTOp.row_dimensions()) x0.setOnes() sigma, q = tt_jacobi_davidson(TTOp, x0, symmetric=True, eps=1.e-8) r = pitts_py.TensorTrain_double(x0.dimensions()) pitts_py.apply(TTOp, q, r) sigma_ref = pitts_py.dot(q, r) r_nrm = pitts_py.axpby(-sigma, q, 1, r) print("Residual norm: %g" % r_nrm) print("Est. eigenvalue: %g, real Ritz value: %g, error: %g" % (sigma, sigma_ref, np.abs(sigma-sigma_ref))) pitts_py.finalize()
return x, nrm_x if __name__ == '__main__': pitts_py.initialize() TTOp = pitts_py.TensorTrainOperator_double([2, 3, 3, 2, 4, 10, 7], [2, 3, 3, 2, 4, 10, 7]) TTOp.setTTranks(1) pitts_py.randomize(TTOp) TTOpEye = pitts_py.TensorTrainOperator_double([2, 3, 3, 2, 4, 10, 7], [2, 3, 3, 2, 4, 10, 7]) TTOpEye.setEye() pitts_py.axpby(1, TTOpEye, 0.1, TTOp) xref = pitts_py.TensorTrain_double(TTOp.row_dimensions()) xref.setOnes() b = pitts_py.TensorTrain_double(TTOp.row_dimensions()) pitts_py.apply(TTOp, xref, b) nrm_b = pitts_py.normalize(b) def AOp(x, y, eps): pitts_py.apply(TTOp, x, y) y_nrm = pitts_py.normalize(y, eps) return y_nrm x, nrm_x = tt_gmres(AOp, b, nrm_b, maxIter=30, eps=1.e-4) print("nrm_x %g" % nrm_x) r = pitts_py.TensorTrain_double(b.dimensions()) pitts_py.apply(TTOp, x, r)
def tt_jacobi_davidson(A, x0, symmetric, eps=1.e-6, maxIter=20, arnoldiIter=5, gmresTol=0.01, gmresIter=10, verbose=True): """ Simplistic tensor-train Jacobi-Davidson algorithm """ assert(x0.dimensions() == A.col_dimensions()) # create empty search space W = list() AW = list() H = np.zeros((0,0)) for j in range(maxIter): if j == 0: # initial search direction v = pitts_py.TensorTrain_double(A.col_dimensions()) pitts_py.copy(x0, v) pitts_py.normalize(v, 0.01*eps) elif j < arnoldiIter: # start with some Arnoldi iterations v = r else: # calculate new search direction def JDop(x, y, eps): # Jacobi-Davidson operator with projections # y = (I - q q^T) (A - sigma I) (I - q q^T) x # we only do # y = (I - q q^T) (A - sigma I) x # because we assume that x is already orthogonal to q pitts_py.apply(A, x, y) y_nrm = pitts_py.normalize(y, 0.01*eps) y_nrm = pitts_py.axpby(-sigma, x, 1, y, 0.01*eps) qTy = y_nrm * pitts_py.dot(q, y) y_nrm = pitts_py.axpby(-qTy, q, y_nrm, y, 0.01*eps) return y_nrm v, _ = tt_gmres(JDop, r, r_nrm, eps=gmresTol, maxIter=gmresIter, verbose=False) # orthogonalize new search direction wrt. previous vectors for iOrtho in range(5): max_wTv = 0 for i in range(j): wTv = pitts_py.dot(W[i], v) max_wTv = max(max_wTv, abs(wTv)) if abs(wTv) > 0.01*eps: pitts_py.axpby(-wTv, W[i], 1., v, 0.01*eps) if max_wTv < 0.01*eps: break Av = pitts_py.TensorTrain_double(A.row_dimensions()) pitts_py.apply(A, v, Av) W = W + [v] AW = AW + [Av] # # calculate orthogonality error # WtW = np.zeros((j+1,j+1)) # for i in range(j+1): # for k in range(j+1): # WtW[i,k] = pitts_py.dot(W[i], W[k]) # WtW_err = np.linalg.norm(WtW - np.eye(j+1,j+1)) # print('WtW_err', WtW_err) # update H = W^T AW H = np.pad(H, ((0,1),(0,1))) for i in range(j): H[i,-1] = pitts_py.dot(W[i], Av) H[-1,i] = H[i,-1] if symmetric else pitts_py.dot(v,AW[i]) H[-1,-1] = pitts_py.dot(v, Av) # compute Schur decomposition H QH = QH RH if symmetric: (RH,QH) = np.linalg.eigh(H) eigIdx = 0 else: (RH,QH) = np.linalg.eig(H) eigIdx = np.argmin(RH) # compute Ritz value and vector sigma = RH[eigIdx] q = pitts_py.TensorTrain_double(A.col_dimensions()) q_nrm = 0 for i in range(j+1): q_nrm = pitts_py.axpby(QH[i,eigIdx], W[i], q_nrm, q, 0.01*eps) # calculate residual r = A*q-sigma*q r = pitts_py.TensorTrain_double(A.col_dimensions()) pitts_py.apply(A, q, r) r_nrm = pitts_py.axpby(-sigma, q, 1, r, 0.01*eps) # explicitly orthogonalize r wrt. q (we have approximation errors) qTr = pitts_py.dot(q,r) pitts_py.axpby(-qTr, q, 1, r, 0.01*eps) if verbose: print("TT-JacobiDavidson: Iteration %d, approximated eigenvalue: %g, residual norm: %g (orthog. error %g)" %(j, sigma, r_nrm, qTr)) # abort if accurate enough if r_nrm < eps: break # return resulting eigenvalue and eigenvector approximation return sigma, q
def test_copy_dimensionMismatch(self): tt = pitts_py.TensorTrain_double([2, 4, 3]) tt2 = pitts_py.TensorTrain_double([2, 4, 2]) with self.assertRaises(ValueError): pitts_py.copy(tt, tt2)
def test_createTensorTrain_double(self): tt = pitts_py.TensorTrain_double([3, 4, 5]) self.assertEqual([3, 4, 5], tt.dimensions()) self.assertEqual([1, 1], tt.getTTranks())
def test_axpby_dimensionMismatch(self): tt = pitts_py.TensorTrain_double([2, 4, 3]) tt2 = pitts_py.TensorTrain_double([2, 4, 2]) with self.assertRaises(ValueError): nrm = pitts_py.axpby(1, tt, 2, tt2)
def test_setGetTTranks(self): tt = pitts_py.TensorTrain_double([2, 2, 2, 2, 2]) self.assertEqual([1, 1, 1, 1], tt.getTTranks()) tt.setTTranks([1, 3, 5, 2]) self.assertEqual([1, 3, 5, 2], tt.getTTranks())