对Tensor Ring Decomposition with Rank Minimization onLatent Space一文求解算法部分的python实现

下面是TRLRF算法的函数部分

from matplotlib.image import imread
import numpy as np
from numpy import int8, linalg
from numpy import array, random
import matplotlib.pyplot as plt

def RSE_fun(X,X_hat,W):
    pos_test = np.where((X != 0) & (W == 0))
    rse = (np.linalg.norm(X_hat[pos_test] - X[pos_test], 2) 
                  / np.linalg.norm(X[pos_test], 2))#计算的是张量分解后的结果和原先结果的对比，采用rse来衡量
    print(rse)

# through G get the X_hat
def coreten2tr(Z):
    N=len(Z)
    S=[]
    for i in range(N):
        S.append(Z[i].shape[1])
    P=Z[0]

    for i in range(1,N):
        L=P.reshape(int(P.size/Z[i-1].shape[2]),Z[i-1].shape[2],order='F')
        R=Z[i].reshape(Z[i].shape[0],S[i]*Z[i].shape[2],order='F')
        P=L@R

    P=P.reshape(Z[0].shape[0],np.prod(S),Z[N-1].shape[2],order='F')
    P=np.moveaxis(P,0,-1)
    P=P.reshape(np.prod(S),Z[0].shape[0]*Z[0].shape[0],order='F')
    temp=np.eye(Z[0].shape[0],Z[0].shape[0])
    P=P@temp.reshape(np.prod(temp.shape))
    X=P.reshape(S,order='F')
    return X

# SVT algorithm
def Pro2TraceNorm(Z,tau):
    m,n=Z.shape
    if 2*m < n:
        AAT = Z@Z.T#求AA.T
        S, Sigma2, D =linalg.svd(AAT)#实对称矩阵的svd求解的是特征值矩阵
        V = np.sqrt(Sigma2)#V是奇异值，是特征值根号
        tol = max(Z.shape) * np.spacing(max(V)) #
        n = sum(V > max(tol, tau))#统计奇异值大于tau的个数
        mid = np.maximum(V[0:n]-tau, 0) / V[0:n] #
        X = S[:, 0:n] @ np.diag(mid) @ S[:, 0:n].T @ Z
        return X
    if m>2*n:
        X= Pro2TraceNorm(Z.T, tau)
        X = X.T
        return X
        
    S,V,D = linalg.svd(Z)
    D=D.T
    Sigma2 = np.diag(V) 
    n = sum(V > tau)
    X = S[0:,0:n] @ np.maximum(Sigma2[0:n,0:n]-tau, 0) @ D[:, 0:n].T
    return X


# mode-n unfold ,input-mode is belong to 0-2
def Gunfold(GT,mode):
    return np.reshape(np.moveaxis(GT, mode, 0), (GT.shape[mode], -1), order = 'F')
    
# mode-n unfold ,input-mode is belong to 0-2
def Gfold(mat, tensor_size, mode):
    index = list()
    index.append(mode)
    for i in range(len(tensor_size)):
        if i != mode:
            index.append(int(i))
    size = []
    for i in index:
        size.append(int(tensor_size[i]))
    return np.moveaxis(np.reshape(mat, size, order = 'F'), 0, mode)

# get G^{(\neq n)}
# input \in 0-2
def Z_neq(Y,n):
    # Z=circshift(Z,-n); % arrange Z{n} to the last core, so we only need to multiply the first N-1 core
    # 上述的代码python中没有，但是可以自己手写一个
    G=[]
    for i in range(n+1,len(Y)):
        G.append(Y[i])
    for i in range(n+1):
        G.append(Y[i])
    Z=G
    N=len(Z)
    P=Z[0]
    for i in range(N-2):
        zl=P.reshape(int(np.size(P)/(Z[i].shape[2])),int(Z[i].shape[2]),order='F')  
        zr=Z[i+1].reshape(int(Z[i+1].shape[0]),int(np.size(Z[i+1])/Z[i+1].shape[0]),order='F') 
        P=zl@zr
    Z_neq_out=P.reshape(int(Z[0].shape[0]),int(np.size(P)/(Z[0].shape[0]*Z[N-2].shape[2])),int(Z[N-2].shape[2]),order='F') 
    return Z_neq_out

# mode-k unfolding of X (square bracket unfolding)
# Ik x Ik+1 ... IN I1...Ik-1
# X=8*768* 8
# input \in 0-2
def tenmat_sb(X,k):
    S=X.shape
    N=len(S)
    if k==0:
        X_sb_k=X.reshape(int(S[0]),int(X.size/S[0]),order='F')
    elif k==N-1:
        X_sb_k=X.reshape(int(X.size/S[N-1]),int(S[N-1]),order='F')
        X_sb_k=np.transpose(X_sb_k)  # necessary???????????????
    else:
        X=X.reshape(int(np.prod(S[0:k])),int(X.size/np.prod(S[0:k])),order='F')
        X=np.transpose(X)
        X_sb_k=X.reshape(int(S[k]),int(X.size/S[k]),order='F')
    return X_sb_k

