make gfs input

make_era5.py

+import os
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def make_era5(init_time, data_dir):
+
+    init_time = pd.to_datetime(init_time)
+    print(f"process {init_time} ...")
+
+    pl_file = os.path.join(data_dir, init_time.strftime('P%Y%m%d%H.nc'))
+    pl = xr.open_dataset(pl_file)
+
+    sfc_file = os.path.join(data_dir, init_time.strftime('S%Y%m%d%H.nc'))
+    sfc = xr.open_dataset(sfc_file)
+
+    tp_file = os.path.join(data_dir, init_time.strftime('R%Y%m%d.nc'))
+    tp = xr.open_dataarray(tp_file).fillna(0)
+    tp = tp.rolling(time=6).sum() * 1000
+    tp = tp.sel(time=tp.time[::6])
+    tp = tp.clip(min=0, max=1000)
+    sfc['tp'] = tp
+
+    pl_names = ['z', 't', 'u', 'v', 'r']
+    sfc_names = ['t2m', 'u10', 'v10', 'msl', 'tp']
+    levels = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
+
+    channel = [f'{n.upper()}{l}' for n in pl_names for l in levels]
+    channel +=[n.upper() for n in sfc_names]
+
+    ds = []
+    for name in pl_names + sfc_names:
+        if name in ['z', 't', 'u', 'v', 'r']:
+            v = pl[name]
+
+        if name in ['t2m', 'u10', 'v10', 'msl', 'tp']:
+            v = sfc[name]
+            level = xr.DataArray([1], coords={'level': [1]}, dims=['level'])
+            v = v.expand_dims({'level': level}, axis=1)             
+
+        if np.isnan(v).sum() > 0:
+            print(f"{name} has nan value")
+            raise ValueError
+
+        v.name = "data"
+        v.attrs = {}                
+        print(f"{name}: {v.shape}, {v.min().values} ~ {v.max().values}")
+        ds.append(v)
+     
+    ds = xr.concat(ds, 'level')
+    ds = ds.assign_coords(level=channel)
+    ds = ds.rename({'longitude': 'lon', 'latitude': 'lat'})
+    ds = ds.astype(np.float32)
+    return ds
+    
+ds12 = make_era5('20230725-12', 'ERA520230725')
+ds18 = make_era5('20230725-18', 'ERA520230725')
+ds = xr.concat([ds12, ds18], 'time')
+ds.to_netcdf('input.nc')
+

make_gfs.py

+import os
+import numpy as np
+import pandas as pd
+import pygrib as pg
+import xarray as xr
+
+def make_gfs(src_name):
+    assert os.path.exists(src_name)
+
+    levels = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
+    pl_names = ['gh', 't', 'u', 'v', 'r']
+    sf_names = ['2t', '10u', '10v', 'mslet']
+        
+    try:
+        ds = pg.open(src_name)
+    except:
+        print(f"{src_name} not found")
+        return
+
+    input = []
+    level = []
+
+    for name in pl_names + sf_names + ["tp"]:
+        
+        if name in pl_names:
+            try:
+                data = ds.select(shortName=name, level=levels)
+            except:
+                print("pl wrong")
+                return
+
+            data = data[:len(levels)]
+
+            if len(data) != len(levels):
+                print("pl wrong")
+                return 
+
+            if name == "gh":
+                name = "z"
+            
+            for v in data:
+                init_time = f'{v.date}-{v.time//100:02d}'
+
+                lat = v.distinctLatitudes
+                lon = v.distinctLongitudes    
+
+                img, _, _ = v.data()
+
+                if name == "z":
+                    img = img * 9.8
+
+                input.append(img)
+                level.append(f'{name}{v.level}')
+                print(f"{v.name}: {v.level}, {img.shape}, {img.min()} ~ {img.max()}")
+
+        if name in sf_names:
+            try:
+                data = ds.select(shortName=name)
+            except:
+                print('sfc wrong')
+                return
+
+            name_map = {'2t': 't2m', '10u': 'u10', '10v': 'v10', 'mslet': 'msl'}
+            name = name_map[name]
+
+            for v in data:
+                img, _, _ = v.data()
+                input.append(img)
+                level.append(name)
+                print(f"{v.name}: {img.shape}, {img.min()} ~ {img.max()}")
+
+        if name == "tp":
+            tp = img * 0
+            input.append(tp)
+            level.append("tp")
+
+    input = np.stack(input)
+    assert input.shape[-3:] == (70, 721, 1440)
+    assert input.max() < 1e10
+
+    times = [pd.to_datetime(init_time)]
+    input = xr.DataArray(
+        data=input[None],
+        dims=['time', 'level', 'lat', 'lon'],
+        coords={'time': times, 'level': level, 'lat': lat, 'lon': lon},
+    )
+
+    if np.isnan(input).sum() > 0:
+        print("Field has nan value")
+        return 
+    
+    return ds
+
+
+
+def test_make_gfs():
+    d1 = make_gfs('30/gfs.t06z.pgrb2.0p25.f000')
+    d2 = make_gfs('30/gfs.t12z.pgrb2.0p25.f000')
+
+    if d1 and d2:
+        ds = xr.concat([d1, d2], 'time')
+        ds = ds.assign_coords(time=ds.time.astype(np.datetime64))
+        ds.to_netcdf('input.nc')
\ No newline at end of file