testing for new module and fix bug in loading data

NOAA-CEFI-Portal · Apr 22, 2024 · 455cd12 · 455cd12
1 parent be09848
commit 455cd12
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 30 deletions.
diff --git a/mom6/stats/mom6_tercile_probability.py b/mom6/stats/mom6_tercile_probability.py
@@ -21,8 +21,7 @@
 import matplotlib.pyplot as plt
 import cartopy.crs as ccrs
 from scipy.stats import norm as normal
-from mom6 import DATA_PATH
-from mom6.mom6_module.mom6_io import MOM6Static, MOM6Forecast
+from mom6.mom6_module.mom6_io import MOM6Forecast
 
 
 if __name__=='__main__':
@@ -32,17 +31,21 @@
     ini_month = 3
     lead_season_index = 3
     varname = 'tos'
-    data_grid = 'regrid'
+    data_grid = 'raw'
 
-    # getting the ocean mask on mom grid
-    da_lmask = MOM6Static.get_mask(DATA_PATH+'static/',mask='wet')
 
     # loaded the mom6 raw field
-    mom6Forecast = MOM6Forecast(varname,'hindcast/','static/','tercile_calculation/',grid=data_grid)
+    if data_grid == 'regrid':
+        data_dir = 'hindcast/regrid/'
+        tercile_dir = 'tercile_calculation/regrid/'
+    elif data_grid == 'raw':
+        data_dir = 'hindcast/'
+        tercile_dir = 'tercile_calculation/'
+    mom6Forecast = MOM6Forecast(varname,data_dir,'static/',tercile_dir,grid=data_grid)
     ds_data = mom6Forecast.get_single(iyear=ini_year,imonth=ini_month)
 
     # load variable to memory
-    da_data = ds_data[varname].isel(init=0)
+    da_data = ds_data[varname]
 
     # setup lead bins to average during forecast lead time
     # (should match lead bins used for the historical data
@@ -70,6 +73,7 @@
     # load the predetermined hindcast/forecast tercile value
     #  this is based on 30 years statistic 1993-2023
     ds_tercile = mom6Forecast.get_tercile('grid')
+    ds_tercile = ds_tercile.sel(init=ini_month)
 
     # average the forecast over the lead bins
     ds_tercile_binned = (
@@ -116,7 +120,6 @@
     ds_tercile_prob=xr.Dataset()
     ds_tercile_prob['tercile_prob'] = da_tercile_prob
     ds_tercile_prob['tercile_prob_max'] = da_tercile_prob_max
-    ds_tercile_prob['mask'] = da_lmask
 
     ############# plotting the max tercile probability
     # get time format
@@ -178,17 +181,17 @@
     cbar.ax.tick_params(labelsize=12,rotation=0)
     cbar.set_label(label=colorbar_labelname,size=12, labelpad=15)
 
-    imm=(ds_tercile_prob['mask']
-        .plot.pcolormesh(
-            x=lon,
-            y=lat,
-            ax=ax2,
-            levels=np.arange(0,1+1),
-            cmap='grey',
-            transform=ccrs.PlateCarree(central_longitude=0)
-        )
-    )
-    imm.colorbar.remove()
+    # imm=(ds_tercile_prob['mask']
+    #     .plot.pcolormesh(
+    #         x=lon,
+    #         y=lat,
+    #         ax=ax2,
+    #         levels=np.arange(0,1+1),
+    #         cmap='grey',
+    #         transform=ccrs.PlateCarree(central_longitude=0)
+    #     )
+    # )
+    # imm.colorbar.remove()
 
     abovearrow = ax2.text(1.31, 0.65, "Upper tercile %",color='white',transform=ax2.transAxes,
                 ha="center", va="bottom", rotation=90, size=12,
@@ -216,4 +219,4 @@
         )
         )
     )
-    fig.savefig(f'{proj_dir}/figures/tercile_regrid.png')
+    fig.savefig(f'{proj_dir}/figures/tercile_{data_grid}.png')
diff --git a/mom6/stats/mom6_tercile_probability_regional.py b/mom6/stats/mom6_tercile_probability_regional.py
@@ -31,7 +31,12 @@
     lead_season_index = 3
     varname = 'tos'
     data_grid = 'raw'
