inverseGLM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "jax.config.update('jax_enable_x64',True)\n",
    "import jax.numpy as jnp\n",
    "import pandas as pd\n",
    "from jax.random import PRNGKey as pkey\n",
    "from jax.scipy.special import expit\n",
    "from jaxopt import LBFGS\n",
    "import statsmodels.api as sm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_true = 1.\n",
    "num_covs = 1\n",
    "num_subjects_per_group = 1000\n",
    "num_groups = 3\n",
    "total_subjects = num_subjects_per_group*num_groups\n",
    "num_exposures = num_groups -1 \n",
    "intercept_val = 0.2\n",
    "confounding_constant = 1.\n",
    "assignment = jnp.vstack([jnp.arange(num_groups)]*num_subjects_per_group).T.flatten().reshape(-1,1)\n",
    "covs = jnp.hstack([jax.random.normal(pkey(12),(total_subjects,num_covs))]) + confounding_constant*assignment\n",
    "beta_covs = jax.random.normal(pkey(13),(num_covs,))\n",
    "full_beta = jnp.hstack([beta_true,intercept_val,beta_covs])\n",
    "\n",
    "\n",
    "#This is just for simulating\n",
    "X_sim = jnp.hstack([assignment,jnp.ones((total_subjects,1)),covs])\n",
    "\n",
    "Y_obs = jax.random.binomial(pkey(1),1,expit(X_sim@full_beta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neg_log_lik(z,p,n):\n",
    "    loss = -1*n*(p*jnp.log(expit(z)) + (1-p)*jnp.log(1-expit(z)))\n",
    "    return jnp.sum(loss)\n",
    "\n",
    "@jax.jit\n",
    "def loss(beta,X,Y):\n",
    "    return neg_log_lik(X@beta,Y,1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummies = pd.get_dummies(pd.Series(assignment[:,0]),prefix='treat',drop_first=True).astype('float64')\n",
    "X_reg = jnp.hstack([jnp.ones((total_subjects,1)),covs,dummies.values])\n",
    "\n",
    "df_reg = pd.DataFrame(X_reg,columns = ['intercept']+['covariate']+list(dummies.columns))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = LBFGS(loss,jit = True)\n",
    "beta_init = jnp.zeros(len(X_reg[0]))\n",
    "result = solver.run(beta_init,X_reg,Y_obs)\n",
    "beta_reg = result.params\n",
    "L = beta_reg[-num_exposures:]\n",
    "cov_reg = jnp.linalg.inv(jax.hessian(loss,argnums=0)(result.params,X_reg,Y_obs))\n",
    "\n",
    "cov_L = cov_reg[-num_exposures:,-num_exposures:]\n",
    "V = jnp.diag(cov_reg)[-num_exposures:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#It matches statsmodels\n",
    "# model = sm.Logit(Y_obs, df_reg)\n",
    "# results = model.fit()\n",
    "# results.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up inverse GLM\n",
    "\n",
    "df_summary = pd.DataFrame({\n",
    "    \"exposure\":assignment[:,0],\n",
    "    \"Y\":Y_obs,\n",
    "    \"cov\":covs[:,0]\n",
    "})\n",
    "\n",
    "P = df_summary.groupby('exposure')['Y'].mean().values\n",
    "N = df_summary.groupby(\"exposure\").size().values\n",
    "\n",
    "inv_GLM_design = jnp.hstack(\n",
    "    [\n",
    "        jnp.ones((num_groups,1)),#Intercept for everyone\n",
    "        jnp.identity(num_groups)[:,1:]# Dummies for exposure\n",
    "        ])\n",
    "L_offset = jnp.hstack([jnp.zeros(1),L])\n",
    "\n",
    "@jax.jit\n",
    "def loss_offset(beta,X,Y,offset,N):\n",
    "    return neg_log_lik(X@beta+offset,Y,N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = LBFGS(loss_offset,jit = True,tol = 1e-5)\n",
    "beta_init = jnp.zeros(len(inv_GLM_design[0]))\n",
    "result = solver.run(beta_init,inv_GLM_design,P,L_offset,N)\n",
    "inv_glm_sol = result.params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Flip back over\n",
    "\n",
    "CC_recover_beta = jnp.hstack([inv_glm_sol[0],L])\n",
    "CC_recover_offset = jnp.hstack([jnp.zeros(1),inv_glm_sol[1:]])\n",
    "H = jax.hessian(loss_offset,argnums = 0)(CC_recover_beta,inv_GLM_design,P,CC_recover_offset,N)\n",
    "\n",
    "cov_iglm = jnp.linalg.inv(H)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.01169197, 0.00515849],\n",
       "       [0.00515849, 0.02400847]], dtype=float64)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sd_iglm = jnp.sqrt(jnp.diag(cov_iglm[-num_exposures:,-num_exposures:]))\n",
    "CC_iglm = jnp.diag(jnp.sqrt(V)/sd_iglm)@cov_iglm[-num_exposures:,-num_exposures:]@jnp.diag(jnp.sqrt(V)/sd_iglm)\n",
    "CC_iglm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.01169197, 0.00742085],\n",
       "       [0.00742085, 0.02400847]], dtype=float64)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cov_L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.00471194, 0.14893243, 0.29368091], dtype=float64)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.hstack([inv_glm_sol[0] - intercept_val,inv_glm_sol[1:]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.00890422, 0.18378993, 0.35276649])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(df_summary.groupby('exposure')['cov'].mean()*beta_covs).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.20471194, 0.93548627, 1.91708523], dtype=float64)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CC_recover_beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "pseudocov = inv_GLM_design[:,1:]@inv_glm_sol[1:]\n",
    "pseudo_design = jnp.hstack([jnp.ones((num_groups,1))*inv_glm_sol[0]+pseudocov.reshape(-1,1),inv_GLM_design[:,1:]])\n",
    "iglm_cc_beta = jnp.hstack([jnp.ones(1),L])\n",
    "\n",
    "\n",
    "H = jax.hessian(loss_offset,argnums = 0)(iglm_cc_beta,pseudo_design,P,jnp.zeros(3),N)\n",
    "alt_cov_iglm = jnp.linalg.inv(H)[-2:,-2:]\n",
    "sd_iglm_alt = jnp.sqrt(jnp.diag(alt_cov_iglm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.01169197, 0.01101055],\n",
       "       [0.01101055, 0.02400847]], dtype=float64)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CC_iglm = jnp.diag(jnp.sqrt(V)/sd_iglm_alt)@alt_cov_iglm@jnp.diag(jnp.sqrt(V)/sd_iglm_alt)\n",
    "CC_iglm"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sindy",
   "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": "undefined.undefined.undefined"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}