{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from scipy.io import loadmat\n",
    "import cupy\n",
    "import xlrd\n",
    "import xlwt\n",
    "\n",
    "height=95\n",
    "width=95\n",
    "N=height*width\n",
    "max_iter=3000\n",
    "L=156\n",
    "P=3\n",
    "lr_decay=0.9\n",
    "unfreeze_time=0.9\n",
    "max_iter2=7\n",
    "\n",
    "#load data\n",
    "gt = loadmat('./dataset/end3.mat') \n",
    "M=gt['M'] #true endmember \n",
    "A=gt['A'] #true abundance\n",
    "im = loadmat('./dataset/samson_1.mat')['V']\n",
    "reference_abundance=A.reshape(height,width,-1)\n",
    "for i in range(3):\n",
    "    reference_abundance[:,:,i]=reference_abundance[:,:,i].T\n",
    "reference_abundance=reference_abundance.reshape(-1,P)\n",
    "abd=np.hstack((reference_abundance[:,0].reshape(N,1),reference_abundance[:,1].reshape(N,1),reference_abundance[:,2].reshape(N,1)))\n",
    "abd=abd.reshape(height,width,-1)\n",
    "im=im.T\n",
    "im=torch.from_numpy(im).float().cuda()\n",
    "\n",
    "def excel_to_matrix(path):\n",
    "\ttable = xlrd.open_workbook(path).sheets()[0]\n",
    "\trow = table.nrows\n",
    "\tcol = table.ncols\n",
    "\tdatamatrix = np.zeros((row, col))\n",
    "\tfor x in range(col):\n",
    "\t\tcols = np.matrix(table.col_values(x))\n",
    "\t\tdatamatrix[:, x] = cols\n",
    "\treturn datamatrix\n",
    "edm=excel_to_matrix('./dataset/VCA_edm_samson.xls')# endmembers by VCA\n",
    "edm=torch.from_numpy(edm).float().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neighbor_x(height, width):\n",
    "    \"\"\"Return smooth matrix h (numpy array)\"\"\"\n",
    "    N = height * width\n",
    "    x0 = np.zeros([N,N])\n",
    "    for i in range(N):\n",
    "        if i - width >= 0 and i + width <= N - 1 and i % width != 0 and (i + 1) % width != 0:\n",
    "            x0[i + 1, i] = 1\n",
    "            x0[i - 1, i] = 1\n",
    "            x0[i + width, i] = 1\n",
    "            x0[i - width, i] = 1\n",
    "        elif i - width < 0:\n",
    "            if i == 0 or i == width - 1:\n",
    "                x0[1, 0] = 1\n",
    "                x0[width, 0] = 1\n",
    "                x0[width - 2, width - 1] = 1\n",
    "                x0[2 * width - 1, width - 1] = 1\n",
    "            else:\n",
    "                x0[i - 1, i] = 1\n",
    "                x0[i + 1, i] = 1\n",
    "                x0[i + width, i] = 1\n",
    "        elif i + width > N - 1:\n",
    "            if i == N - 1 or i == N - width:\n",
    "                x0[N - 2, N - 1] = 1\n",
    "                x0[N - width, N - 1] = 1\n",
    "                x0[N - width + 1, N - width] = 1\n",
    "                x0[N - 2 * width, N - width] = 1\n",
    "            else:\n",
    "                x0[i - 1, i] = 1\n",
    "                x0[i + 1, i] = 1\n",
    "                x0[i - width, i] = 1\n",
    "        elif i % width == 0:\n",
    "            x0[i - width, i] = 1\n",
    "            x0[i + width, i] = 1\n",
    "            x0[i + 1, i] = 1\n",
    "        elif (i + 1) % width == 0:\n",
    "            x0[i - width, i] = 1\n",
    "            x0[i + width, i] = 1\n",
    "            x0[i - 1, i] = 1\n",
    "    return x0\n",
    "#define the neighbor pixels\n",
    "def weight_update(N, P):\n",
    "    \"\"\"Return smooth matrix h (numpy array)\"\"\"\n",
    "    weight_update = np.ones([N,P])\n",
    "    for i in range(N):\n",
    "        if i==0:\n",
    "            weight_update[i,0]=0\n",
    "            weight_update[i,1]=0\n",
    "        elif i==width-1:\n",
    "            weight_update[i,0]=0\n",
    "            weight_update[i,2]=0\n",
    "        elif i==(height-1)*width:\n",
    "            weight_update[i,1]=0\n",
    "            weight_update[i,3]=0\n",
    "        elif i==height*width-1:\n",
    "            weight_update[i,2]=0\n",
    "            weight_update[i,3]=0\n",
    "        elif i//width==0 and i%width>0:\n",
    "            weight_update[i,0]=0\n",
    "        elif i//width==height-1 and i//width>0:\n",
    "            weight_update[i,3]=0\n",
    "        elif i%height==0:\n",
    "            weight_update[i,1]=0\n",
    "        elif i%height==width-1:\n",
    "            weight_update[i,2]=0\n",
    "    return weight_update\n",
    "#the LRS weight of neighbor pixels\n",
    "def creat_weight_crf():\n",
    "    x0=neighbor_x(height=height, width=width)\n",
    "    index=np.transpose(np.nonzero(x0.reshape(-1)))\n",
    "    weight_crf=x0\n",
    "    weight0=weight_update(N, 4)\n",
    "    weight=torch.randn(N,4).numpy()\n",
    "    weight_new=weight*weight0\n",
    "    weight_ele=weight_new.reshape(-1)[np.transpose(np.nonzero(weight0.reshape(-1)))]\n",
    "    weight_crf.reshape(-1)[index]=weight_ele.reshape(-1,1)\n",
    "    return weight_crf\n",
    "weight_crf=creat_weight_crf()\n",
    "weight_crf=torch.from_numpy(weight_crf).float()\n",
    "weight_crf=weight_crf.cuda()\n",
    "weight_crf.requires_grad = True\n",
    "\n",
    "#the abundance estimation by FCLS\n",
    "d0=excel_to_matrix('./dataset/FCLS_abd_samson.xls')\n",
    "d0=torch.from_numpy(d0).float()\n",
    "d0=d0.cuda()\n",
    "\n",
    "#define LCS model\n",
    "def CRFI(code,d0,weight_crf):\n",
    "    step1=code[:,0].reshape(-1,1)+2*torch.mm(weight_crf,torch.unsqueeze(d0[:,0],dim=1))\n",
    "    step2=torch.ones((len(code[:,0]),1)).cuda()+2*torch.mm(weight_crf,torch.ones((len(code[:,0]),1)).cuda())\n",
    "    d1=step1/step2\n",
    "    step1=code[:,1].reshape(-1,1)+2*torch.mm(weight_crf,torch.unsqueeze(d0[:,1],dim=1))\n",
    "    step2=torch.ones((len(code[:,1]),1)).cuda()+2*torch.mm(weight_crf,torch.ones((len(code[:,1]),1)).cuda())\n",
    "    d2=step1/step2\n",
    "    step1=code[:,2].reshape(-1,1)+2*torch.mm(weight_crf,torch.unsqueeze(d0[:,2],dim=1))\n",
    "    step2=torch.ones((len(code[:,2]),1)).cuda()+2*torch.mm(weight_crf,torch.ones((len(code[:,2]),1)).cuda())\n",
    "    d3=step1/step2\n",
    "    d=torch.cat((d1,d2,d3),dim=1)\n",
    "    return d\n",
    "\n",
    "\n",
    "class Stream:\n",
    "\tptr = torch.cuda.current_stream().cuda_stream\n",
    "\n",
    "IRNNForward = open('./model/IRNN_Forward_cuda.cu','r').read()\n",
    "\n",
    "IRNNBackward = open('./model/IRNN_Backward_cuda.cu','r').read()\n",
    "\n",
    "@cupy.memoize(for_each_device=True)\n",
    "def cunnex(strFunction):\n",
    "\treturn cupy.cuda.compile_with_cache(globals()[strFunction]).get_function(strFunction)\n",
    "# end\n",
    "\n",
    "#define GRS model\n",
    "def irnn(input_feature, weight_up, weight_right, weight_down, weight_left, bias_up, bias_right, bias_down, bias_left):\n",
    "    input_feature=input_feature.cuda()\n",
    "    weight_up=weight_up.cuda()\n",
    "    weight_right=weight_right.cuda()\n",
    "    weight_down=weight_down.cuda()\n",
    "    weight_left=weight_left.cuda()\n",
    "    bias_up=bias_up.cuda()\n",
    "    bias_right=bias_right.cuda()\n",
    "    bias_down=bias_down.cuda()\n",
    "    bias_left=bias_left.cuda()\n",
    "    \n",
    "    assert(input_feature.is_contiguous() == True)\n",
    "    assert(weight_left.is_contiguous() == True)\n",
    "    assert(weight_right.is_contiguous() == True)\n",
    "    assert(weight_down.is_contiguous() == True)\n",
    "\n",
    "    assert(weight_up.is_contiguous() == True)\n",
    "    assert(bias_left.is_contiguous() ==True)\n",
    "    assert(bias_right.is_contiguous() == True)\n",
    "    assert(bias_up.is_contiguous() == True)\n",
    "    assert(bias_down.is_contiguous() == True)\n",
    "\n",
    "    output_left = input_feature.clone()\n",
    "    output_right = input_feature.clone()\n",
    "    output_up = input_feature.clone()\n",
    "    output_down = input_feature.clone()\n",
    "    \n",
    "\n",
    "    if input_feature.is_cuda == True:\n",
    "        n = input_feature.nelement()\n",
    "        cuda_num_threads = 1024\n",
    "        cunnex('IRNNForward')(\n",
    "            grid=tuple([ int((n +  cuda_num_threads - 1) / cuda_num_threads ), 1, 1 ]),\n",
    "            block=tuple([ cuda_num_threads , 1, 1 ]),\n",
    "            args=[\n",
    "                input_feature.data_ptr(),\n",
    "\t\t\t\n",
    "\t\t\t\tweight_up.data_ptr(), \n",
    "\t\t\t\tweight_right.data_ptr(),\n",
    "\t\t\t\tweight_down.data_ptr(),\n",
    "\t\t\t\tweight_left.data_ptr(),\n",
    "                \n",
    "\t\t\t\tbias_up.data_ptr(),\n",
    "\t\t\t\tbias_right.data_ptr(),\n",
    "\t\t\t\tbias_down.data_ptr(),\n",
    "\t\t\t\tbias_left.data_ptr(),\n",
    "\t\n",
    "\t\t\t\toutput_up.data_ptr(), \n",
    "\t\t\t\toutput_right.data_ptr(),\n",
    "\t\t\t\toutput_down.data_ptr(), \n",
    "\t\t\t\toutput_left.data_ptr(),\n",
    "\t\n",
    "\t\t\t\tinput_feature.size(1),\n",
    "\t\t\t\tinput_feature.size(2),\n",
    "\t\t\t\tinput_feature.size(3),\n",
    "\t\t\t\tn],\n",
    "\t\t\tstream=Stream\n",
    "\t\t)\n",
    "\t\t\n",
    "\t\t\n",
    "    return output_up,output_right,output_down,output_left\n",
    "\t# end\n",
    "\n",
    "#----------------------------------the global model----------------------------------------#\n",
    "\n",
    "class Spacial_IRNN(nn.Module):\n",
    "    def __init__(self,in_channels,alpha=1.0):\n",
    "        super(Spacial_IRNN,self).__init__()\n",
    "        self.left_weight = nn.Conv2d(in_channels,in_channels,kernel_size=1,stride=1,groups=in_channels,padding=0)\n",
    "        self.right_weight = nn.Conv2d(in_channels,in_channels,kernel_size=1,stride=1,groups=in_channels,padding=0)\n",
    "        self.up_weight = nn.Conv2d(in_channels,in_channels,kernel_size=1,stride=1,groups=in_channels,padding=0)\n",
    "        self.down_weight = nn.Conv2d(in_channels,in_channels,kernel_size=1,stride=1,groups=in_channels,padding=0)\n",
    "        self.left_weight.weight  = nn.Parameter(torch.tensor([[[[alpha]]]]*in_channels))\n",
    "        self.right_weight.weight = nn.Parameter(torch.tensor([[[[alpha]]]]*in_channels))\n",
    "        self.up_weight.weight    = nn.Parameter(torch.tensor([[[[alpha]]]]*in_channels))\n",
    "        self.down_weight.weight  = nn.Parameter(torch.tensor([[[[alpha]]]]*in_channels))\n",
    "        \n",
    "\n",
    "    def forward(self,input):\n",
    "        return irnn(input,self.up_weight.weight,self.right_weight.weight,self.down_weight.weight,self.left_weight.weight, self.up_weight.bias,self.right_weight.bias,self.down_weight.bias,self.left_weight.bias)\n",
    "\n",
    "\n",
    "class TS_IRNN(nn.Module):\n",
    "    def __init__(self, L, P):  #L bands   P  the number of endmembers\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(L, L//2, bias=True), \n",
    "            nn.Tanh(),\n",
    "            nn.Linear(L//2, L//4, bias=True),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(L//4, P, bias=True),\n",
    "            nn.LeakyReLU(1e-2,inplace=False),\n",
    "        )\n",
    "        self.irnn1 = Spacial_IRNN(P)\n",
    "        self.irnn2 = Spacial_IRNN(P)\n",
    "        self.BN1=nn.BatchNorm2d(4*P,affine=True)\n",
    "        self.BN2=nn.BatchNorm2d(4*P,affine=True)\n",
    "        self.decoder = nn.Linear(P, L, bias=False)\n",
    "        self.conv1=nn.Conv2d(4*P, P, 1)\n",
    "        self.conv2=nn.Conv2d(4*P, P, 1)\n",
    "        self.conv3=nn.Conv2d(P, P, 1)\n",
    "        self.act=nn.Softmax(dim=1)\n",
    "        self.tanh=nn.Tanh()\n",
    "        self.relu=nn.ReLU()\n",
    "        self.leakyrelu=nn.LeakyReLU(1e-2,inplace=False)\n",
    "\n",
    "    def forward(self, y):#y N×L\n",
    "        code_AAS = self.encoder(y) #abundance  the shape of code_AAS is N*P\n",
    "        d=CRFI(code_AAS,d0,torch.abs(weight_crf))\n",
    "        for i in range(max_iter2):\n",
    "            d=CRFI(code_AAS,d,torch.abs(weight_crf))\n",
    "        code=code_AAS.T.reshape(P,height,width)\n",
    "        code=code.unsqueeze(dim=0).contiguous()\n",
    "        top_up,top_right,top_down,top_left=self.irnn1(code) \n",
    "        code_g = torch.cat([top_up,top_right,top_down,top_left],dim=1)\n",
    "        code_g = self.leakyrelu(code_g)\n",
    "        code_g = self.BN1(code_g)\n",
    "        code_g = self.conv1(code_g)\n",
    "        top_up,top_right,top_down,top_left=self.irnn2(code_g) \n",
    "        code_g = torch.cat([top_up,top_right,top_down,top_left],dim=1)\n",
    "        code_g = self.leakyrelu(code_g)\n",
    "        code_g = self.BN2(code_g)\n",
    "        code_g = self.conv2(code_g)\n",
    "        code_g_out = code_g.unsqueeze(dim=0).reshape(P,-1).T\n",
    "        code=(code_g+d.T.reshape(-1,height,width).unsqueeze(dim=0))/2\n",
    "        code=code.squeeze(dim=0).reshape(P,-1).T\n",
    "        code=self.act(code)\n",
    "        output = self.decoder(code) #reconstruction\n",
    "        return code_AAS,code_g_out,code,output\n",
    "\n",
    "    \n",
    "#define matrix H\n",
    "def smooth_matrix(height, width):\n",
    "    \"\"\"Return smooth matrix h (numpy array)\"\"\"\n",
    "    N = height * width\n",
    "    h = np.eye(N)\n",
    "    for i in range(N):\n",
    "        if i - width >= 0 and i + width <= N - 1 and i % width != 0 and (i + 1) % width != 0:\n",
    "            h[i + 1, i] = -0.25\n",
    "            h[i - 1, i] = -0.25\n",
    "            h[i + width, i] = -0.25\n",
    "            h[i - width, i] = -0.25\n",
    "        elif i - width < 0:\n",
    "            if i == 0 or i == width - 1:\n",
    "                h[1, 0] = -0.5\n",
    "                h[width, 0] = -0.5\n",
    "                h[width - 2, width - 1] = -0.5\n",
    "                h[2 * width - 1, width - 1] = -0.5\n",
    "            else:\n",
    "                h[i - 1, i] = -1 / 3\n",
    "                h[i + 1, i] = -1 / 3\n",
    "                h[i + width, i] = -1 / 3\n",
    "        elif i + width > N - 1:\n",
    "            if i == N - 1 or i == N - width:\n",
    "                h[N - 2, N - 1] = -0.5\n",
    "                h[N - width, N - 1] = -0.5\n",
    "                h[N - width + 1, N - width] = -0.5\n",
    "                h[N - 2 * width, N - width] = -0.5\n",
    "            else:\n",
    "                h[i - 1, i] = -1 / 3\n",
    "                h[i + 1, i] = -1 / 3\n",
    "                h[i - width, i] = -1 / 3\n",
    "        elif i % width == 0:\n",
    "            h[i - width, i] = -1 / 3\n",
    "            h[i + width, i] = -1 / 3\n",
    "            h[i + 1, i] = -1 / 3\n",
    "        elif (i + 1) % width == 0:\n",
    "            h[i - width, i] = -1 / 3\n",
    "            h[i + width, i] = -1 / 3\n",
    "            h[i - 1, i] = -1 / 3\n",
    "    return h\n",
    "\n",
    "def norm2squ(x):  #compute l-2 norm\n",
    "    \"\"\"Return L2norm^2\"\"\"\n",
    "    return torch.sum(torch.pow(x, exponent=2))\n",
    "\n",
    "h1=torch.from_numpy(smooth_matrix(height=height, width=width)).float()\n",
    "h1=h1.cuda()\n",
    "h1.requires_grad = True\n",
    "h2=torch.from_numpy(smooth_matrix(height=height, width=width)).float()\n",
    "h2=h2.cuda()\n",
    "h2.requires_grad = True\n",
    "h3=torch.from_numpy(smooth_matrix(height=height, width=width)).float()\n",
    "h3=h3.cuda()\n",
    "h3.requires_grad = True\n",
    "\n",
    "\n",
    "model=TS_IRNN(L,P).cuda()\n",
    "model.train()\n",
    "\n",
    "#Initialize the decoder weight\n",
    "layer_name='decoder.weight'\n",
    "def _load_endmember(edm, layer_name):\n",
    "    edm = torch.from_numpy(edm)\n",
    "    model_dict =model.state_dict()\n",
    "    model_dict[layer_name] = edm\n",
    "    model.load_state_dict(model_dict)\n",
    "_load_endmember(np.array(edm.cpu()), layer_name)\n",
    "\n",
    "\n",
    "frozen_layer='decoder.weight'\n",
    "def _freeze(): #freeze self.frozen_layer\n",
    "        if frozen_layer is None:\n",
    "            return\n",
    "        for name, value in model.named_parameters():\n",
    "            if name in frozen_layer:\n",
    "                value.requires_grad = False\n",
    "_freeze()\n",
    "\n",
    "\n",
    "def _unfreeze(): #unfreeze self.frozen_layer\n",
    "        if frozen_layer is None:\n",
    "            return\n",
    "        for name, para in model.named_parameters():\n",
    "            if name in frozen_layer:\n",
    "                para.requires_grad = True\n",
    "                \n",
    "                        \n",
    "#define regularization_loss\n",
    "def _regularization_loss():\n",
    "    r_loss = 0\n",
    "    reg_layer: list = ['encoder.0.weight', 'encoder.2.weight', 'encoder.4.weight']\n",
    "    reg_norm: list = [2, 2, 2] \n",
    "    p=reg_norm\n",
    "    for name, param in model.named_parameters():\n",
    "        if name in reg_layer:\n",
    "            norm_num = p[reg_layer.index(name)]\n",
    "            if norm_num == 1 or norm_num == 2:\n",
    "                r_loss += 1e-5 * torch.norm(param, p=norm_num)\n",
    "    return r_loss\n",
    "\n",
    "\n",
    "loss_fn = torch.nn.MSELoss(reduction='mean')\n",
    "s0=torch.abs(h1) > 0.1\n",
    "s1=h1 < -0.1\n",
    "s2=torch.eye(height*width)\n",
    "t0=weight_crf<0\n",
    "s0=s0.cuda()\n",
    "s1=s1.cuda()\n",
    "s2=s2.cuda()\n",
    "t0=t0.cuda()\n",
    "def loss_total(y,height,width):\n",
    "    \n",
    "    pix = y.T \n",
    "    code_AAS,code_g,code,output= model(pix)\n",
    "    weight_loss=1e-6*norm2squ(weight_crf*t0)\n",
    "    params=list(model.named_parameters())\n",
    "    spacial_loss2=1e-6 * norm2squ(torch.mm(h2*s0, code))+5 * (norm2squ(torch.sum(h2*s1, dim=0)+torch.tensor(1)) +norm2squ(h2*s2-s2))\n",
    "    code_loss=1e-5 * torch.mean(torch.sum(torch.sqrt(torch.abs(code)+1e-4), dim=-1))\n",
    "    base_loss1=torch.mean(torch.acos(torch.sum(output * pix, dim=-1) /(torch.norm(output, dim=-1, p=2)*torch.norm(pix, dim=-1, p=2)+1e-4)))\n",
    "    regularization_loss = _regularization_loss()\n",
    "    loss=spacial_loss2+code_loss+base_loss1+regularization_loss+weight_loss\n",
    "    params=list(model.named_parameters())\n",
    "    return loss,spacial_loss2,code_loss,base_loss1,regularization_loss,params\n",
    "\n",
    "opt = torch.optim.Adam([{'params': model.encoder.parameters(), 'lr': 1e-3},\n",
    "                                         {'params': model.decoder.parameters(), 'lr': 1e-2},\n",
    "                                         {'params': h1, 'lr': 1e-1},\n",
    "                                         {'params': h2, 'lr': 1e-1},\n",
    "                                         {'params': h3, 'lr': 1e-1},\n",
    "                                         {'params':weight_crf,'lr':1e-3},\n",
    "                                         {'params': model.conv1.parameters(), 'lr': 1e-2},\n",
    "                                         {'params': model.BN1.parameters(), 'lr': 5e-3},\n",
    "                                         {'params': model.irnn1.parameters(), 'lr': 1e-2},\n",
    "                                         {'params': model.conv2.parameters(), 'lr': 1e-2},\n",
    "                                         {'params': model.BN2.parameters(), 'lr': 5e-3},\n",
    "                                         {'params': model.irnn2.parameters(), 'lr': 1e-2}])\n",
    "\n",
    "def Unmixing(y,height,width,epoch):\n",
    "    loss,spacial_loss2,code_loss,base_loss1,regularization_loss,params=loss_total(y,height,width)\n",
    "    assert torch.isnan(loss).sum() == 0, print(loss)\n",
    "    opt.zero_grad()\n",
    "    #nn.utils.clip_grad_norm_(model.parameters(), max_norm=1, norm_type=1)  #clip grad\n",
    "    loss.backward()\n",
    "    assert torch.isnan(model.BN1.weight).sum() == 0, print(model.BN1.weight)\n",
    "    assert torch.isnan(model.BN1.weight.grad).sum() == 0, print(model.BN1.weight.grad)\n",
    "    opt.step()\n",
    "    assert torch.isnan(model.BN1.weight.grad).sum() == 0, print(model.BN1.weight.grad)\n",
    "    assert torch.isnan(model.BN1.weight).sum() == 0, print(model.BN1.weight)\n",
    "    if (epoch+1) % (max_iter//10) == 0:\n",
    "        print('epoch [{}/{}], loss:{:.8f}, '.format(epoch + 1, max_iter, loss.item()), end=' ')\n",
    "        print('epoch [{}/{}], spacial_loss2:{:.8f}, '.format(epoch + 1, max_iter, spacial_loss2.item()), end=' ')\n",
    "        print('epoch [{}/{}], code_loss:{:.8f}, '.format(epoch + 1, max_iter, code_loss.item()), end=' ')\n",
    "        print('epoch [{}/{}], base_loss1:{:.8f}, '.format(epoch + 1, max_iter, base_loss1.item()), end=' ')\n",
    "        print('epoch [{}/{}], regularization_loss:{:.8f}, '.format(epoch + 1, max_iter, regularization_loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch [300/3000], loss:0.05655465,  epoch [300/3000], spacial_loss2:0.00003190,  epoch [300/3000], code_loss:0.00001435,  epoch [300/3000], base_loss1:0.04442556,  epoch [300/3000], regularization_loss:0.00016825, \n",
      "epoch [600/3000], loss:0.04914539,  epoch [600/3000], spacial_loss2:0.00003613,  epoch [600/3000], code_loss:0.00001441,  epoch [600/3000], base_loss1:0.04053157,  epoch [600/3000], regularization_loss:0.00017602, \n",
      "epoch [900/3000], loss:0.04591200,  epoch [900/3000], spacial_loss2:0.00003924,  epoch [900/3000], code_loss:0.00001437,  epoch [900/3000], base_loss1:0.03950562,  epoch [900/3000], regularization_loss:0.00017742, \n",
      "epoch [1200/3000], loss:0.04394728,  epoch [1200/3000], spacial_loss2:0.00004105,  epoch [1200/3000], code_loss:0.00001433,  epoch [1200/3000], base_loss1:0.03902149,  epoch [1200/3000], regularization_loss:0.00017703, \n",
      "epoch [1500/3000], loss:0.04265584,  epoch [1500/3000], spacial_loss2:0.00004184,  epoch [1500/3000], code_loss:0.00001431,  epoch [1500/3000], base_loss1:0.03871992,  epoch [1500/3000], regularization_loss:0.00017724, \n",
      "epoch [1800/3000], loss:0.04172334,  epoch [1800/3000], spacial_loss2:0.00004204,  epoch [1800/3000], code_loss:0.00001431,  epoch [1800/3000], base_loss1:0.03849358,  epoch [1800/3000], regularization_loss:0.00017884, \n",
      "epoch [2100/3000], loss:0.04105126,  epoch [2100/3000], spacial_loss2:0.00004209,  epoch [2100/3000], code_loss:0.00001429,  epoch [2100/3000], base_loss1:0.03832112,  epoch [2100/3000], regularization_loss:0.00018176, \n",
      "epoch [2400/3000], loss:0.04053621,  epoch [2400/3000], spacial_loss2:0.00004224,  epoch [2400/3000], code_loss:0.00001428,  epoch [2400/3000], base_loss1:0.03818199,  epoch [2400/3000], regularization_loss:0.00018556, \n",
      "epoch [2700/3000], loss:0.04015298,  epoch [2700/3000], spacial_loss2:0.00004247,  epoch [2700/3000], code_loss:0.00001426,  epoch [2700/3000], base_loss1:0.03807964,  epoch [2700/3000], regularization_loss:0.00018962, \n",
      "epoch [3000/3000], loss:0.03689106,  epoch [3000/3000], spacial_loss2:0.00003841,  epoch [3000/3000], code_loss:0.00001460,  epoch [3000/3000], base_loss1:0.03503936,  epoch [3000/3000], regularization_loss:0.00018808, \n"
     ]
    }
   ],
   "source": [
    "for epoch in range(max_iter): #train\n",
    "    if (epoch+1) % 200 == 0:\n",
    "        for i, dic in enumerate(opt.param_groups): #opt.param_groups \n",
    "            dic['lr'] *= lr_decay  #decay lr\n",
    "    if epoch >= max_iter * unfreeze_time:\n",
    "        _unfreeze() \n",
    "    Unmixing(im.T,height,width,epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62976 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 62812 characters]",
      "text/plain": [
       "<Figure size 432x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "pix = im #test\n",
    "code_AAS,code_g,code,output= model(pix)\n",
    "output=output.cpu().data.detach().numpy()\n",
    "output=np.reshape(output, (height, width, -1))\n",
    "plt.subplot(121)\n",
    "plt.title('Reconstructed')\n",
    "plt.imshow(output[:, :, [90, 60, 30]])\n",
    "plt.subplot(122)\n",
    "plt.title('Real')\n",
    "y_plot = np.reshape(im.cpu(), (height, width, -1))\n",
    "plt.imshow(y_plot[:, :, [90, 60, 30]])\n",
    "plt.show()\n",
    "code=code.reshape(-1,P)\n",
    "#code=torch.cat((code[:,1].reshape(N,1),code[:,2].reshape(N,1),code[:,3].reshape(N,1),code[:,0].reshape(N,1)),1)\n",
    "code=code.cpu().data.numpy() #abundance\n",
    "code = np.reshape(code, (height, width, -1))\n",
    "for i in range(3):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.imshow(code[:,:,i]) \n",
    "\n",
    "params=list(model.named_parameters())\n",
    "print(params.__len__())\n",
    "m_hat=params[26][1].data #endmember"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 60440 characters]",
      "text/plain": [
       "<Figure size 432x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "reference_abundance=A.T\n",
    "reference_abundance=np.hstack((reference_abundance[:,2].reshape(N,1),reference_abundance[:,1].reshape(N,1),reference_abundance[:,0].reshape(N,1)))\n",
    "reference_abundance=reference_abundance.reshape(height,width,-1)\n",
    "for i in range(3):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.imshow(reference_abundance[:,:,i]) \n",
    "params=list(model.named_parameters())\n",
    "print(params.__len__())\n",
    "m_hat=params[26][1].data \n",
    "#eSAD\n",
    "#torch.mean(torch.acos(torch.sum(m_hat.cpu().T * torch.from_numpy(M).float().T, dim=-1) /(torch.norm(m_hat.cpu().T, dim=-1, p=2)*torch.norm(torch.from_numpy(M).float().T, dim=-1, p=2))))\n",
    "#aSAD\n",
    "code=code.reshape(-1,P)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0302)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "P=3\n",
    "temp=torch.from_numpy(reference_abundance.reshape(-1,P)).float()-torch.from_numpy(code.reshape(-1,P)).float()\n",
    "torch.sum(temp**2)/height/width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0793)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#eSAD\n",
    "m_hat=m_hat.cpu()\n",
    "M[:,0]=M[:,0]/np.max(M[:,0])\n",
    "M[:,1]=M[:,1]/np.max(M[:,1])\n",
    "M[:,2]=M[:,2]/np.max(M[:,2])\n",
    "m_hat[:,0]=m_hat[:,0]/torch.max(m_hat[:,0])\n",
    "m_hat[:,1]=m_hat[:,1]/torch.max(m_hat[:,1])\n",
    "m_hat[:,2]=m_hat[:,2]/torch.max(m_hat[:,2])\n",
    "reference_M=np.hstack((M[:,2].reshape(L,1),M[:,1].reshape(L,1),M[:,0].reshape(L,1)))\n",
    "torch.mean(torch.acos(torch.sum(m_hat.T * torch.from_numpy(reference_M).float().T, dim=-1) /(torch.norm(m_hat.T, dim=-1, p=2)*torch.norm(torch.from_numpy(reference_M).float().T, dim=-1, p=2))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.11477678172660011"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#rRMSE\n",
    "pix = im.cpu()\n",
    "loss_fn = torch.nn.MSELoss(reduction='mean')\n",
    "loss_fn(pix,torch.from_numpy(output.reshape(-1,L)).float()).item()**0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "abd=code.reshape(-1,P)\n",
    "def save(data, path):\n",
    "    f = xlwt.Workbook()  \n",
    "    sheet1 = f.add_sheet(u'sheet1', cell_overwrite_ok=True)  # create sheet\n",
    "    [h, l] = data.shape  \n",
    "    for i in range(h):\n",
    "        for j in range(l):\n",
    "            sheet1.write(i, j, data[i, j])\n",
    "    f.save(path)\n",
    "save(abd.astype(np.float64),'abd_CRF_10.xls') \n",
    "save(m_hat.cpu().numpy().astype(np.float64),'edm_CRF_10.xls') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myconda",
   "language": "python",
   "name": "myconda"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
