KAKEHASHI Tech Blog

カケハシのEngineer Teamによるブログです。

TensorFlow Probabilityを使ってドメイン知識を組み込んだ独自の確率分布を定義してみた

カケハシでデータサイエンティストとして働いている蓑田です。

こちらの記事は カケハシ Advent Calendar 2023 の8日目の記事になります。 今回はTensorFlow Probability(以降、TFPと呼ぶ)を使って独自の確率分布を定義するための方法について書いていこうと思います。

みなさんさまざまな領域でデータを活用されていると思いますが、得られたデータの背後には色々な事象が存在しています。たとえばカケハシでは薬局での処方データや採用している医薬品のデータが得られていますが、同じ効用でも複数のメーカーの医薬品を扱っていたり1、毎年メーカーが切り替わっていたりします。医薬品を受け取る側の立場だったときには考えたこともなかったのですが、その背後にはメーカー指定の患者さんが存在したり、メーカー毎に医薬品の価格が異なっているといった事象が想定されます。

このように得られたデータを再現するために背景の事象をドメイン知識を活用して確率的にモデリングすることで、背景事象の深い理解に繋げられます。 そこで今回はTFPを活用して、ドメイン知識に基づく仮説から独自の確率分布を定義し、その分布からサンプリングデータを得るための方法について書きます。

TFPとは

TFPとはTensorFlow 1.12から公式サポートされている、ベイズ統計をはじめとした確率的モデリングのためのパッケージです。従来ではこれらのモデリングにはStanが使用されることが多かったですが、TFPではTensorFlowライクに確率分布を定義でき、MCMCなどの実行も簡単にできます。PythonやTensorFlowに慣れているが、Stanにはあまり馴染みのない方にとっては、とても便利なパッケージだと思います。

独自の確率分布定義

今回は薬局が利用している医薬品のメーカーを切り替える事象に着目します。薬局では4月の薬価改定に伴って価格が安いメーカーの医薬品に切り替えることが珍しくありません。もちろんメーカーを切り替える原因は価格だけではなく、供給不安の影響やメーカー指定の患者さんが増減することなど他にもあります。本ブログでは分かりやすくメーカー切り替えと医薬品の価格の違いについてのみに着目し、価格が安いメーカーが採用されやすいという仮説を置きます。この仮説を置いた上で薬局が各医薬品を購入するか否かの独自の確率分布を定義します。

まず、独自の確率分布をクラスとして下記のように定義します。

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.internal import reparameterization
tfd = tfp.distributions

class PurchaseDecisionForChains:  # カスタム確率分布関数
    def __init__(self, now_price, new_price, demand, purchase_judge_param, cost_merit, chain_num, product_num):
        self.now_price = now_price
        self.new_price = new_price
        self.demand = demand
        self.purchase_judge_param = purchase_judge_param
        self.cost_merit = cost_merit
        self.chain_num = chain_num
        self.product_num = product_num
    
        self.change_merit = tf.reduce_sum(
            (self.now_price - tf.minimum(self.new_price, self.now_price)) * self.demand,
            axis=0)
        self.now_sales = tf.reduce_sum(self.now_price * self.demand, axis=0)
        self.relative_change_merit = self.change_merit / self.now_sales
        self.probs = tf.clip_by_value(self.relative_change_merit, 0.0, 1.0) ** self.purchase_judge_param

        # 以降確率分布で定義すべき変数
        self.batch_shape = [0, self.chain_num, self.product_num] # 薬局数×医薬品数
        self.event_shape = [self.chain_num, self.product_num]
        self.dtype = tf.float64
        self.experimental_shard_axis_names = [] 
        self.parameters = {  
            'now_price': now_price, 
            'new_price': new_price, 
            'demand': demand,
            'purchase_judge_param': purchase_judge_param,
            'cost_merit': cost_merit
        }    
        self.reparameterization_type = reparameterization.NOT_REPARAMETERIZED

まず必要なパッケージをインポートします。now_priceが現在採用しているメーカーの価格で、new_priceが切り替え候補となるメーカーの価格です。これらの差額に需要量を掛け合わせたものを切り替えメリットの値として定義します。この値と現在の売上との比が相対的な切り替えメリット値になります。最後に購買判断パラメーターpurchage_judge_paramを準備します。これは価格だけで左右されないことを再現するために用意した変数です。実際には安くなったとしても何らかの理由で切り替えができないケースも存在します。この値と相対的な切り替えメリット値を掛け合わせて切り替え確率値を算出します。以降ではバッチサイズやパラメーターなど、TFPに対してこのクラスを確率分布として認識させるために必要な変数を定義しています。

