如何对齐两个值相同但其中一个含NaN的Python数组，解决statsmodels索引丢失问题

我有两个数值序列的有效取值完全一致，其中一个穿插了NaN值（实际场景是带索引的pandas DataFrame，用数组举例效果一致）。

举例如下：

• 带NaN的序列：a=[1,2,3,np.nan,4,5,np.nan,6,np.nan,7,np.nan,8,np.nan,9,10]
• 无NaN的序列：b=[1,2,3,4,5,6,7,8,9,10]

需要找到简便的向量化操作方法，调整b的结构让它和a完全对齐：即a对应为NaN的位置，b也插入NaN。不想要循环遍历的实现，需要更优雅的方案。

需求背景补充 #

该需求是为了解决statsmodels.api的配套使用问题：

• .get_prediction()返回置信区间数据和估计值时会自动剔除NaN值，丢失原始索引
• .predict()返回的估计值会保留包含NaN的原始索引
  需要将两个结果合并，得到每个原始索引位置对应的估计值置信区间。

如果有开发者在查找statsmodels中.get_prediction方法自动移除NaN的相关问题，可直接使用下文提供的适配代码，目前网上相关解决资料较少。

基础Numpy实现（通用数组对齐） #

def nan_merge(a,b):
    indices = np.where(np.isnan(a))[0]
    ans = np.insert(arr=b, obj=indices-list(range(len(indices))), values=np.nan)
    return ans

def test_nan_merge():
    a = np.array([1,2,3,np.nan,4,5,np.nan,6,np.nan,7,np.nan,8,np.nan,9,10])        
    b = np.array([1,2,3,4,5,6,7,8,9,10]).astype(float)
    assert (nan_merge(a,b)==a).any()

Pandas适配封装（针对statsmodels索引修复场景） #

该实现相比普通pandas merge方案可以处理重复数值的场景，是普通merge方案无法覆盖的：

def get_true_indexes(data,model):
    '''
    Predict returns all indexes
    Get prediction removes nans so destroys indexing.
    But provides Low and high CI.
    So lets merge them. 
    '''
    
    def nan_merge(a,b):        
        indices = np.where(np.isnan(a))[0]
        ans = np.insert(arr=b, obj=indices-list(range(len(indices))), values=np.nan)
        return ans
    
    def test_nan_merge():
         
        a=np.array([1,2,3,np.nan,4,5,np.nan,6,np.nan,7,np.nan,8,np.nan,9,10])
        b=np.array([1,2,3,4,5,6,7,8,9,10]).astype(float)
        assert (nan_merge(a,b)==a).any()
        
    
    df1=model.predict(data).copy(deep=True)
    df1.name='predict'
    df1=df1.to_frame()
    df1=df1.reset_index()
    
    df2= model.get_prediction(data).summary_frame()
   
    real=df1['predict'].values
    mea=df2['mean'].values
    upp=df2['obs_ci_upper'].values
    low=df2['obs_ci_lower'].values
    
    if len(real)!=len(mea):
        mea = nan_merge(real,mea)
        upp= nan_merge(real,upp)
        low= nan_merge(real,low)
    
    ans=pd.DataFrame({'mean':mea,'obs_ci_upper':upp,'obs_ci_lower':low},index=df1['index'].to_list())
               
    return ans

内容的提问来源于stack exchange，提问作者Ocean Scientist