numpyro_version.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2023-09-29T09:32:55.076098Z",
     "start_time": "2023-09-29T09:32:48.319503Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1695979972.307008       1 tfrt_cpu_pjrt_client.cc:349] TfrtCpuClient created.\n",
      "sample: 100%|██████████| 1500/1500 [00:01<00:00, 1008.53it/s, 63 steps of size 5.22e-02. acc. prob=0.94]\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import requests\n",
    "\n",
    "from rpy2.robjects import pandas2ri\n",
    "from rpy2.robjects.conversion import rpy2py\n",
    "import rpy2.robjects as ro\n",
    "import json\n",
    "import jax\n",
    "from collections import OrderedDict\n",
    "\n",
    "import jax.numpy as jnp\n",
    "import numpyro\n",
    "from numpyro import distributions as dist\n",
    "from numpyro.infer import MCMC, NUTS\n",
    "import requests\n",
    "from numpyro.infer import util\n",
    "\n",
    "\n",
    "\n",
    "def model(x1, x2, x3, y):\n",
    "    # Priors\n",
    "    alpha = numpyro.sample(\"alpha\", dist.Cauchy(0., 1.))\n",
    "    beta1 = numpyro.sample(\"beta1\", dist.Cauchy(0., 1.))\n",
    "    beta2 = numpyro.sample(\"beta2\", dist.Cauchy(0., 1.))\n",
    "    beta3 = numpyro.sample(\"beta3\", dist.Cauchy(0., 1.))\n",
    "\n",
    "    # Likelihood\n",
    "    mu = alpha + beta1 * x1 + beta2 * x2 + beta3 * x3\n",
    "    numpyro.sample(\"y\", dist.StudentT(df=5., loc=mu, scale=1.), obs=y)\n",
    "\n",
    "\n",
    "# Load and prepare the dataset\n",
    "url = \"https://github.com/faosorios/heavy/blob/master/data/creatinine.rda?raw=true\"\n",
    "with requests.get(url) as resp:\n",
    "    with open(\"creatinine.rda\", \"wb\") as f:\n",
    "        f.write(resp.content)\n",
    "\n",
    "# Load RDA file into Python\n",
    "ro.r['load'](\"creatinine.rda\")\n",
    "df = pandas2ri.rpy2py_dataframe(ro.r['creatinine'])\n",
    "\n",
    "data_df = pd.DataFrame(columns=['log_SC', 'log_WT', 'log_140_minus_A', 'log_CR'])\n",
    "\n",
    "# Apply transformations following https://openreview.net/pdf?id=HltJfwwfhX\n",
    "data_df['log_SC'] = np.log(df['SC'])\n",
    "data_df['log_WT'] = np.log(df['WT'])\n",
    "data_df['log_CR'] = np.log(df['CR'])\n",
    "data_df['log_140_minus_A'] = np.log(140 - df['Age'])\n",
    "data_df = data_df.dropna() # remove any rows with NaN values after transformation\n",
    "\n",
    "# Convert data to JAX array\n",
    "data_for_numpyro = {\n",
    "    'x1': jnp.array(data_df['log_SC'].values),\n",
    "    'x2': jnp.array(data_df['log_WT'].values),\n",
    "    'x3': jnp.array(data_df['log_140_minus_A'].values),\n",
    "    'y': jnp.array(data_df['log_CR'].values),\n",
    "}\n",
    "\n",
    "# Initialize NUTS sampler\n",
    "nuts_kernel = NUTS(model)\n",
    "\n",
    "# Initialize MCMC method\n",
    "mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000)\n",
    "\n",
    "# Run MCMC\n",
    "mcmc.run(jax.random.PRNGKey(0), **data_for_numpyro)\n",
    "\n",
    "# Extract samples\n",
    "mcmc_samples = mcmc.get_samples()\n",
    "\n",
    "\n",
    "# Define a function that computes the log density for a single sample\n",
    "def single_sample_log_density(sample):\n",
    "    log_density, _ = util.log_density(model, sample, **data_for_numpyro)\n",
    "    return log_density\n",
    "\n",
    "\n",
    "# Vectorize the function using vmap\n",
    "vectorized_log_density = jax.vmap(single_sample_log_density)\n",
    "\n",
    "# Convert the dictionary of samples to a suitable structure for vmap\n",
    "# samples_array_dict = {k: jnp.stack(v) for k, v in mcmc_samples.items()}\n",
    "# stacked_samples = {k: v for k, v in zip(samples_array_dict.keys(), jax.tree_multimap(lambda *x: jnp.stack(x), *samples_array_dict.values()))}\n",
    "\n",
    "# # Compute the log posterior densities for all samples\n",
    "# log_posterior_densities = vectorized_log_density(stacked_samples)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}