xr.map_blocks calculation stalling

[rsx23ser]

I am trying to reproduce the map_blocks example given here (link), which finds the mixed-layer depth. I am using output files from the POP ocean model but cut up into smaller lon/lat boxes (using ncks), as opposed to reading in the global fields. The only change I have made to the subroutine in that link is that I am directly feeding in the density to the mld_dsigma routine instead of computing the density from salinity and temperature.

`def mld_dsigma(RHO, dsigma=0.03, rho_chunks='auto'):

# determine dimensionality
dims_in = RHO.dims
assert 'z_t' in RHO.coords, 'z_t not found in RHO coords'

# drop ancillary coordinates
rho = RHO.reset_coords(drop=True)

if 'nlat' in rho.dims and 'nlon' in rho.dims:
    rho = rho.assign_coords({
        'nlat': xr.DataArray(np.arange(len(RHO.nlat)), dims=('nlat')),
        'nlon': xr.DataArray(np.arange(len(RHO.nlon)), dims=('nlon')),
    })

# Compute density - this will help the workflow from being "bogged down" with too many tasks
rho = rho.chunk('auto').persist()

# Setup a template
template = rho.isel(z_t=0).drop('z_t')
template.attrs['long_name'] = 'MLD'
template.attrs['units'] = RHO.z_t.attrs['units']
template.name = 'MLD'

# compute and return MLD - this is where the parallelization comes in
return xr.map_blocks(
    interp_mld_dsigma, rho,
    kwargs=dict(dsigma=dsigma),
    template=template,
)`

The other routine interp_mld_dsigma is unchanged from what is in that link. When I execute the above routine, for some reason the calculation just stops after a point. I have uploaded a screenshot of the dask dashboard (link) which shows the point beyond which the code does not appear to proceed. Eventually I just kill the job. Has anybody experienced a similar issue before? Thanks for any help.

[max-sixty]

Hi @rsx23ser , I appreciate that's frustrating. We'd need much more information though — what's a simpler example that works? What's happens if you run over smaller data? Does it work if the arrays are numpy arrays?

See the links in #5404 for more

[rsx23ser]

@max-sixty Thanks a lot for your response! I will try to create a simpler example and repost.

[tlsw231]

@max-sixty Just wanted to follow up on this. I am pasting below some code from a jupyter notebook that should run by itself. My original question concerns output from an ocean model (POP) but some of the behavior in that case appears to be occurring in the example below. The code also includes the two functions (with very slight modifications) referred to in the link provided in the first message on this thread.

The execution appears to proceed in fits and bursts. I find that for smaller (nlat,nlon) values the code runs but the tasks do not appear to be distributed evenly among all the workers throughout the computation. It might start out that way but after some time, one worker ends up with most of the tasks which seems suboptimal. I am uploading two videos at these links which show the dask dashboard for two cases: (nlat, nlon) = (25,25) and (nlat, nlon) = (50,50). In both cases, you can see what I am referring to above.

https://www.dropbox.com/s/fz1qbk6z9l7zjkx/nlat_nlon_25x25.mov?dl=0

https://www.dropbox.com/s/cqb23oxk8rv1ef1/nlat_nlon_50x50.mov?dl=0

import xarray as xr               
import numpy as np              
import scipy
from scipy import interpolate
from dask.distributed import progress

from distributed import Client
client=Client(n_workers=10,memory_limit='16.0GB')
client

# Create time and depth arrays
nt=len(np.arange(204))
z_t=np.arange(62); nz=len(z_t);

# Define number of lat/lon points.
nlat=50;nlon=50;
#nlon=601;nlat=251;

# Create chunk sizes for the four dimensions.
# Do not chunk along the depth (z_t) coordinate as it is a core dimension for interpolation.
chunks=dict(z_t=-1,time=12,nlat=-1,nlon=-1)

