Repository: AIH-SGML/mixmil (language: python)

Issue:
Add experiments to v0.0.1



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 68bc17f..87f3989 100644
--- a/.gitignore
+++ b/.gitignore
@@ -158,3 +158,7 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+.DS_Store
+data
+*.out
diff --git a/README.md b/README.md
index 73aaaa0..1b7610a 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,40 @@
 # MixMIL
-Code coming soon.
+Code for the paper: [Attention-based Multi-instance Mixed Models](https://arxiv.org/abs/2311.02455)
+
+Please raise an issue for questions and bug-reports.
+## Installation
+Install with:
+```
+pip install mixmil
+```
+or if you want to adapt the code:
+```
+git clone https://github.com/AIH-SGML/mixmil.git
+cd mixmil
+pip install -e .
+```
+To enable computations on GPU please follow the installation instructions of [PyTorch](https://pytorch.org/) and [PyTorch Scatter](https://github.com/rusty1s/pytorch_scatter).
+MixMIL works e.g. with PyTorch 2.1.
+## Experiments
+See the notebooks in the `experiments` folder for examples on how to run the simulation and histopathology experiments.
+
+### Histopathology
+Install `anndata` (`pip install anndata`) to run the notebook. 
+
+#### Download Data
+To download the embeddings provided by the DSMIL authors, either:
+- Full embeddings: `python scripts/dsmil_data_download.py`
+- PCA reduced embeddings: [Google Drive](https://drive.google.com/drive/folders/1X9ho1_W5ixyHSw_2hCfQsBb5nzkjMviA?usp=sharing)
+
+
+## Citation
+```
+@misc{engelmann2023attentionbased,
+      title={Attention-based Multi-instance Mixed Models}, 
+      author={Jan P. Engelmann and Alessandro Palma and Jakub M. Tomczak and Fabian J Theis and Francesco Paolo Casale},
+      year={2023},
+      eprint={2311.02455},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG}
+}
+```
\ No newline at end of file
diff --git a/experiments/histopathology_camelyon16.ipynb b/experiments/histopathology_camelyon16.ipynb
new file mode 100644
index 0000000..fb0a77a
--- /dev/null
+++ b/experiments/histopathology_camelyon16.ipynb
@@ -0,0 +1,351 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Histopathology Experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mixmil.paths import DATA\n",
+    "import pandas as pd\n",
+    "from tqdm import tqdm\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "import numpy as np\n",
+    "import anndata as ad\n",
+    "from mixmil import MixMIL\n",
+    "import torch\n",
+    "from sklearn.metrics import roc_auc_score\n",
+    "import matplotlib.pyplot as plt\n",
+    "import scipy.stats as st"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Data Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>label</th>\n",
+       "      <th>split</th>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>file</th>\n",
+       "      <th></th>\n",
+       "      <th></th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>1-tumor/test_033.csv</th>\n",
+       "      <td>1</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_148.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/test_095.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_025.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/test_087.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_006.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1-tumor/tumor_003.csv</th>\n",
+       "      <td>1</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_018.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1-tumor/tumor_017.csv</th>\n",
+       "      <td>1</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/test_067.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>399 rows × 2 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                         label  split\n",
+       "file                                 \n",
+       "1-tumor/test_033.csv         1   test\n",
+       "0-normal/normal_148.csv      0  train\n",
+       "0-normal/test_095.csv        0   test\n",
+       "0-normal/normal_025.csv      0  train\n",
+       "0-normal/test_087.csv        0   test\n",
+       "...                        ...    ...\n",
+       "0-normal/normal_006.csv      0  train\n",
+       "1-tumor/tumor_003.csv        1  train\n",
+       "0-normal/normal_018.csv      0  train\n",
+       "1-tumor/tumor_017.csv        1  train\n",
+       "0-normal/test_067.csv        0   test\n",
+       "\n",
+       "[399 rows x 2 columns]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dataset_index_file = DATA / \"camelyon16\" / \"Camelyon16.csv\"\n",
+    "bagdf = pd.read_csv(dataset_index_file)\n",
+    "bagdf.columns = [\"file\", \"label\"]\n",
+    "bagdf[\"file\"] = bagdf[\"file\"].str.replace(\"datasets/Camelyon16/\", \"\")\n",
+    "bagdf[\"split\"] = bagdf[\"file\"].apply(lambda x: \"test\" if \"test\" in x else \"train\")\n",
+    "bagdf = bagdf.set_index(\"file\")\n",
+    "bagdf"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dtype = \"float32\"\n",
+    "\n",
+    "if not (DATA / \"camelyon16\" / \"full_test.h5ad\").exists() or not (DATA / \"full_camelyon16\" / \"full_train.h5ad\").exists():\n",
+    "    train_data = []\n",
+    "    train_bag_indices = []\n",
+    "    for _, row in tqdm(list(bagdf[bagdf[\"split\"] == \"train\"].iterrows())):\n",
+    "        train_data.append(pd.read_csv(dataset_index_file.parent / row.name).values.astype(dtype))\n",
+    "        train_bag_indices.extend([row.name] * len(train_data[-1]))\n",
+    "\n",
+    "    test_data = []\n",
+    "    test_bag_indices = []\n",
+    "    for _, row in tqdm(list(bagdf[bagdf[\"split\"] == \"test\"].iterrows())):\n",
+    "        test_data.append(pd.read_csv(dataset_index_file.parent / row.name).values.astype(dtype))\n",
+    "        test_bag_indices.extend([row.name] * len(test_data[-1]))\n",
+    "\n",
+    "    i_train = np.array([idx for idx, x in enumerate(train_data) for _ in range(len(x))])\n",
+    "    X_train = np.concatenate(train_data, 0).astype(dtype)\n",
+    "    X_test = np.concatenate(test_data, 0).astype(dtype)\n",
+    "    scaler = StandardScaler()\n",
+    "    X_train = scaler.fit_transform(X_train)\n",
+    "    X_test = scaler.transform(X_test)\n",
+    "\n",
+    "    train_obs = pd.DataFrame(\n",
+    "        {\"bag\": train_bag_indices, \"label\": bagdf.loc[train_bag_indices][\"label\"].values, \"split\": \"train\"}\n",
+    "    )\n",
+    "    train_adata = ad.AnnData(X_train, obs=train_obs)\n",
+    "    test_obs = pd.DataFrame(\n",
+    "        {\"bag\": test_bag_indices, \"label\": bagdf.loc[test_bag_indices][\"label\"].values, \"split\": \"test\"}\n",
+    "    )\n",
+    "    test_adata = ad.AnnData(X_test, obs=test_obs)\n",
+    "\n",
+    "    test_adata.write(DATA / \"camelyon16\" / \"full_test.h5ad\")\n",
+    "    train_adata.write(DATA / \"camelyon16\" / \"full_train.h5ad\")\n",
+    "else:\n",
+    "    print(\"Loading precomputed anndatas\")\n",
+    "    train_adata = ad.read_h5ad(DATA / \"camelyon16\" / \"full_train.h5ad\")\n",
+    "    test_adata = ad.read_h5ad(DATA / \"camelyon16\" / \"full_test.h5ad\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def to_device(el, device):\n",
+    "    \"\"\"Move tensor, list of tensors, or dict of tensors to target device.\"\"\"\n",
+    "    if isinstance(el, dict):\n",
+    "        return {k: to_device(v, device) for k, v in el.items()}\n",
+    "    elif isinstance(el, (list, tuple)):\n",
+    "        return [to_device(x, device) for x in el]\n",
+    "    elif isinstance(el, (torch.Tensor, torch.nn.Module)):\n",
+    "        return el.to(device)\n",
+    "    else:\n",
+    "        return el"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prepare train data\n",
+    "train_bags = train_adata.obs[\"bag\"].unique().tolist()\n",
+    "Xs = [torch.Tensor(train_adata[train_adata.obs[\"bag\"] == bag].X) for bag in train_bags]\n",
+    "F = torch.ones((len(train_bags), 1))\n",
+    "Y = torch.Tensor(train_adata.obs[[\"bag\", \"label\"]].drop_duplicates().set_index(\"bag\").loc[train_bags].values)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prepare test data, following official train-test split\n",
+    "test_bags = test_adata.obs[\"bag\"].unique().tolist()\n",
+    "test_Xs = [torch.Tensor(test_adata[test_adata.obs[\"bag\"] == bag].X) for bag in test_bags]\n",
+    "test_Y = torch.Tensor(test_adata.obs[[\"bag\", \"label\"]].drop_duplicates().set_index(\"bag\").loc[test_bags].values)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "50435294e98d46cc8aedd71c558b0566",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "GLMM Init:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test AUC: 0.698 Spearman: 0.333\n"
+     ]
+    }
+   ],
+   "source": [
+    "# initialize model with mean model and Bernoulli likelihood\n",
+    "model = MixMIL.init_with_mean_model(Xs, F, Y, likelihood=\"binomial\", n_trials=1)\n",
+    "y_pred_mean = model.predict(test_Xs)\n",
+    "print(\n",
+    "    \"Test AUC:\",\n",
+    "    round(roc_auc_score(test_Y, y_pred_mean), 3),\n",
+    "    \"Spearman:\",\n",
+    "    round(st.spearmanr(test_Y, y_pred_mean).correlation, 3),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test AUC: 0.977 Spearman: 0.802\n"
+     ]
+    }
+   ],
+   "source": [
+    "# train model for 1000 epochs\n",
+    "device = \"cuda:0\"\n",
+    "model, Xs, F, Y, test_Xs, test_Y = to_device((model, Xs, F, Y, test_Xs, test_Y), device)\n",
+    "model.train(Xs, F, Y, n_epochs=1000)\n",
+    "model.to(\"cpu\")\n",
+    "test_Xs = [x.cpu() for x in test_Xs]\n",
+    "y_pred = model.predict(test_Xs).cpu().numpy()\n",
+    "y_true = test_Y.cpu().numpy()\n",
+    "print(\n",
+    "    \"Test AUC:\",\n",
+    "    round(roc_auc_score(y_true, y_pred), 3),\n",
+    "    \"Spearman:\",\n",
+    "    round(st.spearmanr(y_true, y_pred).correlation, 3),\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gt",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/experiments/reduced_histopathology_camelyon16.ipynb b/experiments/reduced_histopathology_camelyon16.ipynb
new file mode 100644
index 0000000..f370d40
--- /dev/null
+++ b/experiments/reduced_histopathology_camelyon16.ipynb
@@ -0,0 +1,373 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Histopathology Experiment (Reduced Dimension)\n",
+    "> In this version we run PCA to half the data dimensionality to enable in-memory training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mixmil.paths import DATA\n",
+    "import pandas as pd\n",
+    "from tqdm import tqdm\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "from sklearn.decomposition import PCA\n",
+    "import numpy as np\n",
+    "import anndata as ad\n",
+    "from mixmil import MixMIL\n",
+    "import torch\n",
+    "from sklearn.metrics import roc_auc_score\n",
+    "import scipy.stats as st"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Data Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>label</th>\n",
+       "      <th>split</th>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>file</th>\n",
+       "      <th></th>\n",
+       "      <th></th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>1-tumor/test_033.csv</th>\n",
+       "      <td>1</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_148.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/test_095.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_025.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/test_087.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_006.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1-tumor/tumor_003.csv</th>\n",
+       "      <td>1</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/normal_018.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1-tumor/tumor_017.csv</th>\n",
+       "      <td>1</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0-normal/test_067.csv</th>\n",
+       "      <td>0</td>\n",
+       "      <td>test</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>399 rows × 2 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                         label  split\n",
+       "file                                 \n",
+       "1-tumor/test_033.csv         1   test\n",
+       "0-normal/normal_148.csv      0  train\n",
+       "0-normal/test_095.csv        0   test\n",
+       "0-normal/normal_025.csv      0  train\n",
+       "0-normal/test_087.csv        0   test\n",
+       "...                        ...    ...\n",
+       "0-normal/normal_006.csv      0  train\n",
+       "1-tumor/tumor_003.csv        1  train\n",
+       "0-normal/normal_018.csv      0  train\n",
+       "1-tumor/tumor_017.csv        1  train\n",
+       "0-normal/test_067.csv        0   test\n",
+       "\n",
+       "[399 rows x 2 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dataset_index_file = DATA / \"camelyon16\" / \"Camelyon16.csv\"\n",
+    "bagdf = pd.read_csv(dataset_index_file)\n",
+    "bagdf.columns = [\"file\", \"label\"]\n",
+    "bagdf[\"file\"] = bagdf[\"file\"].str.replace(\"datasets/Camelyon16/\", \"\")\n",
+    "bagdf[\"split\"] = bagdf[\"file\"].apply(lambda x: \"test\" if \"test\" in x else \"train\")\n",
+    "bagdf = bagdf.set_index(\"file\")\n",
+    "bagdf"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loading precomputed anndatas\n"
+     ]
+    }
+   ],
+   "source": [
+    "dtype = \"float32\"\n",
+    "n_pcs = 128\n",
+    "\n",
+    "if not (DATA / \"camelyon16\" / \"test.h5ad\").exists() or not (DATA / \"camelyon16\" / \"train.h5ad\").exists():\n",
+    "    train_data = []\n",
+    "    train_bag_indices = []\n",
+    "    for _, row in tqdm(list(bagdf[bagdf[\"split\"] == \"train\"].iterrows())):\n",
+    "        train_data.append(pd.read_csv(dataset_index_file.parent / row.name).values.astype(dtype))\n",
+    "        train_bag_indices.extend([row.name] * len(train_data[-1]))\n",
+    "\n",
+    "    test_data = []\n",
+    "    test_bag_indices = []\n",
+    "    for _, row in tqdm(list(bagdf[bagdf[\"split\"] == \"test\"].iterrows())):\n",
+    "        test_data.append(pd.read_csv(dataset_index_file.parent / row.name).values.astype(dtype))\n",
+    "        test_bag_indices.extend([row.name] * len(test_data[-1]))\n",
+    "\n",
+    "    i_train = np.array([idx for idx, x in enumerate(train_data) for _ in range(len(x))])\n",
+    "    X_train = np.concatenate(train_data, 0).astype(dtype)\n",
+    "    X_test = np.concatenate(test_data, 0).astype(dtype)\n",
+    "    pca = PCA(n_components=n_pcs)\n",
+    "    scaler = StandardScaler()\n",
+    "    X_train = scaler.fit_transform(pca.fit_transform(X_train))\n",
+    "    X_test = scaler.transform(pca.transform(X_test))\n",
+    "\n",
+    "    train_obs = pd.DataFrame(\n",
+    "        {\"bag\": train_bag_indices, \"label\": bagdf.loc[train_bag_indices][\"label\"].values, \"split\": \"train\"}\n",
+    "    )\n",
+    "    train_adata = ad.AnnData(X_train, obs=train_obs, var=pd.DataFrame(index=[f\"PC{i}\" for i in range(n_pcs)]))\n",
+    "    test_obs = pd.DataFrame(\n",
+    "        {\"bag\": test_bag_indices, \"label\": bagdf.loc[test_bag_indices][\"label\"].values, \"split\": \"test\"}\n",
+    "    )\n",
+    "    test_adata = ad.AnnData(X_test, obs=test_obs, var=pd.DataFrame(index=[f\"PC{i}\" for i in range(n_pcs)]))\n",
+    "\n",
+    "    test_adata.write(DATA / \"camelyon16\" / \"test.h5ad\")\n",
+    "    train_adata.write(DATA / \"camelyon16\" / \"train.h5ad\")\n",
+    "else:\n",
+    "    print(\"Loading precomputed anndatas\")\n",
+    "    train_adata = ad.read_h5ad(DATA / \"camelyon16\" / \"train.h5ad\")\n",
+    "    test_adata = ad.read_h5ad(DATA / \"camelyon16\" / \"test.h5ad\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def to_device(el, device):\n",
+    "    if isinstance(el, dict):\n",
+    "        return {k: to_device(v, device) for k, v in el.items()}\n",
+    "    elif isinstance(el, (list, tuple)):\n",
+    "        return [to_device(x, device) for x in el]\n",
+    "    elif isinstance(el, (torch.Tensor, torch.nn.Module)):\n",
+    "        return el.to(device)\n",
+    "    else:\n",
+    "        return el"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prepare train data\n",
+    "train_bags = train_adata.obs[\"bag\"].unique().tolist()\n",
+    "Xs = [torch.Tensor(train_adata[train_adata.obs[\"bag\"] == bag].X) for bag in train_bags]\n",
+    "F = torch.ones((len(train_bags), 1))\n",
+    "Y = torch.Tensor(train_adata.obs[[\"bag\", \"label\"]].drop_duplicates().set_index(\"bag\").loc[train_bags].values)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prepare test data, following official train-test split\n",
+    "test_bags = test_adata.obs[\"bag\"].unique().tolist()\n",
+    "test_Xs = [torch.Tensor(test_adata[test_adata.obs[\"bag\"] == bag].X) for bag in test_bags]\n",
+    "test_Y = torch.Tensor(test_adata.obs[[\"bag\", \"label\"]].drop_duplicates().set_index(\"bag\").loc[test_bags].values)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "66eb150f90564e659021f5bcb6cafdd3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "GLMM Init:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test AUC: 0.706 Spearman: 0.346\n"
+     ]
+    }
+   ],
+   "source": [
+    "# initialize model with mean model and Bernoulli likelihood\n",
+    "model = MixMIL.init_with_mean_model(Xs, F, Y, likelihood=\"binomial\", n_trials=1)\n",
+    "y_pred_mean = model.predict(test_Xs)\n",
+    "print(\n",
+    "    \"Test AUC:\",\n",
+    "    round(roc_auc_score(test_Y, y_pred_mean), 3),\n",
+    "    \"Spearman:\",\n",
+    "    round(st.spearmanr(test_Y, y_pred_mean).correlation, 3),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1620a6afc2d94f0e9202c7cebd220877",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Epoch:   0%|          | 0/1000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test AUC: 0.969 Spearman: 0.789\n"
+     ]
+    }
+   ],
+   "source": [
+    "# train model for 1000 epochs\n",
+    "device = \"cuda:0\"\n",
+    "model, Xs, F, Y, test_Xs, test_Y = to_device((model, Xs, F, Y, test_Xs, test_Y), device)\n",
+    "model.train(Xs, F, Y, n_epochs=1000)\n",
+    "y_pred = model.predict(test_Xs).cpu().numpy()\n",
+    "y_true = test_Y.cpu().numpy()\n",
+    "print(\n",
+    "    \"Test AUC:\",\n",
+    "    round(roc_auc_score(y_true, y_pred), 3),\n",
+    "    \"Spearman:\",\n",
+    "    round(st.spearmanr(y_true, y_pred).correlation, 3),\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gt",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/experiments/simulation.ipynb b/experiments/simulation.ipynb
new file mode 100644
index 0000000..b063c23
--- /dev/null
+++ b/experiments/simulation.ipynb
@@ -0,0 +1,188 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Simulation Experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import scipy.stats as st\n",
+    "import torch\n",
+    "from sklearn.metrics import roc_auc_score\n",
+    "\n",
+    "from mixmil import MixMIL\n",
+    "from mixmil.data import load_data\n",
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Utility Functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _calc_metrics(u_pred, w_pred, u, w):\n",
+    "    rho_bag = st.spearmanr(u_pred, u).correlation  # bag level correlation\n",
+    "    is_top_instance = (w > np.quantile(w, 0.90)).long().ravel()\n",
+    "    auc_instance = roc_auc_score(is_top_instance, w_pred)  # instance-retrieval AUC\n",
+    "    return rho_bag, auc_instance\n",
+    "\n",
+    "\n",
+    "def calc_metrics(model, X, u, w):\n",
+    "    u_pred = model.predict(X[\"test\"]).cpu().numpy()\n",
+    "    w_pred = model.get_weights(X[\"test\"])[0].cpu().numpy()\n",
+    "\n",
+    "    P = u_pred.shape[1]\n",
+    "    if P > 0:\n",
+    "        rho_bag, auc_instance = [], []\n",
+    "        for i in range(P):\n",
+    "            _rho_bag, _auc_instance = _calc_metrics(\n",
+    "                u_pred[..., i], w_pred[..., i].ravel(), u[\"test\"][..., i], w[\"test\"][..., i].ravel()\n",
+    "            )\n",
+    "            rho_bag.append(_rho_bag)\n",
+    "            auc_instance.append(_auc_instance)\n",
+    "\n",
+    "        res_dict = {\n",
+    "            \"rho_bag\": np.mean(rho_bag),\n",
+    "            \"rho_bag_err\": np.std(rho_bag) / np.sqrt(P),\n",
+    "            \"auc_instance\": np.mean(auc_instance),\n",
+    "            \"auc_instance_err\": np.std(auc_instance) / np.sqrt(P),\n",
+    "        }\n",
+    "    else:\n",
+    "        rho_bag, auc_instance = _calc_metrics(u_pred, w_pred.ravel(), u[\"test\"], w[\"test\"].ravel())\n",
+    "        res_dict = {\"rho_bag\": rho_bag, \"auc_instance\": auc_instance}\n",
+    "    return res_dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 52,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def print_metrics(prefix, metrics):\n",
+    "    print(f\"{prefix} metrics:\")\n",
+    "    for k, v in metrics.items():\n",
+    "        print(f\"{k}: {v:.4f}\")\n",
+    "    print()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 56,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c23b2a79444348288359e79609a4d2d6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "GLMM Init:   0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[START] metrics:\n",
+      "rho_bag: 0.6004\n",
+      "rho_bag_err: 0.0131\n",
+      "auc_instance: 0.5000\n",
+      "auc_instance_err: 0.0000\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c85d52a1411b4eb6a64860a0f2e575cf",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Epoch:   0%|          | 0/2000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[END] metrics:\n",
+      "rho_bag: 0.8318\n",
+      "rho_bag_err: 0.0112\n",
+      "auc_instance: 0.9361\n",
+      "auc_instance_err: 0.0078\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
+    "\n",
+    "# Simulate data as described in the paper\n",
+    "# embeddings, fixed effects, labels, sim bag predictions, sim instance weights\n",
+    "# P: number of outputs, simulated from the same embeddings X\n",
+    "X, F, Y, u, w = load_data(P=10, seed=0)\n",
+    "model = MixMIL.init_with_mean_model(X[\"train\"], F[\"train\"], Y[\"train\"], likelihood=\"binomial\", n_trials=2).to(device)\n",
+    "X, F, Y = [{key: val.to(device) for key, val in el.items()} for el in [X, F, Y]]\n",
+    "\n",
+    "print_metrics(\"[START]\", calc_metrics(model, X, u, w))\n",
+    "# Fit model in parallel to each output separately\n",
+    "model.train(X[\"train\"], F[\"train\"], Y[\"train\"], n_epochs=2_000)\n",
+    "print_metrics(\"[END]\", calc_metrics(model, X, u, w))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gt",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/mixmil/__init__.py b/mixmil/__init__.py
new file mode 100644
index 0000000..abca57e
--- /dev/null
+++ b/mixmil/__init__.py
@@ -0,0 +1,4 @@
+from mixmil.model import MixMIL
+
+__all__ = ["MixMIL", "utils", "likelihood", "posterior", "data", "simulation", "paths"]
+__version__ = "0.1.0"
diff --git a/mixmil/data.py b/mixmil/data.py
new file mode 100644
index 0000000..b051a36
--- /dev/null
+++ b/mixmil/data.py
@@ -0,0 +1,66 @@
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from mixmil.simulation import load_simulation
+
+
+def setup_scatter(Xs):
+    device = Xs[0].device
+    x = torch.cat(Xs, dim=0)
+    i = torch.cat([torch.full((x.shape[0],), idx) for idx, x in enumerate(Xs)]).to(device)
+    i_ptr = torch.cat([torch.tensor([0], device=device), i.bincount().cumsum(0)])
+    return x, i, i_ptr
+
+
+def xgower_factor(X):
+    a = np.power(X, 2).sum()
+    b = X.dot(X.sum(0)).sum()
+    return np.sqrt((a - b / X.shape[0]) / (X.shape[0] - 1))
+
+
+class MILDataset(Dataset):
+    def __init__(self, Xs, F, Y):
+        self.Xs = Xs
+        self.F = F
+        self.Y = Y
+
+    def __len__(self):
+        return len(self.Xs)
+
+    def __getitem__(self, idx):
+        X = self.Xs[idx]
+        F = self.F[idx]
+        Y = self.Y[idx]
+        return X, F, Y
+
+
+def mil_collate_fn(batch):
+    X = [item[0] for item in batch]
+    F = torch.stack([item[1] for item in batch])
+    Y = torch.stack([item[2] for item in batch])
+    return X, F, Y
+
+
+def normalize_feats(X, norm_factor="std_sqrt"):
+    assert norm_factor in ["std", "std_sqrt"]
+    train_data = (
+        torch.cat(X["train"], dim=0) if isinstance(X["train"], list) else X["train"].reshape(-1, X["train"].shape[2])
+    )
+    mean = train_data.mean(0, keepdims=True)
+    std = train_data.std(0, keepdims=True)
+    factor = std * np.sqrt(train_data.shape[1]) if norm_factor == "std_sqrt" else std
+
+    for key in X:
+        X[key] = [(x - mean) / factor for x in X[key]] if isinstance(X[key], list) else (X[key] - mean) / factor
+
+    return X
+
+
+def load_data(dataset="simulation", norm_factor="std_sqrt", **kwargs):
+    if dataset == "simulation":
+        X, F, Y, u, w = load_simulation(**kwargs)
+        X = normalize_feats(X, norm_factor)
+    else:
+        raise ValueError(f"Unknown dataset: {dataset}")
+    return X, F, Y, u, w
diff --git a/mixmil/model.py b/mixmil/model.py
new file mode 100644
index 0000000..bfad2a7
--- /dev/null
+++ b/mixmil/model.py
@@ -0,0 +1,208 @@
+import numpy as np
+import torch
+from torch.distributions import Binomial, Categorical, LowRankMultivariateNormal
+from torch.distributions.kl import kl_divergence
+from torch.utils.data import DataLoader
+from torch_scatter import scatter_softmax, segment_add_csr
+from tqdm.auto import trange
+
+from mixmil.data import MILDataset, mil_collate_fn, setup_scatter
+from mixmil.posterior import GaussianVariationalPosterior
+from mixmil.utils import get_init_params
+
+
+class MixMIL(torch.nn.Module):
+    """Attention-based Multi-instance Mixed Models
+
+    https://arxiv.org/abs/2311.02455
+
+    """
+
+    def __init__(self, Q, K, P=1, likelihood="binomial", n_trials=2, mean_field=False, init_params=None):
+        r"""Initialize the MixMil class.
+
+        Parameters:
+        - Q (int): The dimension of the latent space.
+        - K (int): The number of fixed effects.
+        - P (int): The number of outputs.
+        - likelihood (str, optional): The likelihood to use. Either "binomial" or "categorical". Default is "binomial".
+        - n_trials (int, optional): Number of trials for binomial likelihood. Not used for categorical. Default is 2.
+        - mean_field (bool, optional): Toggle mean field approximation for the posterior. Default is False.
+        - init_params (tuple, optional): Tuple of (mean, var, var_z, alpha) to initialize the model. Default is None.
+            mean (torch.Tensor): The mean of the posterior. Shape: (Q, P). d
+            var (torch.Tensor): The variance of the posterior. Shape: (Q, P).
+            var_z (torch.Tensor): The $\sigma_{\beta}^2$ hparam of the prior.
+                Shape: (1, P) with separate and (1, 1) with shared priors .
+            alpha (torch.Tensor): The fixed effect parameters. Shape: (K, P).
+        """
+        super().__init__()
+        self.Q = Q
+
+        alpha = torch.zeros((K, P))
+        log_sigma_u = torch.full((1, P), 0.5 * np.log(0.5))
+        log_sigma_z = torch.full((1, P), 0.5 * np.log(0.5))
+
+        if init_params is not None:
+            *_, var_z, alpha = init_params
+            log_sigma_z = 0.5 * torch.log(var_z)
+
+        self.alpha = torch.nn.Parameter(alpha)
+        self.log_sigma_u = torch.nn.Parameter(log_sigma_u)
+        self.log_sigma_z = torch.nn.Parameter(log_sigma_z)
+
+        self.posterior = GaussianVariationalPosterior(2 * Q, P, mean_field, init_params)
+
+        self.likelihood_name = likelihood
+        self.n_trials = n_trials if likelihood == "binomial" else None
+        self.is_trained = False
+
+    def init_with_mean_model(Xs, F, Y, likelihood="binomial", n_trials=None, mean_field=False):
+        assert (likelihood == "binomial" and n_trials is not None and 0 < n_trials <= 2) or (
+            likelihood == "categorical" and n_trials is None
+        ), f"n_trials must be 1 or 2 to initialize with binomial mean model, got {n_trials=} and {likelihood=}"
+        init_params = get_init_params(Xs, F, Y, likelihood, n_trials)
+        Q, K, P = Xs[0].shape[1], F.shape[1], init_params[0].shape[1]
+        return MixMIL(Q, K, P, likelihood, n_trials, mean_field, init_params)
+
+    @property
+    def prior_distribution(self):
+        device = self.log_sigma_u.device
+        scale_u = self.log_sigma_u.T * torch.ones([1, self.Q], device=device)
+        scale_z = self.log_sigma_z.T * torch.ones([1, self.Q], device=device)
+        cov_logdiag = 2 * torch.cat([scale_u, scale_z], 1)
+        cov_factor = torch.zeros_like(cov_logdiag)[:, :, None]
+        mu = torch.zeros_like(cov_logdiag)
+        return LowRankMultivariateNormal(mu, cov_factor, torch.exp(cov_logdiag))
+
+    @property
+    def posterior_distribution(self):
+        return self.posterior.distribution
+
+    @property
+    def qu_mu(self):
+        return self.posterior.q_mu[: self.Q]
+
+    @property
+    def qz_mu(self):
+        return self.posterior.q_mu[self.Q :]
+
+    def likelihood(self, logits, y):
+        if self.likelihood_name == "binomial":
+            return Binomial(total_count=self.n_trials, logits=logits).log_prob(y[:, :, None]).sum(1).mean()
+        elif self.likelihood_name == "categorical":
+            logits = logits.permute(0, 2, 1)
+            if logits.shape[-1] == 1:
+                logits = torch.cat([-logits, logits], 2)
+            return Categorical(logits=logits).log_prob(y).mean()
+
+    def loss(self, u, f, y, kld_w=1.0, return_dict=False):
+        logits = f.mm(self.alpha)[:, :, None] + u
+
+        ll = self.likelihood(logits, y)
+        kld = kl_divergence(self.posterior_distribution, self.prior_distribution)
+        kld_term = kld_w * kld.sum() / y.shape[0]
+        loss = -ll + kld_term
+        if return_dict:
+            return loss, dict(loss=loss.item(), ll=ll.item(), kld=kld_term.item())
+        return loss
+
+    def get_betas(self, n_samples=None, predict=False):
+        assert not (n_samples and predict)
+        if n_samples:
+            beta = self.posterior.sample(n_samples)
+            beta_u = beta[: self.Q, :, :]
+            beta_z = beta[self.Q :, :, :]
+        else:
+            beta_u = self.qu_mu[:, :, None]
+            beta_z = self.qz_mu[:, :, None]
+        return beta_u, beta_z
+
+    def forward(self, Xs, n_samples=8, scaling=None, predict=False):
+        beta_u, beta_z = self.get_betas(n_samples, predict)
+        b = torch.sqrt((beta_z**2).mean(0, keepdim=True))
+        eta = beta_z / b
+
+        if torch.is_tensor(Xs):
+            u = self._calc_bag_emb_effect_tensor(beta_u, eta, Xs)
+        else:
+            u = self._calc_bag_emb_effect_scatter(beta_u, eta, Xs)
+
+        mean, std = (u.mean(0), u.std(0)) if scaling is None else scaling
+        if std.isnan().any():
+            std = 1
+        u = b * (u - mean) / std
+        return u
+
+    def _calc_bag_emb_effect_tensor(self, beta_u, eta, Xs):
+        _w = torch.einsum("niq,qps->nips", Xs, beta_u)
+        w = torch.softmax(_w, dim=1)
+        t = torch.einsum("niq,qps->nips", Xs, eta)
+        u = torch.einsum("nips,nips->nps", w, t)
+        return u
+
+    def _calc_bag_emb_effect_scatter(self, beta_u, eta, Xs):
+        x, i, i_ptr = setup_scatter(Xs)
+
+        _w = torch.einsum("iq,qps->ips", x, beta_u)
+        w = scatter_softmax(_w, i, dim=0)
+        t = torch.einsum("iq,qps->ips", x, eta)
+        u = segment_add_csr(w * t, i_ptr)
+        return u
+
+    def train(self, X, F, Y, n_epochs=2_000, batch_size=64, lr=1e-3, verbose=True):
+        train_loader = DataLoader(
+            MILDataset(X, F, Y),
+            shuffle=True,
+            batch_size=batch_size,
+            collate_fn=None if torch.is_tensor(X) else mil_collate_fn,
+        )
+        optim = torch.optim.Adam(lr=lr, params=self.parameters())
+
+        history = []
+        for epoch in trange(1, n_epochs + 1, desc="Epoch", disable=not verbose):
+            for step, (xs, f, y) in enumerate(train_loader):
+                u = self(xs)
+                loss, ldict = self.loss(u, f, y, kld_w=len(xs) / len(Y), return_dict=True)
+                ldict["epoch"], ldict["step"] = epoch, step
+                history.append(ldict)
+                optim.zero_grad()
+                loss.backward()
+                optim.step()
+
+        self.is_trained = True
+        return history
+
+    @torch.inference_mode()
+    def predict(self, Xs, scaling=None):
+        return self(Xs, n_samples=None, predict=True, scaling=scaling).squeeze(2)
+
+    @torch.inference_mode()
+    def get_weights(self, Xs, ravel=False):
+        """Get instance weights after and before softmax"""
+        beta_u, _ = self.get_betas(predict=True)
+        beta_u = beta_u.squeeze(2)  # not taking mcmc samples
+        if torch.is_tensor(Xs):
+            _w = torch.einsum("niq,qp->nip", Xs, beta_u)
+            w = torch.softmax(_w, dim=1)
+
+        else:
+            x, i, _ = setup_scatter(Xs)
+            _w = torch.einsum("iq,qp->ip", x, beta_u)
+            w = scatter_softmax(_w, i, dim=0)
+
+        if ravel:
+            w, _w = w.ravel(), _w.ravel()
+        elif not torch.is_tensor(Xs):
+            _w = [_w[i == idx] for idx in range(len(Xs))]
+            w = [w[i == idx] for idx in range(len(Xs))]
+        return w, _w
+
+    def extra_repr(self):
+        string = f"Q={self.Q}, K={self.alpha.shape[0]}, P={self.alpha.shape[1]}, likelihood={self.likelihood_name}"
+        if self.likelihood_name == "binomial":
+            string += f", n_trials={self.n_trials}"
+        string += f", device={self.alpha.device}, trained={self.is_trained}"
+        string += f"\n(alpha): Parameter(shape={tuple(self.alpha.shape)})\n"
+        string += f"(log_sigma_u): Parameter(shape={tuple(self.log_sigma_u.shape)})\n"
+        string += f"(log_sigma_z): Parameter(shape={tuple(self.log_sigma_z.shape)})"
+        return string
diff --git a/mixmil/paths.py b/mixmil/paths.py
new file mode 100644
index 0000000..6d9fce3
--- /dev/null
+++ b/mixmil/paths.py
@@ -0,0 +1,5 @@
+from pathlib import Path
+
+ROOT = Path(__file__).parent.parent
+
+DATA = ROOT / "data"
diff --git a/mixmil/posterior.py b/mixmil/posterior.py
new file mode 100644
index 0000000..1f9398e
--- /dev/null
+++ b/mixmil/posterior.py
@@ -0,0 +1,61 @@
+import numpy as np
+import torch
+from torch.distributions import LowRankMultivariateNormal
+
+
+def get_params(vc_mean, vc_sd, n_outs, n_vars, mean_field):
+    mu_z = torch.Tensor(vc_mean)
+    mu_u = torch.zeros_like(mu_z)
+    mu = torch.cat([mu_u, mu_z], 1)
+    sd_z = torch.Tensor(vc_sd)
+    sd_u = torch.sqrt(0.1 * torch.ones_like(sd_z))
+
+    if mean_field:
+        cov_factor = torch.zeros(n_outs, n_vars, 1)
+        cov_logdiag = 2.0 * torch.log(torch.cat([sd_u, sd_z], 1))
+    else:
+        diag = torch.diag_embed(torch.cat([sd_u, sd_z], 1))
+        cov_factor = 1e-4 * torch.randn(diag.shape) + diag
+        cov_logdiag = np.log(1e-4) * torch.ones(n_outs, n_vars)
+
+    return mu, cov_factor, cov_logdiag
+
+
+class GaussianVariationalPosterior(torch.nn.Module):
+    def __init__(self, n_vars, n_outs, mean_field=True, init_params=None):
+        super().__init__()
+        self.n_vars = n_vars
+        self.n_outs = n_outs
+        self.mean_field = mean_field
+
+        if init_params is not None:
+            mu_z, sd_z, *_ = init_params
+            mu_z = mu_z.T
+            sd_z = sd_z.T
+        else:
+            mu_z = 1e-3 * torch.randn(n_outs, n_vars // 2)
+            sd_z = 1e-3 * torch.randn(n_outs, n_vars // 2)
+
+        mu, cov_factor, cov_logdiag = get_params(mu_z, sd_z, n_outs, n_vars, mean_field)
+
+        self.mu = torch.nn.Parameter(mu)
+        if mean_field:
+            self.register_buffer("cov_factor", cov_factor)
+            self.cov_logdiag = torch.nn.Parameter(cov_logdiag)
+        else:
+            self.cov_factor = torch.nn.Parameter(cov_factor)
+            self.register_buffer("cov_logdiag", cov_logdiag)
+
+    @property
+    def distribution(self):
+        return LowRankMultivariateNormal(self.mu, self.cov_factor, torch.exp(self.cov_logdiag))
+
+    @property
+    def q_mu(self):
+        return self.mu.T
+
+    def sample(self, n_samples):
+        return self.distribution.rsample([n_samples]).permute([2, 1, 0])
+
+    def extra_repr(self) -> str:
+        return f"n_vars=2*{self.n_vars//2}, n_outs={self.n_outs}, mean_field={self.mean_field}"
diff --git a/mixmil/simulation.py b/mixmil/simulation.py
new file mode 100644
index 0000000..89dd554
--- /dev/null
+++ b/mixmil/simulation.py
@@ -0,0 +1,82 @@
+import numpy as np
+import torch
+from sklearn.model_selection import train_test_split
+
+
+def get_X(N=1_000, I=50, Q=30, N_test=200):
+    N = N + N_test
+
+    X = torch.randn([N * I, Q], dtype=torch.float32)
+
+    X = (X - X.mean(0)) / X.std(0)
+    X = X / np.sqrt(X.shape[1])
+    X = X.reshape([N, I, Q])
+
+    return X
+
+
+def simulate(X, v_beta=0.5, v_gamma=0.8, b=-1, F=None, P=1):
+    if F is None:
+        F = torch.ones([X.shape[0], 1])
+
+    # simulate single phenotype
+    K = F.shape[1]
+    b = b * torch.ones([K, P])
+    v_beta = v_beta * torch.ones(P)
+    v_gamma = v_gamma * torch.ones(P)
+
+    # sample weights
+    gamma = torch.randn((X.shape[2], v_gamma.shape[0]))
+    _w = torch.einsum("nik,kp->nip", X, gamma)
+    _scale_w = torch.sqrt(v_gamma / _w.var([0, 1]))
+    gamma = _scale_w[None, :] * gamma
+    _w = _scale_w[None, None, :] * _w
+
+    w = torch.softmax(_w, dim=1)
+
+    # sample z
+    beta = torch.randn((X.shape[2], v_beta.shape[0]))
+    beta = beta / torch.sqrt((beta**2).mean(0, keepdim=True))
+    z = torch.einsum("nik,kp->nip", X, beta)
+    u = torch.einsum("nip,nip->np", w, z)
+    u = (u - u.mean(0)) / u.std(0)
+    u = torch.sqrt(v_beta) * u
+    beta = torch.sqrt(v_beta) * beta
+
+    # compute rates
+    logits = F.mm(b) + u
+
+    # sample Y
+    Y = torch.distributions.Binomial(2, logits=logits).sample()
+
+    return F, Y, u, w
+
+
+def split_data(Xs, test_size=200, val_size=0.0, test_rs=127, val_rs=412):
+    idxs_all = np.arange(Xs[0].shape[0])
+    idxs = {}
+    idxs["train_val"], idxs["test"] = train_test_split(idxs_all, test_size=test_size, random_state=test_rs)
+
+    if not np.isclose(val_size, 0):
+        idxs["train"], idxs["val"] = train_test_split(idxs["train_val"], test_size=val_size, random_state=val_rs)
+    else:
+        idxs["train"] = idxs["train_val"]
+    del idxs["train_val"]
+
+    outs = []
+    for X in Xs:
+        out = {}
+        for key in idxs.keys():
+            out[key] = X[idxs[key]]
+        outs.append(out)
+    return outs
+
+
+def load_simulation(seed=42, N=1_000, I=50, Q=30, N_test=200, P=1, v_beta=0.5, v_gamma=0.8, b=-1, F=None):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    X = get_X(N, I, Q, N_test)
+    data = simulate(X, v_beta, v_gamma, b, F, P)
+    X, F, Y, u, w = split_data((X, *data))
+    return X, F, Y, u, w
diff --git a/mixmil/utils.py b/mixmil/utils.py
new file mode 100644
index 0000000..5a6598e
--- /dev/null
+++ b/mixmil/utils.py
@@ -0,0 +1,108 @@
+import numpy as np
+import scipy.linalg as la
+import statsmodels.api as sm
+import torch
+from sklearn.linear_model import LogisticRegressionCV
+from tqdm.auto import trange
+
+from mixmil.data import xgower_factor
+
+
+def regressOut(Y, X, return_b=False, return_pinv=False):
+    """
+    regresses out X from Y
+    """
+    Xd = la.pinv(X)
+    b = Xd.dot(Y)
+    Y_out = Y - X.dot(b)
+    out = [Y_out]
+    if return_b:
+        out.append(b)
+    if return_pinv:
+        out.append(Xd)
+    return out if len(out) > 1 else out[0]
+
+
+def _get_single_binomial_init_params(X, F, y):
+    ident = np.zeros((X.shape[1]), dtype=int)
+    model = sm.BinomialBayesMixedGLM(y, F, X, ident).fit_vb()
+
+    u = np.dot(X, model.vc_mean)[::2]
+
+    _scale = u.std() / np.sqrt((model.vc_mean**2).mean(0))
+    mu_beta = _scale * model.vc_mean
+    sd_beta = _scale * model.vc_sd
+    var_z = (mu_beta**2 + sd_beta**2).mean().reshape(1)
+    alpha = model.fe_mean
+
+    return mu_beta, sd_beta, var_z, alpha
+
+
+def get_binomial_init_params(X, F, Y):
+    results = [_get_single_binomial_init_params(X, F, Y[:, p]) for p in trange(Y.shape[1], desc="GLMM Init")]
+
+    mu_beta, sd_beta, var_z, alpha = [_list2tensor(listo) for listo in zip(*results)]
+
+    return mu_beta, sd_beta, var_z, alpha
+
+
+def get_lr_init_params(X, Y, b, Fiv):
+    model = LogisticRegressionCV(
+        Cs=10,
+        fit_intercept=True,
+        penalty="l2",
+        multi_class="multinomial",
+        solver="lbfgs",
+        n_jobs=8,
+        verbose=0,
+        random_state=42,
+        max_iter=1000,
+        refit=True,
+    )
+
+    model.fit(X, Y.ravel())
+
+    alpha = model.intercept_[None]
+    beta = model.coef_.T
+
+    # Compute bag prediction u and reparametrize
+    u = X.dot(beta)
+    um = u.mean(0)[None]
+    us = u.std(0)[None]
+    alpha = alpha + um
+    mu_beta = us * beta / np.sqrt((beta**2).mean(0)[None])
+    sd_beta = np.sqrt(0.1 * (mu_beta**2).mean()) * np.ones_like(mu_beta)
+
+    alpha = Fiv.dot(np.ones((Fiv.shape[1], 1))).dot(alpha) - b.dot(mu_beta)
+
+    # init prior
+    var_z = (mu_beta**2 + sd_beta**2).mean().reshape(1, 1)
+
+    return [torch.Tensor(el) for el in (mu_beta, sd_beta, var_z, alpha)]
+
+
+def _list2tensor(_list):
+    return torch.Tensor(np.stack(_list, axis=1))
+
+
+def get_init_params(Xs, F, Y, likelihood, n_trials):
+    Xm = np.concatenate([x.mean(0, keepdims=True) for x in Xs], axis=0)
+    Fe, Ye = F.numpy(), Y.long().numpy()
+
+    if likelihood == "binomial":
+        Xm = (Xm - Xm.mean(0, keepdims=True)) / xgower_factor(Xm)
+
+        if n_trials == 2:
+            Xm, Fe = Xm.repeat(2, axis=0), Fe.repeat(2, axis=0)
+            to_expanded = np.array(([[0, 0], [1, 0], [1, 1]]))
+            Ye = to_expanded[Y.long().numpy().T].transpose(1, 2, 0).reshape(-1, Y.shape[1])
+
+        mu_z, sd_z, var_z, alpha = get_binomial_init_params(Xm, Fe, Ye)
+
+    elif likelihood == "categorical":
+        Xm, b, Fiv = regressOut(Xm, Fe, return_b=True, return_pinv=True)
+        Xm = (Xm - Xm.mean(0, keepdims=True)) / (Xm.std(0, keepdims=True) * np.sqrt(Xm.shape[-1]))
+
+        mu_z, sd_z, var_z, alpha = get_lr_init_params(Xm, Ye, b, Fiv)
+
+    return mu_z, sd_z, var_z, alpha
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..2c38ca0
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,60 @@
+[project]
+name = "mixmil"
+dynamic = ["version"]
+description = "Code for the paper: Attention-based Multi-instance Mixed Models"
+readme = "README.md"
+license.file = "LICENSE"
+authors = [
+    { name = "Jan Engelmann", email = "jan.engelmann@helmholtz-munich.de" },
+    { name = "Alessandro Palma", email = "alessandro.palma@helmholtz-munich.de" },
+    { name = "Paolo Casale", email = "francescopaolo.casale@helmholtz-munich.de" },
+]
+
+dependencies = [
+    "numpy>=1.22.0",
+    "torch>=1.4.0",
+    "torch_scatter>=2.0.1",
+    "scipy>=1.8.0",
+    "scikit-learn>=1.3.0",
+    "tqdm>=4.0.0",
+    "statsmodels>=0.11.0",
+]
+requires-python = ">=3.9"
+
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Natural Language :: English",
+    "Operating System :: MacOS :: MacOS X",
+    "Operating System :: Microsoft :: Windows",
+    "Operating System :: POSIX :: Linux",
+]
+[project.urls]
+Homepage = "https://github.com/AIH-SGML/MixMIL"
+"Bug Tracker" = "https://github.com/AIH-SGML/MixMIL/issues"
+Discussions = "https://github.com/AIH-SGML/MixMIL/discussions"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["mixmil"]
+
+[tool.hatch]
+version.path = "mixmil/__init__.py"
+
+[tool.black]
+line-length = 120
+
+[tool.ruff]
+line-length = 120
+ignore = ["E741"]
diff --git a/scripts/dsmil_data_download.py b/scripts/dsmil_data_download.py
new file mode 100644
index 0000000..47ebb2f
--- /dev/null
+++ b/scripts/dsmil_data_download.py
@@ -0,0 +1,68 @@
+"""
+Script adapted from
+https://github.com/binli123/dsmil-wsi/blob/master/download.py
+
+Script and data provided by the DSMIL authors:
+
+@inproceedings{li2021dual,
+  title={Dual-stream multiple instance learning network for whole slide image classification with self-supervised contrastive learning},
+  author={Li, Bin and Li, Yin and Eliceiri, Kevin W},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={14318--14328},
+  year={2021}
+}
+"""
+
+import argparse
+import os
+import urllib.request
+import zipfile
+
+from tqdm import tqdm
+
+
+class DownloadProgressBar(tqdm):
+    def update_to(self, b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            self.total = tsize
+        self.update(b * bsize - self.n)
+
+
+def download_url(url, output_path):
+    with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=url.split("/")[-1]) as t:
+        urllib.request.urlretrieve(url, filename=output_path, reporthook=t.update_to)
+
+
+def unzip_data(zip_path, data_path):
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(data_path)
+
+
+DATASET_URLS = {
+    "camelyon16": "https://uwmadison.box.com/shared/static/l9ou15iwup73ivdjq0bc61wcg5ae8dwe.zip",
+}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, default="camelyon16", help="Dataset to be downloaded: camelyon16")
+    parser.add_argument("--keep-zip", action="store_true", help="Keep the downloaded zip file")
+    args = parser.parse_args()
+
+    assert args.dataset in DATASET_URLS, f"Dataset {args.dataset} not found"
+
+    print(f"downloading dataset: {args.dataset}")
+    unzip_dir = f"data/{args.dataset}"
+    zip_file_path = f"data/{args.dataset}-dataset.zip"
+    os.makedirs(unzip_dir, exist_ok=True)
+    download_url(DATASET_URLS[args.dataset], zip_file_path)
+    unzip_data(zip_file_path, unzip_dir)
+
+    if not args.keep_zip:
+        os.remove(f"{args.dataset}-dataset.zip")
+
+    print("done!")
+
+
+if __name__ == "__main__":
+    main()


Test patch:
diff --git a/tests/test_model.py b/tests/test_model.py
new file mode 100644
index 0000000..23a2bce
--- /dev/null
+++ b/tests/test_model.py
@@ -0,0 +1,61 @@
+import numpy as np
+import pytest
+import torch
+
+from mixmil import MixMIL
+
+
+@pytest.fixture
+def mock_data_binomial():
+    Xs = [torch.randn(10, 3) for _ in range(5)]  # List of tensors
+    F = torch.randn(5, 4)  # Fixed effects
+    Y = torch.randint(0, 2, (5, 1))  # Labels for binomial
+    return Xs, F, Y
+
+
+@pytest.fixture
+def mock_data_categorical():
+    N, Q, K = 50, 10, 4
+    bag_sizes = torch.randint(5, 15, (N,))
+    Xs = [torch.randn(bag_sizes[n], Q) for n in range(N)]  # List of tensors
+    F = torch.randn(N, K)  # Fixed effects
+    Y = torch.randint(0, 5, (N, 1))  # Labels for categorical
+    return Xs, F, Y
+
+
+def test_init_with_mean_model_binomial(mock_data_binomial):
+    Xs, F, Y = mock_data_binomial
+    model = MixMIL.init_with_mean_model(Xs, F, Y, likelihood="binomial", n_trials=2)
+    assert isinstance(model, MixMIL)
+    assert model.likelihood_name == "binomial"
+    assert model.n_trials == 2
+    assert model.log_sigma_u.numel() == 1
+
+
+def test_init_with_mean_model_categorical(mock_data_categorical):
+    Xs, F, Y = mock_data_categorical
+    model = MixMIL.init_with_mean_model(Xs, F, Y, likelihood="categorical")
+    assert isinstance(model, MixMIL)
+    assert model.likelihood_name == "categorical"
+    assert model.n_trials is None
+    assert model.log_sigma_u.numel() == len(np.unique(Y))  # separate prior for each class
+
+
+def test_initialization():
+    model = MixMIL(Q=10, K=5, P=2, likelihood="binomial", n_trials=2)
+    assert model.Q == 10
+    assert model.alpha.shape == (5, 2)
+
+
+@pytest.mark.parametrize(
+    "Q, K, P, likelihood, n_trials, mean_field",
+    [
+        (10, 5, 2, "categorical", None, True),
+        (10, 5, 2, "categorical", None, False),
+        (10, 5, 2, "binomial", 1, True),
+        (10, 5, 2, "binomial", 2, False),
+        (10, 5, 1, "binomial", 2, False),
+    ],
+)
+def test_init_model(Q, K, P, likelihood, n_trials, mean_field):
+    MixMIL(Q=Q, K=K, P=P, likelihood=likelihood, n_trials=n_trials, mean_field=mean_field)
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
new file mode 100644
index 0000000..128ae17
--- /dev/null
+++ b/tests/test_simulation.py
@@ -0,0 +1,28 @@
+import numpy as np
+import scipy.stats as st
+from sklearn.metrics import roc_auc_score
+
+from mixmil import MixMIL
+from mixmil.data import load_data
+
+
+def calc_metrics(model, X, u, w):
+    u_pred = model.predict(X["test"]).cpu().numpy().ravel()
+    w_pred = model.get_weights(X["test"])[0].cpu().numpy().ravel()
+    rho_bag = st.spearmanr(u_pred, u["test"]).correlation  # bag level correlation
+    is_top_instance = (w["test"] > np.quantile(w["test"], 0.90)).long().ravel()
+    auc_instance = roc_auc_score(is_top_instance, w_pred)  # instance-retrieval AUC
+    return rho_bag, auc_instance
+
+
+def test_simulation():
+    X, F, Y, u, w = load_data(P=1, seed=0)
+    model = MixMIL.init_with_mean_model(X["train"], F["train"], Y["train"], likelihood="binomial", n_trials=2)
+
+    start_rho_bag, start_auc_instance = calc_metrics(model, X, u, w)
+    model.train(X["train"], F["train"], Y["train"], n_epochs=40)
+    end_rho_bag, end_auc_instance = calc_metrics(model, X, u, w)
+
+    # assert that both metrics improved by at least 10%
+    assert end_rho_bag > start_rho_bag * 1.1
+    assert end_auc_instance > start_auc_instance * 1.1