Zenn
🚀

PyMC：カウントデータに対する回帰

2023/07/09に公開
Python
#
ppl
pymc
tech

はじめに

今回はカウントデータに対する回帰を紹介します。

ライブラリのインポート

import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import arviz as az
from copy import copy
from scipy import stats

import pymc as pm
import pytensor.tensor as pt


# Initialize random number generator
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

デモデータの作成

def get_nb_vals(mu, alpha):
    """Generate negative binomially distributed samples by
    """
    g = stats.gamma.rvs(alpha, scale=mu / alpha)
    return stats.poisson.rvs(g)

alpha = 1
intercept = 1
slope = np.array([1, 1])

n, d = 100, 2
X = np.abs(np.random.randn(n,d))
mu = np.dot(X, slope) + intercept

y = get_nb_vals(mu, alpha)

print(y.mean())
print(y.std())

plt.hist(y, bins=50)
plt.title("Observed data")

回帰モデルの構築

カウントデータの回帰モデルとしてはポアソン回帰がよく使用されますが、ポアソン分布は平均値と分散が等しいという強い制約を持っているので、分散が大きい、小さいようなカウントデータの分布に対して適用するのは適切でありません。

そこで、今回は以下の４つのモデルを適用して結果を確認していきます。それぞれ過分散、過小分散に対応可能な分布ですが、個人的に触った感じだと上から順番に収束性が悪くなるようです。

  • Poisson distribution
  • Negative Binomial distribution：過分散に対応可能な分布.パラメータαが無限大の時にポアソン分布に近似
  • The Generalized Poisson distribution：過分散、過小分散に対応可能.パラーメータλが0より大きい場合は過分散、0より小さい場合は過小分散.0の場合はポアソン分布と等しい
  • Conway-Maxwell Poisson distribution：過分散、過小分散に対応可能.パラメータnu=1の時ポアソン分布に等しい

Poisson distribution

結果を見ると、ポアソン分布では過分散を表現できないためうまくフィットできてないように見えます。

with pm.Model() as poisson_model:
    intercept = pm.Normal("intercept", mu=1.0, sigma=1.0)
    slope = pm.Normal("slope", mu=1.0, sigma=1.0, shape=(X.shape[1]))
    
    mu = pm.Deterministic("mu", pt.exp(intercept + pt.dot(X, slope)))
    
    obs = pm.Poisson("obs", mu=mu, observed=y)

with poisson_model:
    poisson_idata = pm.sample()
    
az.summary(poisson_idata)


with poisson_model:
    pp = pm.sample_posterior_predictive(poisson_idata)
    
az.plot_ppc(pp, num_pp_samples=200);

Negative Binomial distribution

過分散をうまく表現できているように見えます。

with pm.Model() as negative_binomial_model:
    intercept = pm.Normal("intercept", mu=1.0, sigma=1.0)
    slope = pm.Normal("slope", mu=1.0, sigma=1.0, shape=(X.shape[1]))
    alpha = pm.Exponential("alpha", 0.5)
    
    mu = pm.Deterministic("mu", pt.exp(intercept + pt.dot(X, slope)))
    
    obs = pm.NegativeBinomial("obs", mu=mu, alpha=alpha, observed=y)

with negative_binomial_model:
    negative_binomial_idata = pm.sample()
    
az.summary(negative_binomial_idata)


with negative_binomial_model:
    pp = pm.sample_posterior_predictive(negative_binomial_idata)
    
az.plot_ppc(pp, num_pp_samples=200);

The Generalized Poisson distribution

The Generalized Poisson distributionは通常のPyMCのdistributionsに実装されていないので、今回はcustom distributionを自分で作ります。（pymc_experimental.distributions.GeneralizedPoissonがあるようなので、実際に使う際はこちらを使用してください）

Custom Distribution

pm.CustomDist.distを使用して独自の分布を作成します。必要な関数としては、以下のようなlogpを計算する関数とサンプリングする関数です。また、以下のようにlogp関数の中身はpytensorの関数で記述する必要があります。