# Create a 4D function that is basically just the vertical level in the model and is broadcast to dimensions (time,z_t,nlat,nlon).
# We will test the interpolation function on this synthetic field (not really density).
field=np.transpose(np.broadcast_to(np.arange(nz),(nlon,nlat,nz)));
field=np.broadcast_to(field,(nt,nz,nlat,nlon));
field = xr.DataArray(data=field,dims=['time','z_t','nlat','nlon'], \
                    coords=dict(time=('time',np.arange(nt)),z_t=('z_t',z_t), \
                    lon=(['nlon'],np.arange(nlon)),lat=(['nlat'],np.arange(nlat)) ))
field.z_t.attrs['units']='kg/m3'

def mld_dsigma(RHO, dsigma=30.34, rho_chunks={'z_t':-1,'time':6}):
    """
    Compute MLD based on ∆σ criterion. Uses xarray.map_blocks.

    Parameters
    ----------

    RHO : xarray.DataArray
      potential density
    dsigma : float, optional
      The value for ∆σ.
    rho_chunks: dictionary, optional
      Dimension chunk parameters to chunk rho by

    Returns
    -------

    MLD : xarray.DataArray
      The MLD (m) defined as the point in the water column where
      density exceeds rho[0] + dsigma.
    """

    # determine dimensionality
    dims_in = RHO.dims
    assert 'z_t' in RHO.coords, 'z_t not found in RHO coords'

    # drop ancillary coordinates
    rho = RHO.reset_coords(drop=True)

    # add these coordinates which creep in somewhere, I think when xarray does the unstack
    # without these, I get an error: not expecting {'nlat', 'nlon'}
    # chunking here is arbitrary and maybe suitable for the global POP_gx1v7 grid
    if 'nlat' in rho.dims and 'nlon' in rho.dims:
        rho = rho.assign_coords({
            'nlat': xr.DataArray(np.arange(len(RHO.nlat)), dims=('nlat')),
            'nlon': xr.DataArray(np.arange(len(RHO.nlon)), dims=('nlon')),
        })

    # Compute density - this will help the workflow from being "bogged down" with too many tasks
    rho = rho.chunk(rho_chunks).persist()
#    rho = rho.chunk(rho_chunks)

    # Setup a template
    template = rho.isel(z_t=0).drop('z_t')
    template.attrs['long_name'] = 'MLD'
    template.attrs['units'] = RHO.z_t.attrs['units']
    template.name = 'MLD'
    
    # compute and return MLD - this is where the parallelization comes in
    return xr.map_blocks(
        interp_mld_dsigma, rho,
        kwargs=dict(dsigma=dsigma),
        template=template,
    )


def interp_mld_dsigma(rho_in, dsigma=30.34):
    """compute MLD at point using interpolation"""

    non_vertical_dims = [d for d in rho_in.dims if d not in ['z_t']]
    rho_stack = rho_in.stack(non_vertical_dims=non_vertical_dims)
    mld_stack = xr.full_like(rho_stack.isel(z_t=slice(0, 1)), fill_value=np.nan)
    N = rho_stack.sizes["non_vertical_dims"]
    z_t = rho_in.z_t

    for i in range(N):
        # if all NaN, skip this column
        if rho_stack[:, i].isnull().all():
            continue

        # if the whole column has density less than the threshold, set MLD to deepest point
        if (rho_stack[:, i] < rho_stack[0, i] + dsigma).all():
            k = np.where(~np.isnan(rho_stack[:, i]))[0]
            mld_stack[:, i] = z_t[k[-1]]
            
        # linearly interpolate
        else:
            f = interpolate.interp1d(
                rho_stack[:, i], z_t,
                assume_sorted=False,
            )
            mld_stack[:, i] = f(rho_stack[0, i] + dsigma)

    return mld_stack.unstack().isel(z_t=0, drop=True).transpose(*non_vertical_dims)

# Compute the final result
result=mld_dsigma(field,rho_chunks=chunks).persist() 

[tlsw231]

Oops, forgot to mention I am the OP on this thread. It was posted under a different username.