example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "julia EM model fitting example, Nathaniel Daw 8/2020\n",
    "\n",
    "####### NOTE NOTE NOTE: PARALLEL COMPUTATION IS NOW AUTOMATIC IN THIS VERSION \n",
    "####### BUT TO RUN PARALLEL YOU MUST SET ENVIRONMENT VARIABLE JULIA_NUM_THREADS  \n",
    "####### BEFORE STARTING JULIA OR JUPYTER-NOTEBOOK\n",
    "\n",
    "eg in linux/bash:\n",
    "      export JULIA_NUM_THREADS=`nproc`; julia\n",
    "\n",
    "or just run julia with --threads=auto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m    Updating\u001b[22m\u001b[39m git-repo `https://github.com/ndawlab/em.git`\n",
      "\u001b[32m\u001b[1m   Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `C:\\Users\\ndaw\\.julia\\environments\\v1.9\\Project.toml`\n",
      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `C:\\Users\\ndaw\\.julia\\environments\\v1.9\\Manifest.toml`\n"
     ]
    }
   ],
   "source": [
    "# to install the package (only need to do this once):\n",
    "# (note this fails with a weird error if you run this notebook from this directory)\n",
    "import Pkg\n",
    "Pkg.add(url=\"https://github.com/ndawlab/em.git/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "###### setup \n",
    "\n",
    "full = false    # Maintain full covariance matrix (vs a diagional one) at the group level\n",
    "emtol = 1e-3    # stopping condition (relative change) for EM\n",
    "\n",
    "# load EM package\n",
    "using EM\n",
    "\n",
    "# this loads additional packages used in examples below\n",
    "# install them with install\n",
    "\n",
    "using Statistics\n",
    "using Random\n",
    "using GLM\n",
    "using DataFrames\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "###### Q learning example\n",
    "\n",
    "# simulate some  qlearning data\n",
    "\n",
    "Random.seed!(1234); # (for repeatability)\n",
    "\n",
    "NS = 250;\n",
    "NT = 200;\n",
    "NP = 2;\n",
    "\n",
    "params = zeros(NS,NP);\n",
    "\n",
    "cov = randn(NS); # simulated between-subject variable, e.g. age or IQ\n",
    "cov = cov .- mean(cov);\n",
    "\n",
    "cov2 = randn(NS); # simulated between-subject variable, e.g. age or IQ\n",
    "cov2 = cov2 .- mean(cov2);\n",
    "\n",
    "# subject level parameters\n",
    "\n",
    "params[:,1] = 1 .+ 0.5 * randn(NS) + cov; # softmax  temp: mean 1, effect of cov\n",
    "params[:,2] = 0 .+ 1 * randn(NS) + cov2;  # learning rate (pre sigmoidal transform): mean 0, effect of cov2\n",
    "\n",
    "c = zeros(Int64,NS*NT);\n",
    "r = zeros(Int64,NS*NT);\n",
    "s = zeros(Int64,NS*NT);\n",
    "\n",
    "for i = 1:NS\n",
    "\t(c[(i-1)*NT+1:i*NT],r[(i-1)*NT+1:i*NT]) = simq(params[i,:],NT);\n",
    "\ts[(i-1)*NT+1:i*NT] .= i;\n",
    "end\n",
    "\n",
    "data = DataFrame(sub=s,c=c,r=r);\n",
    "subs = 1:NS;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up the fit\n",
    "\n",
    "# design matrix specifying the group level model\n",
    "# this is replicated once for each model parameter\n",
    "#\n",
    "# in particular for each subject-level parameter x_ij  (subject i, parameter j)\n",
    "#\n",
    "# x_ij ~ Normal(X beta_j, Sigma)\n",
    "#\n",
    "# thus X has a row for each subject and a column for each predictor\n",
    "# in the simplest case where the only predictor is an intercept, X = ones(NS)\n",
    "# then beta_j specifies the group-level mean for parameter j\n",
    "#\n",
    "# but in this example we have two covariates that vary by subject\n",
    "# so x_ij = beta_1j + beta_2j * cov_i + beta_3j * cov2_i\n",
    "# and we infer the slopes beta for each parameter j as well as the intercept\n",
    "#\n",
    "# so we have a design matrix with 3 columns, and a row per subject:\n",
    "\n",
    "X = [ones(NS) cov cov2];\n",
    "\n",
    "# note: when you have no covariates (only intercepts) omit the brackets to get a column vector\n",
    "\n",
    "# X = ones(NS)\n",
    "\n",
    "# starting points for group level parameters\n",
    "# betas: one column for each parameter, one row for each regressor (so here: 3 rows, 2 columns)\n",
    "# make sure these are floats\n",
    "# note: if you have a single predictor you need a row vector (length: # params)\n",
    "# eg betas = [0. 0.];\n",
    "# and if there is also only a single model parameter and no covariates, then betas is a scalar\n",
    "# eg betas = 0.\n",
    "\n",
    "startbetas = [1. 0; 0 0; 0 0]\n",
    "\n",
    "# sigma: one element starting variance for each model parameter (this is really variance not SD)\n",
    "# if there is only one model parameter it needs to be a length-one vector eg. sigma = [5.]\n",
    "\n",
    "startsigma = [5., 1];\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Estimation & significance tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "iter: 12\n",
      "betas: [0.95 -0.08; 0.94 0.01; -0.02 0.85]\n",
      "sigma: [0.29, 0.64]\n",
      "free energy: -22912.710021\n",
      "change: [4.0e-5, -0.000122, 6.3e-5, 0.00053, -0.000175, 2.0e-6, 0.000114, 9.0e-6]\n",
      "max: 0.00053\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "3×2 Matrix{Float64}:\n",
       "  0.947158  -0.0813925\n",
       "  0.940341   0.0115734\n",
       " -0.016061   0.850036"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# fit the model\n",
    "# (this takes: a data frame, a list of subjects, a group level design matrix, \n",
    "#  starting group level betas, starting group-level variance or covariance, a likelihood function\n",
    "#  and some optional options)\n",
    "#\n",
    "# (return values: betas are the group level means and slopes\n",
    "#  sigma is the group level *variance* or covariance\n",
    "#  x is a matrix of MAP/empirical Bayes per-subject parameters\n",
    "#  l is the per-subject negative log likelihoods \n",
    "#  h is the *inverse* per subject hessians) \n",
    "\n",
    "(betas,sigma,x,l,h) = em(data,subs,X,startbetas,startsigma,qlik; emtol=emtol, full=full);\n",
    "\n",
    "betas # ground truth would be [1 0 ; 1 0; 0 1] so this is closei"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×2 adjoint(::Matrix{Float64}) with eltype Float64:\n",
       " 4.8045e-87   0.16772\n",
       " 2.30126e-78  0.827098\n",
       " 0.658954     3.8568e-36"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# standard errors on the subject-level means, based on an asymptotic Gaussian approx \n",
    "# (these may be inflated for small n)\n",
    "# returns standard errors, pvalues, and a covariance matrix \n",
    "# these are a vector ordered as though the betas matrix were read out column-wise\n",
    "# eg parameter 1, (intercept covariate covariate) then parameter 2\n",
    "\n",
    "(standarderrors,pvalues,covmtx) = emerrors(data,subs,x,X,h,betas,sigma,qlik)\n",
    "reshape(pvalues,size(betas'))'  # cov1 (2nd row) is significant for beta (first column)\n",
    "# & cov2 (3rd row) is significant for alpha (second column)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "iter: 12\n",
      "betas: [0.9 -0.04]\n",
      "sigma: [1.17, 1.11]\n",
      "Warning: Omitting from LML 5 subjects with non-invertible Hessian\n",
      "free energy: -23783.131705\n",
      "change: [0.0, -4.0e-6, 1.0e-6, 0.0]\n",
      "max: 4.0e-6\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "StatsModels.TableRegressionModel{LinearModel{GLM.LmResp{Vector{Float64}}, GLM.DensePredChol{Float64, LinearAlgebra.CholeskyPivoted{Float64, Matrix{Float64}, Vector{Int64}}}}, Matrix{Float64}}\n",
       "\n",
       "beta ~ 1 + cov + cov2\n",
       "\n",
       "Coefficients:\n",
       "────────────────────────────────────────────────────────────────────────────\n",
       "                   Coef.  Std. Error      t  Pr(>|t|)   Lower 95%  Upper 95%\n",
       "────────────────────────────────────────────────────────────────────────────\n",
       "(Intercept)   0.898618     0.0364278  24.67    <1e-67   0.82687    0.970367\n",
       "cov           0.82229      0.034146   24.08    <1e-66   0.755036   0.889545\n",
       "cov2         -0.00558035   0.0357861  -0.16    0.8762  -0.0760652  0.0649045\n",
       "────────────────────────────────────────────────────────────────────────────"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "StatsModels.TableRegressionModel{LinearModel{GLM.LmResp{Vector{Float64}}, GLM.DensePredChol{Float64, LinearAlgebra.CholeskyPivoted{Float64, Matrix{Float64}, Vector{Int64}}}}, Matrix{Float64}}\n",
       "\n",
       "alpha ~ 1 + cov + cov2\n",
       "\n",
       "Coefficients:\n",
       "────────────────────────────────────────────────────────────────────────────\n",
       "                   Coef.  Std. Error      t  Pr(>|t|)   Lower 95%  Upper 95%\n",
       "────────────────────────────────────────────────────────────────────────────\n",
       "(Intercept)  -0.0436486    0.0493965  -0.88    0.3778  -0.140941   0.0536435\n",
       "cov          -0.00676074   0.0463024  -0.15    0.8840  -0.0979586  0.0844371\n",
       "cov2          0.572827     0.0485264  11.80    <1e-25   0.477249   0.668405\n",
       "────────────────────────────────────────────────────────────────────────────"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# another way to get a p value for a covariate, by omitting it from the model and regressing\n",
    "# this seems to work better when full=false\n",
    "# in general not super well justified and can clearly be biased in some cases\n",
    "# but works well in practice as long as you avoid the bias cases (which are pretty obvious)\n",
    "\n",
    "X2 = ones(NS);\n",
    "startbetas2 = [0. 0.];\n",
    "startsigma2 = [5., 1];\n",
    "(betas2,sigma2,x2,l2,h2) = em(data,subs,X2,startbetas2,startsigma2,qlik; emtol=1e-5, full=full);\n",
    "\n",
    "display(lm(@formula(beta~cov+cov2),DataFrame(beta=x2[:,1],cov=cov,cov2=cov2)))\n",
    "display(lm(@formula(alpha~cov+cov2),DataFrame(alpha=x2[:,2],cov=cov,cov2=cov2)))\n",
    "\n",
    "# again the first covariate is significant for beta and the second for alpha"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model comparison metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "22331.093777006565"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# laplace approximation to the aggregate log marginal likelihood of the whole dataset\n",
    "# marginalized over the individual params\n",
    "\n",
    "ll1 = lml(x,l,h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "22374.372890144205"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# to compare these between models you need to correct for the group level free parameters\n",
    "# either aic or bic (this is Quentin Huys' IBIC or IAIC, i.e. the subject level\n",
    "# params are marginalized by laplace approx, and aggregated, and the group level\n",
    "# params are corrected by AIC or BIC)\n",
    "\n",
    "ibic(x,l,h,betas,sigma,NS*NT)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "22339.093777006565"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iaic(x,l,h,betas,sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "iter: 4\n",
      "betas: [0.94 -0.07; 0.94 0.02; -0.02 0.85]\n",
      "sigma: [0.29, 0.62]\n",
      "free energy: -22894.045173\n",
      "change: [9.9e-5, -0.000891, 0.000159, 7.4e-5, -0.000796, 7.0e-6, 0.000266, 0.000403]\n",
      "max: 0.000891\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "22341.82677834569"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# or by computing unbiased per subject marginal likelihoods via cross validation.\n",
    "# you can do paired t tests on these between models\n",
    "# these are also appropriate for SPM_BMS etc\n",
    "\n",
    "liks = loocv(data,subs,x,X,betas,sigma,qlik;emtol=emtol, full=full)\n",
    "sum(liks)\n",
    "\n",
    "# note that iaic does an excellent job of predicting the aggregate held out likelihood\n",
    "# but importantly these are per subject scores that you can compare in paired tests\n",
    "# across models as per Stephan et al. random effects model comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.1",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}