How can I replace missing values in a dataset with a value that depends on a coordinate?

[DaniJonesOcean]

I am working with a set of vectors (i.e. profiles) that have a number of missing values. I have used linear interpolation to fill some of the missing values, but one problem remains: there are still missing values where one cannot interpolate, and extrapolating is not especially sensible in this case.

I would like to replace the remaining missing values with a "mean" vector (i.e. mean profile) that changes with one of the coordinates (specifically, it changes with depth). How can I do this? @AndrewWilliams3142 made a good suggestion on Twitter yesterday, but I haven't been able to implement it. Here are some details that I hope will make my case clear.

Here is my xarray dataset (named df):

where I would like to replace NaN values in ![p2](https://user-images.githubusercontent.com/11757453/117964556-a8a6f380-b319-11eb-89ef-06e162fc932f.jpg) with the values at the corresponding location in mean_S, which only depends on the depth coordinate. (This is a bit of an odd thing to do, but I want to try it for a toy example.) Here is what the field interp_S looks like:

Now, here is the set of commands that I tried to use in order to replace the values in interp_S' with the values in mean_S', based on the suggestion from @AndrewWilliams3142:

da = df.interp_S
clean_S = da.where(da != np.nan, df.mean_S)
df = df.assign({'clean_S':clean_S})

However, the field clean_S looks identical to interp_S. If I change the != to a ==, then the da.where() command replaces every values with the mean value. What should I try next?

EDIT: a little more info

• I am using xarray 0.16.2
• I am using python 3.7.4

Thanks in advance for any help/guidance you all can provide.

[keewis]

I think the issue is that np.nan == np.nan returns False. This will work:

da.where(lambda arr: arr.notnull(), df.mean_S)

However, you might also want to try .fillna:

In [10]: arr = xr.DataArray(
    ...:     [[nan, 0, 5, 3], [7, 5, 3, nan], [6, 9, 1, 4], [8, 7, 3, nan], [nan, 8, 1, 4]],
    ...:     dims=("x", "z"),
    ...:     coords={
    ...:         "x": [-1, 0, 1, 2, 3],
    ...:         "z": [10, 11, 12, 13],
    ...:     }
    ...: )
    ...: arr
Out[10]: 
<xarray.DataArray (x: 5, z: 4)>
array([[nan,  0.,  5.,  3.],
       [ 7.,  5.,  3., nan],
       [ 6.,  9.,  1.,  4.],
       [ 8.,  7.,  3., nan],
       [nan,  8.,  1.,  4.]])
Coordinates:
  * x        (x) int64 -1 0 1 2 3
  * z        (z) int64 10 11 12 13

In [11]: values = xr.DataArray(np.arange(88, 104), dims="z", coords={"z": range(16)})
    ...: values
Out[11]: 
<xarray.DataArray (z: 16)>
array([ 88,  89,  90,  91,  92,  93,  94,  95,  96,  97,  98,  99, 100,
       101, 102, 103])
Coordinates:
  * z        (z) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15

In [12]: arr.fillna(values)
Out[12]: 
<xarray.DataArray (x: 5, z: 4)>
array([[ 98.,   0.,   5.,   3.],
       [  7.,   5.,   3., 101.],
       [  6.,   9.,   1.,   4.],
       [  8.,   7.,   3., 101.],
       [ 98.,   8.,   1.,   4.]])
Coordinates:
  * z        (z) int64 10 11 12 13
  * x        (x) int64 -1 0 1 2 3

[DaniJonesOcean]

Thanks! Both of these approaches worked. Here's the code for my particular problem.

First solution:

# replace remaining missing values with mean profile : salt
da = df.interp_S
clean_S = da.where(lambda da: da.notnull(), df.mean_S)
df = df.assign({'clean_S':clean_S})

Second solution:

# replace remaining missing values with mean profile : salt
da = df.interp_S
clean_S = da.fillna(df.mean_S)
df = df.assign({'clean_S':clean_S})

I appreciate the quick reply. : )