diff --git a/tools/analysis/MOM6_refineDiag.csh b/tools/analysis/MOM6_refineDiag.csh
index c849e0a06b..fe881eb8d6 100755
--- a/tools/analysis/MOM6_refineDiag.csh
+++ b/tools/analysis/MOM6_refineDiag.csh
@@ -101,6 +101,9 @@ chmod +x $script_dir/*.py
 set ocean_static_file = $yr1.ocean_static.nc
 if ( -e $yr1.ocean_static_no_mask_table.nc ) set ocean_static_file = $yr1.ocean_static_no_mask_table.nc
 
+set basin_codes_file = basin_codes.nc
+if ( -e $yr1.basin_codes.nc ) set basin_codes_file = $yr1.basin_codes.nc
+
 set basin_codes_d2_file = ocean_static_d2.nc
 
 #echo '====annual mean Eddy Kinetic Energy======'
@@ -109,13 +112,13 @@ set basin_codes_d2_file = ocean_static_d2.nc
 
 echo '==== Offline Diagnostics ===='
 if ( -f $yr1.ocean_month.nc ) then
-  $script_dir/refineDiag_ocean_month.py -b basin_codes.nc -r $refineDiagDir $yr1.ocean_month.nc
+  $script_dir/refineDiag_ocean_month.py -b $basin_codes_file -r $refineDiagDir $yr1.ocean_month.nc
 endif
 if ( -f $yr1.ocean_month_z.nc ) then
-  $script_dir/refineDiag_ocean_month_z.py -b basin_codes.nc -r $refineDiagDir -s ./ $yr1.ocean_month_z.nc
+  $script_dir/refineDiag_ocean_month_z.py -b $basin_codes_file -r $refineDiagDir -s ./ $yr1.ocean_month_z.nc
 endif
 if ( -f $yr1.ocean_month_rho2.nc ) then
-  $script_dir/refineDiag_ocean_month_rho2.py -b basin_codes.nc -r $refineDiagDir $yr1.ocean_month_rho2.nc
+  $script_dir/refineDiag_ocean_month_rho2.py -b $basin_codes_file -r $refineDiagDir $yr1.ocean_month_rho2.nc
 endif
 
 echo '==== Offline Diagnostics downsampled ===='
diff --git a/tools/analysis/refineDiag_ocean_month.py b/tools/analysis/refineDiag_ocean_month.py
index f83e747d24..90ac10e4e0 100755
--- a/tools/analysis/refineDiag_ocean_month.py
+++ b/tools/analysis/refineDiag_ocean_month.py
@@ -40,7 +40,10 @@ def run():
 
 def heat_trans_by_basin(x,mask=None,lat=None,minlat=None):
     if mask is not None:
-        if not (x.shape[0] == 1+mask.shape[0]): #symmetric case
+        assert len(mask.shape) == 2, 'mask should be 2 dimensions'
+        assert len(x.shape) == 2 or len(x.shape) == 3, 'data should be 2 or 3 dimensions'
+        # symmetric case
+        if x.shape[-2:][1] == 1+mask.shape[1] and x.shape[-2:][0] == 1+mask.shape[0]:
             mask=np.append(mask,np.zeros((1,mask.shape[1])),axis=0)
         varmask = np.sum(mask,axis=-1)
         varmask = np.expand_dims(varmask,0)