06_ifs.py 9.58 KB
Newer Older
1
2
3
4
5
6
#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
description:    Extract lat-lon box of data from model IFS.
author:         Christoph Heim
date created:   05.07.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
25
def sellatlon_IFS(inp_file, out_file, dt, box, options, var_dict,
                  var_name, res):
26
27
    
    TM = Timer()
28
    file_code = '{}km_{}_{:%Y%m%d%H%M}'.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
        TM.stop('cdo')
    else:
35
        print('\t\t'+file_code)
36
37
38
39
40
41

        split = os.path.split(out_file)
        tmp_file = os.path.join(split[0],'tmp_'+split[1])

        # cdo
        TM.start('cdo')
42
43
44
45
46
47
        if not os.path.exists(tmp_file): 
            cdo.select('name='+var_dict['file'],
                        input=(inp_file),
                        options='--eccodes',
                        output=tmp_file)

48
49
        if var_dict['dim'] == '3D':
            input = ('-sellevel,'+str(box['vert0'])+'/'+
50
51
                           str(box['vert1'])+
                           ' -setgridtype,regular'+
52
53
54
55
56
57
58
59
60
                           ' '+tmp_file)
        else:
            input = (' -setgridtype,regular'+
                     ' '+tmp_file               )

        ofile = cdo.sellonlatbox(
                    box['lon'].start, box['lon'].stop,
                    box['lat'].start, box['lat'].stop,
                    input=input,
61
                    output=out_file, options='-f nc')
62

63
64
65
66
        # delete tmp_file
        if options['rm_tmp_files']:
            os.remove(tmp_file)

67
68
69
70
71
72
73
74
75
76
77
78
        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',
79
                                'christoph_heim','newdata')
80

81
82
    # box to subselect
    box = domain
83
84
    #box.update({'vert0':105,'vert1':137}) #top at 3km
    box.update({'vert0':94,'vert1':137}) #top at 6km
Christoph Heim's avatar
Christoph Heim committed
85
86
    box['lon'] = slice(box['lon'].start - padding, box['lon'].stop + padding)
    box['lat'] = slice(box['lat'].start - padding, box['lat'].stop + padding)
87
88
89
90
91

    # name of model 
    model_name = 'IFS'

    # variables to extract
92
93
94
95
96
97
98
99
    var_names = ['HSURF']
    var_names = ['QV', 'QC', 'T', 'W', 'U', 'V',
                 'PS', 'U10M', 'V10M', 'T2M',
                 'LWUTOA', 'SWDSFC', 'SWNDTOA',
                 'SST', 'SLHFLX', 'SSHFLX',
                 'TQC', 'TQI', 'TQV',
                 'CLCL', 'CLCT', 'PPCONV', 'PPGRID']
    # missing: P (derive with PS and vertical grid)
100
101
102
103
104
105
106
107
108
109

    """
    parallel tasks:
    9km
    2D: ok: 12,18
    3D:
    4km
    2D: ok:
    3D: ok: 6
    """
110
111
    
    # model resolutions [km] of simulations
112
113
114
    #ress = [9,4]
    ress = [9]
    ress = [4]
115
116
    
    # date range
Christoph Heim's avatar
Christoph Heim committed
117
    first_date = datetime(2016,8,1)
118
    last_date = datetime(2016,9,9)
119

