DDSC/evaluate.py at main · NUIST-GenAI-Lab/DDSC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import os

import numpy as np
import torch

from torch import device, cuda

from main import Task
from train import test_model

if __name__ == '__main__':

    torch.manual_seed(42)
    torch.cuda.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    np.random.seed(42)

    validate_year_all = [[2018, 2019, 2020],
                         [1993, 1994, 1995],
                         [1998, 1999, 2000],
                         [2003, 2004, 2005],
                         [2008, 2009, 2010],
                         [2013, 2014, 2015]]
    test_year_all = [[1991, 1992, 1993, 1994, 1995],
                     [1996, 1997, 1998, 1999, 2000],
                     [2001, 2002, 2003, 2004, 2005],
                     [2006, 2007, 2008, 2009, 2010],
                     [2011, 2012, 2013, 2014, 2015],
                     [2016, 2017, 2018, 2019, 2020]]

    result_dir = f'../results'
    os.makedirs(result_dir, exist_ok=True)

    init_list = ['dec','jun']
    var_list = ['t2m', 'tp']

    for init_mon in init_list:
        for var in var_list:
            result_all = []
            result_save_path = os.path.join(result_dir, f'corrected_{init_mon}_{var}')

            for i in range(6):

                validation_year = validate_year_all[i]
                test_year = test_year_all[i]
                dev = device("cuda" if cuda.is_available() else "cpu")
                device_ids = [0]
                nwp_improving = Task(init_mon,var,validation_year=validation_year, test_year=test_year)
                print(i, dev, nwp_improving, sep='\n')

                model_save_path = os.path.join(f'../models/model_{nwp_improving.info[0]}_{nwp_improving.info[1]}', f'{i}-model.pth')
                model_trained = torch.load(model_save_path, map_location=dev, weights_only=False)

                if hasattr(model_trained, "module"):
                    model_trained = model_trained.module

                result = test_model(model_trained, nwp_improving.data[2], dev, nwp_improving.data[0].diff_std,
                                    nwp_improving.data[0].diff_mean)
                result_all.append(result)

            result_all = np.array(result_all)
            result_all = result_all.squeeze()
            result_all = result_all.reshape(30, 6, 25, 32, 64)
            result_all = result_all.transpose(2, 0, 1, 3, 4)
            np.save(result_save_path, result_all)