GTCTV-DPC_traffic.py

import os
import time
import torch 
import logging

import numpy as np

import sys
sys.path.append("./utils/")
from utils_log import logger_info, logger_close
from utils_data import initial_seed
from utils_data import read_traffic_data, missing_pattern
from utils_data import compute_rmse, compute_mape
from utils_lowrank import tensor_svd, shrinkage
from utils_lowrank import dct_unit, idct_unit
from utils_lowrank import diff_mdi, diff_mdi_T, KKT_mdi

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

def svt_lrtc(tensor, tau, lamb, gamma, 
             unitary_operation=dct_unit, 
             inv_unitary_operation=idct_unit,
             ):
    u, s, v = tensor_svd(unitary_operation(tensor))
    ss = shrinkage(s, [tau, lamb, gamma], mode="scad")
    # ss = shrinkage(s, tau, mode="soft")
    try:
        idx = torch.max(torch.where(ss > 0)[-1]) + 1
        # if the unitary_operation is DFT, the dtype of u, v is complexfloat32, while the dtype of ss is float32.
        return inv_unitary_operation(u[..., :idx] @ torch.diag_embed(ss[..., :idx]).to(u.dtype) @ v[..., :idx, :])
    except:
        # when the shrinkage threshold is too large, the result is all zeros
        # if the shrinkage result is all zeros, we return a zero tensor with the same shape as the input tensor
        return torch.zeros_like(tensor)

def deep_pnp_solver(tensor, net, sigma):
    # tensor: (T, C, H, W)
    T, _, H, W = tensor.shape
    max_value = torch.max(tensor)

    noise_level_map = torch.ones(T, 1, H, W).mul_(sigma).float().to(tensor.device)
    input = torch.cat([tensor / max_value, noise_level_map], dim=1)
    with torch.no_grad():
        return net(input).clamp(0., 1.) * max_value

def fidelity_prox(tensor, Y, mask):
    # proximal operator for completion problem
    return (1 - mask) * tensor + Y

def dys_restore(
        observe_tensor, mask, # fidelity term
        tau, beta, rho, lamb, gamma, orders, inner_iter, # tctv-scad term
        model, sigma_init, sigma_min, alpha, # deep denoiser
        check_iter=1, max_iter=200, tole=1e-4,
        ori_tensor=None, logger=None, 
        ):
    
    initial_seed(1000)
    # variable initialization
    z = observe_tensor.clone()
    pos_missing = torch.where(mask == 0)
    z[pos_missing] = torch.mean(observe_tensor[mask != 0])
    out_tensor = z.clone()

    # here sigma is the noise level, which divided by 255
    sigma = sigma_init

    # variable initialization for proximal solution in Low-rank with dim-1 TV
    n = len(orders)
    Gs = [diff_mdi(z, axis=order) for order in orders]
    Bs = [torch.zeros_like(z) for _ in orders]
    [T, C, H, W] = z.shape
    mu = 4 / (gamma - 1)  # relaxation parameter for SCAD-TCTV

    # prepare the KKT tensor for the total variation
    # only use the shape of the tensor to compute the KKT tensor
    KKT_fft = sum(KKT_mdi(z, axis=order) for order in orders)
    
    # variable initialization for proximal solution in fidelity term
    used_time = 0
    rmses = []
    mapes = []

    # iterative optimization
    for i in range(max_iter):
        if ori_tensor is not None:
            if i % check_iter == 0:
                # compute the MAPE, RMSE
                pos_test = torch.where((ori_tensor != 0) & (observe_tensor == 0))
                mape_value = compute_mape(ori_tensor[pos_test], out_tensor[pos_test]) * 100
                rmse_value = compute_rmse(ori_tensor[pos_test], out_tensor[pos_test])
                mapes.append([i, mape_value.cpu()])
                rmses.append([i, rmse_value.cpu()])
                logger.info(f"Iter: {i}, MAPE: {mape_value:.6f}, RMSE: {rmse_value:.6f}")
                
        # record the time and update the input sigma value
        start_time = time.time()

        # DYS step
        # proximal operator for Low-rank with dim-{orders} TV
        for inner_it in range(inner_iter):
            ### tctv
            # update x_B
            H = sum(diff_mdi_T(beta * Gs[index] - Bs[index], axis=orders[index]) for index in range(n))
            Nominator = torch.fft.fftn(tau * H + z)
            Denominator = tau * beta * KKT_fft + 1 + mu * tau
            x_B = torch.real(torch.fft.ifftn(Nominator / Denominator))

            # update the gradient tensors and lagrange multipliers
            for index in range(n):
                temp = diff_mdi(x_B, axis=orders[index]) + Bs[index] / beta
                Gs[index] = svt_lrtc(temp, 1 / (n * beta), lamb, gamma)
                Bs[index] += beta * (diff_mdi(x_B, axis=orders[index]) - Gs[index])
            
            beta = min(beta * rho, 1e10)
            if inner_it > 0:
                inner_tol = torch.linalg.norm(x_B - x_B_old) / torch.linalg.norm(x_B_old)
                if inner_tol < 1e-5:
                    logger.info("Inner iteration converged at iter: {}".format(inner_it+1))
                    break
            
            x_B_old = x_B.clone()

        x_C = deep_pnp_solver(x_B, model, sigma)
        temp_x_A = 2 * x_B - z - tau * alpha * (x_B - x_C) # 2 * J_{tau * B} - I - tau * C(J_{tau * B})
        x_A = fidelity_prox(temp_x_A, observe_tensor, mask) # A(x_B) fidelity term of completion problem

        sigma = max(sigma / rho, sigma_min)
        del x_C, temp_x_A

        # update the z
        # z += lamb_t * (x_A - x_B)
        if i < 100:
            z += (x_A - x_B) # set lamb_t to 1.
        else:
            z += 100 / i * (x_A - x_B) # set lamb_t to 1 / t.

        # record the time
        used_time += time.time() - start_time
        
        tol = torch.linalg.norm(out_tensor - x_A) / torch.linalg.norm(out_tensor)
        if tol < tole:
            break
        out_tensor = x_A.clone()
    
    # compute the PSNR, SSIM
    if ori_tensor is not None:
        pos_test = torch.where((ori_tensor != 0) & (observe_tensor == 0))
        mape = compute_mape(ori_tensor[pos_test], out_tensor[pos_test]) * 100
        rmse = compute_rmse(ori_tensor[pos_test], out_tensor[pos_test])
        mapes.append([i+1, mape_value.cpu()])
        rmses.append([i+1, rmse_value.cpu()])
        logger.info(f"Total iteration: {i + 1}, MAPE: {mape:.6f}, RMSE: {rmse:.6f}")
        logger.info("Total time: {:.5f}s".format(used_time))
        
    return out_tensor, mapes, rmses, used_time

