UND_Code/pull_variables_from_qtrack_output.py at main · kelliemthrower/UND_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
import numpy as np
import xarray as xr
from pathlib import Path

MODEL = 'VHR'
RESOLUTION = '0.25'
YEAR = 1986
BASE = Path(f"/data/kellie.thrower/IPSL_{MODEL}_QTRACK/{YEAR}_JunOct_{RESOLUTION}_deg")

track = xr.open_dataset(BASE / f"AEW_tracks_post_processed_{YEAR}_junOct.nc")
wind = xr.open_dataset(BASE / f"adjusted_data_{YEAR}_junOct.nc")
curv = xr.open_dataset(BASE / f"curv_vort_ipsl_{YEAR}_junOct.nc")

outdir = BASE / "AEW_wave_files"
outdir.mkdir(exist_ok=True)

lat = wind.latitude.values
lon = wind.longitude.values

dlat = abs(lat[1]-lat[0])
dlon = abs(lon[1]-lon[0])

half_lat = 15.0
half_lon = 15.0

ny = int(round(half_lat/dlat))
nx = int(round(half_lon/dlon))

for system in range(track.sizes["system"]):

    u_list=[]
    v_list=[]
    c_list=[]
    t_list=[]
    clon=[]
    clat=[]
    strength=[]

    for ti in range(track.sizes["time"]):

        lo=float(track.AEW_lon[system,ti])
        la=float(track.AEW_lat[system,ti])

        if np.isnan(lo) or np.isnan(la):
            continue

        iy=np.argmin(np.abs(lat-la))
        ix=np.argmin(np.abs(lon-lo))

        if iy-ny<0 or iy+ny>=len(lat):
            continue
        if ix-nx<0 or ix+nx>=len(lon):
            continue

        u=wind.u.isel(time=ti,
                      latitude=slice(iy-ny,iy+ny+1),
                      longitude=slice(ix-nx,ix+nx+1)).values

        v=wind.v.isel(time=ti,
                      latitude=slice(iy-ny,iy+ny+1),
                      longitude=slice(ix-nx,ix+nx+1)).values

        cv=curv.curv_vort.isel(time=ti,
                      latitude=slice(iy-ny,iy+ny+1),
                      longitude=slice(ix-nx,ix+nx+1)).values

        u_list.append(u)
        v_list.append(v)
        c_list.append(cv)
        t_list.append(track.time.values[ti])
        clon.append(lo)
        clat.append(la)
        strength.append(float(track.AEW_strength[system,ti]))

    if len(u_list)==0:
        print(f"Skipping wave {system+1}")
        continue

    u=np.stack(u_list)
    v=np.stack(v_list)
    cv=np.stack(c_list)

    ds=xr.Dataset(
        {
            "u700":(("time","y","x"),u),
            "v700":(("time","y","x"),v),
            "wind_speed":(("time","y","x"),np.sqrt(u*u+v*v)),
            "curv_vort":(("time","y","x"),cv),
            "AEW_strength":("time",strength),
            "center_lon":("time",clon),
            "center_lat":("time",clat),
        },
        coords={
            "time":t_list,
            "x":np.arange(-nx,nx+1)*dlon,
            "y":np.arange(-ny,ny+1)*dlat
        }
    )

    ds.to_netcdf(outdir/f"wave_{system+1:03d}.nc")
    print(f"Saved wave {system+1}")

print("Done.")