th_train_codec_mixed_residual.py

"""Physics-constraint surrogates.
Convolutional Encoder-decoder networks for surrogate modeling of darcy flow.
Assume the PDEs and boundary conditions are known.
Train the surrogate with mixed residual loss, instead of maximum likelihood.

5 runs per setup
setup:
    - training with different number of input: 
        512, 1024, 2048, 4096, 8192
    with mini-batch size 8, 8, 16, 32, 32, correpondingly.
    - metric: 
        relative L2 error, i.e. NRMSE 
        R^2 score
    - Other default hyperparameters in __init__ of Parser class
"""
import torch
from tqdm import tqdm
from torch.nn.modules import utils
import torch.optim as optim
from models.codec import DenseED
from models.cppn import CPPN
from models.th import th_conv_constitutive_constraint as constitutive_constraint
#from models.th import th_conv_continuity_constraint as continuity_constraint
from models.th import th_conv_boundary_condition as boundary_condition
from utils.image_gradient import SobelFilter
from utils.load import load_data_vtk_train, load_data_vtk_test
from utils.misc import mkdirs, to_numpy
from utils.plot import plot_prediction_det, save_stats, plot_velocity_temperature_ax, plot_temperature_ax, plot_velocity_temperature
from utils.practices import OneCycleScheduler, adjust_learning_rate, find_lr
import time
import argparse
import random
from pprint import pprint
import json
import sys
import matplotlib.pyplot as plt
plt.switch_backend('agg')
import meshio
import numpy as np
# from FyeldGenerator import generate_field

# Helper that generates power-law power spectrum
def Pkgen(n):
    def Pk(k):
        return np.power(k, -n)

    return Pk

# Draw samples from a normal distribution
def distrib(shape):
    a = np.random.normal(loc=0, scale=1, size=shape)
    b = np.random.normal(loc=0, scale=1, size=shape)
    c = np.random.normal(loc=0, scale=1, size=shape)
    return a*c + 1j * b

class Parser(argparse.ArgumentParser):
    def __init__(self): 
        super(Parser, self).__init__(description='Learning surrogate with mixed residual norm loss')
        self.add_argument('--exp-name', type=str, default='codec/mixed_residual', help='experiment name')
        self.add_argument('--exp-dir', type=str, default="./experiments", help='directory to save experiments')      
        # codec
        self.add_argument('--blocks', type=list, default=[3, 4, 3], help='list of number of layers in each dense block')
        self.add_argument('--growth-rate', type=int, default=24, help='number of output feature maps of each conv layer within each dense block')
        self.add_argument('--init-features', type=int, default=32, help='number of initial features after the first conv layer')        
        self.add_argument('--drop-rate', type=float, default=0.2, help='dropout rate')
        self.add_argument('--upsample', type=str, default='nearest', choices=['nearest', 'bilinear'])
        # data 
        self.add_argument('--data-dir', type=str, default="./datasets", help='directory to dataset')
        self.add_argument('--data', type=str, default='grf_kle512', choices=['grf_kle512', 'channelized'])
        self.add_argument('--ntrain', type=int, default=58, help="number of training data")
        self.add_argument('--ntest', type=int, default=10, help="number of validation data")
        self.add_argument('--imsize', type=int, default=64)
        # training
        self.add_argument('--run', type=int, default=1, help='run instance')
        self.add_argument('--epochs', type=int, default=1000, help='number of epochs to train')
        self.add_argument('--lr', type=float, default=1e-3, help='learning rate')
        self.add_argument('--lr-div', type=float, default=2., help='lr div factor to get the initial lr')
        self.add_argument('--lr-pct', type=float, default=0.3, help='percentage to reach the maximun lr, which is args.lr')
        self.add_argument('--weight-decay', type=float, default=0., help="weight decay")
        self.add_argument('--weight-bound', type=float, default=10, help="weight for boundary loss")
        self.add_argument('--batch-size', type=int, default=10, help='input batch size for training')
        self.add_argument('--test-batch-size', type=int, default=2, help='input batch size for testing')
        self.add_argument('--seed', type=int, default=1, help='manual seed used in Tensor')
        self.add_argument('--cuda', type=int, default=1, choices=[0, 1, 2, 3], help='cuda index')
        # logging
        self.add_argument('--debug', action='store_true', default=False, help='debug or verbose')
        self.add_argument('--ckpt-epoch', type=int, default=None, help='which epoch of checkpoints to be loaded')
        self.add_argument('--ckpt-freq', type=int, default=10, help='how many epochs to wait before saving model')
        self.add_argument('--log-freq', type=int, default=1, help='how many epochs to wait before logging training status')
        self.add_argument('--plot-freq', type=int, default=10, help='how many epochs to wait before plotting test output')
        self.add_argument('--plot-fn', type=str, default='imshow', choices=['contourf', 'imshow'], help='plotting method')
        self.add_argument('--in-channels', type=int, default='2', choices=['contourf', 'imshow'], help='plotting method')

    def parse(self):
        args = self.parse_args()

        hparams = f'{args.data}_ntrain{args.ntrain}_run{args.run}_bs{args.batch_size}_lr{args.lr}_epochs{args.epochs}_output'
        if args.debug:
            hparams = 'debug/' + hparams
        args.run_dir = args.exp_dir + '/' + args.exp_name + '/' + hparams
        args.ckpt_dir = args.run_dir + '/checkpoints'
        mkdirs(args.run_dir, args.ckpt_dir)

        #assert args.ntrain % args.batch_size == 0 and \
        #    args.ntest % args.test_batch_size == 0

        if args.seed is None:
            args.seed = random.randint(1, 10000)
        print("Random Seed: ", args.seed)
        random.seed(args.seed)
        torch.manual_seed(args.seed)

        print('Arguments:')
        pprint(vars(args))
        with open(args.run_dir + "/args.txt", 'w') as args_file:
            json.dump(vars(args), args_file, indent=4)

        return args


