add split param for save_like

data_util.py

@@ -25,13 +25,13 @@ def split_variable(ds, name):
     return v
 
 
-def save_like(output, input, step, save_dir="", freq=6):
+def save_like(output, input, step, save_dir="", freq=6, split=False):
     if save_dir:
         os.makedirs(save_dir, exist_ok=True)
         dtime = (step+2) * freq # 
         init_time = pd.to_datetime(input.time.values[0])
 
-        data = xr.DataArray(
+        ds = xr.DataArray(
             output[None, None],
             dims=['member', 'time', 'dtime', 'level', 'lat', 'lon'],
             coords=dict(
@@ -44,6 +44,7 @@ def save_like(output, input, step, save_dir="", freq=6):
             )
         ).astype(np.float32)  
         
+        if split:
             def rename(name):
                 if name == "tp":
                     return "TP06"
@@ -51,13 +52,12 @@ def save_like(output, input, step, save_dir="", freq=6):
                     return "RH"
                 return name.upper()
             
-        ds = []
+            new_ds = []
             for k in pl_names + sfc_names:
-            v = split_variable(data, k)
+                v = split_variable(ds, k)
                 v.name = rename(k)
-            # print(f"{k}: {v.shape} {v.values.min()} ~ {v.values.max()}")
-            ds.append(v)
-        ds = xr.merge(ds, compat="no_conflicts")
+                new_ds.append(v)
+            ds = xr.merge(new_ds, compat="no_conflicts")
 
         print(f'Save to {save_name} ...')
         save_name = os.path.join(save_dir, f'{dtime:03d}.nc')