Transitioning from pm.sample() to pm.fit()

Questions v5
16

Sure, for some context, each record in my dataset is a customer transaction. Each customer will have 1+ records, and each record has a value and a transaction type. I have implemented (or, have tried to implement)

  • a random effect per customer (epsilon)
  • a transaction type effect (tau)

I used beta functions for these to ensure positivity of the gamma likelihood’s parameters. I multiplied those each by 2 to get distributions that somewhat resemble normal distributions centered at 1. Those get multiplied by the alpha parameter. This felt a but hacky but I got it to work–I am relatively new to this stuff and certainly open to suggestions.

The model code is here:

customer_idx, customers = factorize(train, "customer_id")
transaction_idx, transactions = factorize(train, "transaction_id")
transaction_type_idx, transaction_types = factorize(train, "transaction_type")

values = train.select(pl.col("values")).to_numpy().squeeze()

coords = {
    "customers": customers,
    "transactions": transactions,
    "transaction_types": transaction_types
}

with pm.Model(coords=coords) as model:

    # population level params
    pop_mu_1 = pm.Normal("pop_mu_1", mu=8.25, sigma=0.25)
    pop_mu_2 = pm.Normal("pop_mu_2", mu=2.5, sigma=0.25)
    
    # customer level beta
    customer_beta = pm.Normal(
        'customer_beta',
        mu=pop_mu_2,
        sigma=0.2,
        dims="customers"
    )
    
    # Random effect
    alpha_e = pm.Normal("alpha_e", mu=2.5, sigma=0.2)
    beta_e = pm.Normal("beta_e", mu=4.7, sigma=0.3)
    epsilon = pm.Beta("epsilon", alpha=alpha_e, beta=beta_e, dims="customers")
    
    # transaction type level
    # non-centering alpha_t
    m_alpha_t = pm.Normal('m_alpha_t', mu=4, sigma=0.4)
    s_alpha_t = pm.Normal('s_alpha_t', mu=4, sigma=0.4)
    z = pm.Beta('z', alpha=3.5, beta=3.5, dims='transaction_types')
    alpha_t = pm.Deterministic('alpha_t', m_alpha_t + s_alpha_t * z, dims="transaction_types")
    
    beta_t = pm.Normal('beta_t', mu=5.6, sigma=0.25)
    
    tau = pm.Beta(
        'tau',
        alpha=alpha_t,
        beta=beta_t,
        dims=("customers", "transaction_types")
    )  
    
    transaction_alpha = pm.Deterministic(
        'transaction_alpha',
        pop_mu_1 * (2*epsilon[customer_idx]) * (2*tau[customer_idx, transaction_type_idx]),
        dims="transactions"
    )
    
    # Likelihood
    y_gamma = pm.Gamma(
        'y_gamma',
        alpha=transaction_alpha[transaction_idx],
        beta=customer_beta[customer_idx],
        observed=values,
        dims="transactions"
    )
1 Like