arithmetic.jl (2956B)
1 #/usr/bin/julia 2 3 using LinearAlgebra 4 using Plots 5 using Plots.Measures 6 using Distributions 7 using LaTeXStrings 8 using TSVD: tsvd 9 using TensorToolbox: tenmat, ttm, contract # todo:implement 10 11 @doc raw""" 12 Addition of two tensors in the TT-MPS format of $A$ and $B$ of the same size 13 $n_1 × ⋯ × n_d$ with different TT-MPS decompositions U with ranks $p_1, … , p_{d-1}$ 14 and V with ranks $q_1, … ,q_{d-1}$ respectively. 15 16 Output is the TT-MPS representation of $A+B$ with ranks $(q_1 + p_1), ⋯ , (q_{d-1} + p_{d-1})$. 17 """ 18 function tt_add(U, V, α, β) 19 W = [] 20 d = length(U) 21 for k ∈ 1:d 22 U_k = U[k]; V_k = V[k] 23 p_k, n_k, p_k1 = size(U_k) 24 q_k, n_k, q_k1 = size(V_k) 25 if k == 1 26 W_1 = zeros(1, n_k, q_k1 + p_k1) 27 for i=1:n_k 28 W_1[:, i, :] += hcat(α * U_k[:, i, :], β * V_k[:, i, :]) 29 end 30 append!(W, [W_1]) 31 elseif k == d 32 W_d = zeros(q_k + p_k, n_k, 1) 33 for i=1:n_k 34 W_d[:, i, :] += vcat(U_k[:, i, :], V_k[:, i, :]) 35 end 36 append!(W, [W_d]) 37 else 38 W_k = zeros(p_k + q_k, n_k, q_k1 + p_k1) 39 for i=1:n_k 40 W_k[:, i, :] += hcat(vcat(U_k[:, i, :], zeros(q_k, p_k1)), 41 vcat(zeros(p_k, q_k1), V_k[:, i, :])) 42 end 43 append!(W, [W_k]) 44 end 45 end 46 return W 47 end 48 49 @doc raw""" 50 Hadamard product of two tensors in the TT-MPS format of $A$ and $B$ of the same size 51 $n_1 × ⋯ × n_d$ with different TT-MPS decompositions U with ranks $p_1, … , p_{d-1}$ 52 and V with ranks $q_1, … ,q_{d-1}$ respectively. 53 54 Output is the representation of $A ⊙ B$ with ranks $(q_1 ⋅ p_1), ⋯ , (q_{d-1} ⋅ p_{d-1})$. 55 """ 56 function tt_mult(U, V) 57 W = [] 58 d = length(U) 59 for k ∈ 1:d 60 p_k, n_k, p_k1 = size(U[k]) 61 q_k, n_k, q_k1 = size(V[k]) 62 W_k = zeros(p_k*q_k, n_k, p_k1 * q_k1) 63 for i_k ∈ 1:n_k 64 W_k[:, i_k, :] += kron(U[k][:, i_k, :], V[k][:, i_k, :]) 65 end 66 append!(W, [W_k]) 67 end 68 return W 69 end 70 71 @doc raw""" 72 Matrix Vector product of two tensors in the TT-MPS format of $A$, size $(m_1⋯m_d×n_1⋯n_d)$ and $u$ size $(n_1⋯n_d)$ 73 with different TT-MPS decompositions U with ranks $p_1, … , p_{d-1}$ 74 and V with ranks $q_1, … ,q_{d-1}$ respectively. 75 76 Output is the TT-MPS representation of $A ⋅ u$ with ranks $(q_1 ⋅ p_1), ⋯ , (q_{d-1} ⋅ p_{d-1})$. 77 """ 78 function tt_matvec(A, U) 79 W = [] 80 d = length(U) 81 for k ∈ 1:d 82 α_k, n_k, m_k, α_k1 = size(A[k]) 83 β_k, m_k, β_k1 = size(U[k]) 84 W_k = zeros(α_k * β_k, n_k, α_k1 * β_k1) 85 for i_k ∈ 1:n_k 86 for j_k ∈ 1:m_k 87 W_k[:, i_k, :] += kron(A[k][:, i_k, j_k, :], U[k][:, j_k, :]) 88 end 89 end 90 append!(W, [W_k]) 91 end 92 return W 93 end