# M is the cell of size Nx3, add all the row together!
def Msum_fun(M):
    N=len(M)
    Msum_out=[]
    for i in range(N):
        Msum_out.append(M[i][0]+M[i][1]+M[i][2])
    return Msum_out

# generate tensor W by size and missing rate
def gen_W(S,mr):
    random_tensor = np.random.rand(S[0], S[1], S[2])#随机生成的矩阵，元素值在(0,1)之间
    p = mr#需要舍弃的数据量的比例
    W=np.round(random_tensor + 0.5 - p)#本实验直接舍弃90%的数据
    return W

# this is for simulation data, randomly generate tensor cores Z 
def TR_initcoreten(S,r):
    N=len(S)#计算张量阶数
    Z=[]#存放核心张量
    for i in range(N-1):
        Z.append(np.random.randn(r[0][i],S[i],r[0][i+1]))#创建了三个张量4*256*4，4*256*4   
    Z.append(np.random.randn(r[0][N-1],S[N-1],r[0][0]))#创建了一个张量4*3*4
    return Z

# TRLRF algorithm
def TRLRF(data,W,r,maxiter,mu,ro,Lamda,tol):
    Truth_data=data.copy()
    T=data*W  #创建不完全的张量
    N=len(T.shape) #不完全张量的维度
    S=T.shape #不完全张量的大小
    X= np.random.rand(S[0], S[1], S[2]) #随机生成的张量，作为输出张量
    G=TR_initcoreten(S,r) #初始化核心张量,G为产生后的张量
    M=[[] for _ in range(N)]
    Y=[[] for _ in range(N)]
    for i in range(N):
        G[i]=1*G[i]  #??????有什么意义
        for j in range(N):
            M[i].append(np.zeros(G[i].shape))
            Y[i].append(np.sign(G[i]))
    mu_max=10^2
    Convergence_rec=np.zeros((1,maxiter))
    iter=0
    while iter<maxiter:
        iter=iter+1
        # update G
        for n in range(N):
            Msum=Msum_fun(M)
            Ysum=Msum_fun(Y)
            Q=tenmat_sb(Z_neq(G,n),1)  
            Q=Q.T # Q is the right part of the right part of the relation equation
            G[n]=Gfold((Lamda*tenmat_sb(X,n)@Q.T+mu*Gunfold(Msum[n],1)+Gunfold(Ysum[n],1))@linalg.pinv((Lamda*(Q@Q.T)+3*mu*np.eye(Q.shape[0],Q.shape[0]))),G[n].shape,1)
            #j \in 0-2
            for j in range(3):
                Df=Gunfold(G[n]-Y[n][j]/mu,j)   
                M[n][j]=Gfold(Pro2TraceNorm(Df,1/mu),G[n].shape,j)
        # update X
        lastX=X
        X_hat=coreten2tr(G)
        X=X_hat
        X[W==1]=T[W==1]

        #update Y
        for n in range(N):
            for j in range(3):
                Y[n][j]=Y[n][j]+mu*(M[n][j]-G[n])
        mu=min(mu*ro,mu_max)

        # evaluation
        print('Iter: {}'.format(iter + 1))
        pos_test = np.where((Truth_data != 0) & (W == 0))
        rse = (np.linalg.norm(X_hat[pos_test] - Truth_data[pos_test], 2) 
                  / np.linalg.norm(Truth_data[pos_test], 2))#计算的是张量分解后的结果和原先结果的对比，采用rse来衡量
        print(rse)
        print()
    return X

下面是对于随机缺失80%图片的恢复代码

image=imread(r'C:\Users\wh\Desktop\TRLRF\TRLRF-Python\lena.bmp')
X=np.array((image)/255)
mr=0.8
W=gen_W(X.shape,mr);#W是位置算子

r=8*np.ones((1,3),dtype=int); ## TR-rank 
maxiter=1000 # maxiter 300~500
tol=1e-8 # 1e-6~1e-8
Lambda=5 # usually 1~10 \lambda 是F-2范数的参数
ro=1.1 # 1~1.5 \rho,调节惩罚参数的大小
mu=1e0 # 1e-1~1e0 
X_hat=TRLRF(X,W,r,maxiter,mu,ro,Lambda,tol)


# Evaluation
fig = plt.figure(figsize = (15, 8))
ax = fig.add_subplot(1, 3, 1)
plt.imshow(image)


ax = fig.add_subplot(1, 3, 2)
plt.imshow(np.uint8(W*image))


ax = fig.add_subplot(1, 3, 3)
image_hat=np.uint8(255*X_hat)
plt.imshow(image_hat)

plt.show()