CustomDistを作る際の詳細は以下を参考にしてください。
https://www.pymc.io/projects/docs/en/stable/api/distributions/generated/pymc.CustomDist.html
https://www.pymc.io/projects/docs/en/stable/contributing/implementing_distribution.html

from pymc.distributions.dist_math import check_parameters, factln, logpow

def _generalized_poisson_logp(value, theta, lam):
    theta_lam_value = theta + lam * value
    log_prob = pt.log(theta) + logpow(theta_lam_value, value - 1) - theta_lam_value - factln(value)
    
    log_prob = pt.switch(
        pt.or_(
            theta_lam_value < 0,
            value < 0,
        ),
        -np.inf,
        log_prob,
    )
    
    return check_parameters(log_prob, value >= 0, theta > 0, pt.abs(lam) <= 1, -theta / 4 <= lam)

def _generalized_poisson_random(theta, lam, rng=None, size=None):
    if size is not None:
        assert size == theta.shape
    else:
        size = theta.shape
        
    #lam = lam[0]
    omega = np.exp(-lam)
    X = np.full(size, 0)
    S = np.exp(-theta)
    P = np.copy(S)
    for i in range(size[0]):
        U = np.random.uniform()
        while U > S[i]:
            X[i] += 1
            C = theta[i] - lam + lam * X[i]
            P[i] = omega * C * (1 + lam / C) ** (X[i] - 1) * P[i] / X[i]
            S[i] += P[i]
    return X

def _generalized_poisson_moment(rv, size, theta, lam):
    return pt.ones(size)

作成した関数がうまくいっているか確認します。thetaが負の時は過小分散、正の時は過分散、0の時は通常のポアソン回帰と同等であることがわかります。

fig, ax = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True, figsize=(10, 8))
plt.setp(ax, xlim=(0, 20))

d = pm.Poisson.dist(mu=np.full(1000, 5))
samples = pm.draw(d)
ax[0][0].hist(samples, bins=np.arange(21), label="poisson", alpha=0.3)
ax[0][0].set_title("poisson")

d = pm.CustomDist.dist(np.full(1000, 5), 0, random=_generalized_poisson_random, logp=_generalized_poisson_logp, dtype="int64")
samples = pm.draw(d)
ax[0][1].hist(samples, bins=np.arange(21), label="lam=5, theta=0", alpha=0.3)
ax[0][1].set_title("lam=5, theta=0")

d = pm.CustomDist.dist(np.full(1000, 5), -0.5, random=_generalized_poisson_random, logp=_generalized_poisson_logp, dtype="int64")
samples = pm.draw(d)
ax[1][0].hist(samples, bins=np.arange(21), label="lam=5, theta=-0.5", alpha=0.3)
ax[1][0].set_title("lam=5, theta=-0.5")

d = pm.CustomDist.dist(np.full(1000, 5), 0.3, random=_generalized_poisson_random, logp=_generalized_poisson_logp, dtype="int64")
samples = pm.draw(d)
ax[1][1].hist(samples, bins=np.arange(21), label="lam=5, theta=0.3", alpha=0.3)
ax[1][1].set_title("lam=5, theta=0.3")

実行

pm.CustomDistのndim_suppはsupport dimensions、ndims_paramsはパラメータのdimensions、dtypeは出力の型を指定しています。

結果を見ると、the generalized poissonでもうまく推定できていることがわかります。

with pm.Model() as generalized_poisson_model:
    intercept = pm.Normal("intercept", mu=1.0, sigma=1.0)
    slope = pm.Normal("slope", mu=1.0, sigma=1.0, shape=(X.shape[1]))
    lam = pm.TruncatedNormal("lam", mu=0, sigma=0.1, lower=-1, upper=1)
    #lam = 0.1
    
    theta = pm.Deterministic("mu", pt.exp(intercept + pt.dot(X, slope)))
    
    obs = pm.CustomDist("obs", 
                        theta, 
                        lam, 
                        random=_generalized_poisson_random, 
                        logp=_generalized_poisson_logp, 
                        moment=_generalized_poisson_moment,
                        ndim_supp=0,
                        ndims_params=[0, 0],
                        dtype="int64",
                        observed=y)
        
    idata = pm.sample(initvals={"intercept": 1, "slope": np.ones((X.shape[1])), "lam": 0.1})
    
