04_um.py 6.84 KB
Newer Older
1
2
3
4
5
6
#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
description:    Extract lat-lon box of data from model UM.
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
def sellatlon_UM(inp_file, out_file, dt, box, options, var_name, res):
24
25
    
    TM = Timer()
26
    file_code = '{}km_{}_{:%Y%m%d}'.format(res, var_name, dt)
27
    
28
    if os.path.exists(out_file) and not options['recompute']:
29
        TM.start('cdo')
30
        print('\t\t'+file_code+' already computed')
31
32
33
        TM.stop('cdo')
    else:
        TM.start('cdo')
34
35
        print('\t\t'+file_code)

36
        ofile = cdo.sellonlatbox(
37
38
                    box['lon'].start, box['lon'].stop,
                    box['lat'].start, box['lat'].stop,
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
                    input=('-sellevidx,'+str(box['vert0'])+'/'+
                           str(box['vert1'])+' '+inp_file),
                    output=out_file)
        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',
54
                                'christoph_heim','newdata')
55

56
57
    # box to subselect
    box = domain
58
59
    #box.update({'vert0':1,'vert1':21}) #3km top
    box.update({'vert0':1,'vert1':30}) #6km top
Christoph Heim's avatar
Christoph Heim committed
60
61
    box['lon'] = slice(box['lon'].start - padding, box['lon'].stop + padding)
    box['lat'] = slice(box['lat'].start - padding, box['lat'].stop + padding)
62
63
64
65
66

    # name of model 
    model_name = 'UM'

    # variables to extract
67
68
69
70
71
72
73
74
75
76
77
    var_names = ['QV', 'QC', 'T', 'W', 'U', 'V', 'P',
                 'PS', 'U10M', 'V10M', 'T2M',
                 'LWUTOA', 'SWDSFC', 'SWUTOA',
                 'SLHFLX', 'SSHFLX',
                 'TQC', 'TQI', 'TQV',
                 'CLCT', 'PP']
    # missing: SST, CLCL
    #var_names = ['U', 'V']
    #var_names = ['P']
    #var_names = ['TQV', 'CLCT']
    #var_names = ['PS']
78
79
80
81
82
83

    """
    parallel tasks:
    2D: ok: 6
    3D: ok: 4
    """
84
85
86
87
88
    
    # model resolutions [km] of simulations
    ress = [5]
    
    # date range
Christoph Heim's avatar
Christoph Heim committed
89
    first_date = datetime(2016,8,1)
90
    last_date = datetime(2016,9,8) # 3D fields only go to 2016,9,8
91
    #last_date = datetime(2016,9,9)
92

93
94
    # options for computation
    options = {}
95
96
    options['recompute']        = 0
    options['rm_tmp_folder']    = 0
97
98
99
100
101
102
    ###########################################################################


    # UM SPECIFIC SETTINGS
    ###########################################################################
    var_dict = {
Christoph Heim's avatar
Christoph Heim committed
103
        'QV'    :{'file':'hus',}, 
104
105
        'QC'    :{'file':'clw',}, 
        'T'     :{'file':'ta',}, 
Christoph Heim's avatar
Christoph Heim committed
106
        'W'     :{'file':'wa',}, 
107
108
109
        'U'     :{'file':'ua',}, 
        'V'     :{'file':'va-NEW',}, 
        'P'     :{'file':'pfull',}, 
110

111
        'PS'    :{'file':'ps',}, 
Christoph Heim's avatar
Christoph Heim committed
112
113
114
        'U10M'  :{'file':'uas',}, 
        'V10M'  :{'file':'vas',}, 
        'T2M'   :{'file':'tas',}, 
115
        'LWUTOA':{'file':'rlut',}, 
Christoph Heim's avatar
Christoph Heim committed
116
        'SWDSFC':{'file':'rsds',}, 
Christoph Heim's avatar
Christoph Heim committed
117
        'SWUTOA':{'file':'rsut',}, 
Christoph Heim's avatar
Christoph Heim committed
118
119
        'SLHFLX':{'file':'hfls',}, 
        'SSHFLX':{'file':'hfss',}, 
120
        'TQC'   :{'file':'clwvi',}, 
121
122
123
        'TQI'   :{'file':'clivi',}, 
        'TQV'   :{'file':'prw',}, 
        'CLCT'  :{'file':'clt',}, 
124
        'PP'    :{'file':'pr',}, 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
    }
    ###########################################################################

    ## 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
143
144
    # args used as input to parallel executed function
    args = []
145
146
147
148
149
150
151
152
153
    for var_name in var_names:
        #var_name = 'T'
        print('############## var ' + var_name + ' ##################')
        for res in ress:
            print('############## res ' + str(res) + ' ##################')
            #res = 4

            sim_name = model_name + '-' + str(res) + 'km'
            inp_dir	= os.path.join(raw_data_dir, sim_name)
154
155
156
157
158
159

            # 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)
160
            Path(out_dir).mkdir(parents=True, exist_ok=True)
161
162
163
            # remove temporary fils if desired
            if options['rm_tmp_folder'] and os.path.exists(out_tmp_dir):
                shutil.rmtree(out_tmp_dir)
164
165
166
167
168
            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:
169
170
171
172
                #print(os.path.join(inp_dir,
                #                            var_dict[var_name]['file'],
                #                            '*{:%Y%m%d}*.nc'.format(dt)))
                #quit()
173
174
175
                inp_file = glob.glob(os.path.join(inp_dir,
                                            var_dict[var_name]['file'],
                                            '*{:%Y%m%d}*.nc'.format(dt)))[0]
176
177
                out_file = os.path.join(out_tmp_dir,
                            var_name+'_{:%Y%m%d}'.format(dt)+'.nc')
178
179
                args.append( (inp_file, out_file, dt, box, options,
                              var_name, res) )
180
            #quit()
181

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
    # run function serial or parallel
    if n_tasks > 1:
        with Pool(processes=n_tasks) as pool:
            results = pool.starmap(sellatlon_UM, args)
    else:
        results = []
        for arg in args:
            results.append(sellatlon_UM(*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)
201
            #cdo_mergetime(out_tmp_dir, out_dir, var_name)
202
    TM.stop('merge')
203

204
    TM.print_report()