傾向スコアマッチングの非復元抽出

概要

因果推論における傾向スコアマッチング(非復元抽出)を実装します。
コードの核となる部分は、最近傍法による全データ間の距離をとった後、距離の昇順で逐次ペアマッチングをし、一度マッチさせたデータは探索先から外すようにしています。

コード

ライブラリとサンプルデータ

import pandas as pd
import numpy as np
from scipy import stats
from sklearn.neighbors import NearestNeighbors

#　標準偏差を1、データサイズを50に
std = 1
n = 50

#　探索元サンプルデータ
target = pd.DataFrame(
    data=stats.norm.rvs(loc=0, scale=std, size=n, random_state=0),
    index=range(0, n),
    columns=['col']
)
#　探索先サンプルデータ
population = pd.DataFrame(
    data=stats.norm.rvs(loc=0, scale=std, size=n, random_state=1),
    index=range(0, n),
    columns=['col']
)

ペアマッチング

def matching_data(target, population, caliper, n_neighbors=1):
    """
        最近傍法により近似した値を探索し、見つかったらそのindexと距離をDataFrame化して返す

        Parameters
        ----------
        target : pd.Series
            検索対象データ（インデックス,検索値）
        population : pd.Series
            検索先データ（インデックス,検索値）
        caliper : float
            マッチング時に許容する誤差範囲

        Returns
        -------
        result_arr : pd.DataFrame
            target_index : int
                検索対象インデックス番号
            population_index : int
                検索先インデックス番号
            distance : float
                ヒットしたデータ間の距離
            distance_rank : int
                distanceの昇順順位
    """

    # 最近傍探索
    neigh = NearestNeighbors(n_neighbors=n_neighbors, metric='euclidean')
    neigh.fit(population.to_numpy().reshape(1, -1).T)

    #　ユークリッド距離と探索先のインデックス
    distances, indexes = neigh.kneighbors(target.to_numpy().reshape(1, -1).T, 
return_distance=True)
    distances, indexes = distances.reshape(-1), indexes.reshape(-1)

    # ユークリッド距離の昇順順位を取得
    distance_ranks = np.argsort(np.argsort(distances))

    # 結合する配列を初期化
    result_arr = np.empty([0, 4])

    # ユークリッド距離の小さい順から紐づけ
    # 非復元抽出のため一度紐づけたdistances,indexesはskipする
    # 各配列(distance_ranks,distances,indexes)の添え字は同一であるため、distance_ranksの添え字に基づいて処理
    
    for i in range(len(distance_ranks)):
        idx = np.where(i==distance_ranks)[0].item()

        # result_arrの1列目(population_idx)に値が格納されていない場合,かつcaliper以内のdistancesに限定
        if (~np.isin(result_arr[:, 1], indexes[idx]).any()) and (distances[idx].item()<=caliper):
            target_idx = idx
            population_idx = indexes[idx].item()
            distance = distances[idx].item()
            distance_rank = distance_ranks[idx].item()
            new_row = np.array([target_idx, population_idx, distance, distance_rank])

            result_arr = np.vstack([result_arr, new_row])

        else: pass

    return pd.DataFrame(
        data=result_arr,
        columns=['target_index', 'population_index', 'distance', 'distance_rank']
    ).astype({'target_index': int, 'population_index': int, 'distance': float, 'distance_rank': int})

mapping_df = matching_data(target['col'], population['col'], std*0.2) #　キャリパーは標準偏差の0.2倍

マッチング結果

result_df = pd.merge(mapping_df, target, how='left', left_on='target_index', right_index=True)
result_df = pd.merge(result_df, population, how='left', left_on='population_index', right_index=True)
result_df.rename(columns={'col_x': 'target_col', 'col_y': 'population_col'}, inplace=True)
result_df

	target_index	population_index	distance	distance_rank	target_col	population_col
0	32	41	0.000157	0	-0.887786	-0.887629
1	22	4	0.000971	1	0.864436	0.865408
2	13	48	0.001516	2	0.121675	0.120159
3	38	13	0.003272	3	-0.387327	-0.384054
4	26	18	0.003545	4	0.045759	0.042214
5	27	40	0.004652	5	-0.187184	-0.191836
6	23	42	0.004993	6	-0.742165	-0.747158
7	18	8	0.005971	7	0.313068	0.319039
8	29	10	0.007251	8	1.469359	1.462108
9	19	33	0.008890	10	-0.854096	-0.845206
10	44	2	0.018520	13	-0.509652	-0.528172
11	12	39	0.018994	14	0.761038	0.742044
12	0	6	0.019241	15	1.764052	1.744812
13	8	26	0.019671	16	-0.103219	-0.122890
14	39	12	0.020114	17	-0.302303	-0.322417
15	7	16	0.021071	19	-0.151357	-0.172428
16	40	3	0.024416	21	-1.048553	-1.072969
17	35	46	0.034567	24	0.156349	0.190915
18	21	49	0.036415	27	0.653619	0.617203
19	45	31	0.041321	28	-0.438074	-0.396754
20	5	27	0.041508	29	-0.977278	-0.935769
21	6	22	0.048498	30	0.950088	0.901591
22	37	21	0.057656	31	1.202380	1.144724
23	14	23	0.058631	32	0.443863	0.502494
24	33	11	0.079344	36	-1.980796	-2.060141
25	46	36	0.135485	41	-1.252795	-1.117310
26	3	47	0.140638	42	2.240893	2.100255