03_icon.py 10.3 KB
Newer Older
1
2
3
4
5
6
#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
description:    Extract lat-lon box of data from model ICON.
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, write_grid_file, 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 comp_weights_file(target_grid, weights_file, inp_file, grid_def_file,
25
                      res, box, options):
26
27
    """
    """
Christoph Heim's avatar
Christoph Heim committed
28
29
    #if (not os.path.exists(target_grid)) or (options['recompute']):
    #    write_grid_file(box, target_grid, res)
30

31
    print('Compute weights file')
32
33
    ofile = cdo.gennn(target_grid, input=(' -setgrid,'+grid_def_file+
                           ' '+inp_file), output=weights_file,
34
                      options='-P 1')
35
36
37
38



def sellatlon_ICON(inp_file, out_file, grid_def_file, weights_file,
39
40
                   target_grid, dt, box, options, var_dict, var_name,
                   res):
41
42
    """
    """
43
44
    
    TM = Timer()
45
    file_code = '{}km_{}_{:%Y%m%d}'.format(res, var_name, dt)
46
    
47
    if os.path.exists(out_file) and not options['recompute']:
48
        TM.start('cdo')
49
        print('\t\t'+file_code+' already computed')
50
51
52
53
        TM.stop('cdo')
    else:
        # cdo
        TM.start('cdo')
54
        print('\t\t'+file_code)
55

56
57
58
59
60
61
62
        if var_dict['dim'] == '3d':
            ofile = cdo.remap(target_grid, weights_file,
                        input=(
                               ' -sellevidx,'+
                               str(box['vert0'])+'/'+str(box['vert1'])+
                               ' -setgrid,'+grid_def_file+
                               ' '+inp_file),
63
                        output=out_file, options='-f nc')
64
65
66
67
        elif var_dict['dim'] == '2d':
            ofile = cdo.remap(target_grid, weights_file,
                        input=(
                               ' -setgrid,'+grid_def_file+
68
                               ' -selname,'+var_dict['key']+
69
                               ' '+inp_file),
70
                        output=out_file, options='-f nc')
71
72
73
74
75
76
77
78
79
80
81
82
83

        TM.stop('cdo')

    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',
84
                                'christoph_heim','newdata')
85

86
    # vert box to subselect
87
    box = domain
88
89
    #box.update({'vert0':73-14,'vert1':91-14}) #3km top
    box.update({'vert0':64-14,'vert1':91-14}) #6km top
Christoph Heim's avatar
Christoph Heim committed
90
91
    box['lon'] = slice(box['lon'].start - padding, box['lon'].stop + padding)
    box['lat'] = slice(box['lat'].start - padding, box['lat'].stop + padding)
92
93
94
95
96

    # name of model 
    model_name = 'ICON'

    # variables to extract
97
98
99
100
101
102
103
104
105
106
107
108
    var_names = ['QV', 'QC', 'T', 'W', 'U', 'V', 'P',
                 'PS', 'U10M', 'V10M', 'T2M',
                 'LWUTOA', 'SWNDTOA',
                 'SLHFLX', 'SSHFLX',
                 'TQC', 'TQI', 'TQV',
                 'CLCT', 'PP']
    # missing: SST, CLCL
    #var_names = ['U', 'V', 'P']
    #var_names = ['TQV', 'CLCT', 'PS']
    #var_names = ['V']
    #var_names = ['P']
    
109
110
111
112
113
114
115
116
117


    """
    parallel tasks:
    10km
    2D: ok:     - fail:
    3D: ok: 6    - fail:
    2.5km
    2D: ok: 3 - fail: 4
118
    3D: ok: 2 - fail: 2,6,
119
    """
120
    
121
    ress = [10,2.5]
122
123
    ress = [10]
    ress = [2.5]
124
125
    
    # date range
Christoph Heim's avatar
Christoph Heim committed
126
    first_date = datetime(2016,8,1)
127
    last_date = datetime(2016,9,9)
128
    #last_date = datetime(2016,8,1)
129

130
131
    # options for computation
    options = {}
132
133
    options['recompute']        = 0
    options['rm_tmp_folder']    = 0
134
135
136
137
138
139
    ###########################################################################


    # ICON SPECIFIC SETTINGS
    ###########################################################################
    grid_def_base_dir = os.path.join('/work','bk1040','experiments', 'input')
140

