Towards model agnosticism + pint functionality in Meridional Overturning Circulation 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})

    ...

Also this PR adds pint functionality to deal with units and conversions from, e.g., kg/s to Sv in a systematic manner.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[review-notebook-app]

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!

[dhruvbhagtani]

Hi, I just made a few minor enhancements to the example.

Towards the end of the notebook, it is mentioned that sub mesoscale contributions are not considered. I thought that ty_trans_rho_GM includes the sub mesoscale contribution?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

anton-seaice commented on 2024-07-03T01:42:56Z
----------------------------------------------------------------

Line #7.            "chunks": {},

Can we just remove these chunks objects ?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

anton-seaice commented on 2024-07-03T01:42:57Z
----------------------------------------------------------------

Line #19.            psiGM = psiGM.sel(time=slice(start_time, end_time))

We probably don't need to specify start_time / end_time twice (its already in the getvar)

dhruvbhagtani commented on 2024-07-03T02:11:38Z
----------------------------------------------------------------

Good point -- I need to check it again, but my understanding is that using start and end times within getvar only selects the time roughly. To exactly be within the start and end time bounds, we need to write this line.

[dhruvbhagtani]

View / edit / reply to this conversation on ReviewNB

anton-seaice commented on 2024-07-03T01:42:56Z ----------------------------------------------------------------

Line #7. "chunks": {},
Can we just remove these chunks objects ?

We could but then we get a bunch of warnings every time we load a variable. Perhaps we can comment the use of this line?

[dhruvbhagtani]

Good point -- I need to check it again, but my understanding is that using start and end times within getvar only selects the time roughly. To exactly be within the start and end time bounds, we need to write this line.

---

View entire conversation on ReviewNB

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

anton-seaice commented on 2024-07-03T02:24:32Z
----------------------------------------------------------------

Line #1.    %config InlineBackend.figure_format='retina'

I think we can remove this line too.

[navidcy]

Line #30. for ii in range(nmin, nmax):
if we can replace this for loop with xarray operations it'd be nice!

Yes. I'm happy to open an issue pointing this out. I think there might be one already? #321?
In this PR I tried to reduce the if model=='mom5': do_this() elif model='mom6': do_that() statements.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:31Z
----------------------------------------------------------------

Delete this

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:32Z
----------------------------------------------------------------

Can we add some documentation on what the GM transport is, and why we need to include it in this calculation? What are we missing if we don't include it?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:33Z
----------------------------------------------------------------

Now let's remap the overturning into depth coordinates. First we define the method that does the job. define a function to remap the vertical dimension into depth coordinates.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:33Z
----------------------------------------------------------------

Line #3.        experiment = psi_args[model]["expt"]

We are using here psi_args to open something that is not psi . It would be worth making a more general model_args dictionary, which has psi_variable information, density,etc.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:34Z
----------------------------------------------------------------

Line #28.        z1 = rho2_zonal_mean[remap_dens[model]["depth"]].values

I think remap_dens[model]["depth"] is replaceable by rho2_zonal_mean.cf['vertical'] . That's one less argument needed in the remap_dens dictionary.

Although z1 is never used. Why is it here?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:35Z
----------------------------------------------------------------

        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)

I feel all this would be simpler by:

xr.DataArray(data = rho2_zonal_mean.cf['vertical'].values,
             dims = ['rho_ref'],
             coords = {'rho_ref': rho2_zonal_mean.cf.isel(latitude=ii).values})

