rhopot0

rhopot0#

# Parameters
variable = "rhopot0"
stream = "z"
long_name = "Potential density referenced to surface"

from IPython.display import display, Markdown

# Dynamically generate markdown content
markdown_text = f" This notebook compares area-weighted maps, in some cases, vertical profiles for {variable} in different basins."

# Display the updated markdown content
display(Markdown(markdown_text))

This notebook compares area-weighted maps, in some cases, vertical profiles for rhopot0 in different basins.

%load_ext autoreload
%autoreload 2

%%capture 
# comment above line to see details about the run(s) displayed
import sys, os
sys.path.append(os.path.abspath(".."))
from misc import *
import glob
print("Last update:", date.today())
%matplotlib inline

months = ['January', 'February', 'March', 'April', 
          'May', 'June', 'July', 'August', 'September', 
          'October', 'November', 'December']

# load data
ds = []
for c, p in zip(casename, climo_path):
  file = glob.glob(p+'{}.{}.{}.??????-??????.nc'.format(c, stream, variable))[0]
  ds.append(xr.open_dataset(file))

def identify_xyz_dims(dims):
    dims = tuple(dims)

    z_options = ['zl', 'z_l', 'zi', 'z_i']
    y_options = ['yh', 'yq']
    x_options = ['xh', 'xq']

    z_dim = next((dim for dim in dims if dim in z_options), None)
    y_dim = next((dim for dim in dims if dim in y_options), None)
    x_dim = next((dim for dim in dims if dim in x_options), None)

    # Set default values for coordinates and area
    x_coord = y_coord = area_var = None

    if y_dim == 'yh' and x_dim == 'xh':
        x_coord = 'geolon'
        y_coord = 'geolat'
        area_var = 'areacello'
    elif y_dim == 'yq' and x_dim == 'xh':
        x_coord = 'geolon_v'
        y_coord = 'geolat_v'
        area_var = 'areacello_cv'
    elif y_dim == 'yh' and x_dim == 'xq':
        x_coord = 'geolon_u'
        y_coord = 'geolat_u'
        area_var = 'areacello_cu'

    return x_dim, y_dim, z_dim, x_coord, y_coord, area_var

dims = identify_xyz_dims(ds[0][variable+'_annual_mean'].dims)

def annual_plot(variable, dims, label):
    area = grd_xr[0][dims[5]].fillna(0)
    x = dims[0]; y = dims[1]; z = dims[2]
    lon = dims[3]; lat = dims[4] 
    model = []
    for i in range(len(label)):
        if z is None:
            model.append(np.ma.masked_invalid(ds[i][variable+'_annual_mean'].values))
        else:
            model.append(np.ma.masked_invalid(ds[i][variable+'_annual_mean'].isel({z: 0}).values))

        if i == 0:
            xyplot(model[i], 
                grd_xr[i].geolon.values, grd_xr[i].geolat.values, area.values,
                title = 'Annual mean '+str(variable)+ ' ('+str(ds[0].units)+')', 
                suptitle= label[i]+', '+ str(start_date) + ' to ' + str(end_date), 
                extend='max')
        else:
            xyplot((model[i]-model[0]), 
                grd_xr[i].geolon.values, grd_xr[i].geolat.values, area.values,
                title = 'Annual mean '+str(variable)+ ' ('+str(ds[0].units)+')', 
                suptitle= label[i]+' - '+label[0]+', '+ str(start_date) + ' to ' + str(end_date), 
                extend='max')
            
    fig, ax = plt.subplots(figsize=(8,4))
    for i in range(len(label)):
        if z is None:
            ds[i][variable+'_annual_mean'].weighted(area).mean(x).plot(y=y, 
                                            ax=ax, label=label[i])
        else:
            ds[i][variable+'_annual_mean'].isel({z: 0}).weighted(area).mean(x).plot(y=y, 
                                            ax=ax, label=label[i])
            
    ax.set_title('Zonally averaged '+str(variable)+' ('+str(ds[0].units)+'), annual mean')
    ax.grid()
    ax.legend();
    return

Annual mean#

annual_plot(variable, dims, label)
../_images/8581cf2ad648004c8ef87e70b4b0f7bf98fd6d9deeeef0f529144688b0c9a691.png ../_images/a04328c56825371388cab48b5075534c98880ee74849ea376a1329c13290b4ce.png

Monthly climatology#

area = grd_xr[0][dims[5]].fillna(0)
x = dims[0]; y = dims[1]; z = dims[2]
lon = dims[3]; lat = dims[4]
model = []
for i in range(len(label)):
    if z is None:
        model.append(ds[i][variable+'_monthly_climatology'])
    else:
        model.append(ds[i][variable+'_monthly_climatology'].isel({z: 0}))
        
    if i == 0:
        g = model[i].plot(x='geolon', y='geolat', col='month', col_wrap=3,
            figsize=(12,12), robust=True,
            cbar_kwargs={"label": variable + ' ({})'.format(str(ds[0].units)),
                        "orientation": "horizontal", 'shrink': 0.8, 'pad': 0.05})
        
        plt.suptitle(label[i]+ ', ' +str(start_date) + ' to ' + str(end_date), y=1.02, fontsize=17)  

    else:
        g = (model[i]-model[0]).plot(x='geolon', y='geolat', col='month', col_wrap=3,
            figsize=(12,12), robust=True,
            cbar_kwargs={"label": variable + ' ({})'.format(str(ds[0].units)),
                        "orientation": "horizontal", 'shrink': 0.8, 'pad': 0.05})
        plt.suptitle(label[i] + ' - ' + label[0]+ ', ' +str(start_date) + ' to ' + str(end_date), 
                     y=1.02, fontsize=17)
