adaptation_specialization_material/code/etm_ei.py

from socketserver import ThreadingUnixStreamServer
import pandas as pd
import numpy as np

from cmdstanpy import CmdStanModel

import argparse
import pickle

from os.path import join as opj

import seaborn as sns
from matplotlib import pyplot as plt

parser = argparse.ArgumentParser()
parser.add_argument("--input")
parser.add_argument("--n", type=int, default=200)
parser.add_argument("--min-pop", type=int, default=100)
parser.add_argument("--stack-rows", type=int, default=1)

parser.add_argument("--chains", type=int, default=1)
parser.add_argument("--threads-per-chain", type=int, default=4)
parser.add_argument("--samples", type=int, default=500)
parser.add_argument("--warmup", type=int, default=1000)

args = parser.parse_args()

n_topics = len(pd.read_csv(opj(args.input, "topics.csv")))

df = pd.read_csv(opj(args.input, "aggregate.csv"))
df = df[df[[f"start_{k+1}" for k in range(n_topics)]].sum(axis=1) >= args.min_pop]
df = df.sample(n=args.n)

resources = pd.read_parquet(opj(args.input, "pooled_resources.parquet"))
df = df.merge(resources, left_on="bai", right_on="bai")

data = {
    "NR": np.stack(df[[f"start_{k+1}" for k in range(n_topics)]].values).astype(int),
    "NC": np.stack(df[[f"end_{k+1}" for k in range(n_topics)]].values).astype(int),
    "R": n_topics,
    "C": n_topics,
    "n_units": len(df),
    "threads": args.threads_per_chain
}

data["cov"] = np.stack(df["pooled_resources"])

junk = np.sum(data["NR"] + data["NC"], axis=0) == 0

for col in ["NR", "NC", "cov"]:
    data[col] = data[col][:, ~junk]

data["R"] -= junk.sum()
data["C"] -= junk.sum()

data["cov"] = data["cov"] / np.maximum(data["cov"].sum(axis=1)[:, np.newaxis], 1)
sns.heatmap(
    np.corrcoef(data["NC"].T, data["cov"].T), vmin=-0.5, vmax=0.5, cmap="RdBu"
)
plt.show()


print(data["cov"].shape)


model = CmdStanModel(
    stan_file=f"code/ei_cov_softmax_control.stan",
    cpp_options={"STAN_THREADS": "TRUE"},
    compile="force",
)

fit = model.sample(
    data=data,
    chains=args.chains,
    threads_per_chain=args.threads_per_chain,
    iter_sampling=args.samples,
    iter_warmup=args.warmup,
)

vars = fit.stan_variables()
samples = {}
for (k, v) in vars.items():
    samples[k] = v

np.savez_compressed(opj(args.input, "ei_samples.npz"), **samples)