And similar in other places in this function. Seems to me some of the steps are redundant (like the naming and renaming of latitude_b latitude_bsect etc. Apologies if these steps are essential and I'm failing to understand them

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:35Z
----------------------------------------------------------------

Line #61.        del psi_avg2

Not needed

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:36Z
----------------------------------------------------------------

I am getting this error here:

RuntimeError: NetCDF: Not a valid ID

And then my ARE session collapses. I am using xxlarge (28cpus), so not sure what the problem is, but I think there is something amiss in the remapping function. I've tried running outside the function to identify what's the problematic line, but the iteration through latitudes is very slow. Maybe ufuncs iwould be good for this?

adele-morrison commented on 2024-08-21T10:40:15Z
----------------------------------------------------------------

I got the same error Julia, but was able to fix it using client = Client(threads_per_worker=1), see https://forum.access-hive.org.au/t/netcdf-not-a-valid-id-errors/389

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:37Z
----------------------------------------------------------------

Line #1.    stretching_factor = 6 # A powvalues set the stretching

what's a powvalues?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

julia-neme commented on 2024-08-05T03:54:37Z
----------------------------------------------------------------

Does this mean the ty_trans_gm that is not available for these experiments?

[julia-neme]

My main issue is that I can't run this notebook. I get a runtime error when trying to do the remapping and sometimes my ARE session crashes. I am using 24.04 conda-environment, and an xxlarge ARE session (28 cpus). I've heard that this notebook sometimes run, sometimes not, in a bit of a stochastic manner.

My second questions is whether potrho coordinate in mom5 is potrho2 and not potrho0? This would matter for the re-mapping. The metadata does not say and I personally don't know.

I think there are a number of simplifications/edits to do (see my comments), but I might not have understood stuff so apologies if the suggestions are not appropriate. I think we also need a bit of an explanation as to when/why do we need a GM transport.

A minor personal opinion is to do mom5 and mom6 one after the other and plot them together to compare.

[adele-morrison]

Confirming that the potrho coord in MOM5 is potrho2.

[navidcy]

I know I opened this PR but I'm out of steam. I'm tempted to close it because I don't have capacity to deal with it. But if somebody else feels that this is useful and want to take over please do.

[adele-morrison]

Regarding this:

Towards the end of the notebook, it is mentioned that sub mesoscale contributions are not considered. I thought that ty_trans_rho_GM includes the sub mesoscale contribution?

Well, GM supposedly captures the mesoscale... Is that where the comment on submesoscale is based on?

There is a separate submesoscale parameterisation, based on Fox-Kemper et al. 2008 and Fox-Kemper et al. 2011, that is different to GM and is still switched on even in 1/10deg and higher res configs. To include the advection from this parameterisation, you would need to save the diagnostic ty_trans_submeso. This figure from Fox-Kemper et al. 2011 gives you an idea of the magnitude and distribution of this term in a couple of low-res GFDL and CCSM models:

Perhaps we could add a comment to the notebook with the name of the required diagnostic if someone did want to include it and a reference to Fox-Kemper et al. 2011?

[julia-neme]

I have a curious questions about this notebook - unrelated to the reviews coming from personal interest.

Most of the mom5 experiments include a transport in depth coordinates (ty_trans), and I don't understand why we remap from density to get depth instead of simply using ty_trans. I suspect it is because there is something conceptually wrong in doing this but I don't understand why.

[adele-morrison]

@julia-neme if you do the time average in depth space first, then you only get the mean overturning component, not the residual (i.e. total = mean + eddy). So you have to do the time average in density space (i.e. use the ty_trans_rho diagnostic) to get the full overturning. But then often people want to see what it looks like in depth space not density space, so that's why the recipe maps back to depth space for visualisation purposes.

[adele-morrison]

@navidcy I will have a go at getting it running and incorporating @julia-neme's review comments.

[adele-morrison]

I got the same error Julia, but was able to fix it using client = Client(threads_per_worker=1), see https://forum.access-hive.org.au/t/netcdf-not-a-valid-id-errors/389

---

View entire conversation on ReviewNB

[adele-morrison]

I included nearly all of @julia-neme's review comments, and it runs again! There is definitely still a bunch more cleaning up that could be done (e.g. I didn't make these changes @julia-neme suggested), but I reckon we should merge for now, and can restart further cleanup in a new PR if someone has energy.

[adele-morrison]

@navidcy or @julia-neme are you able to sort the conflicts? I have no idea how to do that.

[julia-neme]

Ohh thanks Adele! It ran quite fast for me now... seems like the threads_per_worker needs to be included, not only when using intake!

I have no idea how to resolve conflicts so might leave that to Navid.

[anton-seaice]

I did the merge from main , if you can sanity check its not broken before merging this PR that would be great (there were no conflicts, I don't understand why it need the commandline !)

[navidcy]

Thanks!!!

[adele-morrison]

It ran quite fast for me now...

I think mostly it runs faster, because I reduced it to using only 1 year of output. I got sick of waiting for it!