../_images/9f7ff36de5fd803324346996b26cb0f36c468cfc5ecda09e831140f65a3ee3c8.png 
def monthly_plot(variable, dims, label, m):
    area = grd_xr[0][dims[5]].fillna(0)
    x = dims[0]; y = dims[1]; z = dims[2]
    lon = dims[3]; lat = dims[4]
          
    fig, ax = plt.subplots(figsize=(8,4))
    for i in range(len(label)):
        if z is None:
            ds[i][variable+'_monthly_climatology'].isel(month=m).weighted(area).mean(x).plot(y=y, 
                                               ax=ax, label=label[i])
        else:
            ds[i][variable+'_monthly_climatology'].isel({z: 0, 'month': m}).weighted(area).mean(x).plot(y=y, 
                                                ax=ax, label=label[i])
    ax.set_title(str(months[m])+', zonally averaged '+str(variable)+' ('+str(ds[0].units)+')')
    ax.grid()
    ax.legend();
    return

January#

m=0
monthly_plot(variable, dims, label, m)
../_images/6df0fc8999dc43d1117bd397580a3f70c1b969d64428c6d6133f339898230b41.png

February#

m=1
monthly_plot(variable, dims, label, m)
../_images/c8a318f322e2704d63f52e15c131b85b66e7bd93100c6bbd1d0138c8991aa6dc.png

March#

m=2
monthly_plot(variable, dims, label, m)
../_images/284f80fe4fa2b9847ed2a701ce6393ee178047a5e8152afba8be112c53e5baad.png

April#

m=3
monthly_plot(variable, dims, label, m)
../_images/777b55604d5f0be552fda43e8d958690215b4dea80a4fa5c862a5c2befda1bd4.png

May#

m=4
monthly_plot(variable, dims, label, m)
../_images/46d81ffaef973dfa04fd6d9c33c317315a48d6e2f28fa85ccc21fd5d6be6eb4a.png

June#

m=5
monthly_plot(variable, dims, label, m)
../_images/acbcf360f5fa7f2664dc06ab7e18a701a3ee9f41455bc9500ad4cd07dccee560.png

July#

m=6
monthly_plot(variable, dims, label, m)
../_images/c47606d976d372e91def7f93656092f5a330da8696896c17d03e3549decb2b32.png

August#

m=7
monthly_plot(variable, dims, label, m)
../_images/8687a1b9f132e111ce4e8bb1a74c3cfc97b0cdaffd4b429c64ac432888093aab.png

September#

m=8
monthly_plot(variable, dims, label, m)
../_images/2bde682eb97cfa0082a66f6a74b69756cf15899fca8f471fd95e8a3166b7d1c9.png

October#

m=9
monthly_plot(variable, dims, label, m)
../_images/cb8ec9a393364d12d5c7dbf3eed4dd4b563975b712aaa1e0b54054d8246612a2.png

November#

m=10
monthly_plot(variable, dims, label, m)
../_images/2eb2f668ab03f27850cb272e91a79a89b9a901ac2c74ae74e1b1e6a4bafd8e1c.png

December#

m=11
monthly_plot(variable, dims, label, m)
../_images/4b51f322acaad088e5a24d6ae30ffd0010974eeb506f4fc5f855767b2fb3c578.png

By basins#

Monthly climo @ surface#

# GMM, update this
basin_code = xr.open_dataset('/glade/work/gmarques/cesm/tx2_3/basin_masks/basin_masks_tx2_3v2_20250318.nc')['basin_masks']

area = grd_xr[0][dims[5]].fillna(0)
x = dims[0]; y = dims[1]; z = dims[2]
model_mean_wgt = []
    
for i in range(len(label)):
    basin_code_dummy = basin_code.rename({'yh': y, 'xh': x})
    if z is None:
        model = ds[i][variable+'_monthly_climatology']
    else:
        model = ds[i][variable+'_monthly_climatology'].isel({z: 0})
    
    model_mean_wgt.append((model * basin_code_dummy).weighted(area*basin_code_dummy).mean(dim=[y, x]))
        
for i in range(len(label)):
    g = model_mean_wgt[i].plot(x="month", yincrease=False, col="region", col_wrap=5, label=label[i])
    
fig = g.fig  # not g.figure
fig.suptitle(str(variable)+' ('+str(ds[0].units)+')', fontsize=16)
fig.tight_layout()
fig.subplots_adjust(top=0.9)
for ax in g.axes.flat:
    ax.grid(True);
ax.legend()

<matplotlib.legend.Legend at 0x14aa0d055610>
../_images/2908540adf63fafe528775367bb2fba7ef12f7b9f575843147cfd8cc65dc0812.png

Vertical profiles#

Averaged over annual means

z_max=1000 # change this to 6000 to see full profile

if stream == 'z' and (z == 'z_l' or z == 'z_i'):

    model_mean_wgt = []
    
    for i in range(len(label)):
        basin_code_dummy = basin_code.rename({'yh': y, 'xh': x})
        model = ds[i][variable+'_annual_mean']
        
        model_mean_wgt.append((model * basin_code_dummy).weighted(area*basin_code_dummy).mean(dim=[y, x]))
            
    for i in range(len(label)):
        g = model_mean_wgt[i].sel(**{z: slice(0., z_max)}).plot(y=z, yincrease=False, col="region", col_wrap=5, label=label[i])
    
    fig = g.fig  # not g.figure
    fig.suptitle(str(variable)+' ('+str(ds[0].units)+')', fontsize=16)
    fig.tight_layout()
    fig.subplots_adjust(top=0.9)
    plt.legend()
    for ax in g.axes.flat:
        ax.grid(True);
../_images/80334977bf43729bed1e7548376ae7a6821df5b008fe29aac3e82a26e7fd6094.png
Contents