02_sam.py 9.88 KB
Newer Older
v99 Workshop's avatar
v99 Workshop committed
1
2
#!/usr/bin/python
# -*- coding: utf-8 -*-
v99 Workshop's avatar
v99 Workshop committed
3
"""
v99 Workshop's avatar
v99 Workshop committed
4
5
6
description:    Extract lat-lon box of data from model SAM.
author:         Christoph Heim
date created:   20.06.2019
7
date changed:   27.02.2020
v99 Workshop's avatar
v99 Workshop committed
8
9
usage:          arguments:
                1st:    n jobs for multiprocessing pool
v99 Workshop's avatar
v99 Workshop committed
10
"""
v99 Workshop's avatar
v99 Workshop committed
11
###############################################################################
12
import os, glob, subprocess, sys, time, shutil
v99 Workshop's avatar
v99 Workshop committed
13
14
15
import numpy as np
from datetime import datetime, timedelta
from multiprocessing import Pool
v99 Workshop's avatar
v99 Workshop committed
16
from pathlib import Path
v99 Workshop's avatar
v99 Workshop committed
17
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
v99 Workshop's avatar
v99 Workshop committed
21
22
###############################################################################

23
def sellatlon_SAM(inp_file, out_file, dt, box, options, var_name, res):
v99 Workshop's avatar
v99 Workshop committed
24
    
25
    TM = Timer()
26
27

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

v99 Workshop's avatar
v99 Workshop committed
38
39
40
41
42
        split = os.path.split(out_file)
        nco_file = os.path.join(split[0],'nco_'+split[1])

        # nco
        if not os.path.exists(nco_file):
43
            TM.start('nco')
v99 Workshop's avatar
v99 Workshop committed
44
45
46
47
48
49
            bash_command = ('ncatted -O -a units,lon,o,c,degrees_east ' + 
                            '-a units,lat,o,c,degrees_north '+inp_file+
                            ' '+nco_file)
            process = subprocess.Popen(bash_command.split(),
                                        stdout=subprocess.PIPE)
            output, error = process.communicate()
50
51
52
53
            TM.stop('nco')
        else:
            TM.start('nco')
            TM.stop('nco')
v99 Workshop's avatar
v99 Workshop committed
54
55

        # cdo
56
        TM.start('cdo')
Christoph Heim's avatar
Christoph Heim committed
57
58
59
60
        if var_dict[var_name]['loc'] == 'OUT_3D':
            time_fmt = '{:%Y-%m-%d,%H:%M:%S,3hour}'
        elif var_dict[var_name]['loc'] == 'OUT_2D':
            time_fmt = '{:%Y-%m-%d,%H:%M:%S,30min}'
v99 Workshop's avatar
v99 Workshop committed
61
        ofile = cdo.sellonlatbox(
62
63
                    box['lon'].start,box['lon'].stop,
                    box['lat'].start,box['lat'].stop,
v99 Workshop's avatar
v99 Workshop committed
64
                    input=('-sellevidx,'+str(box['vert0'])+'/'+
Christoph Heim's avatar
Christoph Heim committed
65
66
67
68
                           str(box['vert1'])+
                           ' -setreftime,2016-08-01,00:00:00,minutes'+
                           ' -settaxis,'+time_fmt.format(dt)+
                           ' '+nco_file),
v99 Workshop's avatar
v99 Workshop committed
69
                    output=out_file)
70
71
        TM.stop('cdo')

72
73
74
75
        # delete tmp_file
        if options['rm_tmp_files']:
            os.remove(nco_file)

76
    return(TM)
v99 Workshop's avatar
v99 Workshop committed
77
78
79
80


if __name__ == '__main__':

v99 Workshop's avatar
v99 Workshop committed
81
82
83
84
85
    # 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',
86
                                'christoph_heim','newdata')
v99 Workshop's avatar
v99 Workshop committed
87
88

    # lat lon vert box to subselect
89
    box = domain
90
91
    #box.update({'vert0':1,'vert1':28}) #3km
    box.update({'vert0':1,'vert1':35}) #6km
Christoph Heim's avatar
Christoph Heim committed
92
93
    box['lon'] = slice(box['lon'].start - padding, box['lon'].stop + padding)
    box['lat'] = slice(box['lat'].start - padding, box['lat'].stop + padding)
v99 Workshop's avatar
v99 Workshop committed
94
95
96
97
98

    # name of model 
    model_name = 'SAM'

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

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

122
123
124
    #first_date = datetime(2016,8,2)
    #last_date = datetime(2016,9,2)

125
126
    # options for computation
    options = {}
127
128
    options['recompute']        = 0
    options['rm_tmp_files']     = 1
129
    options['rm_tmp_folder']    = 0