az.summary(idata)


with generalized_poisson_model:
    pp = pm.sample_posterior_predictive(idata)
    
az.plot_ppc(pp, num_pp_samples=200);

Conway-Maxwell Poisson Distribution

上記と同様に独自の関数を定義して実行してみます。

Custom Distribution

from scipy.special import gammaln
from pymc.distributions.dist_math import check_parameters, factln, logpow

def _cmpoisson_logp(value, lam, nu):
    pi = pt.constant(np.pi)
        
    alpha = pt.power(lam, 1.0/nu)
    log_Z = nu * alpha - (nu - 1.0) * 0.5 * pt.log(2.0 * pi * alpha) - 0.5 * pt.log(nu)
    log_Z += pt.log(1.0 + (nu**2 - 1.0) / 24.0 * pt.power(nu * alpha, -1.0) + (nu**2 - 1.0) * (nu**2 + 23.0) / 1152.0 * pt.power(nu * alpha, -2.0))
    log_prob = value * pt.log(lam) - nu * factln(value + 1) - log_Z 
    
    log_prob = pt.switch(
        pt.lt(value, 0),
        -np.inf,
        log_prob,
    )
    
    return check_parameters(log_prob, value >= 0, lam > 0, nu > 0)

def _cmpoisson_random(lambda_, nu, rng=None, size=None):
    if size is None:
        if isinstance(lambda_, float):
            size = 1
        else:
            size = lambda_.shape
        
    def logp_numpy(value, lambda_, nu):
        alpha = np.power(lambda_, 1.0/nu)
        log_Z = nu * alpha - (nu - 1.0) * 0.5 * np.log(2.0 * np.pi * alpha) - 0.5 * np.log(nu)
        log_Z += np.log(1.0 + (nu**2 - 1.0) / 24.0 * np.power(nu * alpha, -1.0) + (nu**2 - 1.0) * (nu**2 + 23.0) / 1152.0 * np.power(nu * alpha, -2.0))
        res = value * np.log(lambda_) - nu * gammaln(value + 1) - log_Z      
        return res
    
    u = rng.uniform(0, 1, size=size).ravel()

    values = np.empty(u.shape)
    for i in range(len(u)):
        value = 0
        cdf = np.sum(np.exp(logp_numpy(np.arange(value+1), lambda_, nu)))
        while u[i] > cdf:
            value += 1
            cdf = np.sum(np.exp(logp_numpy(np.arange(value+1), lambda_, nu)))
        values[i] = value
        
    return values.reshape(size)

def _cmpoisson_moment(rv, size, lam, nu):
    return pt.ones(size)

実行

実装が悪いせいかかなり計算が不安定です。収束もしていません。。原因が分かる方がいればコメント等で教えてください。

with pm.Model() as cmpoisson_model:
    intercept = pm.Normal("intercept", mu=0.0, sigma=5.0)
    slope = pm.Normal("slope", mu=0.0, sigma=5.0, shape=(X.shape[1],))
    nu = pm.HalfNormal("nu", 5)
    
    mu = pt.log(pt.exp(intercept + pt.dot(X, slope)) + 1)
    
    obs = pm.CustomDist("obs", 
                        mu, 
                        nu, 
                        logp=_cmpoisson_logp, 
                        random=_cmpoisson_random,
                        moment=_cmpoisson_moment,
                        ndim_supp=0,
                        ndims_params=[0, 0],
                        dtype="int64", 
                        observed=y)
    
    idata = pm.sample(draws=3000, tune=3000, target_accept=0.90, initvals={"intercept": 1, "slope": np.ones((X.shape[1])), "nu": 1})

az.summary(idata)


with generalized_poisson_model:
    pp = pm.sample_posterior_predictive(idata)
    
az.plot_ppc(pp, num_pp_samples=200);

最後に

今回はカウントデータを扱いました。カウントデータに対してはポアソン回帰を適用するケースが多いですが、ポアソン分布は制約が強い分布になるので必要に応じて上記で紹介したようなアプローチを適用してみてください。

Discussion