120
121
    exceptions = ['4_T_201608070900']

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


    # IFS SPECIFIC SETTINGS
    ###########################################################################
    var_dict = {
133
        'QV'    :{'file':'q',       'dim':'3D',
Christoph Heim's avatar
Christoph Heim committed
134
                  'group':'mars_out_ml_moist',}, 
135
136
        'QC'    :{'file':'clwc',    'dim':'3D',
                  'group':'mars_out_ml_moist',}, 
137
        'T'     :{'file':'t',       'dim':'3D',
Christoph Heim's avatar
Christoph Heim committed
138
139
                  'group':'gg_mars_out_ml_upper_sh',}, 
        'W'     :{'file':'param120.128.192','dim':'3D',
140
                  'group':'gg_mars_out_ml_upper_sh',}, 
141
142
143
144
        'U'     :{'file':'u',       'dim':'3D',
                  'group':'gg_uv_mars_out_ml_vor_div_sh',}, 
        'V'     :{'file':'v',       'dim':'3D',
                  'group':'gg_uv_mars_out_ml_vor_div_sh',}, 
145

146
147
148
149
        'HSURF' :{'file':'z',     'dim':'2D',
                  'group':'mars_out',}, 
        'PS'    :{'file':'msl',     'dim':'2D',
                  'group':'mars_out',}, 
Christoph Heim's avatar
Christoph Heim committed
150
151
152
153
154
155
156
157
158
159
        'U10M'  :{'file':'10u',     'dim':'2D',
                  'group':'mars_out',}, 
        'V10M'  :{'file':'10v',     'dim':'2D',
                  'group':'mars_out',}, 
        'T2M'   :{'file':'2t',      'dim':'2D',
                  'group':'mars_out',}, 
        'LWUTOA':{'file':'ttr',     'dim':'2D',
                  'group':'mars_out',}, 
        'SWDSFC':{'file':'ssrd',    'dim':'2D',
                  'group':'mars_out',}, 
Christoph Heim's avatar
Christoph Heim committed
160
161
        'SWNDTOA':{'file':'tsr',    'dim':'2D',
                  'group':'mars_out',}, 
162
163
        'SST'   :{'file':'sst',     'dim':'2D',
                  'group':'mars_out',}, 
Christoph Heim's avatar
Christoph Heim committed
164
165
166
        'SLHFLX':{'file':'slhf',    'dim':'2D',
                  'group':'mars_out',}, 
        'SSHFLX':{'file':'sshf',    'dim':'2D',
167
168
169
                  'group':'mars_out',}, 
        'TQC'   :{'file':'tclw',    'dim':'2D',
                  'group':'mars_out',}, 
170
171
172
173
174
175
176
177
        'TQI'   :{'file':'tciw',    'dim':'2D',
                  'group':'mars_out',}, 
        'TQV'   :{'file':'tcwv',    'dim':'2D',
                  'group':'mars_out',}, 
        'CLCL'  :{'file':'lcc',     'dim':'2D',
                  'group':'mars_out',}, 
        'CLCT'  :{'file':'tcc',     'dim':'2D',
                  'group':'mars_out',}, 
178
179
180
181
        'PPCONV':{'file':'crr',     'dim':'2D',
                  'group':'mars_out',}, 
        'PPGRID':{'file':'lsrr',    'dim':'2D',
                  'group':'mars_out',}, 
182
183
184
    }

    base_time = datetime(2016,8,1)
185
186
187

    os.environ['GRIB_DEFINITION_PATH'] = ('/work/ka1081/programs/'+
                                    'eccodes-2.9.0/share/eccodes/definitions')
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
    ###########################################################################

    ## PREPARING STEPS
    TM = Timer()
    TM.start('real')
    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 + ' ##################')
        for res in ress:
            print('############## res ' + str(res) + ' ##################')

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

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

            # find times and files that should be extracted
            inp_files_glob = glob.glob(os.path.join(inp_dir,
223
224
225
226
227
                                var_dict[var_name]['group']+'.*'))
            
            inp_files_glob = [f for f in inp_files_glob if not '.tar.gz' in f]
            times = [base_time + timedelta(hours=int(
                        f.split('.')[1])) 
228
229
230
231
232
233
234
235
236
237
                        for f in inp_files_glob]
            use_times = [dt for dt in times if dt >= first_date and
                                           dt < last_date+timedelta(days=1)]
            use_files = [inp_files_glob[i] 
                        for i in range(len(inp_files_glob)) if 
                         times[i] in use_times]

            # prepare arguments for function
            args = []
            for i in range(len(use_times)):
238
239
240
241
242
243
244
245
                except_str = '{:g}_{}_{:%Y%m%d%H%M}'.format(res, var_name,
                                                            use_times[i])
                if except_str not in exceptions:
                    inp_file = use_files[i]
                    out_file = os.path.join(out_tmp_dir,
                                var_name+'_{:%Y%m%d%H%M}'.format(use_times[i])+'.nc')
                    args.append( (inp_file, out_file, use_times[i], box,
                                  options, var_dict[var_name], var_name, res) )
246
247
248
249
250
251
252
253
254
255
256
257
258
259

            # run function serial or parallel
            if n_tasks > 1:
                with Pool(processes=n_tasks) as pool:
                    results = pool.starmap(sellatlon_IFS, args)
            else:
                results = []
                for arg in args:
                    results.append(sellatlon_IFS(*arg))

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

260
            # merge all time step files to one
261
            TM.start('merge')
262
            #cdo_mergetime(out_tmp_dir, out_dir, var_name)
263
            TM.stop('merge')
264

265
    print('ran IFS at resolutions {} for variables {}'.format(ress, var_names))
266
    TM.print_report()