Improve efficiency on Surface Water-Mass Transformation computation

[josuemtzmo]

Hello,

I recently started computing the Surface Water-Mass Transformation and I found a more efficient and quicker way to do the binning in density space. It requires flox a library that allows to aggregate xr.DataArray. See flox.xarray.xarray_reduce.

The required change on the notebook would simplify the code from:

bin_bottoms = isopycnal_bins[:-1]
binned_salt_transformation = xr.DataArray(np.zeros((len(bin_bottoms), len(yt_ocean), len(xt_ocean))), coords = [bin_bottoms, yt_ocean,xt_ocean], dims = ['isopycnal_bins', 'yt_ocean', 'xt_ocean'], name = 'salt transformation in isopycnal bins summed over time')
binned_salt_transformation.chunk({'isopycnal_bins':1})
for i in range(len(isopycnal_bins)-1):
    bin_mask = pot_rho_1.where(pot_rho_1 <= isopycnal_bins[i+1]).where(pot_rho_1 > isopycnal_bins[i]) * 0 + 1
    masked_transform = (salt_transformation * bin_mask).sum(dim = 'time')
    masked_transform = masked_transform.where(masked_transform != 0)
    masked_transform = masked_transform.load()
    binned_salt_transformation[i, :, :] = masked_transform
print('salt_transformation binning done')

to:

aligned_salt_transformation, aligned_pot_rho_1 = xr.align( salt_transformation, pot_rho_1)
binned_salt_transformation = xarray_reduce(
    aligned_salt_transformation,
    aligned_pot_rho_1,
    func="sum",
    expected_groups=isopycnal_bins,
    dim="time"
)

This same library can be use to more efficiently compute the overturning circulation and other "complex" aggregations. If I have time, I will submit a PR, but I wanted to leave this here in case it's useful for somebody else.

[navidcy]

Nice! I just don't see where this flox is used in the bottom code you wrote up. Did you import it at the top and used an alias?

[josuemtzmo]

You import from flox.xarray import xarray_reduce

[josuemtzmo]

I'm still having problems in my data since it's a 6000x6000 grid, but I think for a 1/4° sim should be easier to implement.

[navidcy]

@aidanheerdegen, @angus-g do we have this flox.xarray in the conda analysis environment?

[aidanheerdegen]

@aidanheerdegen, @angus-g do we have this flox.xarray in the conda analysis environment?

This is unknowable ...

$ module load conda/analysis3-23.01
$ python -c 'import flox.xarray'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
ModuleNotFoundError: No module named 'flox'

[navidcy]

I guess one needs to install it? I was asking if it's installed -- you answered :)

[aidanheerdegen]

I was being snarky. But I only do it to people who I know have an excellent sense of humour .. and you @navidcy

Yeah, you'll need email CMS help to get that installed (cws_help@nci.org.au)