param_file = "./params/SPC09_gray.pth"
from utils_drunet.network_unet import UNetRes as net
model = net(
    in_nc=2, 
    out_nc=1, 
    nc=[64, 128, 256, 512], 
    nb=4, 
    act_mode="R", 
    downsample_mode="strideconv", 
    upsample_mode="convtranspose", 
    bias=False,
    interpolation_mode="bilinear"
)
state_dict = torch.load(param_file, weights_only=True)
model.load_state_dict(state_dict, strict=True)
model.eval()

data_root = './test_datasets/traffic_datas/'

beta, rho = 1e-4, 1.02
hypers = {
    'Guangzhou': (5.00, 2e3, 0.85, 1.50), 
    'Seattle': (3.00, 3e3, 0.95, 2.00),
    'PeMS': (3.00, 2e2, 0.65, 2.00),
}
inner_iter = 5
sigma_min = 1e-3
tau = 1.

datasets = ['Guangzhou', 'Seattle', 'PeMS']
sampling_rates = [0.3, 0.5, 0.7]
res_dir_root = "./exp_results/traffics"

for dataset in datasets:
    # create the log and result directories
    temp_res_dir_root = res_dir_root + f'_{dataset}'
    log_dir = os.path.join(temp_res_dir_root, "logs")
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    res_dir = os.path.join(temp_res_dir_root, "results")
    if not os.path.exists(res_dir):
        os.makedirs(res_dir)
    lamb, gamma, sigma_init, alpha = hypers[dataset]
    
    rmses_sum = 0
    mapes_sum = 0
        
    for sampling_rate in sampling_rates:
        dense_tensor = read_traffic_data(
            dataset_name = dataset,
            dataroot = data_root,
        )
        binary_tensor = missing_pattern(dense_tensor, 1 - sampling_rate, seed=1000)
        sparse_tensor = dense_tensor * binary_tensor
        
        # (H, W, T) to (T, 1, H, W), where H means the number of locations, W means the number of time intervals in one day, T means the number of days.
        dense_tensor = dense_tensor.permute(1, 0, 2).unsqueeze(1)
        sparse_tensor = sparse_tensor.permute(1, 0, 2).unsqueeze(1)
        mask = torch.ones_like(sparse_tensor)
        mask[sparse_tensor == 0] = 0
            
        logger_info(log_dir, log_path=os.path.join(log_dir, f"sampling_rate_{sampling_rate}.log"))
        result_logger = logging.getLogger(log_dir)
        result_logger.info("--------------------------------------------------------------------------------------------------------")
                
        out_tensor, mapes, rmses, used_time = dys_restore(
            sparse_tensor.to(device), mask=mask.to(device), 
            tau=tau, beta=beta, rho=rho, lamb=lamb, gamma=gamma, orders=[0, 2, 3], inner_iter=inner_iter,
            model=model.to(device), sigma_init=sigma_init, sigma_min=sigma_min, alpha=alpha, 
            ori_tensor=dense_tensor.to(device), logger=result_logger,
            check_iter=100, max_iter=200, tole=1e-4,
            )
        
        rmses_sum += rmses[-1][-1]
        mapes_sum += mapes[-1][-1]
        # save the results
        save_dir_name = os.path.join(res_dir, f"sampling_rate_{sampling_rate}")
        if not os.path.exists(save_dir_name):
            os.makedirs(save_dir_name)
        # (T, 1, H, W) to (H, W, T)
        save_data = out_tensor.cpu().detach().squeeze(1).permute(1, 0, 2).numpy()
        np.save(os.path.join(save_dir_name, "output.npy"), save_data)
        np.save(os.path.join(save_dir_name, "rmses.npy"), rmses)
        np.save(os.path.join(save_dir_name, "mapes.npy"), mapes)
        result_logger.info("--------------------------------------------------------------------------------------------------------")
        logger_close(result_logger)
    
    # save the average results
    avg_rmse = rmses_sum / len(sampling_rates)
    avg_mape = mapes_sum / len(sampling_rates)
    logger_info(log_dir, log_path=os.path.join(log_dir, f"avg-rmse_{avg_rmse:.3f}_mape_{avg_mape:.4f}.log"))
    temp_logger = logging.getLogger(log_dir)
    temp_logger.info("--------------------------------------------------------------------------------------------------------")
    temp_logger.info("Average rmse: {:.3f}, mape: {:.4f}".format(avg_rmse, avg_mape))
    temp_logger.info("--------------------------------------------------------------------------------------------------------")
    logger_close(temp_logger)