長さの異なる多重リストをnumpy.arrayライクに扱う

結論

疎行列を利用する．scipy.sparseは2d arrayまでしか対応していないので，今回は3d以上のarrayにも対応しているsparseを利用する．

pydata/sparse: Sparse multi-dimensional arrays for the PyData ecosystem

多重リストからsparse.COOを生成する関数を作成する．

import sparse

def from_list(l, fill_value=None):
    # sparseにはリストからsparse.COOを生成する関数がないので自分で書く

    def coords_tr_builder(l):
        # sparse.COOの座標の転置を生成する関数
        if isinstance(l, list):
            return [[ix] + coords for ix, subl in enumerate(l) for coords in coords_tr_builder(subl)]
        else:
            return [[]]

    def data_builder(l):
        # sparse.COOのデータを生成する関数
        if isinstance(l, list):
            return [data for subl in l for data in data_builder(subl)]
        else:
            return [l]

    return sparse.COO(
        coords=list(zip(*coords_tr_builder(l))),
        data=data_builder(l),
        fill_value=fill_value
    )

Operations on COO arrays — sparse 0.5.0+16.ge80efc7.dirty documentationによると，scipy.sparseやnumpy.arrayと同じ感覚で利用できる．
長さの異なる多重リストを表現するという目的上，fill_valueはデフォルト値0.ではなくnp.nanとするとよい．

### 簡単な例 ###

A = from_list([[1., 2.], [3., 4., 5.], [6.]])
A # <COO: shape=(3, 3), dtype=float64, nnz=6, fill_value=0.0>

# fill_value=0.だと全体がずれてfill_value=1.になってしまう
A + 1. # <COO: shape=(3, 3), dtype=float64, nnz=6, fill_value=1.0>

# fill_value=np.nanだと意図した挙動になる
import numpy as np
A = from_list([[1., 2.], [3., 4., 5.], [6.]], fill_value=np.nan)
A # <COO: shape=(3, 3), dtype=float64, nnz=6, fill_value=nan>

A + 1. # <COO: shape=(3, 3), dtype=float64, nnz=6, fill_value=nan>

### 複雑な例 ###

df = pd.read_csv('https://raw.githubusercontent.com/mwaskom/seaborn-data/master/titanic.csv')
df['survived'].value_counts()
# 0    549
# 1    342
# Name: survived, dtype: int64

A = from_list(df['fare'].groupby(df['survived']).agg(lambda srs: srs.tolist()).tolist(), fill_value=np.nan)
A # <COO: shape=(2, 549), dtype=float64, nnz=891, fill_value=nan>

A - np.nanmean(A, axis=1, keepdims=True) # <COO: shape=(2, 549), dtype=float64, nnz=891, fill_value=nan>

動機（雑説明）

Beta-Binomial compositionにおけるBeta分布のパラメータを経験Bayesで決定するにあたり，Bayes推定とモーメント法を反復して雑にハイパーパラメータを探索する．
K個のBeta分布からそれぞれL_k個のBinomial分布が生成され，それぞれのBinomial分布から観測値が生成される．

beta_binomial_pd.py

import pandas as pd
import numpy as np
import scipy.stats as st

# データ作成
num_total = 100
alpha, beta = np.asarray([2., 3., 4.]), np.asarray([4., 3., 2.])
Nmin, Nmax = 10, 1000
gix = np.random.choice(range(len(alpha)), size=num_total)
p = st.beta.rvs(alpha[gix], beta[gix])
N = np.random.randint(Nmin, Nmax, size=num_total)
n = np.random.binomial(N, p)
df = pd.DataFrame({'gix': gix, 'N': N, 'n': n})

# (μ_k, M_k) for k=1,...,K
# → (ベータ分布) p_kl for l=1,...,L_k
#   → (二項分布) n_kl positive from N_kl trial

# 二項分布の初期パラメータは最尤推定量／ベータ分布の初期パラメータは一様分布
p = df['n'] / df['N']
mu, M = np.full_like(p, 0.5), np.full_like(p, 2.)

# 反復
maxiter = 10
for _ in range(maxiter):
    p = (df['n'] + M * mu) / (df['N'] + M)
    mean, std = p.groupby(df['gix']).transform('mean'), p.groupby(df['gix']).transform('std')
    mu, M = mean, mean * (1. - mean) / std - 1.
    # 停止条件などは省略

transformでSQLのウィンドウ関数みたいなことをしているが，transformは死ぬほど遅い．
gixによるグルーピングは反復を通して変わらないので，グルーピングしたまま処理したい．
そこでnumpy.arrayを利用する．

beta_binomial_np.py

# 横持ちのnとNを作る．str.split(expand=True)が使いたかったので一度文字列にしている
DELIM = ':'
df_tmp = df[['N', 'n']].groupby(df['gix']).agg(lambda x: DELIM.join(x.map(str)))
n = df_tmp['n'].str.split(DELIM, expand=True).fillna(np.nan).astype(np.float64).values
N = df_tmp['N'].str.split(DELIM, expand=True).fillna(np.nan).astype(np.float64).values

# 二項分布の初期パラメータは最尤推定量／ベータ分布の初期パラメータは一様分布
p = n / N
mu, M = np.full(len(p), 0.5)[:, np.newaxis], np.full(len(p), 2.)[:, np.newaxis]

# 反復
maxiter = 10
for _ in range(maxiter):
    p = (n + M * mu) / (N + M)
    mean, std = np.nanmean(p, axis=1, keepdims=True), np.nanstd(p, axis=1, keepdims=True)
    mu, M = mean, mean * (1. - mean) / std - 1.
    # 停止条件などは省略

かなり高速化したが，密行列で保持するためグループの大きさがバラバラのときにメモリが足りなくなることがある．
そこでsparse.COOを利用する．

beta_binomial_sp.py

import sparse

def from_list(l, fill_value=None):
    # sparseにはリストからsparse.COOを生成する関数がないので自分で書く

    def coords_tr_builder(l):
        # sparse.COOの座標の転置を生成する関数
        if isinstance(l, list):
            return [[ix] + coords for ix, subl in enumerate(l) for coords in coords_tr_builder(subl)]
        else:
            return [[]]

    def data_builder(l):
        # sparse.COOのデータを生成する関数
        if isinstance(l, list):
            return [data for subl in l for data in data_builder(subl)]
        else:
            return [l]

    return sparse.COO(
        coords=list(zip(*coords_tr_builder(l))),
        data=data_builder(l),
        fill_value=fill_value
    )

# 横持ちのnとNを作る
df_tmp = df[['N', 'n']].groupby(df['gix']).agg(lambda x: x.astype(np.float64).tolist())
ns = from_list(df_tmp['n'].tolist(), fill_value=np.nan)
Ns = from_list(df_tmp['N'].tolist(), fill_value=np.nan)

# 二項分布の初期パラメータは最尤推定量／ベータ分布の初期パラメータは一様分布
p = ns / Ns
mu, M = np.full(len(p), 0.5)[:, np.newaxis], np.full(len(p), 2.)[:, np.newaxis]

# 反復
maxiter = 10
for _ in range(maxiter):
    p = (n + M * mu) / (N + M)
    mean, std = np.nanmean(p, axis=1, keepdims=True), np.nanstd(p, axis=1, keepdims=True)
    mu, M = mean, mean * (1. - mean) / std - 1.
    # 停止条件などは省略

初めに作成したfrom_listを使って実装することができた．

余談

Beta-Binomial compositionにおけるBeta分布のパラメータの，必ず負にならない点推定量ってどんなのがありますかね？　ご存知の方がいたら教えていただきたいです．