01_nicam.py 8.5 KB
Newer Older
1
2
3
4
5
6
#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
description:    Extract lat-lon box of data from model NICAM.
author:         Christoph Heim
date created:   27.06.2019
7
date changed:   27.02.2020
8
9
10
11
usage:          arguments:
                1st:    n jobs for multiprocessing pool
"""
###############################################################################
12
import os, glob, subprocess, sys, time, shutil
13
14
15
16
17
import numpy as np
from datetime import datetime, timedelta
from multiprocessing import Pool
from pathlib import Path
from cdo import Cdo
18
from package.utilities import Timer, cdo_mergetime
Christoph Heim's avatar
Christoph Heim committed
19
from namelist import domain, padding
20
from functions import paste_dir_names
21
22
###############################################################################

23
24
def sellatlon_NICAM(inp_file, out_file, dt, box, options, var_dict,
                    var_name, res):
25
26
    
    TM = Timer()
27
28

    file_code = '{}km_{}_{:%Y%m%d}'.format(res, var_name, dt)
29
    
30
    if os.path.exists(out_file) and not options['recompute']:
31
        TM.start('cdo')
32
        print('\t\t'+file_code+' already computed')
33
34
35
36
        TM.stop('cdo')
    else:
        # cdo
        TM.start('cdo')
37
38
        print('\t\t'+file_code)

39
        if var_dict['dim'] == '3D':
40
41
            input = ('-sellevidx,'+str(box['vert0'])+'/'+
                               str(box['vert1'])+' '+inp_file)
42
        elif var_dict['dim'] == '2D':
43
            input=inp_file
44

45
46
47
48
49
        ofile = cdo.sellonlatbox(
                    box['lon'].start,box['lon'].stop,
                    box['lat'].start,box['lat'].stop,
                    input=input,
                    output=out_file)
50

51
        TM.stop('cdo')
52
53
54
55
56
57
58
59
60
61
62

    return(TM)


if __name__ == '__main__':

    # GENERAL SETTINGS
    ###########################################################################
    # input and output directories
    raw_data_dir = os.path.join('/work','ka1081','DYAMOND')
    out_base_dir = os.path.join('/work','ka1081','2019_06_Hackathon_Mainz',
63
                                'christoph_heim','newdata')
64
65

    # lat lon vert box to subselect
66
    box = domain
67
68
    #box.update({'vert0':1,'vert1':18}) # 3km
    box.update({'vert0':1,'vert1':26}) # 6km
Christoph Heim's avatar
Christoph Heim committed
69
70
    box['lon'] = slice(box['lon'].start - padding, box['lon'].stop + padding)
    box['lat'] = slice(box['lat'].start - padding, box['lat'].stop + padding)
71
72
73
74
75

    # name of model 
    model_name = 'NICAM'

    # variables to extract
76
77
78
79
80
81
82
83
84
85
    var_names = ['QV', 'QC', 'T', 'W', 'U', 'V', 'P',
                 'MSLP', 'U10M', 'V10M', 'T2M',
                 'LWUTOA', 'SWDSFC', 'SWUTOA',
                 'SST', 'SLHFLX', 'SSHFLX',
                 'TQC', 'TQI',
                 'PP']
    # missing: TQV, CLCL, CLCT, and in NICAM-3.5 also SST
    # missing: PS
    #var_names = ['MSLP']

86
87
88
89

    """
    parallel tasks:
    7km
90
    2D: ok: 12  - fail: 12
