The ordered logistic distribution

import tensorflow as tf
import tensorflow_probability as tfp

print(f"tf: {tf.__version__}")
print(f"tfp: {tfp.__version__}")

import pandas as pd
import numpy as np

data_list = []
for x in range(10, 19):
    season = f"{x}{x+1}"
    url = f"http://www.football-data.co.uk/mmz4281/{season}/E0.csv"
    data_list.append(pd.read_csv(url))

soccer_data = (
    pd.concat(data_list, axis=0)
    .rename(
        columns={
            "HomeTeam": "home_team_name",
            "AwayTeam": "away_team_name",
            "FTHG": "home_goals",
            "FTAG": "away_goals",
        }
    )
    .filter(regex="^home_|^away_")
    .dropna()
    .astype({"home_goals": np.int32, "away_goals": np.int32})
    .assign(
        clipped_score_difference=lambda x: np.clip(
            x["home_goals"] - x["away_goals"], -3, 3
        )
    )
)

team_names = np.unique(
    pd.concat([soccer_data["home_team_name"], soccer_data["away_team_name"]])
)


def team_code(team_name):
    def fn(df):
        codes = pd.Categorical(df[team_name], categories=team_names).codes
        return codes.astype(np.int32)

    return fn


soccer_data = soccer_data.assign(
    home_team=team_code("home_team_name"), away_team=team_code("away_team_name"),
)

print(soccer_data.head())

num_outcomes = 7
num_teams = len(team_names)
home_team = soccer_data["home_team"].to_numpy()
away_team = soccer_data["away_team"].to_numpy()
clipped_score_difference = soccer_data["clipped_score_difference"].to_numpy()

tfb = tfp.bijectors
tfd = tfp.distributions

Root = tfd.JointDistributionCoroutine.Root


def logit(x):
    return -tf.math.log(tf.math.reciprocal(x) - 1.0)


def make_cutpoints(simplex):
    return logit(tf.math.cumsum(simplex[..., tf.newaxis, :-1], axis=-1))


def model():
    cutpoints_simplex = yield Root(tfd.Dirichlet([2.0, 4.0, 6.0, 8.0, 6.0, 4.0, 2.0]))
    team_scale = yield Root(tfd.HalfNormal(1.0))
    unit_team_strength = yield Root(tfd.MultivariateNormalDiag(loc=tf.zeros(num_teams)))

    team_strength = team_scale[..., tf.newaxis] * unit_team_strength
    home_strength = tf.gather(team_strength, home_team, axis=-1)
    away_strength = tf.gather(team_strength, away_team, axis=-1)

    clipped_score_difference = yield tfd.Independent(
        tfd.OrderedLogistic(
            cutpoints=make_cutpoints(cutpoints_simplex),
            loc=home_strength - away_strength,
        ),
        reinterpreted_batch_ndims=1,
    )


joint_dist = tfd.JointDistributionCoroutine(model)

import matplotlib.pyplot as plt

plt.rcParams["figure.figsize"] = (12, 8)

_, _, _, prior_predictive_samples = joint_dist.sample(1_000)

sampled_score_difference_counts = [
    np.bincount(sample, minlength=7) for sample in prior_predictive_samples
]
observed_score_difference_counts = np.bincount(clipped_score_difference + 3)

sampled_mean = np.mean(sampled_score_difference_counts, axis=0)
sampled_quantiles = np.quantile(
    sampled_score_difference_counts, q=[0.025, 0.975], axis=0
)

low_error = sampled_mean - sampled_quantiles[0]
high_error = sampled_quantiles[1] - sampled_mean

plt.bar(
    x=np.arange(-3, 4),
    height=observed_score_difference_counts,
    label="observed data",
    color="blue",
)

plt.errorbar(
    x=np.arange(-3, 4),
    y=sampled_mean,
    yerr=[low_error, high_error],
    fmt="o",
    label="simulated data",
    color="red",
    capsize=5,
)

plt.legend(loc="upper left")
plt.xlabel("clipped score difference")
plt.ylabel("frequency")

import time as tm

num_chains = 10
num_burnin_steps = 1_000
num_results = 500


def target_log_prob_fn(*state):
    return joint_dist.log_prob(list(state) + [clipped_score_difference + 3])


def trace_fn(states, pkr):
    return (
        pkr.inner_results.inner_results.target_log_prob,
        pkr.inner_results.inner_results.leapfrogs_taken,
        pkr.inner_results.inner_results.has_divergence,
        pkr.inner_results.inner_results.energy,
        pkr.inner_results.inner_results.log_accept_ratio,
    )


def step_size_setter_fn(pkr, new_step_size):
    return pkr._replace(
        inner_results=pkr.inner_results._replace(step_size=new_step_size)
    )


def step_size_getter_fn(pkr):
    return pkr.inner_results.step_size


def log_accept_prob_getter_fn(pkr):
    return pkr.inner_results.log_accept_ratio


initial_state = list(joint_dist.sample(num_chains)[:-1])
initial_step_size = [0.25] * len(initial_state)

nuts = tfp.mcmc.NoUTurnSampler(target_log_prob_fn, step_size=initial_step_size)

transformed_nuts = tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=nuts, bijector=[tfb.SoftmaxCentered(), tfb.Softplus(), tfb.Identity()],
)

