Merge pull request #5 from martiningram/add_mcmc_comparison

Add code for MCMC comparison

Author: canyon289

mcmc_comparison/Compare runtimes -- multi-run.ipynb

@@ -0,0 +1,44 @@
+# MCMC benchmarks
+
+![ESS](images/ess_values.png)
+
+This code compares Stan, PyMC, and PyMC + JAX numpyro sampler on a model for
+tennis. It accompanies the blog post available
+[here](https://martiningram.github.io/mcmc-comparison/).
+
+This is a copy of the repository here:
+https://github.com/martiningram/mcmc_runtime_comparison.
+
+### Setup notes
+
+This benchmark uses Jeff Sackmann's tennis data. You can obtain it as follows:
+
+```
+git clone https://github.com/JeffSackmann/tennis_atp.git
+
+# If you want to reproduce the results in the blog post, check out this commit:
+cd tennis_atp && git checkout 89c20f1ef56f69db1b73b5782671ee85203b068a
+```
+
+Requirements that can be installed using pip are listed in
+`requirements.txt`. Please install these first.
+
+Once these are done, here are the steps I followed to setup PyMC v4 with JAX support:
+
+* PyMC v4 installed using the instructions here: https://github.com/pymc-devs/pymc/wiki/Installation-Guide-(Linux)#pymc-v4-installation
+* `blackjax` and `numpyro` were also installed using those instructions.
+
+To run the Stan code, it's best to install `cmdstanpy`. Instructions for
+installing it can be found [here](https://mc-stan.org/cmdstanpy/installation.html).
+
+### How to run
+
+It's easiest to run the benchmarks using the `fit_all.sh` script. Make sure to
+first edit the `target_dir` variable in it and amend it to a directory that
+makes sense for you. All the model runs will be stored in it under
+subdirectories.
+
+Once benchmarks have been run, you can analyse the results and make plots using
+the `Compare_runtimes.ipynb` notebook.
+
+If you run into any problems, please raise an issue!

mcmc_comparison/Compare runtimes.ipynb

@@ -0,0 +1,62 @@
+from sackmann import get_data
+import pymc as pm
+from sklearn.preprocessing import LabelEncoder
+import numpy as np
+
+
+def create_arrays(
+    start_year=1960,
+    data_dir="./tennis_atp",
+    include_qualifying_and_challengers=False,
+    include_futures=False,
+):
+
+    df = get_data(
+        data_dir,
+        include_qualifying_and_challengers=include_qualifying_and_challengers,
+        include_futures=include_futures,
+    )
+
+    rel_df = df[df["tourney_date"].dt.year >= start_year]
+
+    encoder = LabelEncoder()
+
+    encoder.fit(
+        rel_df["winner_name"].values.tolist() + rel_df["loser_name"].values.tolist()
+    )
+
+    winner_ids = encoder.transform(rel_df["winner_name"])
+    loser_ids = encoder.transform(rel_df["loser_name"])
+
+    return {
+        "winner_ids": winner_ids,
+        "loser_ids": loser_ids,
+        "player_encoder": encoder,
+    }
+
+
+def get_pymc_model(start_year=1960, data_dir="./tennis_atp"):
+
+    arrays = create_arrays(start_year=start_year, data_dir=data_dir)
+
+    n_players = len(arrays["player_encoder"].classes_)
+
+    winner_ids = arrays["winner_ids"]
+    loser_ids = arrays["loser_ids"]
+
+    with pm.Model() as model:
+
+        player_sd = pm.HalfNormal("player_sd", sigma=1.0)
+
+        player_skills_raw = pm.Normal(
+            "player_skills_raw", 0.0, sigma=1.0, shape=(n_players,)
+        )
+
+        player_skills = pm.Deterministic("player_skills", player_skills_raw * player_sd)
+        logit_skills = player_skills[winner_ids] - player_skills[loser_ids]
+
+        lik = pm.Bernoulli(
+            "win_lik", logit_p=logit_skills, observed=np.ones(winner_ids.shape[0])
+        )
+
+    return model

mcmc_comparison/README.md

@@ -0,0 +1,24 @@
+# Modify as desired; will be created if it does not exist
+base_target_dir="/media/martin/External Drive/projects/pymc_vs_stan/multi_run/fits"
+n_runs=10
+
+for cur_run in `seq 1 $n_runs`; do
+
+    echo "Running $cur_run"
+
+    random_seed=$cur_run
+    target_dir="$base_target_dir"/"$cur_run"
+
+    for start_year in 2020 2019 2015 2010 2000 1990 1980 1968; do
+	echo "Fitting $start_year"
+	echo "PyMC JAX GPU parallel" && python fit_pymc_numpyro.py $start_year gpu parallel "$target_dir" $random_seed
+	echo "PyMC JAX GPU vectorized" && python fit_pymc_numpyro.py $start_year gpu vectorized "$target_dir" $random_seed
+	echo "PyMC JAX CPU parallel" && python fit_pymc_numpyro.py $start_year cpu parallel "$target_dir" $random_seed
+	echo "PyMC JAX CPU vectorized" && python fit_pymc_numpyro.py $start_year cpu vectorized "$target_dir" $random_seed
+	echo "PyMC BlackJAX CPU" && python fit_pymc_blackjax.py $start_year cpu "$target_dir" $random_seed parallel
+	echo "PyMC BlackJAX GPU" && python fit_pymc_blackjax.py $start_year gpu "$target_dir" $random_seed vectorized
+	echo "PyMC" && python fit_pymc.py $start_year "$target_dir" $random_seed
+	echo "cmdstanpy" && python fit_cmdstanpy.py $start_year "$target_dir" $random_seed
+    done
+
+done

mcmc_comparison/fetch_data.py

@@ -0,0 +1,40 @@
+import sys
+import os
+from fetch_data import create_arrays
+from time import time
+import numpy as np
+import arviz as az
+from cmdstanpy import CmdStanModel
+
+start_year = int(sys.argv[1])
+target_dir = sys.argv[2] + "/cmdstanpy"
+seed = int(sys.argv[3])
+
+os.makedirs(target_dir, exist_ok=True)
+
+arrays = create_arrays(start_year=start_year)
+
+start_time = time()
+
+winner_ids = arrays["winner_ids"]
+loser_ids = arrays["loser_ids"]
+player_encoder = arrays["player_encoder"]
+
+stan_data = {
+    "n_matches": len(winner_ids),
+    "n_players": len(player_encoder.classes_),
+    "winner_ids": winner_ids + 1,
+    "loser_ids": loser_ids + 1,
+}
+
+model = CmdStanModel(stan_file="stan_model_optimised.stan")
+model.compile()
+
+fit = model.sample(data=stan_data, parallel_chains=4, seed=seed)
+
+runtime = time() - start_time
+
+arviz_version = az.from_cmdstanpy(posterior=fit)
+
+az.to_netcdf(arviz_version, os.path.join(target_dir, f"samples_{start_year}.netcdf"))
+print(runtime, file=open(os.path.join(target_dir, f"runtime_{start_year}.txt"), "w"))

mcmc_comparison/fit_all.sh

@@ -0,0 +1,29 @@
+import sys
+import os
+from fetch_data import get_pymc_model
+from time import time
+import pymc as pm
+
+start_year = int(sys.argv[1])
+target_dir = sys.argv[2] + "/pymc"
+seed = int(sys.argv[3])
+
+os.makedirs(target_dir, exist_ok=True)
+
+model = get_pymc_model(start_year=start_year)
+
+start_time = time()
+
+with model:
+    hierarchical_trace = pm.sample(
+        1000,
+        tune=1000,
+        return_inferencedata=True,
+        compute_convergence_checks=False,
+        random_seed=seed,
+    )
+
+runtime = time() - start_time
+
+hierarchical_trace.to_netcdf(os.path.join(target_dir, f"samples_{start_year}.netcdf"))
+print(runtime, file=open(os.path.join(target_dir, f"runtime_{start_year}.txt"), "w"))

mcmc_comparison/fit_cmdstanpy.py

@@ -0,0 +1,37 @@
+import sys
+import os
+import pymc.sampling_jax
+from fetch_data import get_pymc_model
+from time import time
+import pymc as pm
+
+start_year = int(sys.argv[1])
+platform = sys.argv[2]
+base_dir = sys.argv[3]
+seed = int(sys.argv[4])
+chain_method = sys.argv[5]
+
+assert platform in ["cpu", "gpu"]
+
+if platform == "cpu":
+    # Disable GPU
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+target_dir = f"{base_dir}/pymc_blackjax_{platform}_{chain_method}"
+
+os.makedirs(target_dir, exist_ok=True)
+
+model = get_pymc_model(start_year=start_year)
+
+start_time = time()
+
+with model:
+    # No progress bar?
+    hierarchical_trace = pymc.sampling_jax.sample_blackjax_nuts(
+        random_seed=seed, chain_method=chain_method,
+        idata_kwargs={'log_likelihood': False})
+
+runtime = time() - start_time
+
+hierarchical_trace.to_netcdf(os.path.join(target_dir, f"samples_{start_year}.netcdf"))
+print(runtime, file=open(os.path.join(target_dir, f"runtime_{start_year}.txt"), "w"))

mcmc_comparison/fit_pymc.py

@@ -0,0 +1,37 @@
+import sys
+import os
+import pymc.sampling_jax
+from fetch_data import get_pymc_model
+from time import time
+import pymc as pm
+
+start_year = int(sys.argv[1])
+platform = sys.argv[2]
+chain_method = sys.argv[3]
+base_dir = sys.argv[4]
+seed = int(sys.argv[5])
+
+assert platform in ["cpu", "gpu"]
+
+if platform == "cpu":
+    # Disable GPU
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+target_dir = f"{base_dir}/pymc_numpyro_{platform}_{chain_method}"
+
+os.makedirs(target_dir, exist_ok=True)
+
+model = get_pymc_model(start_year=start_year)
+
+start_time = time()
+
+with model:
+    hierarchical_trace = pymc.sampling_jax.sample_numpyro_nuts(
+        chain_method=chain_method, random_seed=seed,
+        idata_kwargs={'log_likelihood': False}
+    )
+
+runtime = time() - start_time
+
+hierarchical_trace.to_netcdf(os.path.join(target_dir, f"samples_{start_year}.netcdf"))
+print(runtime, file=open(os.path.join(target_dir, f"runtime_{start_year}.txt"), "w"))

mcmc_comparison/fit_pymc_blackjax.py

@@ -0,0 +1,39 @@
+import sys
+import os
+from fetch_data import create_arrays
+from time import time
+import stan
+import numpy as np
+import arviz as az
+
+start_year = int(sys.argv[1])
+target_dir = sys.argv[2] + "/stan"
+
+os.makedirs(target_dir, exist_ok=True)
+
+arrays = create_arrays(start_year=start_year)
+
+start_time = time()
+
+winner_ids = arrays["winner_ids"]
+loser_ids = arrays["loser_ids"]
+player_encoder = arrays["player_encoder"]
+
+stan_data = {
+    "n_matches": len(winner_ids),
+    "n_players": len(player_encoder.classes_),
+    "winner_ids": winner_ids + 1,
+    "loser_ids": loser_ids + 1,
+}
+
+with open("./stan_model.stan", "r") as f:
+    posterior = stan.build(program_code=f.read(), data=stan_data)
+
+fit = posterior.sample(num_chains=4, num_samples=1000)
+
+runtime = time() - start_time
+
+arviz_version = az.from_pystan(fit)
+
+az.to_netcdf(arviz_version, os.path.join(target_dir, f"samples_{start_year}.netcdf"))
+print(runtime, file=open(os.path.join(target_dir, f"runtime_{start_year}.txt"), "w"))