91
92
93
94
95
    3D: ok:     - fail:
    3.5km
    2D: ok: 2 - fail: 8,4
    3D: ok: 1
    """
96
97
    
    # model resolutions [km] of simulations
98
    ress = [7, 3.5]
99
    ress = [7]
100
    ress = [3.5]
101
102
    
    # date range
Christoph Heim's avatar
Christoph Heim committed
103
    first_date = datetime(2016,8,1)
104
    last_date = datetime(2016,9,9)
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
    #last_date = datetime(2016,8,1)

    exceptions = ['7_QC_20160801', '7_QC_20160802', '7_QC_20160803', '7_QC_20160804',
                  '7_T_20160815', '7_T_20160831', '7_T_20160901', '7_T_20160902',
                  '7_T_20160903', '7_T_20160904', '7_T_20160905', '7_T_20160906',
                  '7_T_20160907', '7_T_20160908', '7_T_20160909',
                  '7_U_20160824', '7_U_20160831', '7_U_20160901', '7_U_20160902',
                  '7_U_20160903', '7_U_20160904', '7_U_20160905', '7_U_20160906',
                  '7_U_20160907', '7_U_20160908', '7_U_20160909',
                                  '7_V_20160831', '7_V_20160901', '7_V_20160902',
                  '7_V_20160903', '7_V_20160904', '7_V_20160905', '7_V_20160906',
                  '7_V_20160907', '7_V_20160908', '7_V_20160909',
                  '7_P_20160815', '7_P_20160831', '7_P_20160901', '7_P_20160902',
                  '7_P_20160903', '7_P_20160904', '7_P_20160905', '7_P_20160906',
                  '7_P_20160907', '7_P_20160908', '7_P_20160909']
120

121
122
    # options for computation
    options = {}
123
124
    options['recompute']        = 0
    options['rm_tmp_folder']    = 0
125
    options['cdo_merge_days']   = 1
126
127
128
129
130
131
    ###########################################################################


    # NICAM SPECIFIC SETTINGS
    ###########################################################################
    var_dict = {
132
133
134
135
        'QV'    :{'file':'ms_qv',       'dim':'3D',  }, 
        'QC'    :{'file':'ms_qc',       'dim':'3D',  }, 
        'T'     :{'file':'ms_tem',      'dim':'3D',  }, 
        'W'     :{'file':'ms_w',        'dim':'3D',  }, 
136
137
138
        'U'     :{'file':'ms_u',        'dim':'3D',  }, 
        'V'     :{'file':'ms_v',        'dim':'3D',  }, 
        'P'     :{'file':'ms_pres',     'dim':'3D',  }, 
139

140
        'MSLP'  :{'file':'ss_slp',      'dim':'2D',  }, 
141
142
143
144
145
        'U10M'  :{'file':'ss_u10m',     'dim':'2D',  }, 
        'V10M'  :{'file':'ss_v10m',     'dim':'2D',  }, 
        'T2M'   :{'file':'ss_t2m',      'dim':'2D',  }, 
        'LWUTOA':{'file':'sa_lwu_toa',  'dim':'2D',  }, 
        'SWDSFC':{'file':'ss_swd_sfc',  'dim':'2D',  }, 
Christoph Heim's avatar
Christoph Heim committed
146
        'SWUTOA':{'file':'ss_swu_toa',  'dim':'2D',  }, 
147
148

        'SST'   :{'file':'oa_sst',      'dim':'2D',  }, 
149
150
        'SLHFLX':{'file':'ss_lh_sfc',   'dim':'2D',  }, 
        'SSHFLX':{'file':'ss_sh_sfc',   'dim':'2D',  }, 
151

152
        'TQC'   :{'file':'sa_cldw',     'dim':'2D',  }, 
153
154
        'TQI'   :{'file':'sa_cldi',     'dim':'2D',  }, 

155
        'PP'    :{'file':'sa_tppn',     'dim':'2D',  }, 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
    }
    ###########################################################################

    ## PREPARING STEPS
    TM = Timer()

    dt_range = np.arange(first_date, last_date + timedelta(days=1),
                        timedelta(days=1)).tolist()

    cdo = Cdo()

    if len(sys.argv) > 1:
        n_tasks = int(sys.argv[1])
    else:
        n_tasks = 1
    print('Using ' + str(n_tasks) + ' taks.')

    ## EXTRACT VARIABLES FROM SIMULATIONS
174
175
    # args used as input to parallel executed function
    args = []
176
177
178
179
    for var_name in var_names:
        print('############## var ' + var_name + ' ##################')
        for res in ress:
            print('############## res ' + str(res) + ' ##################')
180
181
            
            # data input directory
182
183
            sim_name = model_name + '-' + str(res) + 'km'
            inp_dir	= os.path.join(raw_data_dir, sim_name)
184
185
186
187
188
189

            # out_dir = directory for final model output (after mergetime)
            # out_tmp_dir = directory for output of files in time merge
            # level of raw model output
            out_dir, out_tmp_dir = paste_dir_names(out_base_dir,
                                                   model_name, res, domain)
190
            Path(out_dir).mkdir(parents=True, exist_ok=True)
191
192
193
            # remove temporary fils if desired
            if options['rm_tmp_folder'] and os.path.exists(out_tmp_dir):
                shutil.rmtree(out_tmp_dir)
194
195
196
197
198
            Path(out_tmp_dir).mkdir(parents=True, exist_ok=True)

            # find times and files that should be extracted
            # and prepare arguments for function
            for dt in dt_range:
199
200
201
202
203
204
205
206
207
208
209
                except_str = '{:g}_{}_{:%Y%m%d}'.format(res, var_name, dt)
                if except_str not in exceptions:
                    inp_files_glob = glob.glob(os.path.join(inp_dir,
                                                '{:%Y%m%d}*.000000'.format(dt)))
                    inp_file = os.path.join(inp_files_glob[0],
                                        var_dict[var_name]['file'] + '.nc')

                    out_file = os.path.join(out_tmp_dir,
                                var_name+'_{:%Y%m%d}'.format(dt)+'.nc')
                    args.append( (inp_file, out_file, dt, box, options,
                                  var_dict[var_name], var_name, res) )
210

211
212
213
214
215
216
217
218
    # run function serial or parallel
    if n_tasks > 1:
        with Pool(processes=n_tasks) as pool:
            results = pool.starmap(sellatlon_NICAM, args)
    else:
        results = []
        for arg in args:
            results.append(sellatlon_NICAM(*arg))
219

220
221
222
    # collect timings from subtasks
    for task_TM in results:
        TM.merge_timings(task_TM)
223

224
225
226
227
228
229
    # merge all time step files to one
    TM.start('merge')
    for var_name in var_names:
        for res in ress:
            out_dir, out_tmp_dir = paste_dir_names(out_base_dir,
                                                   model_name, res, domain)
230
            #cdo_mergetime(out_tmp_dir, out_dir, var_name)
231
    TM.stop('merge')
232

233
    TM.print_report()