使用数据中的 nans 计算 3 arrays 的加权平均值 python

Question

I have 3, two dimensional arrays that represent geospatial data. Each array shape is (721,1440) , ie, 721 latitude values and 1440 longitude values. I want to compute a weighted mean of these 3 arrays. Normally that is simple and would generally be sum(array*weight)/sum(weights). This works great except in cases where you have nans in the data.

In my specific case arr1 should have a weight of 0.7, arr2 0.2, and arr3 0.1. However, anytime there is a nan, the mean obviously becomes nan. In my case the only data with nans is arr3 .

What I want though is when there is a nan for the weighted mean to only comprise of the first two arrays, which would be (arr1*0.7 + arr2*0.2)/0.9 . I tried using xr.where() to accomplish this but for some reason it goes hog wild on my RAM and crashes my kernel every single time. Are there any other ways to accomplish this task?

Answer 1

You can use np.nansum() and np.isnan() :

import numpy as np

# Dummy example
x = np.ones((5,5))
y = np.ones((5,5))*2
x[0,0] = np.nan

# Stack your array 
stack  = np.stack((x,y))
# Compute the weight for each value:                 
weight = np.apply_along_axis(np.multiply,0,~np.isnan([x,y]),[0.2,0.8])
# Get the result
res    = np.nansum(stack*weight,axis=0)/weight.sum(axis=0)