141
    var_dict = {
142
143
144
145
    'QV'        :{'file':'qv',          'dim':'3d',  }, 
    'QC'        :{'file':'tot_qc_dia',  'dim':'3d',  }, 
    'T'         :{'file':'t',           'dim':'3d',  }, 
    'W'         :{'file':'w',           'dim':'3d',  }, 
146
147
148
    'U'         :{'file':'u',           'dim':'3d',  }, 
    'V'         :{'file':'v',           'dim':'3d',  }, 
    'P'         :{'file':'pres',        'dim':'3d',  }, 
149

150
    'PS'        :{'file':'atm2_2d_ml',      'dim':'2d',   'key':'PS'}, 
151
152
153
154
    'U10M'      :{'file':'atm3_2d_ml',      'dim':'2d',   'key':'U_10M'}, 
    'V10M'      :{'file':'atm3_2d_ml',      'dim':'2d',   'key':'V_10M'}, 
    'T2M'       :{'file':'atm3_2d_ml',      'dim':'2d',   'key':'T_2M'}, 
    'LWUTOA'    :{'file':'atm_2d_avg_ml',   'dim':'2d',   'key':'ATHB_T'}, 
Christoph Heim's avatar
Christoph Heim committed
155
156
    'SWNDSFC'   :{'file':'atm_2d_avg_ml',   'dim':'2d',   'key':'ASOB_S'}, 
    'SWNDTOA'   :{'file':'atm_2d_avg_ml',   'dim':'2d',   'key':'ASOB_T'}, 
157
158
159
160
    'SWDIFFUSFC':{'file':'atm_2d_avg_ml',   'dim':'2d',   'key':'ASWDIFU_S'}, 
    'SLHFLX'    :{'file':'atm2_2d_ml',      'dim':'2d',   'key':'LHFL_S'}, 
    'SSHFLX'    :{'file':'atm2_2d_ml',      'dim':'2d',   'key':'SHFL_S'}, 
    'TQC'       :{'file':'atm1_2d_ml',      'dim':'2d',   'key':'TQC_DIA'}, 
161
162
163
    'TQI'       :{'file':'atm1_2d_ml',      'dim':'2d',   'key':'TQI_DIA'}, 
    'TQV'       :{'file':'atm1_2d_ml',      'dim':'2d',   'key':'TQV_DIA'}, 
    'CLCT'      :{'file':'atm2_2d_ml',      'dim':'2d',   'key':'CLCT'}, 
164
    'PP'        :{'file':'atm2_2d_ml',      'dim':'2d',   'key':'TOT_PREC'}, 
165
166
    }
    grid_dict = {
167
168
169
        10:  {'grid_def_file':os.path.join(grid_def_base_dir,
                             '10km/icon_grid_0025_R02B08_G.nc'),
            }, 
170
171
172
173
174
175
        5:  {'grid_def_file':os.path.join(grid_def_base_dir,
                             '5km_2/icon_grid_0015_R02B09_G.nc'),
            }, 
        2.5:{'grid_def_file':os.path.join(grid_def_base_dir,
                             '2.5km/icon_grid_0017_R02B10_G.nc'),
            }, 
176
    }
177
178
179

    os.environ['GRIB_DEFINITION_PATH'] = ('/mnt/lustre01/sw/rhel6-x64/eccodes'+
                                          '/definitions')
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
    ###########################################################################

    ## 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
197
198
199
200
201
202
    # args used as input to parallel executed function
    args = []
    for res in ress:
        print('############## res ' + str(res) + ' ##################')
        for var_name in var_names:
            print('\t############## var ' + var_name + ' ##################')
203
204
205

            sim_name = model_name + '-' + str(res) + 'km'
            inp_dir	= os.path.join(raw_data_dir, sim_name)
206
207
208
209
210
211

            # 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)
212
            Path(out_dir).mkdir(parents=True, exist_ok=True)
213
214
215
            # remove temporary fils if desired
            if options['rm_tmp_folder'] and os.path.exists(out_tmp_dir):
                shutil.rmtree(out_tmp_dir)
216
217
            Path(out_tmp_dir).mkdir(parents=True, exist_ok=True)

218
219
220
221
222
            # prepare grid interpolation
            # icon grid definition file
            grid_def_file = grid_dict[res]['grid_def_file']
            # weights file to compute once for a grid
            weights_file = os.path.join(out_tmp_dir,
223
224
                                'weights_{}km_dom_{}'.format(res,
                                domain['code']))
225
            # target grid on which to interpolate the model output
226
227
            target_grid = os.path.join('grids','latlon_{}km_dom_{}'.format(
                                        res, domain['code']))
Christoph Heim's avatar
Christoph Heim committed
228
            write_grid_file(box, target_grid, res)
229

230
231
232
233
234
235
236
237
238
239
            # find times and files that should be extracted
            # and prepare arguments for function
            for dt in dt_range:
                inp_files_glob = glob.glob(os.path.join(inp_dir,
                                            '*_{}_*{:%Y%m%d}*'.format(
                                            var_dict[var_name]['file'], dt)))
                inp_file = os.path.join(inp_files_glob[0])

                out_file = os.path.join(out_tmp_dir,
                            var_name+'_{:%Y%m%d}'.format(dt)+'.nc')
240
241
                args.append( (inp_file, out_file, grid_def_file,
                              weights_file, target_grid,
242
243
                              dt, box, options, var_dict[var_name],
                              var_name, res) )
244

245
            TM.start('grid')
246
247
            if ((not os.path.exists(weights_file)) or 
                (not os.path.exists(target_grid))):
248
                comp_weights_file(target_grid, weights_file,
249
                                  inp_file, grid_def_file,
250
                                  res, box, options)
251
            TM.stop('grid')
252

253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
    # run function serial or parallel
    if n_tasks > 1:
        with Pool(processes=n_tasks) as pool:
            results = pool.starmap(sellatlon_ICON, args)
    else:
        results = []
        for arg in args:
            results.append(sellatlon_ICON(*arg))

    # collect timings from subtasks
    for task_TM in results:
        TM.merge_timings(task_TM)

    # merge all time step files to one
    TM.start('merge')
    for res in ress:
        for var_name in var_names:
            out_dir, out_tmp_dir = paste_dir_names(out_base_dir,
                                                   model_name, res, domain)
272
            #cdo_mergetime(out_tmp_dir, out_dir, var_name)
273
    TM.stop('merge')
274

275
    print('ran ICON at resolutions {} for variables {}'.format(ress, var_names))
276
    TM.print_report()