05_mpas.py 11.9 KB
Newer Older
1
2
3
4
5
6
#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
description:    Extract lat-lon box of data from model MPAS.
author:         Christoph Heim
date created:   05.07.2019
7
date changed:   27.02.2020
8
9
usage:          arguments:
                1st:    n jobs for multiprocessing pool
10
                MPAS_3.75 : 3D var: 2 jobs (sometimes 3 possible)
Christoph Heim's avatar
Christoph Heim committed
11
                MPAS_7.50 : 3D var: 5 jobs
12
13
"""
###############################################################################
14
import os, glob, subprocess, sys, time, shutil
15
16
17
18
19
import numpy as np
from datetime import datetime, timedelta
from multiprocessing import Pool
from pathlib import Path
from cdo import Cdo
20
from package.utilities import Timer, write_grid_file#, cdo_mergetime
Christoph Heim's avatar
Christoph Heim committed
21
from namelist import domain, padding
22
from functions import paste_dir_names
23
24
25
###############################################################################


Christoph Heim's avatar
Christoph Heim committed
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
def fix_time_MPAS(out_file, dt, var_dict, var_name):
    
    if os.path.exists(out_file):
        file_code = '{}km_{}_{:%Y%m%d%H%M}'.format(res, var_name, dt)
        print('\t\t'+file_code)

        domain_str = "{},{},{},{}".format(
                        box['lon'].start, box['lon'].stop,
                        box['lat'].start, box['lat'].stop)
        levels_str = "{}/{}".format(
                        box['vert0'], box['vert1'])

        tmp_file = os.path.join(os.path.split(out_file)[0],
                                'temp_'+file_code+'.nc')
        subprocess.call(['mv', out_file, tmp_file])

        if var_dict['type'] == 'history':
            time_fmt = '{:%Y-%m-%d,%H:%M:%S,3hour}'
        elif var_dict['type'] == 'diag':
            time_fmt = '{:%Y-%m-%d,%H:%M:%S,15min}'
        cdo.setreftime('2016-08-01,00:00:00,minutes',
                        input=(' -settaxis,'+time_fmt.format(dt)+
                               ' '+ tmp_file),
                        output=out_file)

        if var_dict['type'] == 'history':
52
53
54
55
56
            if 'vdim' in var_dict:
                vdim = var_dict['vdim']
            # in case of SST
            else:
                vdim = 'nodim'
Christoph Heim's avatar
Christoph Heim committed
57
58
59
60
61
62
63
64
65
66
67
        elif var_dict['type'] == 'diag':
            vdim = 'nodim'
        subprocess.call(['ncpdq', '-O',
                         '--rdr=time,{},lat,lon'.format(vdim),
                         out_file, out_file])

        os.remove(tmp_file)




68
def sellatlon_MPAS(inp_file, out_file, dt, box, options, var_dict,
Christoph Heim's avatar
Christoph Heim committed
69
                   target_grid, var_name, res):
70
71
    
    TM = Timer()
72
    file_code = '{}km_{}_{:%Y%m%d%H%M}'.format(res, var_name, dt)
73

Christoph Heim's avatar
Christoph Heim committed
74
75
76
77
78
79
80
81
82
83
84
85
86
    if os.path.exists(out_file) and not options['recompute']:
        TM.start('run')
        print('\t\t'+file_code+' already computed')
        TM.stop('run')
    else:
        TM.start('run')
        print('\t\t'+file_code)

        domain_str = "{},{},{},{}".format(
                        box['lon'].start, box['lon'].stop,
                        box['lat'].start, box['lat'].stop)
        levels_str = "{}/{}".format(
                        box['vert0'], box['vert1'])
Christoph Heim's avatar
Christoph Heim committed
87
88
89
90
91
92
93
94
        if var_dict['type'] == 'history':
            time_fmt = '{:%Y-%m-%d,%H:%M:%S,3hour}'.format(dt)
        elif var_dict['type'] == 'diag':
            time_fmt = '{:%Y-%m-%d,%H:%M:%S,15min}'.format(dt)
        if var_dict['type'] == 'history':
            vdim = var_dict['vdim']
        elif var_dict['type'] == 'diag':
            vdim = 'nodim'
Christoph Heim's avatar
Christoph Heim committed
95
96
97
98
99
100

        if i_bash_output:
            subprocess.call(['./run_MPAS.sh', domain_str, levels_str,
                             str(res), os.path.split(out_file)[0],
                             os.path.split(out_file)[1][:-3],
                             var_dict['file'], inp_file, target_grid,
Christoph Heim's avatar
Christoph Heim committed
101
                             time_fmt, vdim])
Christoph Heim's avatar
Christoph Heim committed
102
103
104
105
106
        else:
            subprocess.call(['./run_MPAS.sh', domain_str, levels_str,
                             str(res), os.path.split(out_file)[0],
                             os.path.split(out_file)[1][:-3],
                             var_dict['file'], inp_file, target_grid,
Christoph Heim's avatar
Christoph Heim committed
107
                             time_fmt, vdim
Christoph Heim's avatar
Christoph Heim committed
108
109
110
111
112
                            ], stdout=subprocess.DEVNULL,
                               stderr=subprocess.DEVNULL)
        TM.stop('run')

    TM.print_report(short=True)
113
114
115
116
117
118
119
120
121
122
    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',
123
                                'christoph_heim','newdata')
124

125
126
    # box to subselect
    box = domain
127
128
    #box.update({'vert0':1,'vert1':22}) #3km top
    box.update({'vert0':1,'vert1':31}) #6km top
Christoph Heim's avatar
Christoph Heim committed
129
130
    box['lon'] = slice(box['lon'].start - padding, box['lon'].stop + padding)
    box['lat'] = slice(box['lat'].start - padding, box['lat'].stop + padding)
131
132
133
134
135

    # name of model 
    model_name = 'MPAS'

    # variables to extract
Christoph Heim's avatar
Christoph Heim committed
136
    var_namess = {
137
138
139
140
141
142
143
144
145
146
147
148
        '3D':['W', 'T', 'QV', 'QC', 'U', 'V', 'P'],
        '2D':['MSLP', 'U10M', 'V10M', 'T2M',
              'LWUTOA', 'SWNDTOA',
              'TQC', 'TQI', 'TQV',
              'CLCT', 'PPCONV', 'PPGRID'],

        '3D':['U', 'P', 'V'],
        '3D':['P'],
        #'3D':['U'],
        #'3D':['V'],
        #'3D':['U', 'V'],
        # missing CLCL
Christoph Heim's avatar
Christoph Heim committed
149
150
151
    }

    run_var_type = '3D'
152
153
154
155
156
157
158
159
160
161
162
    #run_var_type = '2D'

    """
    parallel tasks:
    7.5km
    2D: ok: 12
    3D: ok: 5
    3.75km
    2D:         fail: 6
    3D: ok: 2,3 fail: 3,2
    """
Christoph Heim's avatar
Christoph Heim committed
163
164

    var_names = var_namess[run_var_type]
165

166
167
168
    #ress = [7.5, 3.75]
    ress = [7.5]
    ress = [3.75]
Christoph Heim's avatar
Christoph Heim committed
169

Christoph Heim's avatar
Christoph Heim committed
170
    i_bash_output = 0
171
172
    
    # date range
173
    # 7.5
Christoph Heim's avatar
Christoph Heim committed
174
    first_date = datetime(2016,8,1)
175
    last_date = datetime(2016,9,9)
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    ######################

    ## done U, V, P
    #first_date = datetime(2016,8,1)
    #last_date = datetime(2016,8,10)

    ## done U, V, P
    #first_date = datetime(2016,8,11)
    #last_date = datetime(2016,8,20)

    ## done U, V, P
    #first_date = datetime(2016,8,21)
    #last_date = datetime(2016,8,30)

    ## done U, V, P
    #first_date = datetime(2016,8,31)
    #last_date = datetime(2016,9,9)
193

194
195
    # options for computation
    options = {}
196
    options['recompute']        = 0
Christoph Heim's avatar
Christoph Heim committed
197
    options['rm_tmp_files']     = 1
198
    options['rm_tmp_folder']    = 0
199
200
201
202
203
    ###########################################################################


    # MPAS SPECIFIC SETTINGS
    ###########################################################################
204
205
    grid_def_base_dir = os.path.join('/work','ka1081',
                            '2019_06_Hackathon_Mainz', 'falko')
Christoph Heim's avatar
Christoph Heim committed
206

207
    var_dict = {
Christoph Heim's avatar
Christoph Heim committed
208
209
210
211
        'QV'    :{'file':'qv',          'type':'history','vdim':'nVertLevels'}, 
        'QC'    :{'file':'qc',          'type':'history','vdim':'nVertLevels'}, 
        'T'     :{'file':'temperature', 'type':'history','vdim':'nVertLevels'}, 
        'W'     :{'file':'w',           'type':'history','vdim':'nVertLevelsP1'}, 
212
213
214
215
        'U'     :{'file':'uReconstructZonal',       'type':'history','vdim':'nVertLevels'}, 
        'V'     :{'file':'uReconstructMeridional',  'type':'history','vdim':'nVertLevels'}, 
        'P'     :{'file':'pressure',    'type':'history','vdim':'nVertLevels'}, 
        'SWDSFC':{'file':'acswdnb',     'type':'history'}, 
Christoph Heim's avatar
Christoph Heim committed
216

217
        'MSLP'  :{'file':'mslp',        'type':'diag'}, 
Christoph Heim's avatar
Christoph Heim committed
218
219
220
221
        'U10M'  :{'file':'u10',         'type':'diag'}, 
        'V10M'  :{'file':'v10',         'type':'diag'}, 
        'T2M'   :{'file':'t2m',         'type':'diag'}, 
        'LWUTOA':{'file':'olrtoa',      'type':'diag'}, 
222
        'SWNDTOA':{'file':'acswnett',   'type':'diag'}, 
Christoph Heim's avatar
Christoph Heim committed
223
        # missing in model output
224
        #'SST'   :{'file':'sst',         'type':'diag'}, 
Christoph Heim's avatar
Christoph Heim committed
225
226
        #'SLHFLX':{'file':'',     'type':'diag'}, 
        #'SSHFLX':{'file':'',     'type':'diag'}, 
227
        'TQC'   :{'file':'vert_int_qc', 'type':'diag'}, 
228
229
230
        'TQI'   :{'file':'vert_int_qi', 'type':'diag'}, 
        'TQV'   :{'file':'vert_int_qv', 'type':'diag'}, 
        'CLCT'  :{'file':'cldcvr',      'type':'diag'}, 
231
232
        'PPCONV':{'file':'rainc',       'type':'diag'}, 
        'PPGRID':{'file':'rainnc',      'type':'diag'}, 
233
234
    }
    grid_dict = {
Christoph Heim's avatar
Christoph Heim committed
235
236
237
238
239
240
241
242
        #3.75:{'grid_def_file':os.path.join(grid_def_base_dir,
        #                     'MPAS_3.75km_grid.nc'),}, 
        #7.5: {'grid_def_file':os.path.join(grid_def_base_dir,
        #                     'MPAS_7.5km_grid.nc'),}, 
        3.75:{'grid_def_file':os.path.join(out_base_dir, 'MPAS_3.75km',
                             'MPAS_3.75km_grid.nc'),}, 
        7.5: {'grid_def_file':os.path.join(out_base_dir, 'MPAS_7.5km',
                             'MPAS_7.5km_grid.nc'),}, 
243
    }
244
245
    #inc_min = {'history':180, 'diag':15}
    inc_min = {'history':180, 'diag':30}
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
    ###########################################################################

    ## PREPARING STEPS
    TM = Timer()
    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
    for var_name in var_names:
        print('############## var ' + var_name + ' ##################')

        dt_range = np.arange(first_date, last_date + timedelta(days=1),
                        timedelta(minutes=inc_min[
                        var_dict[var_name]['type']])).tolist()

        for res in ress:
            print('############## res ' + str(res) + ' ##################')

            sim_name = model_name + '-' + str(res) + 'km'
            inp_dir	= os.path.join(raw_data_dir, sim_name)
271
272
273
274
275
276

            # 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)
277
            Path(out_dir).mkdir(parents=True, exist_ok=True)
278
279
280
            # remove temporary fils if desired
            if options['rm_tmp_folder'] and os.path.exists(out_tmp_dir):
                shutil.rmtree(out_tmp_dir)
281
282
            Path(out_tmp_dir).mkdir(parents=True, exist_ok=True)

283
284
285
286
            # prepare grid interpolation
            # mpas grid definition file
            grid_def_file = grid_dict[res]['grid_def_file']
            # target grid on which to interpolate the model output
287
288
            target_grid = os.path.join('grids','latlon_{}km_dom_{}'.format(
                                        res, domain['code']))
Christoph Heim's avatar
Christoph Heim committed
289
            write_grid_file(box, target_grid, res)
290

291
292
293
294
            # find times and files that should be extracted
            # and prepare arguments for function
            args = []
            for dt in dt_range:
295
296
297
298
299
300
301
302
303
304
                # missing time step:
                if dt != datetime(2016,9,5,0,0,0):
                    inp_file = glob.glob(os.path.join(
                                    inp_dir,var_dict[var_name]['type']+
                                    '.{:%Y-%m-%d_%H.%M.%S}.nc'.format(dt)))[0]
                    out_file = os.path.join(out_tmp_dir,
                                var_name+'_{:%Y%m%d%H%M}'.format(dt)+'.nc')
                    args.append( (inp_file, out_file, dt, box, options,
                                  var_dict[var_name],
                                  target_grid, var_name, res) )
305

Christoph Heim's avatar
Christoph Heim committed
306
307
                #fix_time_MPAS(out_file, dt, var_dict[var_name], var_name)

308
309
310
311
312
313
314
315
            # run function serial or parallel
            if n_tasks > 1:
                with Pool(processes=n_tasks) as pool:
                    results = pool.starmap(sellatlon_MPAS, args)
            else:
                results = []
                for arg in args:
                    results.append(sellatlon_MPAS(*arg))
Christoph Heim's avatar
Christoph Heim committed
316
            
317
318
319
320
            # collect timings from subtasks
            for task_TM in results:
                TM.merge_timings(task_TM)

321
            # merge all time step files to one
322
            TM.start('merge')
323
            #cdo_mergetime(out_tmp_dir, out_dir, var_name)
324
            TM.stop('merge')
325

326
    print('ran MPAS at resolutions {} for variables {}'.format(ress, var_names))
327
    TM.print_report()