Get a variable (salinity/temperature) but only at specified mixed layer depths

[pecos27]

Apologies in advance for this being not reproducible. This is also more of an advice question, but I have these monthly mixed layer depths (time, lat, and lon), where time in monthly intervals corresponds to the different depth of the mixed layer per month.

However I have 4d salinity (time, depth, lat and lon), is there a way to pick all the salinity values but only at the mixed layer depth for each grid point? This would result in a new data array of salinity but in time, lat and lon.
The mixed layer depth may not match exactly the depth of the salinity, but I was thinking a "nearest neighbor" depth could be taken instead.

Is it possible to do this with xarray?

I was thinking something along the lines of this:

# outputting some variables
print(mld_monthly)
<xarray.DataArray 'depth' (time: 324, lat: 60, lon: 75)>
dask.array<where, shape=(324, 60, 75), dtype=float32, chunksize=(1, 60, 75), chunktype=numpy.ndarray>
Coordinates:
  * lat      (lat) float32 20.0 21.0 22.0 23.0 24.0 ... 75.0 76.0 77.0 78.0 79.0
  * lon      (lon) float32 281.0 282.0 283.0 284.0 ... 352.0 353.0 354.0 355.0
  * time     (time) datetime64[ns] 1993-01-16T12:00:00 ... 2019-12-16T12:00:00

print(salinity)
<xarray.DataArray 'salinity' (time: 324, depth: 42, lat: 60, lon: 75)>
dask.array<where, shape=(324, 42, 60, 75), dtype=float32, chunksize=(1, 42, 60, 75), chunktype=numpy.ndarray>
Coordinates:
  * depth    (depth) float32 5.022 15.08 25.16 ... 4.752e+03 5.051e+03 5.35e+03
  * lat      (lat) float32 20.0 21.0 22.0 23.0 24.0 ... 75.0 76.0 77.0 78.0 79.0
  * lon      (lon) float32 281.0 282.0 283.0 284.0 ... 352.0 353.0 354.0 355.0
  * time     (time) datetime64[ns] 1993-01-16T12:00:00 ... 2019-12-16T12:00:00
Attributes:
    long_name:      salinity
    units:          psu
    standard_name:  sea_water_salinity
    valid_min:      -5.0
    valid_max:      48.0

salinities_at_the_mld = np.zeros((mld_monthly.shape))

for i in range(324): #324 = total months 
    for j in range(42): # 42 depth levels 
        salinities_at_the_mld[i,:,:] = salinity[i,j].where(mld_monthly[i])

The nested loop takes a long time though... and I'm not sure it's really doing the right thing. Is there another recommended way to get the salinity but only at the mixed layer depths using xarray? Thank you in advance!

[mathause]

Have you tried salinity.sel(depth=mld_monthly, method="nearest")? The coords must align (but they seem to do). A small example:

s = xr.DataArray(np.arange(6).reshape(3, 2), dims=("d", "x"), coords=dict(d=[0, -1, -2], x=[11, 12]))
d = xr.DataArray([0, -1.4], name="d", dims=("x",), coords=dict(x=[11, 12]))

s.sel(d=d, method="nearest")

Which returns:

<xarray.DataArray (x: 2)>
array([0, 3])
Coordinates:
    d        (x) int64 0 -1
  * x        (x) int64 11 12