Waveguide_code/Waveguide_preprocess_ERA5.py at main · rhwhite/Waveguide_code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import xarray as xr
import netCDF4
import numpy as np
import sys

#from rhwhitepackages3.waveguides import *
from rhwhitepackages3.waveguides_pre import *
from datetime import date
import warnings

args = (sys.argv)
if len(args) < 3:
    sys.error('need at least 2 arguments: pressure level and hemisphere')
else:
    plev = int(args[1])
    hemi = args[2]


## Read in raw U data
Ufile = 'ERA5_u_850_250_dailymean_1degree_1979-2022.nc'
dirout = # Define directory to write output to

startyear = 1979
endyear = 2022

if hemi == 'NH':
    latstart = 10
    latend = 90
elif hemi == 'SH':
    latstart = -90
    latend = -10
else:
    sys.error('hemisphere arguement must be "NH" or "SH"')

# Filter U data - Northern hemisphere
with xr.open_dataset(Ufile).sel(level=plev) as datain:
    # For this version of ERA5, flip latitudes

    try:
        if datain.latitude[0] > datain.latitude[1]:
            print('flipping latitudes')
            datain = datain.sel(latitude=slice(None, None,-1))
        # resample to daily data in case input is not daily - timefilter assumes daily data
        Udatain = datain.sel(time=slice('01-Dec-' + str(startyear),
                              '01-Jan-' + str(endyear +1))).sel(latitude=slice(latstart,latend)).u.resample(time='D').mean()
    except AttributeError:
        if datain.lat[0] > datain.lat[1]:
            print('flipping latitudes')
            datain = datain.sel(lat=slice(None, None,-1))
        # resample to daily data in case input is not daily - timefilter assumes daily data
        Udatain = datain.sel(time=slice('01-Dec-' + str(startyear),
                              '01-Jan-' + str(endyear +1))).sel(lat=slice(latstart,latend)).u.resample(time='D').mean()
    print('reached here OK')

    for wavelfilter in [2,3]:
        print(wavelfilter)
        ## Zonal filter data

        ERA5_filter_temp = zonal_filter_wind(Udatain,wavelfilter)

        for timefilter in [10,15]:
            print(timefilter)

            if timefilter == 0:
                ERA5_filter_temp2 = ERA5_filter_temp

            else:
                ERA5_filter_temp2 = butter_time_filter_wind(ERA5_filter_temp.u,timefilter)

            print('calculated')
            # Convert to datasets and append

            ds = ERA5_filter_temp2.u.to_dataset(name = 'U_filtered')

            # save file
            ds.attrs['history'] = ('Created on ' + str(date.today()) + ' by Waveguide_preprocess_ERA5.py script')
            ds.attrs['input data'] = ('Input files = ' + Ufile)
            ds.attrs['parameters'] = ('timefilter: ' + str(timefilter) +
                                      ' and zonal wavenumber filter: ' + str(wavelfilter))

            ds.to_netcdf(dirout + '/ERA5_' + hemi + '_u' + str(plev) + '_dailymean_filtered_' + str(startyear) + '-' + str(endyear) +
                         '_tf' + str(timefilter) + '_zfnt' + str(wavelfilter) + '_1x1.nc')

            print('written')

# Calculate Ks

for wavelfilter in [2,3]:
    print(wavelfilter)

    for timefilter in [10,15]:
        print(timefilter)

        filterfilename = ('ERA5_' + hemi + '_u' + str(plev) + '_dailymean_filtered_' + str(startyear) + '-' + str(endyear) +
		 '_tf' + str(timefilter) + '_zfnt' + str(wavelfilter) + '_1x1')

        with xr.open_dataset(dirout + '/' + filterfilename + '.nc') as ERA5_filter_temp:
            (U_temp, Ks_temp, Ks2_temp) = calc_Ks_wg(ERA5_filter_temp.U_filtered)

            print('calculated')
            # Convert to datasets and append
            ds = Ks_temp.to_dataset(name = 'Ks')
            ds['Ks2'] = Ks2_temp
            ds['Ufilter'] = U_temp

            # save file
            ds.attrs['history'] = ('Created on ' + str(date.today()) + ' by Waveguide_preprocess_ERA5.py script')
            ds.attrs['input data'] = filterfilename + ' from ' + (ERA5_filter_temp.attrs['input data'])
            ds.attrs['parameters'] = ('Timefilter: ' + str(timefilter) +' and zonal wavenumber filter: ' + str(wavelfilter))

            ds.to_netcdf(dirout + '/Ks_' + filterfilename + '.nc')

            print('written')