-    region = 'MAB'
+    region = 'SS_LME'
+
+    # RegionalOptions = Literal[
+    # 'MAB','GOM','SS','GB','SS_LME','NEUS_LME','SEUS_LME',
+    # 'GOMEX','GSL','NGOMEX','SGOMEX','Antilles','Floridian'
+    # ]
 
     # getting the regionl mask on mom grid
     ds_lme = MOM6Static.get_regionl_mask(DATA_PATH+'masks/')
@@ -44,7 +49,7 @@
     ds_data = mom6Forecast.get_single(iyear=ini_year,imonth=ini_month)
 
     # load variable to memory (remove the init dimension)
-    da_data = ds_data[varname].isel(init=0)
+    da_data = ds_data[varname]
 
     # Area weighted average to the specific region
     da_data = (
@@ -78,6 +83,7 @@
     # load the predetermined hindcast/forecast tercile value
     #  this is based on 30 years statistic 1993-2023
     da_tercile = mom6Forecast.get_tercile('region').sel(region=region)
+    da_tercile = da_tercile.sel(init=ini_month)
 
     # average the forecast over the lead bins
     da_tercile_binned = (
@@ -212,4 +218,4 @@
         )
         )
     )
-    fig.savefig(f'{proj_dir}/figures/tercile_regional.png')
+    fig.savefig(f'{proj_dir}/figures/tercile_{data_grid}_regional.png')
diff --git a/mom6/stats/mom6_tercile_regional.py b/mom6/stats/mom6_tercile_regional.py
@@ -3,6 +3,9 @@
 avereaged tercile value (EPU) based on the raw gridded
 regional MOM6 forecast output
 
+regrid product need to have EPU mask avaialble if one 
+want to add to the script
+
 """
 # %%
 import os
@@ -35,12 +38,9 @@
     if grid == 'raw':
         MOM6_DIR = os.path.join(DATA_PATH,"hindcast/")
         MOM6_TERCILE_DIR = os.path.join(DATA_PATH,"tercile_calculation/")
-    elif grid == 'regrid':
-        MOM6_DIR = os.path.join(DATA_PATH,"hindcast/regrid/")
-        MOM6_TERCILE_DIR = os.path.join(DATA_PATH,"tercile_calculation/regrid/")
     else:
         print("Usage: python mom6_tercile_regional.py VARNAME GRIDTYPE")
-        raise NotImplementedError('GRIDTYPE can only be "raw" or "regrid"')
+        raise NotImplementedError('GRIDTYPE can only be "raw"')
 
     file_list = glob.glob(MOM6_DIR+'/*.nc')
     var_file_list = []
@@ -95,14 +95,14 @@
         ds_tercile = ds_tercile.drop_vars('quantile')
 
         # output the netcdf file
-        print(f'output {MOM6_TERCILE_DIR}{file[len(MOM6_DIR):-6]}tercile_{file[-6:-4]}.region.nc')
+        print(f'output {MOM6_TERCILE_DIR}{file[len(MOM6_DIR):-6]}tercile_{file[-6:-3]}.region.nc')
         MOM6Misc.mom6_encoding_attr(
                 ds,
                 ds_tercile,
                 var_names=list(ds_tercile.keys()),
                 dataset_name='regional mom6 tercile'
             )
         try:
-            ds_tercile.to_netcdf(f'{MOM6_TERCILE_DIR}{file[len(MOM6_DIR):-6]}tercile_{file[-6:-4]}.region.nc',mode='w')
+            ds_tercile.to_netcdf(f'{MOM6_TERCILE_DIR}{file[len(MOM6_DIR):-6]}tercile_{file[-6:-3]}.region.nc',mode='w')
         except PermissionError:
-            print(f'{MOM6_TERCILE_DIR}{file[len(MOM6_DIR):-6]}tercile_{file[-6:-4]}.region.nc is used by other scripts' )
+            print(f'{MOM6_TERCILE_DIR}{file[len(MOM6_DIR):-6]}tercile_{file[-6:-3]}.region.nc is used by other scripts' )