def sample(self, sample_shape=(), seed=0, name='sample', **kwargs):
    assert sample_shape == ()  
    probs = self.probs
    # X%より大きいメリットであれば切り替えるとする。
    sample = tf.cast(tf.where(probs > self.cost_merit, 1, 0), self.dtype)
    res = tf.broadcast_to(tf.reshape(sample, (1, self.chain_num)), (self.product_num, self.chain_num))
    return res

次に、独自の確率分布のクラスにsampleメソッドを実装します。こちらのメソッドには確率分布から出力されるデータを定義します。今回のPurchaseDecisionForChains分布は薬局がある医薬品を買うか買わないかの2値を出力する確率分布のため、上記のように定義しました。cost_merit変数はこの値以上確率値(相対的な切り替えメリット)があれば、買うものとする閾値であり、この値が高い場合、メーカーを切り替えにくい薬局だと定義できます。

sampleメソッドの各引数はすでに組み込まれている確率分布の関数に合わせています。

これでPurchaseDecisionForChainsという価格によって購買判断を出力する確率分布を定義することができました。

シミュレーション

前項で独自の確率分布を定義することができたので、その確率分布を使ってシミュレーションクラスを定義します。

class SimModel(tf.Module):
    def __call__(self, features):
        model = tfd.JointDistributionNamed(
            {
                'new_price': tfd.Deterministic(features['new_price_tensor']),
                'demand': tfd.Deterministic(features['demand']),
                'purchase_judge_param': tfd.Uniform(
                    features['purchase_judge_min'],
                    features['purchase_judge_max']
                ),
                'cost_merit': tfd.Uniform(features['cost_merit_min'], features['cost_merit_max']),
                'purchase_decision': lambda new_price, demand, purchase_judge_param, cost_merit: 
                    PurchaseDecisionForChains(
                        features['now_price_tensor'],
                        new_price,
                        demand,
                        purchase_judge_param,
                        cost_merit,
                        features['chain_num'],
                        features['product_num']
                    )
            },
            batch_ndims=0,
            use_vectorized_map=True
        )
        model._to_track = self
        return model

まずシミュレーションクラスを作成するには、tf.Moduleをサブクラスとして設定する必要があります。tf.Moduleのcallメソッドを継承しここに確率分布の関係性を書いていきます。今回はPurchaseDecisionForChains分布に必要な4つのパラメーターも事前分布として設定しています。たとえば購買判断パラメーターは実際に観測できないので最小値〜最大値までの一様分布として事前分布を設定しました。このように階層的な関係を表す場合、JointDistributionNamed関数が便利です。この中で各パラメーターの名前と分布を辞書形式で設定することで、階層的な分布の表現が簡単にできます。TFPの便利な機能の一部だと思います。

ここまで完成したら後は必要なサンプリング数を指定して、下記のようにすればサンプリングデータを得ることができます。

my_model = SimModel()
loop_num = 100 # サンプリング数
samples = []
for i in range(loop):
    model_instance=my_model(
        features={
            'purchase_judge_min': 0.1,
            'purchase_judge_max': 0.8,
            'now_price_tensor': now_price_tensor, # 各医薬品の現在の価格をtensorで入れたもの
            'new_price_tensor': new_price_tensor, # 各医薬品の新しい価格をtensorで入れたもの
            'demand': demand,
            'cost_merit_min': 0.2,
            'cost_merit_max': 0.5,
            'chain_num': 10,
            'product_num': 100,
        }
    )
    sample_data = model_instance.sample()
    samples.append(sample_data)

最後に

本来はここから実際のメーカー切り替えのデータを入手して、このデータをより再現するような各パラメーターを見つけるために、変分ベイズ法やMCMCなどを用いる必要があります。これによって得られたパラメーターを解釈することで、メーカー切り替えと医薬品の価格の関係を深く考察できるようになります。 本ブログでは長くなってしまうので、独自確率分布の定義とそこから得られるサンプリングデータまでとしますが、次回機会があれば得られたパラメーターの最適値から考察まで書いていこうと思います。

参考


  1. 特許が切れた医薬品はジェネリック医薬品としてさまざまなメーカーが製造・販売する事ができます。