if __name__ == '__main__':

    args = Parser().parse()
    device = torch.device(f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")

    args.train_dir = args.run_dir + '/training'
    args.pred_dir = args.train_dir + '/predictions'
    input_channels = args.in_channels
    mkdirs(args.train_dir, args.pred_dir)
    # Change the input channel to one with the boundary velocities.
    # Use these velocities for the loss function. Assume these velocities
    # are far away and remain constant.
    model = DenseED(in_channels=input_channels, out_channels=1, 
                    imsize=args.imsize,
                    blocks=args.blocks,
                    growth_rate=args.growth_rate, 
                    init_features=args.init_features,
                    drop_rate=args.drop_rate,
                    out_activation=None,
                    upsample=args.upsample)
    if args.debug:
        print(model)
    # if start from ckpt
    if args.ckpt_epoch is not None:
        ckpt_file = args.run_dir + f'/checkpoints/model_epoch{args.ckpt_epoch}.pth'
        model.load_state_dict(torch.load(ckpt_file, map_location='cpu'))
        print(f'Loaded ckpt: {ckpt_file}')
        print(f'Resume training from epoch {args.ckpt_epoch + 1} to {args.epochs}')
    model = model.to(device)
    # load data
    if args.data == 'grf_kle512':
        train_hdf5_file = args.data_dir + \
            f'/{args.imsize}x{args.imsize}/kle512_lhs10000_train.hdf5'
        test_hdf5_file = args.data_dir + \
            f'/{args.imsize}x{args.imsize}/kle512_lhs1000_val.hdf5'
        ntrain_total, ntest_total = 10000, 1000
    elif args.data == 'channelized':
        train_hdf5_file = args.data_dir + \
            f'/{args.imsize}x{args.imsize}/channel_ng64_n4096_train.hdf5'
        test_hdf5_file = args.data_dir + \
            f'/{args.imsize}x{args.imsize}/channel_ng64_n512_test.hdf5'
        ntrain_total, ntest_total = 4096, 512
    assert args.ntrain <= ntrain_total, f"Only {args.ntrain_total} data "\
        f"available in {args.data} dataset, but needs {args.ntrain} training data."
    assert args.ntest <= ntest_total, f"Only {args.ntest_total} data "\
        f"available in {args.data} dataset, but needs {args.ntest} test data."      
    train_loader, _, x_data_train, y_data_train = load_data_vtk_train(train_hdf5_file, args.imsize, input_channels, args.ntrain, args.batch_size, 
        only_input=True, return_stats=False)
    test_loader, test_stats, x_data_test, y_data_test = load_data_vtk_test(train_hdf5_file, args.imsize, input_channels, args.ntest, 
        args.batch_size, only_input=False, return_stats=False)

    postfix_dict = {"loss": "", "lr": "NaN", "pde": "NaN",  "dir": "NaN", "neu": "NaN"}
    progress_bar = tqdm(
            desc="Epoch: ", total=args.epochs, postfix=postfix_dict, delay=0.5
        )


    imsize = args.imsize
    y_data_train = torch.FloatTensor(y_data_train).to(device)
    y_data_test = torch.FloatTensor(y_data_test).to(device)
    x_data_train = torch.FloatTensor(x_data_train).to(device)
    x_data_test = torch.FloatTensor(x_data_test).to(device)

    optimizer = optim.Adam(model.parameters(), lr=args.lr,
                        weight_decay=args.weight_decay)
    # scheduler = OneCycleScheduler(lr_max=args.lr, div_factor=args.lr_div, 
    #                     pct_start=args.lr_pct)
    scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, steps_per_epoch=len(train_loader), epochs=args.epochs)
    sobel_filter = SobelFilter(args.imsize, correct=True, device=device)

    n_out_pixels = test_loader.dataset[0][1].numel()
    print(f'Number of out pixels per image: {n_out_pixels}')

    logger = {}
    logger['loss_train'] = []
    logger['loss_test'] = []
    logger['r2_test'] = []
    logger['nrmse_test'] = []

    printCounter = 0

    def test(epoch,totalPrintCounter):
        print("Plotting prediction")
        model.eval()
        loss_test = 0.
        relative_l2, err2 = [], []
        n_batches = len(test_loader)
        for batch_idx, (input, target) in enumerate(test_loader, start=1):
            #input, target = input.to(device), tart.to(device)
            #print("Input: ", input)
            input = input.to(device)
            target = target.to(device)
            output = model(input)
            '''
            loss_pde = constitutive_constraint(input, output, sobel_filter) \
                #+ continuity_constraint(output, sobel_filter)
            loss_dirichlet, loss_neumann = boundary_condition(output)
            loss_boundary = loss_dirichlet + loss_neumann
            loss = loss_pde + loss_boundary * args.weight_bound
            '''
            loss_pde = 0#constitutive_constraint(input, output, sobel_filter)
            loss_dirichlet, loss_neumann = boundary_condition(output,target)
            loss_boundary = loss_dirichlet + loss_neumann
            loss = loss_pde + loss_boundary * args.weight_bound
            loss_test += loss.item()
            # sum over H, W --> (B, C)
            err2_sum = torch.sum((output - target) ** 2, [-1, -2])
            relative_l2.append(torch.sqrt(err2_sum / (target ** 2).sum([-1, -2])))
            err2.append(err2_sum)
            # plot predictions
            #if (epoch % args.plot_freq == 0 or epoch == args.epochs): #and \
                #batch_idx == len(test_loader) - 1:
            
            n_samples = 1 #if epoch == args.epochs else 2
            idx = torch.randperm(input.size(0))[:n_samples]
            samples_output = output.data.cpu().numpy()
            samples_input = input.data.cpu()[0].numpy()
            samples_target = target.data.cpu()[0].numpy()
            
            fig, axes = plt.subplots(1,3, figsize=(12, 12))
            temperature_pred = np.zeros((imsize,imsize))
            temperature_gt = np.zeros((imsize,imsize))
            temperature_err = np.zeros((imsize,imsize))
            velocities = np.zeros((2, imsize, imsize))  
            for i in range(0,imsize):
                for j in range(0,imsize):
                    #print(samples_output[0,0,i,j])
                    temperature_pred[i,j] = samples_output[0,0,i,j]
                    temperature_gt[i,j] = samples_target[0,i,j]
                    temperature_err[i,j] = np.abs(samples_output[0,0,i,j] - samples_target[0,i,j])
                    velocities = samples_input
            
            im_pred = plot_temperature_ax(axes[0], temperature_pred)
            im_gt = plot_temperature_ax(axes[1], temperature_gt)
            im_err = plot_temperature_ax(axes[2], temperature_err)
            axes[0].set_title("Pred")
            axes[1].set_title("GT")
            axes[2].set_title("Err")
            # plt.colorbar(im_pred)
            cbar = fig.colorbar(im_pred, ax=axes.ravel().tolist(), shrink=0.33)
            fig.savefig('plots/field-output' + str(totalPrintCounter) + '.png',dpi=400)
            plt.close()

            #print(min(samples_output[0,0,:,:]))

        loss_test /= n_batches
        relative_l2 = to_numpy(torch.cat(relative_l2, 0).mean(0))
        r2_score = 1 - to_numpy(torch.cat(err2, 0).sum(0)) #/ y_test_variation
        progress_bar.write(f"Epoch: {epoch}, test r2-score:  {r2_score}")
        progress_bar.write(f"Epoch: {epoch}, test relative-l2:  {relative_l2}")
        progress_bar.write(f'Epoch {epoch}: test loss: {loss_train:.6f}, loss_pde: {loss_pde:.6f}, '\
                f'dirichlet {loss_dirichlet:.6f}, nuemann {loss_neumann:.6f}')

        if epoch % args.log_freq == 0:
           logger['loss_test'].append(loss_test)
           logger['r2_test'].append(r2_score)
           logger['nrmse_test'].append(relative_l2)

    print('Start training...................................................')
    totalPrintCounter = 0
    start_epoch = 1 if args.ckpt_epoch is None else args.ckpt_epoch + 1
    tic = time.time()
    # step = 0
    total_steps = args.epochs * len(train_loader)
    # print(f'total steps: {total_steps}')
    for epoch in range(start_epoch, args.epochs + 1):
        model.train()
        # print("============== Print counter: ", printCounter )
        printCounter += 1
        # if epoch == 30:
        #     print('begin finding lr')
        #     logs,losses = find_lr(model, train_loader, optimizer, loss_fn, 
        #           args.weight_bound, init_value=1e-8, final_value=10., beta=0.98)
        #     plt.plot(logs[10:-5], losses[10:-5])
        #     plt.savefig(args.train_dir + '/find_lr.png')
        #     sys.exit(0)
        loss_train, mse = 0., 0.
            
        n_batches = len(train_loader)
        for batch_idx, (input, target) in enumerate(train_loader, start=1):
            # print("batchID: ", batch_idx)
            input = input.to(device)
            target = target.to(device)
            #print("Input: ", input)
            model.zero_grad()
            output = model(input)
            '''
            loss_pde = constitutive_constraint(input, output, sobel_filter) \
                + continuity_constraint(output, sobel_filter)
            loss_dirichlet, loss_neumann = boundary_condition(output)
            loss_boundary = loss_dirichlet + loss_neumann
            loss = loss_pde + loss_boundary * args.weight_bound
            '''
            #energy = constitutive_constraint(input, output, 
                    #sobel_filter) #+ continuity_constraint(output, sobel_filter)
        
            loss_pde = 0#constitutive_constraint(input, output, sobel_filter)
            loss_dirichlet, loss_neumann = boundary_condition(output,target)
            loss_boundary = loss_dirichlet + loss_neumann
            loss = loss_pde + loss_boundary * args.weight_bound
            loss.backward()
            # lr scheduling
            # step = (epoch - 1) * len(train_loader) + batch_idx
            # pct = step / total_steps
            optimizer.step()
            scheduler.step()
            # adjust_learning_rate(optimizer, lr)
            loss_train += loss.item()

        loss_train /= n_batches

        lr = scheduler.get_last_lr()[0]

        postfix_dict['lr'] = f"{lr:.5f}"
        postfix_dict['loss'] = f"{loss_train:.5f}"
        postfix_dict['dir'] = f"{loss_dirichlet:.4f}"
        postfix_dict['pde'] = f"{loss_pde:.4f}"
        postfix_dict['neu'] = f"{loss_neumann:.4f}"
        progress_bar.set_postfix(postfix_dict)
        progress_bar.update(1)
        # print(f'Epoch {epoch}, lr {lr:.6f}')
        # print(f'Epoch {epoch}: training loss: {loss_train:.6f}, pde: {loss_pde:.6f}, '\
            # f'dirichlet {loss_dirichlet:.6f}, nuemann {loss_neumann:.6f}')
        if epoch % args.log_freq == 0:
            logger['loss_train'].append(loss_train)
        if epoch % args.ckpt_freq == 0:
            torch.save(model.state_dict(), args.ckpt_dir + "/model_epoch{}.pth".format(epoch))
        
        if ((epoch+1) % 100 == 0):
            test(epoch,totalPrintCounter)
            totalPrintCounter += 1
           

    tic2 = time.time()
    print(f'Finished training {args.epochs} epochs with {args.ntrain} data ' \
        f'using {(tic2 - tic) / 60:.2f} mins')
    metrics = ['loss_train', 'loss_test', 'nrmse_test', 'r2_test']
    save_stats(args.train_dir, logger, *metrics)
    args.training_time = tic2 - tic
    args.n_params, args.n_layers = model.model_size
    with open(args.run_dir + "/args.txt", 'w') as args_file:
        json.dump(vars(args), args_file, indent=4)