Enhances "Zonally-averaged overturning" example

[navidcy]

This PR makes this example truly model-agnostic.

Previously there were a lot of if model=='mom5' do this, elseif model=='mom6' do the other in the remap_depth() method. This is not "model agnostic" programming, but rather two set of codes packaged in one notebook. The PR uses cf-xarray functionality to replace the if model='this' do that... with one code that works for all.

In particular:

before this PR

ef remap_depth(remap_dens,psi_args,psi_avg,session,model):

    ...

    #Mask the Mediteranean
    if model=='mom5': rho2=rho2.where(((rho2.xt_ocean<0) | (rho2.xt_ocean>45) ) | ((rho2.yt_ocean<10) | (rho2.yt_ocean>48)))
    if model=='mom6': rho2=rho2.where(((rho2.xh<0) | (rho2.xh>45) ) | ((rho2.yh<10) | (rho2.yh>48)))

    ...

    if model=='mom5':
        nmax=np.size(rho2_zonal_mean.yt_ocean)
        nmin=59
    elif model=='mom6':
        nmax=np.size(rho2_zonal_mean.yh)
        nmin=int(list(rho2_zonal_mean.yh.values).index(rho2_zonal_mean.sel(yh=-78, method='nearest').yh)) #locates lat=-78
        
    ...

    for ii in range(nmin,nmax):
        if model=='mom5':
            rho1 = rho2_zonal_mean.isel(yt_ocean=ii); rho1v=rho1; z=rho1.st_ocean
            rho1 = rho1.rename({'st_ocean' : 'rho_ref'}); rho1['rho_ref']=np.array(rho1v)
            rho1.name='st_ocean'; rho1.values=np.array(z)
            rho1 = rho1.isel(rho_ref = ~np.isnan(rho1.rho_ref)).drop_duplicates(dim='rho_ref', keep='first')
            rho1 = rho1.interp(rho_ref = psi_avg.potrho.values, kwargs={"bounds_error": False, "fill_value": (0, 6000)})
            psi_depth[:, ii] = rho1.rename({'rho_ref' : 'potrho'})
        elif model=='mom6':
            rho1 = rho2_zonal_mean.isel(yh=ii); rho1v=rho1; z=rho1.z_l
            rho1 = rho1.rename({'z_l': 'rho_ref'}); rho1['rho_ref']=np.array(rho1v)
            rho1.name='z_l'; rho1.values=np.array(z)
            rho1 = rho1.isel(rho_ref = ~np.isnan(rho1.rho_ref)).drop_duplicates(dim='rho_ref', keep='first')
            rho1 = rho1.interp(rho_ref = psi_avg.rho2_l.values, kwargs={"bounds_error": False, "fill_value": (0, 6000)})
            psi_depth[:, ii] = rho1.rename({'rho_ref': 'rho2_l'})

    ...

after this PR

def remap_depth(remap_dens, psi_args, psi_avg, session, model):
    
    ...
    
    #Mask the Mediteranean
    rho2 = rho2.cf.where(((rho2.cf['longitude'] < 0) | (rho2.cf['longitude'] > 45) ) |
                         ((rho2.cf['latitude'] < 10) | (rho2.cf['latitude'] > 48))
                        )

    ...
    
    # nmin is the latitude index that corresponds to 78S
    nmin = int(list(rho2_zonal_mean.cf['latitude'].values).index(rho2_zonal_mean.cf.sel(latitude=-78, method='nearest').cf['latitude']))
    nmax = np.size(rho2_zonal_mean.cf['latitude'])
        
   ...

    for ii in range(nmin, nmax):
        rho1 = rho2_zonal_mean.cf.isel(latitude=ii); rho1v = rho1.copy(); z = rho1.cf['vertical']
        rho1 = rho1.rename({rho1.cf['vertical'].name: 'rho_ref'}); rho1['rho_ref'] = np.array(rho1v)
        rho1.name = rho2_zonal_mean.cf['vertical'].name; rho1.values = np.array(z)
        rho1 = rho1.isel(rho_ref = ~np.isnan(rho1.rho_ref)).drop_duplicates(dim='rho_ref', keep='first')
        rho1 = rho1.interp(rho_ref = psi_avg.cf['vertical'].values,
                           kwargs={"bounds_error": False, "fill_value": (0, 6000)})
        psi_depth[:, ii] = rho1.rename({'rho_ref': psi_avg.cf['vertical'].name})

    ...

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[review-notebook-app[bot]]

View / edit / reply to this conversation on ReviewNB

navidcy commented on 2024-02-26T15:02:17Z ----------------------------------------------------------------

Line #30.        for ii in range(nmin, nmax):

if we can replace this for loop with xarray operations it'd be nice!