adaptive_transformed_nuts = tfp.mcmc.DualAveragingStepSizeAdaptation(
    inner_kernel=transformed_nuts,
    num_adaptation_steps=int(0.8 * num_burnin_steps),
    target_accept_prob=0.8,
    step_size_setter_fn=step_size_setter_fn,
    step_size_getter_fn=step_size_getter_fn,
    log_accept_prob_getter_fn=log_accept_prob_getter_fn,
)


@tf.function(autograph=False, experimental_compile=False)
def run_mcmc():
    return tfp.mcmc.sample_chain(
        num_results=num_results,
        current_state=initial_state,
        num_burnin_steps=num_burnin_steps,
        kernel=adaptive_transformed_nuts,
        trace_fn=trace_fn,
    )

start_tfp = tm.time()
samples, sample_stats = run_mcmc()
end_tfp = tm.time()

print(f"TFP took {end_tfp - start_tfp:.2f} seconds")

import arviz as az
import numpy as np


sample_names = ["cutpoints_simplex", "team_scale", "unit_team_strength"]
summary_vars = ["mean", "sd", "hpd_3%", "hpd_97%", "ess_bulk", "r_hat"]

az_samples = {
    name: np.swapaxes(sample, 0, 1) for name, sample in zip(sample_names, samples)
}

sample_stats_names = [
    "lp",
    "tree_size",
    "diverging",
    "energy",
    "mean_tree_accept",
]

az_sample_stats = {
    name: np.swapaxes(stat, 0, 1)
    for name, stat in zip(sample_stats_names, sample_stats)
}

tfp_fit = az.from_dict(
    az_samples,
    sample_stats=az_sample_stats,
    coords={"teams": team_names},
    dims={"unit_team_strength": ["teams"]},
)

print(az.summary(tfp_fit).filter(items=summary_vars))

az.plot_forest(tfp_fit, var_names="unit_team_strength", combined=True)

data {
    int<lower = 1> n;
    int<lower = 2> num_teams;
    int<lower = 2> num_outcomes;
    int<lower = 1, upper = num_teams> home_team[n];
    int<lower = 1, upper = num_teams> away_team[n];
    int<lower = 1, upper = num_outcomes> clipped_score_difference[n];
}

parameters {
    simplex[num_outcomes] cutpoints_simplex;
    vector[num_teams] unit_team_strength;
    real<lower = 0.0> team_scale;
}

model {
    vector[num_outcomes - 1] cutpoints = logit(
        cumulative_sum(cutpoints_simplex[1:(num_outcomes - 1)]));
    vector[num_teams] team_strength = team_scale * unit_team_strength;
    vector[n] location = team_strength[home_team] - team_strength[away_team];

    cutpoints_simplex ~ dirichlet([2.0, 4.0, 6.0, 8.0, 6.0, 4.0, 2.0]');
    unit_team_strength ~ std_normal();
    team_scale ~ std_normal();

    clipped_score_difference ~ ordered_logistic(location, cutpoints);
}

import pystan
import arviz as az

go_faster_flags = ["-O3", "-march=native", "-ffast-math"]

stan_model = pystan.StanModel(
    file=stan_file, model_name="ordered_model", extra_compile_args=go_faster_flags
)

stan_data = {
    "n": len(soccer_data),
    "num_teams": len(team_names),
    "num_outcomes": 7,
    "home_team": soccer_data["home_team"] + 1,
    "away_team": soccer_data["away_team"] + 1,
    "clipped_score_difference": soccer_data["clipped_score_difference"] + 4,
}

start_stan = tm.time()
stan_fit = az.from_pystan(
    stan_model.sampling(data=stan_data, chains=5),
    coords={"team": team_names},
    dims={"unit_team_strength": ["team_names"]},
)
end_stan = tm.time()

print(f"stan took {end_stan - start_stan:.2f} seconds")

print(az.summary(stan_fit).filter(items=summary_vars))

az.plot_forest(stan_fit, var_names="unit_team_strength", combined=True)

cutpoints_simplex, team_scale, unit_team_strength = samples

team_strength = team_scale[..., tf.newaxis] * unit_team_strength
home_strength = tf.gather(team_strength, home_team, axis=-1)
away_strength = tf.gather(team_strength, away_team, axis=-1)

posterior_predictive_dist = tfd.OrderedLogistic(
    cutpoints=make_cutpoints(cutpoints_simplex),
    loc=home_strength - away_strength,
)

with tf.device("/CPU:0"):
    posterior_predictive_samples = posterior_predictive_dist.sample()

sampled_score_difference_counts = [
    np.bincount(sample, minlength=7)
    for sample in tf.reshape(posterior_predictive_samples, [-1, len(soccer_data)])
]

sampled_mean = np.mean(sampled_score_difference_counts, axis=0)
sampled_quantiles = np.quantile(
    sampled_score_difference_counts, q=[0.025, 0.975], axis=0
)

low_error = sampled_mean - sampled_quantiles[0]
high_error = sampled_quantiles[1] - sampled_mean

plt.bar(
    x=np.arange(-3, 4),
    height=observed_score_difference_counts,
    label="observed data",
    color="blue",
)

plt.errorbar(
    x=np.arange(-3, 4),
    y=sampled_mean,
    yerr=[low_error, high_error],
    fmt="o",
    label="simulated data",
    color="red",
    capsize=5,
)

plt.legend(loc="upper left")
plt.xlabel("clipped score difference")
plt.ylabel("frequency")