v99 Workshop's avatar
v99 Workshop committed
130
131
132
133
134
135
    ###########################################################################


    # SAM SPECIFIC SETTINGS
    ###########################################################################
    var_dict = {
Christoph Heim's avatar
Christoph Heim committed
136
137
138
139
140
141
142
143
        'QV'        :{'file':'QV',
                      'loc':'OUT_3D','fntime':(-16,-6),}, 
        'QC'        :{'file':'QC',
                      'loc':'OUT_3D','fntime':(-16,-6),}, 
        'T'         :{'file':'TABS',
                      'loc':'OUT_3D','fntime':(-18,-8),}, 
        'W'         :{'file':'W',
                      'loc':'OUT_3D','fntime':(-15,-5),}, 
144
145
146
147
148
149
        'U'         :{'file':'U',
                      'loc':'OUT_3D','fntime':(-15,-5),}, 
        'V'         :{'file':'V',
                      'loc':'OUT_3D','fntime':(-15,-5),}, 
        'P'         :{'file':'PP',
                      'loc':'OUT_3D','fntime':(-16,-6),}, 
Christoph Heim's avatar
Christoph Heim committed
150

151
152
        'PS'        :{'file':'PSFC',
                      'loc':'OUT_2D','fntime':(-21,-11),}, 
Christoph Heim's avatar
Christoph Heim committed
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
        'U10M'      :{'file':'U10m',
                      'loc':'OUT_2D','fntime':(-21,-11),}, 
        'V10M'      :{'file':'V10m',
                      'loc':'OUT_2D','fntime':(-21,-11),}, 
        'T2M'       :{'file':'T2mm',
                      'loc':'OUT_2D','fntime':(-21,-11),}, 
        'LWUTOA'    :{'file':'LWNTA',
                      'loc':'OUT_2D','fntime':(-22,-12),}, 
        'SWDSFC'    :{'file':'SWDSA',
                      'loc':'OUT_2D','fntime':(-22,-12),}, 
        'SWNDTOA'   :{'file':'SWNTA',
                     'loc':'OUT_2D' ,'fntime':(-22,-12),}, 
        'SLHFLX'    :{'file':'LHF',
                      'loc':'OUT_2D','fntime':(-20,-10),}, 
        'SSHFLX'    :{'file':'SHF',
                      'loc':'OUT_2D','fntime':(-20,-10),}, 
        'TQC'       :{'file':'CWP',
                      'loc':'OUT_2D','fntime':(-20,-10),}, 
171
172
        'TQI'       :{'file':'IWP',
                      'loc':'OUT_2D','fntime':(-20,-10),}, 
Christoph Heim's avatar
Christoph Heim committed
173
174
        'PP'        :{'file':'Precac',
                      'loc':'OUT_2D','fntime':(-23,-13),}, 
v99 Workshop's avatar
v99 Workshop committed
175
176
    }

177

v99 Workshop's avatar
v99 Workshop committed
178
179
    dt = 7.5
    base_time = datetime(2016,8,1)
v99 Workshop's avatar
v99 Workshop committed
180
    ###########################################################################
v99 Workshop's avatar
v99 Workshop committed
181

v99 Workshop's avatar
v99 Workshop committed
182
    ## PREPARING STEPS
183
    TM = Timer()
v99 Workshop's avatar
v99 Workshop committed
184
185
    cdo = Cdo()

v99 Workshop's avatar
v99 Workshop committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
    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'
            inp_dir	= os.path.join(raw_data_dir, sim_name,
                                   var_dict[var_name]['loc'])
201
202
203
204
205
            # 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)
v99 Workshop's avatar
v99 Workshop committed
206
            Path(out_dir).mkdir(parents=True, exist_ok=True)
207
208
209
            # remove temporary fils if desired
            if options['rm_tmp_folder'] and os.path.exists(out_tmp_dir):
                shutil.rmtree(out_tmp_dir)
v99 Workshop's avatar
v99 Workshop committed
210
211
212
            Path(out_tmp_dir).mkdir(parents=True, exist_ok=True)

            # find times and files that should be extracted
213
214
215
216
217
218
219
            if var_dict[var_name]['loc'] == 'OUT_3D':
                inp_files_glob = glob.glob(os.path.join(inp_dir,
                                    '*_'+var_dict[var_name]['file']+'.nc'))
            elif var_dict[var_name]['loc'] == 'OUT_2D':
                inp_files_glob = glob.glob(os.path.join(inp_dir,
                                    '*'+var_dict[var_name]['file']+'*.nc'))

220
            #print(inp_files_glob)
v99 Workshop's avatar
v99 Workshop committed
221
222
223
224
            times = [base_time + timedelta(seconds=dt*int(
                        f[var_dict[var_name]['fntime'][0]:
                          var_dict[var_name]['fntime'][1]])) 
                        for f in inp_files_glob]
225

v99 Workshop's avatar
v99 Workshop committed
226
227
            use_times = [dt for dt in times if dt >= first_date and
                                           dt < last_date+timedelta(days=1)]
228
229
230
231
232
233
234
235
236

            # remove files that have problems in some fields (e.g. LWUTOA and SWNDTOA)
            del_vals = []
            for i,dt in enumerate(use_times):
                if '{:%H%M}'.format(dt) in ['0030', '0630', '1230', '1830']:
                    del_vals.append(dt)
            for del_val in del_vals:
                use_times.remove(del_val)

v99 Workshop's avatar
v99 Workshop committed
237
238
239
240
241
242
243
244
245
            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)):
                inp_file = use_files[i]
                out_file = os.path.join(out_tmp_dir,
246
                            var_name+'_{:%Y%m%d%H%M}'.format(use_times[i])+'.nc')
v99 Workshop's avatar
v99 Workshop committed
247
                args.append( (inp_file, out_file, use_times[i], box,
248
                              options, var_name, res) )
v99 Workshop's avatar
v99 Workshop committed
249
250
251
252

            # run function serial or parallel
            if n_tasks > 1:
                with Pool(processes=n_tasks) as pool:
253
                    results = pool.starmap(sellatlon_SAM, args)
v99 Workshop's avatar
v99 Workshop committed
254
255
256
257
258
            else:
                results = []
                for arg in args:
                    results.append(sellatlon_SAM(*arg))

259
260
261
            # collect timings from subtasks
            for task_TM in results:
                TM.merge_timings(task_TM)
262

263
            TM.start('merge')
264
            # merge all time step files to one
265
            #cdo_mergetime(out_tmp_dir, out_dir, var_name)
266
            TM.stop('merge')
267

268
    print('ran SAM at resolutions {} for variables {}'.format(ress, var_names))
269
    TM.print_report()