google-research
1912 строк · 66.3 Кб
1{2"cells": [3{4"cell_type": "markdown",5"metadata": {6"id": "Mf3kOv1YMB5y"7},8"source": [9"# Pharmacokinetics models with TensorFlow Probability\n",10"\n",11"Copyright 2021 Google LLC."12]13},14{15"cell_type": "code",16"execution_count": null,17"metadata": {18"id": "-rOdskBSMfQN"19},20"outputs": [],21"source": [22"#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",23"# you may not use this file except in compliance with the License.\n",24"# You may obtain a copy of the License at\n",25"#\n",26"# https://www.apache.org/licenses/LICENSE-2.0\n",27"#\n",28"# Unless required by applicable law or agreed to in writing, software\n",29"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",30"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",31"# See the License for the specific language governing permissions and\n",32"# limitations under the License."33]34},35{36"cell_type": "markdown",37"metadata": {38"id": "rammI_v-S2LS"39},40"source": [41"This notebook demonstrates how to fit a pharmacokinetic model with TensorFlow probability. This includes defining the relevant joint distribution and working through the basic steps of a Bayesian workflow, e.g. prior and posterior predictive checks, diagnostics for the inference, etc.\n",42"\n",43"There are three main components when building a pharmacokinetic model:\n",44"\n",45"\n",46"1. The pharmacokinetics of the system involves solving ordinary differential equations with varying levels of sophistication.\n",47"2. We need to describe the treatment the patient undergoes, using a _clinical event schedule_.\n",48"3. The data we have comes from multiple patients. To model the heterogeneity between patients, we use a hierarchical model (also termed a population model).\n",49"\n",50"We'll first tackle a one compartment model with a first-order absorption from the gut.\n",51"The ODE describing this system is simple enough to be solved analytically.\n",52"We'll start with a one-dose model for one patient and build our way up to a population model with an event schedule.\n",53"\n",54"Next on the to-do list will be expanding these models to nonlinear ODEs. Examples of ODEs that arise in PK models can be found in this [Stan notebook](https://mc-stan.org/events/stancon2017-notebooks/stancon2017-margossian-gillespie-ode.html).\n",55"\n",56"\n",57"\n",58"\n"59]60},61{62"cell_type": "markdown",63"metadata": {64"id": "wl7S_O0AJYWN"65},66"source": [67"ToDo list:\n",68"- one cpt model (with analytical solution, numerical integrator an option), one patient, one dose $\\checkmark$\n",69"- one cpt model, one patient, multiple doses $\\checkmark$\n",70"- one cpt model, multiple patients, multiple doses $\\checkmark$\n",71"- Michaelis-Mentis PK model, one patient, one dose\n",72"- Michaelis-Mentis PK model, multiple patients, multiple doses.\n",73"- Friberg-Karlsson PKPD model, one patient, multiple doses.\n",74"- Friberg-Karlsson PKPD model, multiple patients, multiple doses."75]76},77{78"cell_type": "code",79"execution_count": null,80"metadata": {81"id": "kd8S5DK8XqIT"82},83"outputs": [],84"source": [85"import tensorflow as tf\n",86"tf.executing_eagerly()\n",87"\n",88"import numpy as np\n",89"from matplotlib.pyplot import *\n",90"%config InlineBackend.figure_format = 'retina'\n",91"matplotlib.pyplot.style.use(\"dark_background\")\n",92"\n",93"import jax\n",94"from jax import random\n",95"from jax import numpy as jnp\n",96"\n",97"from colabtools import adhoc_import\n",98"\n",99"# import tensforflow_datasets\n",100"from inference_gym import using_jax as gym\n",101"\n",102"# import tensorflow as tf\n",103"from tensorflow_probability.python.internal import prefer_static as ps\n",104"from tensorflow_probability.python.internal import unnest\n",105"\n",106"import tensorflow_probability as _tfp\n",107"tfp = _tfp.substrates.jax\n",108"tfd = tfp.distributions\n",109"tfb = tfp.bijectors\n",110"\n",111"tfp_np = _tfp.substrates.numpy\n",112"tfd_np = tfp_np.distributions \n",113"\n",114"from jax.experimental.ode import odeint\n",115"from jax import vmap\n",116"\n",117"import arviz as az\n",118"from tensorflow_probability.python.internal.unnest import get_innermost"119]120},121{122"cell_type": "code",123"execution_count": null,124"metadata": {125"id": "YDoovE217kDx"126},127"outputs": [],128"source": [129"# Define nested Rhat for one parameter.\n",130"# Assume for now the indexed parameter is a scalar.\n",131"def nested_rhat(result_state, num_super_chains, index_param, num_samples,\n",132" warmup_length = 0):\n",133" state_param = result_state[index_param][\n",134" warmup_length:(warmup_length + num_samples), :, :]\n",135" num_samples = state_param.shape[0]\n",136" num_chains = state_param.shape[1]\n",137" num_sub_chains = num_chains // num_super_chains\n",138" \n",139" state_param = state_param.reshape(num_samples, -1, num_sub_chains, 1)\n",140"\n",141" mean_chain = np.mean(state_param, axis = (0, 3))\n",142" between_chain_var = np.var(mean_chain, axis = 1, ddof = 1)\n",143" within_chain_var = np.var(state_param, axis = (0, 3), ddof = 1)\n",144" total_chain_var = between_chain_var + np.mean(within_chain_var, axis = 1)\n",145"\n",146" mean_super_chain = np.mean(state_param, axis = (0, 1, 3))\n",147" between_super_chain_var = np.var(mean_super_chain, ddof = 1)\n",148"\n",149" return np.sqrt(1 + between_super_chain_var / np.mean(total_chain_var))\n",150"\n",151"# WARNING: this is a very poor estimate for ESS, and we shoud note\n",152"# W / B isn't typically used to estimate ESS.\n",153"def ess_per_super_chain(nRhat):\n",154" return 1 / (np.square(nRhat) - 1)"155]156},157{158"cell_type": "markdown",159"metadata": {160"id": "sxh29965zc71"161},162"source": [163"## 1 One compartment model with absoprtion from the gut\n",164"\n",165"A patient orally takes in a drug, which lands in the gut and is then absorbed into a central compartment (e.g. the blood). This process is described by a differential equation:\n",166"\\begin{eqnarray*}\n",167" y_0' \u0026 = \u0026 -k_0 y_0, \\\\\n",168" y_1' \u0026 = \u0026 k_0 y_0 - k_1 y_1,\n",169"\\end{eqnarray*}\n",170"with each state corresponding to the drug mass in the gut and the central compartment.\n",171"This system can be solved analytically for initial conditions $(y_0^I, y_1^I)$ at time $t = 0$:\n",172"\\begin{eqnarray*}\n",173" y_0 \u0026 = \u0026 y_0^I e^{-k_0 t}, \\\\\n",174" y_1 \u0026 = \u0026 \\frac{e^{-k_1 t}}{k_0 - k_1} \\left [ y_0^I k_0(1 - e^{(k_1 - k_0)t}) + (k_0 - k_1) y^I_1 \\right ], \n",175"\\end{eqnarray*}\n",176"provided $k_0 \\neq k_1$.\n",177"\n",178" We can also use on Jax's `odeint` and solve the equation numerically. This will set us up for more complicated problems. The data is noisy observation of $y_1$ (in practice we should use $y_1 / V$ where $V$ is the volume of the central compartment, but I'll omit this for now).\n"179]180},181{182"cell_type": "code",183"execution_count": null,184"metadata": {185"id": "M8brb74_XvoG"186},187"outputs": [],188"source": [189"# NOTE: need to pass the initial time as the first element of t.\n",190"t = np.array([0., 0.5, 0.75, 1, 1.25, 1.5, 2, 3, 4, 5, 6])\n",191"y0 = np.array([100.0, 0.0])\n",192"\n",193"theta = np.array([1.5, 0.25])\n",194"def system(state, time, theta):\n",195" k1 = theta[0]\n",196" k2 = theta[1]\n",197"\n",198" return jnp.array([\n",199" - k1 * state[0] ,\n",200" k1 * state[0] - k2 * state[1]\n",201" ])\n"202]203},204{205"cell_type": "code",206"execution_count": null,207"metadata": {208"id": "IcCdeUYS6m2b"209},210"outputs": [],211"source": [212"use_analytical_sln = True\n",213"\n",214"if (use_analytical_sln):\n",215" def ode_map(k1, k2):\n",216" sln = jnp.exp(- k2 * t) / (k1 - k2) * (y0[0] * k1 * (1 - jnp.exp((k2 - k1) * t)) + (k1 - k2) * y0[1])\n",217" return sln[1:]\n",218"else:\n",219" def ode_map(k1, k2):\n",220" theta = jnp.array([k1, k2])\n",221" return odeint(system, y0, t, theta, mxstep = 1e6)[1:, 1]"222]223},224{225"cell_type": "code",226"execution_count": null,227"metadata": {228"id": "kK0V0YgaBJ8e"229},230"outputs": [],231"source": [232"states = ode_map(k1 = theta[0], k2 = theta[1])\n",233"random.normal(random.PRNGKey(37272710), (states.shape[0],))\n",234"jnp.log(states)"235]236},237{238"cell_type": "markdown",239"metadata": {240"id": "Z_udVycUau-P"241},242"source": [243"## 1.1 Model for one patient recieving a single dose"244]245},246{247"cell_type": "code",248"execution_count": null,249"metadata": {250"id": "lYt8vpauYEv9"251},252"outputs": [],253"source": [254"# Simulate data\n",255"states = ode_map(k1 = theta[0], k2 = theta[1])\n",256"sigma = 0.1\n",257"log_y = sigma * random.normal(random.PRNGKey(37272710), (states.shape[0],)) \\\n",258" + jnp.log(states)\n",259"\n",260"y = jnp.exp(log_y)\n",261"# print(y)\n",262"\n",263"figure(figsize = [6, 6])\n",264"plot(t[1:], states)\n",265"plot(t[1:], y, 'o')\n",266"show()"267]268},269{270"cell_type": "markdown",271"metadata": {272"id": "W3l_LeXPzpBo"273},274"source": [275"### 1.1.1 Run model with TFP\n",276"The model runs faster on a CPU than a GPU, because of the ODE integrator."277]278},279{280"cell_type": "code",281"execution_count": null,282"metadata": {283"id": "v5DB0Yu_Ycy3"284},285"outputs": [],286"source": [287"model = tfd.JointDistributionSequentialAutoBatched([\n",288" # Priors\n",289" tfd.LogNormal(loc = jnp.log(1.), scale = 0.5, name = \"k1\"),\n",290" tfd.LogNormal(loc = jnp.log(1.), scale = 0.5, name = \"k2\"),\n",291" tfd.HalfNormal(scale = 1., name = \"sigma\"),\n",292"\n",293" lambda sigma, k2, k1: (\n",294" tfd.LogNormal(loc = jnp.log(ode_map(k1, k2)),\n",295" scale = sigma[..., jnp.newaxis], name = \"y\"))\n",296"])\n",297"\n",298"def target_log_prob_fn(k1, k2, sigma):\n",299" return model.log_prob((k1, k2, sigma, y))\n"300]301},302{303"cell_type": "code",304"execution_count": null,305"metadata": {306"id": "_LX3pepcZDT4"307},308"outputs": [],309"source": [310"num_dimensions = 3\n",311"def initialize (shape, key = random.PRNGKey(37272709)):\n",312" prior_location = jnp.log(jnp.array([1., 1., 1.]))\n",313" prior_scale = jnp.array([0.5, 0.5, 0.5])\n",314" return jnp.exp(prior_scale * random.normal(key, shape + (num_dimensions,)) + prior_location)\n",315"\n",316"# initial_state = initialize((4, ), key = random.PRNGKey(1954))\n",317"initial_state = model.sample(sample_shape = (4, 1), seed = random.PRNGKey(1954))[:3]"318]319},320{321"cell_type": "code",322"execution_count": null,323"metadata": {324"id": "cPCYbGWOCYh1"325},326"outputs": [],327"source": [328"x = jnp.array(initial_state).reshape(3, 4)\n",329"print(x[0, :])"330]331},332{333"cell_type": "code",334"execution_count": null,335"metadata": {336"id": "WjbHd99uZuqg"337},338"outputs": [],339"source": [340"# TODO: find a wat to do this when the init is a list!! \n",341"# Check call to target_log_prob_fn works\n",342"# target = target_log_prob_fn(initial_state)\n",343"# print(target)"344]345},346{347"cell_type": "code",348"execution_count": null,349"metadata": {350"id": "U4Kj2OBTqbKe"351},352"outputs": [],353"source": [354"# Prior predictive checks\n",355"num_prior_samples = 1000\n",356"*prior_samples, prior_predictive = model.sample(1000, seed = random.PRNGKey(37272709))"357]358},359{360"cell_type": "code",361"execution_count": null,362"metadata": {363"id": "krd1-QxmiB4K"364},365"outputs": [],366"source": [367"figure(figsize = [6, 6])\n",368"plot(t[1:], y, 'o')\n",369"plot(t[1:], np.median(prior_predictive, axis = 0), color = 'yellow')\n",370"plot(t[1:], np.quantile(prior_predictive, q = 0.95, axis = 0), linestyle = ':', color = 'yellow')\n",371"plot(t[1:], np.quantile(prior_predictive, q = 0.05, axis = 0), linestyle = ':', color = 'yellow')\n",372"show()"373]374},375{376"cell_type": "code",377"execution_count": null,378"metadata": {379"id": "sqd_nhqKo8yn"380},381"outputs": [],382"source": [383"# Implement ChEES transition kernel.\n",384"init_step_size = 1\n",385"warmup_length = 1000\n",386"\n",387"kernel = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",388"kernel = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel, warmup_length)\n",389"kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",390" kernel, warmup_length, target_accept_prob = 0.75,\n",391" reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n",392"\n",393"def trace_fn(current_state, pkr):\n",394" return (\n",395" # proxy for divergent transitions\n",396" get_innermost(pkr, 'log_accept_ratio') \u003c -1000\n",397" )"398]399},400{401"cell_type": "code",402"execution_count": null,403"metadata": {404"id": "6gdbPSgLpEuw"405},406"outputs": [],407"source": [408"num_chains = 4\n",409"\n",410"mcmc_states, diverged = tfp.mcmc.sample_chain(\n",411" num_results = 2000,\n",412" current_state = initial_state, \n",413" kernel = kernel,\n",414" trace_fn = trace_fn,\n",415" seed = random.PRNGKey(1954))"416]417},418{419"cell_type": "code",420"execution_count": null,421"metadata": {422"id": "krol-v6AEspw"423},424"outputs": [],425"source": [426"# remove warmup samples\n",427"for i in range(0, len(mcmc_states)):\n",428" mcmc_states[i] = mcmc_states[i][1000:]"429]430},431{432"cell_type": "code",433"execution_count": null,434"metadata": {435"id": "pgAXKnQ_TgXQ"436},437"outputs": [],438"source": [439"# get draws for posterior predictive checks\n",440"*_, posterior_predictive = model.sample(value = mcmc_states, \n",441" seed = random.PRNGKey(37272709))"442]443},444{445"cell_type": "markdown",446"metadata": {447"id": "_mjLAI5RvRNp"448},449"source": [450"### 1.1.2 Analyze results\n"451]452},453{454"cell_type": "code",455"execution_count": null,456"metadata": {457"id": "YEmtNcGERhRu"458},459"outputs": [],460"source": [461"print(\"Divergent transition(s):\", np.sum(diverged[1000:]))"462]463},464{465"cell_type": "markdown",466"metadata": {467"id": "GAFJ-4NK_Xwj"468},469"source": [470"To convert TFP's output to something compatible with Arviz, we'll follow the example in https://jeffpollock9.github.io/bayesian-workflow-with-tfp-and-arviz/."471]472},473{474"cell_type": "code",475"execution_count": null,476"metadata": {477"id": "fYwytRbZONKy"478},479"outputs": [],480"source": [481"parameter_names = model._flat_resolve_names()\n",482"\n",483"az_states = az.from_dict(\n",484" prior = {k: v[tf.newaxis, ...] for k, v in zip(parameter_names, prior_samples)},\n",485" posterior={\n",486" k: np.swapaxes(v, 0, 1) for k, v in zip(parameter_names, mcmc_states)\n",487" },\n",488")"489]490},491{492"cell_type": "code",493"execution_count": null,494"metadata": {495"id": "21J3YlDHQtrm"496},497"outputs": [],498"source": [499"print(az.summary(az_states).filter(items=[\"mean\", \"sd\", \"mcse_sd\", \"hdi_5%\", \n",500" \"hdi_95%\", \"ess_bulk\", \"ess_tail\", \n",501" \"r_hat\"]))"502]503},504{505"cell_type": "code",506"execution_count": null,507"metadata": {508"id": "nPzebaMXSBTO"509},510"outputs": [],511"source": [512"axs = az.plot_trace(az_states, combined = False, compact = False)"513]514},515{516"cell_type": "code",517"execution_count": null,518"metadata": {519"id": "teFRHkjhSUNu"520},521"outputs": [],522"source": [523"# TODO: include potential divergent transitions.\n",524"az.plot_pair(az_states, figsize = (6, 6), kind = 'hexbin', divergences = True);"525]526},527{528"cell_type": "code",529"execution_count": null,530"metadata": {531"id": "e0wmo3pOUTtp"532},533"outputs": [],534"source": [535"ppc_data = posterior_predictive.reshape((4000, 10))"536]537},538{539"cell_type": "code",540"execution_count": null,541"metadata": {542"id": "yBfRKltT8fD2"543},544"outputs": [],545"source": [546"figure(figsize = [6, 6])\n",547"plot(t[1:], y, 'o')\n",548"plot(t[1:], np.median(ppc_data, axis = 0), color = 'yellow')\n",549"plot(t[1:], np.quantile(ppc_data, q = 0.95, axis = 0), linestyle = ':', color = 'yellow')\n",550"plot(t[1:], np.quantile(ppc_data, q = 0.05, axis = 0), linestyle = ':', color = 'yellow')\n",551"show()"552]553},554{555"cell_type": "markdown",556"metadata": {557"id": "xeKKKn7IDYN5"558},559"source": [560"## 1.2 Clinical event schedule\n",561"\n",562"Let's now suppose the patient recieves a bolus dose every $12$ hours for a total of $15$ doses.\n",563"The first dose is administered at time $t = 0$ and the final dose at time $t = 180$ (hours).\n",564"We take many observations during the first, second and fourtennth doses. For all other dosing events, we record the drug plasma concentration at the time of the dosing event (i.e. right before the dosing), and then 6 and 12 hours after the dose is administered.\n"565]566},567{568"cell_type": "code",569"execution_count": null,570"metadata": {571"id": "zvL6F0_sELFw"572},573"outputs": [],574"source": [575"# Construct event times, and identify dosing times (all other times correspond\n",576"# to measurement events).\n",577"time_after_dose = np.array([0.083, 0.167, 0.25, 0.5, 0.75, 1, 1.5, 2, 3, 4, 6, 8])\n",578"\n",579"t = np.append(\n",580" np.append(np.append(np.append(0., time_after_dose),\n",581" np.append(12., time_after_dose + 12)),\n",582" np.linspace(start = 24, stop = 156, num = 12)),\n",583" np.append(jnp.append(168., 168. + time_after_dose),\n",584" np.array([180, 192])))\n",585"\n",586"\n",587"start_event = np.array([], dtype = int)\n",588"dosing_time = range(0, 192, 12)\n",589"\n",590"# Use dosing events to determine times of integration between\n",591"# exterior interventions on the system.\n",592"eps = 1e-4 # hack to deal with some t being slightly offset.\n",593"for t_dose in dosing_time:\n",594" start_event = np.append(start_event, np.where(abs(t - t_dose) \u003c= eps))\n",595"\n",596"amt = jnp.array([1000., 0.])\n",597"n_dose = start_event.shape[0]\n",598"\n",599"start_event = np.append(start_event, t.shape[0] - 1)"600]601},602{603"cell_type": "code",604"execution_count": null,605"metadata": {606"id": "-bwrJoy8mRmf"607},608"outputs": [],609"source": [610"def ode_map (theta, dt, current_state):\n",611" k1 = theta[0]\n",612" k2 = theta[1]\n",613" y0_hat = jnp.exp(- k1 * dt) * current_state[0]\n",614" y1_hat = jnp.exp(- k2 * dt) / (k1 - k2) * (current_state[0] * k1 *\\\n",615" (1 - jnp.exp((k2 - k1) * dt)) + (k1 - k2) * current_state[1])\n",616" return jnp.array([y0_hat, y1_hat])\n"617]618},619{620"cell_type": "code",621"execution_count": null,622"metadata": {623"id": "A5vxTyMDnl__"624},625"outputs": [],626"source": [627"ode_map(theta, np.array([1, 2, 3]), y0)[1, :]"628]629},630{631"cell_type": "markdown",632"metadata": {633"id": "Zgh3JuZda-d4"634},635"source": [636"We now wrap our ODE solver (whehter it be via an analytical solution or a numerical integrator) inside an event schedule handler. For starters, we'll go through the events using a `for` loop. This, it turns out, is fairly inefficient, and we'll later revise this code using `jax.lax.scan`."637]638},639{640"cell_type": "code",641"execution_count": null,642"metadata": {643"id": "xPM7h8KllY50"644},645"outputs": [],646"source": [647"def ode_map_event (theta):\n",648" '''\n",649" Wrapper around the ODE solver, based on the event schedule.\n",650" NOTE: if using the ode integrator, need to adjust the shape of mass.\n",651" '''\n",652" y_hat = jnp.array([])\n",653" current_state = amt\n",654" for i in range(0, n_dose):\n",655" t_integration = jax.lax.dynamic_slice(t, (start_event[i], ), \n",656" (start_event[i + 1] - start_event[i] + 1, ))\n",657" \n",658" mass = ode_map(theta, t_integration - t_integration[0], current_state)\n",659" # mass = odeint(system, current_state, t_integration,\n",660" # theta, rtol = 1e-6, atol = 1e-6, mxstep = 1e3)\n",661"\n",662" y_hat = jnp.append(y_hat, mass[1, 1:])\n",663" current_state = mass[:, mass.shape[1]] + amt\n",664" return y_hat\n",665"\n",666"y_hat = ode_map_event(theta)\n",667"log_y_hat = jnp.log(y_hat[1:])\n",668"\n",669"sigma = 0.5\n",670"# NOTE: no observation at time t = 0.\n",671"log_y = sigma * random.normal(random.PRNGKey(1954), (y_hat.shape[0],)) \\\n",672" + jnp.log(y_hat)\n",673"y_obs = jnp.exp(log_y)\n"674]675},676{677"cell_type": "code",678"execution_count": null,679"metadata": {680"id": "fEYhXQwqbo8h"681},682"outputs": [],683"source": [684"figure(figsize = [6, 6])\n",685"plot(t[1:], y_hat)\n",686"plot(t[1:], y_obs, 'o', markersize = 2)\n",687"show()"688]689},690{691"cell_type": "markdown",692"metadata": {693"id": "l4SpOYIEbvVv"694},695"source": [696"The code above works fine to simulate data but we can do better using `jax.lax.scan`."697]698},699{700"cell_type": "code",701"execution_count": null,702"metadata": {703"id": "ZCdOhAupsRjd"704},705"outputs": [],706"source": [707"t_jax = jnp.array(t)\n",708"amt_vec = np.repeat(0., t.shape[0])\n",709"amt_vec[start_event] = 1000\n",710"amt_vec[amt_vec.shape[0] - 1] = 0.\n",711"amt_vec_jax = jnp.array(amt_vec)\n",712"\n",713"# Overwrite definition of ode_map_event.\n",714"def ode_map_event(theta):\n",715" def ode_map_step (current_state, event_index):\n",716" dt = t_jax[event_index] - t_jax[event_index - 1]\n",717" y_sln = ode_map(theta, dt, current_state)\n",718" return (y_sln + jnp.array([amt_vec_jax[event_index], 0.])), y_sln[1,]\n",719"\n",720" (__, yhat) = jax.lax.scan(ode_map_step, amt, np.array(range(1, t.shape[0])))\n",721" return yhat\n"722]723},724{725"cell_type": "code",726"execution_count": null,727"metadata": {728"id": "8zp45oXI2OKY"729},730"outputs": [],731"source": [732"y_hat = ode_map_event(theta) \n",733"\n",734"figure(figsize = [6, 6])\n",735"plot(t[1:], y_hat)\n",736"plot(t[1:], y_obs, 'o', markersize = 2)\n",737"show()"738]739},740{741"cell_type": "code",742"execution_count": null,743"metadata": {744"id": "Obax1lexpON7"745},746"outputs": [],747"source": [748"# Remark: using more informative priors helps insure the chains mix\n",749"# reasonably well. (Could be interesting to examine with nested-rhat\n",750"# the case where they don't).\n",751"model = tfd.JointDistributionSequentialAutoBatched([\n",752" # Priors\n",753" tfd.LogNormal(loc = jnp.log(1.), scale = 0.5, name = \"k1\"),\n",754" tfd.LogNormal(loc = jnp.log(.5), scale = 0.25, name = \"k2\"),\n",755" tfd.HalfNormal(scale = 1., name = \"sigma\"),\n",756"\n",757" lambda sigma, k2, k1: (\n",758" tfd.LogNormal(loc = jnp.log(ode_map_event(jnp.array([k1, k2]))),\n",759" scale = sigma[..., jnp.newaxis], name = \"y_obs\"))\n",760"])\n",761"\n",762"def target_log_prob_fn(k1, k2, sigma):\n",763" return model.log_prob((k1, k2, sigma, y_obs))\n"764]765},766{767"cell_type": "code",768"execution_count": null,769"metadata": {770"id": "rm-wSAdaDJ0B"771},772"outputs": [],773"source": [774"initial_state = model.sample(sample_shape = (4, 1), seed = random.PRNGKey(1954))[:3]"775]776},777{778"cell_type": "code",779"execution_count": null,780"metadata": {781"id": "K11arQRyqdYV"782},783"outputs": [],784"source": [785"# TODO: find a way to test target_log_prob_fn with init as a list\n",786"# print(initial_state)\n",787"# target_log_prob_fn(initial_state)"788]789},790{791"cell_type": "code",792"execution_count": null,793"metadata": {794"id": "gDQfSMebpv2G"795},796"outputs": [],797"source": [798"# Implement ChEES transition kernel.\n",799"init_step_size = 0.1\n",800"warmup_length = 1000\n",801"\n",802"kernel = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",803"kernel = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel, warmup_length)\n",804"kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",805" kernel, warmup_length, target_accept_prob = 0.75,\n",806" reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n"807]808},809{810"cell_type": "code",811"execution_count": null,812"metadata": {813"id": "ARqxRVdI_tzm"814},815"outputs": [],816"source": [817"def trace_fn(current_state, pkr):\n",818" return (\n",819" # proxy for divergent transitions\n",820" get_innermost(pkr, 'log_accept_ratio') \u003c -1000,\n",821" get_innermost(pkr, 'step_size'),\n",822" get_innermost(pkr, 'max_trajectory_length')\n",823" )"824]825},826{827"cell_type": "code",828"execution_count": null,829"metadata": {830"id": "5TB2NpYjqKYZ"831},832"outputs": [],833"source": [834"num_chains = 4\n",835"\n",836"mcmc_states, diverged = tfp.mcmc.sample_chain(\n",837" num_results = 2000, \n",838" current_state = initial_state, \n",839" kernel = kernel,\n",840" trace_fn = trace_fn,\n",841" seed = random.PRNGKey(1954))"842]843},844{845"cell_type": "code",846"execution_count": null,847"metadata": {848"id": "ayEH0FkvAqk1"849},850"outputs": [],851"source": [852"semilogy(diverged[1], label = \"step size\")\n",853"semilogy(diverged[2], label = \"max_trajectory length\")\n",854"legend(loc = \"best\")\n",855"show()"856]857},858{859"cell_type": "code",860"execution_count": null,861"metadata": {862"id": "ewxRhNdRKrfR"863},864"outputs": [],865"source": [866"# remove warmup samples\n",867"for i in range(0, len(mcmc_states)):\n",868" mcmc_states[i] = mcmc_states[i][1000:]"869]870},871{872"cell_type": "markdown",873"metadata": {874"id": "N1YEdr-fEnuA"875},876"source": [877"We'll only look at some essential diagnostics. For more, we can follow the code in the single-dose example."878]879},880{881"cell_type": "code",882"execution_count": null,883"metadata": {884"id": "-TCIf1yJELrK"885},886"outputs": [],887"source": [888"print(\"Divergent transition(s):\", np.sum(diverged[1000:]))"889]890},891{892"cell_type": "code",893"execution_count": null,894"metadata": {895"id": "NKrEpAZRT8Y4"896},897"outputs": [],898"source": [899"parameter_names = model._flat_resolve_names()\n",900"\n",901"az_states = az.from_dict(\n",902" prior = {k: v[tf.newaxis, ...] for k, v in zip(parameter_names, prior_samples)},\n",903" posterior={\n",904" k: np.swapaxes(v, 0, 1) for k, v in zip(parameter_names, mcmc_states)\n",905" },\n",906")\n",907"\n",908"print(az.summary(az_states).filter(items=[\"mean\", \"sd\", \"mcse_sd\", \"hdi_3%\", \n",909" \"hdi_97%\", \"ess_bulk\", \"ess_tail\", \n",910" \"r_hat\"]))"911]912},913{914"cell_type": "code",915"execution_count": null,916"metadata": {917"id": "vFcX2gsjdQNS"918},919"outputs": [],920"source": [921"# get draws for posterior predictive checks\n",922"*_, posterior_predictive = model.sample(value = mcmc_states, \n",923" seed = random.PRNGKey(37272709))"924]925},926{927"cell_type": "code",928"execution_count": null,929"metadata": {930"id": "PB9XuVUDP-qv"931},932"outputs": [],933"source": [934"# ppc_data = posterior_predictive.reshape(1000, 4, 52)\n",935"\n",936"# az_data = az.from_dict(\n",937"# posterior = dict(x = ppc_data.transpose((1, 0, 2)))\n",938"# )\n",939"# print(az.summary(az_data).filter(items=[\"mean\", \"hdi_3%\", \n",940"# \"hdi_97%\", \"ess_bulk\", \"ess_tail\", \n",941"# \"r_hat\"]))"942]943},944{945"cell_type": "code",946"execution_count": null,947"metadata": {948"id": "RgKRN9XnFBQo"949},950"outputs": [],951"source": [952"# REMARK: hmmm... the ppc's look odd. Not sure why. Everything else looks fine.\n",953"figure(figsize = [6, 6])\n",954"semilogy(t[1:], y_obs, 'o')\n",955"semilogy(t[1:], np.median(posterior_predictive, axis = (0, 1, 2)), color = 'yellow')\n",956"semilogy(t[1:], np.quantile(posterior_predictive, q = 0.95, axis = (0, 1, 2)), linestyle = ':', color = 'yellow')\n",957"semilogy(t[1:], np.quantile(posterior_predictive, q = 0.05, axis = (0, 1, 2)), linestyle = ':', color = 'yellow')\n",958"show()"959]960},961{962"cell_type": "markdown",963"metadata": {964"id": "3u-z8HIDdnIN"965},966"source": [967"## 1.3 Population models\n",968"\n",969"We now model data from multiple patients and use a hierarchical model to describe inter-individual heterogeneity. For simplicity, we assume the patients all undergo the same treatment."970]971},972{973"cell_type": "markdown",974"metadata": {975"id": "X-q70FSapWt5"976},977"source": [978"### 1.3.1 Simulate data"979]980},981{982"cell_type": "code",983"execution_count": null,984"metadata": {985"id": "B5u_7uLf4XM5"986},987"outputs": [],988"source": [989"# (Code from previous cells, rewritten here to make\n",990"# section 1.3 self-contained).\n",991"# TODO: replace this with a function.\n",992"time_after_dose = np.array([0.083, 0.167, 0.25, 0.5, 0.75, 1, 1.5, 2, 3, 4, 6, 8])\n",993"\n",994"t = np.append(\n",995" np.append(np.append(np.append(0., time_after_dose),\n",996" np.append(12., time_after_dose + 12)),\n",997" np.linspace(start = 24, stop = 156, num = 12)),\n",998" np.append(jnp.append(168., 168. + time_after_dose),\n",999" np.array([180, 192])))\n",1000"\n",1001"start_event = np.array([], dtype = int)\n",1002"dosing_time = range(0, 192, 12)\n",1003"\n",1004"# Use dosing events to determine times of integration between\n",1005"# exterior interventions on the system.\n",1006"eps = 1e-4 # hack to deal with some t being slightly offset.\n",1007"for t_dose in dosing_time:\n",1008" start_event = np.append(start_event, np.where(abs(t - t_dose) \u003c= eps))\n",1009"\n",1010"amt = jnp.array([1000., 0.])\n",1011"n_dose = start_event.shape[0]\n",1012"\n",1013"start_event = np.append(start_event, t.shape[0] - 1)"1014]1015},1016{1017"cell_type": "code",1018"execution_count": null,1019"metadata": {1020"id": "mfRo0Ummukam"1021},1022"outputs": [],1023"source": [1024"# NOTE: need to run the first cell under Section 1.2\n",1025"# (Clinical event schedule)\n",1026"\n",1027"n_patients = 100\n",1028"pop_location = jnp.log(jnp.array([1.5, 0.25]))\n",1029"# pop_location = jnp.log(jnp.array([0.5, 1.0]))\n",1030"pop_scale = jnp.array([0.15, 0.35])\n",1031"theta_patient = jnp.exp(pop_scale * random.normal(random.PRNGKey(37272709), \n",1032" (n_patients, ) + (2,)) + pop_location)\n",1033"\n",1034"amt = np.array([1000., 0.])\n",1035"amt_patient = np.append(np.repeat(amt[0], n_patients),\n",1036" np.repeat(amt[1], n_patients))\n",1037"amt_patient = amt_patient.reshape(2, n_patients)\n",1038"\n",1039"# redfine variables from previous section (in case we only run population model)\n",1040"t_jax = jnp.array(t)\n",1041"amt_vec = np.repeat(0., t.shape[0])\n",1042"amt_vec[start_event] = 1000\n",1043"amt_vec[amt_vec.shape[0] - 1] = 0.\n",1044"amt_vec_jax = jnp.array(amt_vec)"1045]1046},1047{1048"cell_type": "markdown",1049"metadata": {1050"id": "4s06cfWJuxBS"1051},1052"source": [1053"We rewrite the ode_map, so that, rather than returning the mass for one patient, it returns the mass across multiple patients. The function `ode_map` now takes in the physiological parameters for all patients, as well as the initial states for each patient.\n",1054"\n",1055"The call to `jax.lax.scan` now takes in an additional argument, `unroll`, which is used to unroll the for loop and make its call on an accelerator more efficient. By default, `unroll = 1` (no unrolling); we observe a major speedup when using `unroll = 10`, and an additional (more minor) speedup when `unroll = 20`."1056]1057},1058{1059"cell_type": "code",1060"execution_count": null,1061"metadata": {1062"id": "7RwFyOggg_dr"1063},1064"outputs": [],1065"source": [1066"# Rewrite ode_map_event for population case.\n",1067"# TODO: remove 'use_second_axis' hack.\n",1068"def ode_map (theta, dt, current_state, use_second_axis = False):\n",1069" if (use_second_axis):\n",1070" k1 = theta[0, :]\n",1071" k2 = theta[1, :]\n",1072" else: \n",1073" k1 = theta[:, 0]\n",1074" k2 = theta[:, 1]\n",1075"\n",1076" y0_hat = jnp.exp(- k1 * dt) * current_state[0, :]\n",1077" y1_hat = jnp.exp(- k2 * dt) / (k1 - k2) * (current_state[0, :] * k1 *\\\n",1078" (1 - jnp.exp((k2 - k1) * dt)) + (k1 - k2) * current_state[1, :])\n",1079" return jnp.array([y0_hat, y1_hat])\n",1080"\n",1081"# @jax.jit # Cannot use jit if function has an IF statement.\n",1082"def ode_map_event(theta, use_second_axis = False):\n",1083" def ode_map_step (current_state, event_index):\n",1084" dt = t_jax[event_index] - t_jax[event_index - 1]\n",1085" y_sln = ode_map(theta, dt, current_state, use_second_axis)\n",1086" dose = jnp.repeat(amt_vec_jax[event_index], n_patients)\n",1087" y_after_dose = y_sln + jnp.append(jnp.repeat(amt_vec_jax[event_index], n_patients),\n",1088" jnp.repeat(0., n_patients)).reshape(2, n_patients)\n",1089" return (y_after_dose, y_sln[1, ])\n",1090"\n",1091" (__, yhat) = jax.lax.scan(ode_map_step, amt_patient, \n",1092" np.array(range(1, t.shape[0])),\n",1093" unroll = 20)\n",1094" return yhat"1095]1096},1097{1098"cell_type": "code",1099"execution_count": null,1100"metadata": {1101"id": "qdAE6KxseIY5"1102},1103"outputs": [],1104"source": [1105"# Simulate some data\n",1106"y_hat = ode_map_event(theta_patient)\n",1107"\n",1108"sigma = 0.1\n",1109"# NOTE: no observation at time t = 0.\n",1110"log_y = sigma * random.normal(random.PRNGKey(1954), y_hat.shape) \\\n",1111" + jnp.log(y_hat)\n",1112"y_obs = jnp.exp(log_y)\n",1113"\n",1114"figure(figsize = [6, 6])\n",1115"plot(t[1:], y_hat)\n",1116"plot(t[1:], y_obs, 'o', markersize = 2)\n",1117"show()"1118]1119},1120{1121"cell_type": "markdown",1122"metadata": {1123"id": "GcgditOnpfOJ"1124},1125"source": [1126"### 1.3.2 Fit the model with TFP "1127]1128},1129{1130"cell_type": "markdown",1131"metadata": {1132"id": "Ypgnnwm0CxKl"1133},1134"source": [1135"This is an adaptation of the previous model, except we're now only working with parameters on the unconstrained scale. This makes it easier for HMC and it is good practice."1136]1137},1138{1139"cell_type": "code",1140"execution_count": null,1141"metadata": {1142"id": "P_N-3eRl6HNa"1143},1144"outputs": [],1145"source": [1146"pop_model = tfd.JointDistributionSequentialAutoBatched([\n",1147" # tfd.LogNormal(loc = jnp.log(1.), scale = 0.25, name = \"k1_pop\"),\n",1148" # tfd.LogNormal(loc = jnp.log(0.3), scale = 0.1, name = \"k2_pop\"),\n",1149" # tfd.Normal(loc = jnp.log(1.), scale = 0.25, name = \"log_k1_pop\"),\n",1150" tfd.Normal(loc = jnp.log(1.), scale = 0.1, name = \"log_k1_pop\"),\n",1151" tfd.Normal(loc = jnp.log(0.3), scale = 0.1, name = \"log_k2_pop\"),\n",1152" tfd.Normal(loc = jnp.log(0.15), scale = 0.1, name = \"log_scale_k1\"),\n",1153" tfd.Normal(loc = jnp.log(0.35), scale = 0.1, name = \"log_scale_k2\"),\n",1154" # tfd.HalfNormal(scale = 1., name = \"sigma\"),\n",1155" tfd.Normal(loc = -1., scale = 1., name = \"log_sigma\"),\n",1156"\n",1157" # non-centered parameterization for hierarchy\n",1158" tfd.Independent(tfd.Normal(loc = jnp.zeros(n_patients),\n",1159" scale = jnp.ones(n_patients),\n",1160" name = \"eta_k1\"),\n",1161" reinterpreted_batch_ndims = 1),\n",1162" \n",1163" tfd.Independent(tfd.Normal(loc = jnp.zeros(n_patients),\n",1164" scale = jnp.ones(n_patients),\n",1165" name = \"eta_k2\"),\n",1166" reinterpreted_batch_ndims = 1),\n",1167"\n",1168" lambda eta_k2, eta_k1, log_sigma, log_scale_k2, log_scale_k1,\n",1169" log_k2_pop, log_k1_pop: (\n",1170" tfd.Independent(tfd.LogNormal(\n",1171" loc = jnp.log(\n",1172" ode_map_event(theta = jnp.array([\n",1173" jnp.exp(log_k1_pop[..., jnp.newaxis] + eta_k1 * jnp.exp(log_scale_k1[..., jnp.newaxis])),\n",1174" jnp.exp(log_k2_pop[..., jnp.newaxis] + eta_k2 * jnp.exp(log_scale_k2[..., jnp.newaxis]))]),\n",1175" use_second_axis = True)),\n",1176" scale = jnp.exp(log_sigma[..., jnp.newaxis]), name = \"y_obs\")))\n",1177"\n",1178" # lambda eta_k2, eta_k1, sigma, log_scale_k2, log_scale_k1,\n",1179" # k2_pop, k1_pop: (\n",1180" # tfd.Independent(tfd.LogNormal(\n",1181" # loc = jnp.log(\n",1182" # ode_map_event(theta = jnp.array(\n",1183" # [jnp.exp(jnp.log(k1_pop[..., jnp.newaxis]) + eta_k1 * jnp.exp(log_scale_k1[..., jnp.newaxis])),\n",1184" # jnp.exp(jnp.log(k2_pop[..., jnp.newaxis]) + eta_k2 * jnp.exp(log_scale_k2[..., jnp.newaxis]))]),\n",1185" # use_second_axis = True)),\n",1186" # scale = sigma[..., jnp.newaxis], name = \"y_obs\")))\n",1187"])\n",1188"\n",1189"def pop_target_log_prob_fn(log_k1_pop, log_k2_pop, log_scale_k1, log_scale_k2,\n",1190" log_sigma, eta_k1, eta_k2):\n",1191" return pop_model.log_prob((log_k1_pop, log_k2_pop, log_scale_k1, log_scale_k2,\n",1192" log_sigma, eta_k1, eta_k2, y_obs))\n",1193" # CHECK -- do we need to parenthesis?\n",1194"\n",1195"\n",1196"\n",1197"# def pop_target_log_prob_fn(k1_pop, k2_pop, log_scale_k1, log_scale_k2,\n",1198"# sigma, eta_k1, eta_k2):\n",1199"# return pop_model.log_prob((k1_pop, k2_pop, log_scale_k1, log_scale_k2,\n",1200"# sigma, eta_k1, eta_k2, y_obs))\n",1201"\n",1202"def pop_target_log_prob_fn_flat(x):\n",1203" k1_pop = x[:, 0]\n",1204" k2_pop = x[:, 1]\n",1205" sigma = x[:, 2]\n",1206" log_scale_k1 = x[:, 3]\n",1207" log_scale_k2 = x[:, 4]\n",1208" eta_k1 = x[:, 5:(5 + n_patients)]\n",1209" eta_k2 = x[:, (5 + n_patients):(5 + 2 * n_patients)]\n",1210"\n",1211" return pop_model.log_prob((k1_pop, k2_pop, log_scale_k1, log_scale_k2,\n",1212" sigma, eta_k1, eta_k2, y_obs))\n"1213]1214},1215{1216"cell_type": "markdown",1217"metadata": {1218"id": "7x630Zb1DHuF"1219},1220"source": [1221"If we want to run many chains in parallel and use $n\\hat R$ (nested $\\hat R$), we need to specify the number of chains and the number of super chains.\n",1222"The number of super chains determined the numbers of distinct starting point, seeing within each super chain, each chain starts at the same location. "1223]1224},1225{1226"cell_type": "code",1227"execution_count": null,1228"metadata": {1229"id": "-R1KuCCyXDVY"1230},1231"outputs": [],1232"source": [1233"# Sample initial states from prior\n",1234"num_chains = 128\n",1235"num_super_chains = 4 # num_chains # 128\n",1236"\n",1237"n_parm = 5 + 2 * n_patients\n",1238"initial_state_raw = pop_model.sample(sample_shape = (num_super_chains, 1),\\\n",1239" seed = random.PRNGKey(37272710))[:7]\n",1240"\n",1241"# QUESTION: does this assignment create a pointer?\n",1242"initial_state = initial_state_raw\n",1243"\n",1244"for i in range(0, len(initial_state_raw)):\n",1245" initial_state[i] = np.repeat(initial_state_raw[i],\n",1246" num_chains // num_super_chains, axis = 0)\n"1247]1248},1249{1250"cell_type": "markdown",1251"metadata": {1252"id": "0J3pKThK1yE5"1253},1254"source": [1255"Some care is required when setting the tuning parameters for ChEES-HMC, in particular the initial step size. In the [ChEES-HMC paper](http://proceedings.mlr.press/v130/hoffman21a/hoffman21a.pdf), the following proceudre is used: \"Initial step sizes were chosen by repeatedly halving the step size (starting from a consistently too-large value of 1.0) until an HMC proposal with a single leapfrog step achieved a harmonic-mean acceptance probability of at least 0.5.\"\n",1256"\n",1257"TODO: implement this.\n",1258"\n",1259"For now we note that an \"appropriate\" initial step size depends on the number chains (for reasons I don't quite understand...)."1260]1261},1262{1263"cell_type": "code",1264"execution_count": null,1265"metadata": {1266"id": "aXS4RkqKr_v8"1267},1268"outputs": [],1269"source": [1270"# Implement ChEES transition kernel. Increase the target acceptance rate\n",1271"# to avoid divergent transitions.\n",1272"# NOTE: increasing the target acceptance probability can lead to poor performance.\n",1273"init_step_size = 0.001 # CHECK -- how to best tune this?\n",1274"warmup_length = 1000 # 1000\n",1275"\n",1276"kernel = tfp.mcmc.HamiltonianMonteCarlo(pop_target_log_prob_fn, \n",1277" step_size = init_step_size, \n",1278" num_leapfrog_steps = 10)\n",1279"kernel = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel, warmup_length)\n",1280"kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",1281" kernel, warmup_length, target_accept_prob = 0.75,\n",1282" reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n",1283"\n",1284"def trace_fn(current_state, pkr):\n",1285" return (\n",1286" # proxy for divergent transitions\n",1287" get_innermost(pkr, 'log_accept_ratio') \u003c -1000,\n",1288" get_innermost(pkr, 'step_size'),\n",1289" get_innermost(pkr, 'max_trajectory_length')\n",1290" )"1291]1292},1293{1294"cell_type": "code",1295"execution_count": null,1296"metadata": {1297"id": "2hEuLZccsdY4"1298},1299"outputs": [],1300"source": [1301"mcmc_states, diverged = tfp.mcmc.sample_chain(\n",1302" num_results = warmup_length + 1000,\n",1303" current_state = initial_state,\n",1304" kernel = kernel,\n",1305" trace_fn = trace_fn,\n",1306" seed = random.PRNGKey(1954))\n"1307]1308},1309{1310"cell_type": "code",1311"execution_count": null,1312"metadata": {1313"id": "s_t77idfg9SW"1314},1315"outputs": [],1316"source": [1317"# Remark: somehow modifying mcmc_states still modifies\n",1318"# mcmc_states_raw.\n",1319"mcmc_states_raw = mcmc_states"1320]1321},1322{1323"cell_type": "markdown",1324"metadata": {1325"id": "E1CCgq0jppah"1326},1327"source": [1328"### 1.3.3 Traditional diagnostics"1329]1330},1331{1332"cell_type": "code",1333"execution_count": null,1334"metadata": {1335"id": "Ura3AS9Nv3EB"1336},1337"outputs": [],1338"source": [1339"# remove warmup samples\n",1340"# NOTE: not a good idea. It's better to store all the states.\n",1341"if False:\n",1342" for i in range(0, len(mcmc_states)):\n",1343" mcmc_states[i] = mcmc_states_raw[i][warmup_length:]"1344]1345},1346{1347"cell_type": "code",1348"execution_count": null,1349"metadata": {1350"id": "_Alb3e0QCd-b"1351},1352"outputs": [],1353"source": [1354"semilogy(diverged[1], label = \"step size\")\n",1355"semilogy(diverged[2], label = \"max_trajectory length\")\n",1356"legend(loc = \"best\")\n",1357"show()"1358]1359},1360{1361"cell_type": "code",1362"execution_count": null,1363"metadata": {1364"id": "JgsVYXdSn0XE"1365},1366"outputs": [],1367"source": [1368"mcmc_states[0].shape"1369]1370},1371{1372"cell_type": "code",1373"execution_count": null,1374"metadata": {1375"id": "6F-lSkMjd86x"1376},1377"outputs": [],1378"source": [1379"# Use this to search for points where the the step size changes\n",1380"# dramatically and divergences that might be happening there.\n",1381"if False:\n",1382" index_l = 219 # 225\n",1383" index_u = index_l + 1 # 235\n",1384" print(\"Max L:\" , diverged[2][index_l:index_u])\n",1385" print(\"Divergence:\", np.sum(diverged[0][index_l:index_u]),\n",1386" \"at\", np.where(diverged[0][index_l:index_u] == 1))\n",1387"\n",1388" chain = 0\n",1389" eta1_state = mcmc_states[5][index_l, chain, :] *\\\n",1390" mcmc_states[2][index_l, chain, 0] + mcmc_states[0][index_l, chain, 0]\n",1391" eta2_state = mcmc_states[6][index_l, chain, :] *\\\n",1392" mcmc_states[3][index_l, chain, 0] + mcmc_states[1][index_l, chain, 0]\n",1393"\n",1394" k0_state = np.exp(eta1_state)\n",1395" k1_state = np.exp(eta2_state) \n",1396" print(k0_state - k1_state)"1397]1398},1399{1400"cell_type": "code",1401"execution_count": null,1402"metadata": {1403"id": "ju1sBMuVX6Nn"1404},1405"outputs": [],1406"source": [1407"print(\"Divergent transition(s):\", np.sum(diverged[0][warmup_length:]))"1408]1409},1410{1411"cell_type": "code",1412"execution_count": null,1413"metadata": {1414"id": "NCAQdcFQYPz9"1415},1416"outputs": [],1417"source": [1418"# NOTE: the last parameter is an 'x': not sure where this comes from...\n",1419"parameter_names = pop_model._flat_resolve_names()[:-1]\n",1420"\n",1421"az_states = az.from_dict(\n",1422" #prior = {k: v[tf.newaxis, ...] for k, v in zip(parameter_names, prior_samples)},\n",1423" posterior={\n",1424" k: np.swapaxes(v, 0, 1) for k, v in zip(parameter_names, mcmc_states)\n",1425" },\n",1426")\n",1427"\n",1428"print(az.summary(az_states).filter(items=[\"mean\", \"sd\", \"mcse_sd\", \"hdi_3%\", \n",1429" \"hdi_97%\", \"ess_bulk\", \"ess_tail\", \n",1430" \"r_hat\"]))"1431]1432},1433{1434"cell_type": "code",1435"execution_count": null,1436"metadata": {1437"id": "PwRaHRUIZBKU"1438},1439"outputs": [],1440"source": [1441"# Only plot the population parameters.\n",1442"axs = az.plot_trace(az_states, combined = False, compact = False,\n",1443" var_names = parameter_names[:5])"1444]1445},1446{1447"cell_type": "code",1448"execution_count": null,1449"metadata": {1450"id": "Edxek3KkHgr0"1451},1452"outputs": [],1453"source": [1454"# posterior predictive checks\n",1455"# NOTE: for 100 patients, running this exhausts memory\n",1456"*_, posterior_predictive = pop_model.sample(value = mcmc_states, \n",1457" seed = random.PRNGKey(37272709))\n",1458"ppc_data = posterior_predictive.reshape(1000 * num_chains, 52, n_patients)"1459]1460},1461{1462"cell_type": "code",1463"execution_count": null,1464"metadata": {1465"id": "kobBLgFmIHHR"1466},1467"outputs": [],1468"source": [1469"# NOTE: unclear why the confidence interval is so small...\n",1470"fig, axes = subplots(n_patients, 1, figsize=(8, 4 * n_patients))\n",1471"\n",1472"for i in range(0, n_patients):\n",1473" patient_ppc = posterior_predictive[:, :, :, :, i]\n",1474" axes[i].semilogy(t[1:], y_obs[:, i], 'o')\n",1475" axes[i].semilogy(t[1:], np.median(patient_ppc, axis = (0, 1, 2)), color = 'yellow')\n",1476" axes[i].semilogy(t[1:], np.quantile(patient_ppc, q = 0.95, axis = (0, 1, 2)), linestyle = ':', color = 'yellow')\n",1477" axes[i].semilogy(t[1:], np.quantile(patient_ppc, q = 0.05, axis = (0, 1, 2)), linestyle = ':', color = 'yellow')\n",1478"show()"1479]1480},1481{1482"cell_type": "markdown",1483"metadata": {1484"id": "0EpU8FfApLKM"1485},1486"source": [1487"### 1.3.3 Diagnostic using $n \\hat R$."1488]1489},1490{1491"cell_type": "markdown",1492"metadata": {1493"id": "4lbjEnANvPsK"1494},1495"source": [1496"For starters, let's examine estimates in the short regime, i.e. using only the first few iterations from each chain. We'll focus on $\\log k_{1,\\text{pop}}$ which seems to have the most difficult expectation value to estimate (given it's relatively low ESS)."1497]1498},1499{1500"cell_type": "code",1501"execution_count": null,1502"metadata": {1503"id": "4KS_r5r6vOqb"1504},1505"outputs": [],1506"source": [1507"# Assumes mcmc_states contains all the samples (including warmup)\n",1508"parameter_index = 0\n",1509"num_samples = 500\n",1510"mc_mean = np.mean(mcmc_states[parameter_index][\n",1511" warmup_length:(warmup_length + num_samples), :, :])\n",1512"\n",1513"print(\"Mean:\", mc_mean)\n",1514"print(\"Estimated squared error:\",\n",1515" np.square(mc_mean -\n",1516" np.mean(mcmc_states[parameter_index][warmup_length:, :, :])))\n",1517"print(\"Upper bound on expected squared error for one iteration:\",\n",1518" np.var(mcmc_states[0]) / num_chains)"1519]1520},1521{1522"cell_type": "code",1523"execution_count": null,1524"metadata": {1525"id": "lQajSPNT7ggR"1526},1527"outputs": [],1528"source": [1529"nRhat = nested_rhat(result_state = mcmc_states, \n",1530" num_super_chains = num_super_chains,\n",1531" index_param = parameter_index, \n",1532" num_samples = num_samples,\n",1533" warmup_length = warmup_length)\n",1534"\n",1535"print(\"num_samples:\", num_samples)\n",1536"print(\"nRhat:\", nRhat)\n",1537"print(\"Rhat:\",\n",1538" tfp.mcmc.potential_scale_reduction(\n",1539" mcmc_states[0][warmup_length:(num_samples + warmup_length), :, :]))"1540]1541},1542{1543"cell_type": "markdown",1544"metadata": {1545"id": "eKflNcfnC4l1"1546},1547"source": [1548"## 2 Michaelis-Menten pharmacokinetics (Incomplete)\n",1549"\n"1550]1551},1552{1553"cell_type": "markdown",1554"metadata": {1555"id": "4AlJ32QWJ5y5"1556},1557"source": [1558"Nonlinear PK model with absorption from the gut.\n",1559"\n",1560"\\begin{eqnarray*}\n",1561" y_0' \u0026 = \u0026 - k_a y_0 \\\\\n",1562" y_1' \u0026 = \u0026 k_a y_0 - \\frac{V_m C}{K_m + C},\n",1563"\\end{eqnarray*}\n",1564"whwre $C = y_1 / V$."1565]1566},1567{1568"cell_type": "code",1569"execution_count": null,1570"metadata": {1571"id": "WvvAVKmyD7OJ"1572},1573"outputs": [],1574"source": [1575"t = np.array([0.0, 0.5, 0.75, 1, 1.25, 1.5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100])\n",1576"y0 = np.array([100.0, 0.0])\n",1577"theta = np.array([0.5, 27, 10, 14])\n",1578"\n",1579"def system(state, time, theta):\n",1580" ka = theta[0]\n",1581" V = theta[1]\n",1582" Vm = theta[2]\n",1583" Km = theta[3]\n",1584" C = state[1] / V\n",1585"\n",1586" return jnp.array([\n",1587" - ka * state[0],\n",1588" ka * state[0] - Vm * C / (Km + C) \n",1589" ])\n",1590"\n",1591"states = odeint(system, y0, t, theta, mxstep = 1000)\n",1592"sigma = 0.5\n",1593"log_y = sigma * random.normal(random.PRNGKey(37272709), (states.shape[0] - 1,)) \\\n",1594" + jnp.log(states[1:, 1])\n",1595"\n",1596"y = jnp.exp(log_y)\n",1597"\n",1598"figure(figsize = [6, 6])\n",1599"plot(t[1:], states[1:, 1])\n",1600"plot(t[1:], y, 'o');"1601]1602},1603{1604"cell_type": "code",1605"execution_count": null,1606"metadata": {1607"id": "ZJDzcZ7sE6Ec"1608},1609"outputs": [],1610"source": [1611"def ode_map(ka, V, Vm, Km):\n",1612" theta = jnp.array([ka, V, Vm, Km])\n",1613" return odeint(system, y0, t, theta, mxstep = 1e3)[1:, 1]\n",1614"\n",1615"model = tfd.JointDistributionSequentialAutoBatched([\n",1616" # Priors\n",1617" tfd.LogNormal(loc = jnp.log(1), scale = 0.5, name = \"ka\"),\n",1618" tfd.LogNormal(loc = jnp.log(35), scale = 0.5, name = \"V\"),\n",1619" tfd.LogNormal(loc = jnp.log(10), scale = 0.5, name = \"Vm\"),\n",1620" tfd.LogNormal(loc = jnp.log(2.5), scale = 1, name = \"Km\"),\n",1621" tfd.HalfNormal(scale = 1., name = \"sigma\"),\n",1622"\n",1623" # Likelihood (TODO: divide location by volume to get concentration)\n",1624" lambda sigma, Km, Vm, V, ka: (\n",1625" tfd.LogNormal(loc = jnp.log(ode_map(ka, V, Vm, Km) / V),\n",1626" scale = sigma[..., jnp.newaxis], name = \"y\"))\n",1627"])\n",1628"\n",1629"def target_log_prob_fn(x):\n",1630" ka = x[:, 0]\n",1631" V = x[:, 1]\n",1632" Vm = x[:, 2]\n",1633" Km = x[:, 3]\n",1634" sigma = x[:, 4]\n",1635" return model.log_prob((ka, V, Vm, Km, sigma, y))\n",1636"\n",1637"num_dimensions = 5\n",1638"def initialize (shape, key = random.PRNGKey(37272709)):\n",1639" prior_location = jnp.log(jnp.array([1.5, 35, 10, 2.5, 0.5]))\n",1640" prior_scale = jnp.array([3, 0.5, 0.5, 3, 1.])\n",1641" return jnp.exp(prior_scale * random.normal(key, shape + (num_dimensions,)) + prior_location)\n",1642"\n",1643"initial_state = initialize((4, ), key = random.PRNGKey(1954))\n"1644]1645},1646{1647"cell_type": "code",1648"execution_count": null,1649"metadata": {1650"id": "lqbZweOXF0cE"1651},1652"outputs": [],1653"source": [1654"# Test target probability density can be computed\n",1655"target = target_log_prob_fn(initial_state)\n",1656"print(target)"1657]1658},1659{1660"cell_type": "code",1661"execution_count": null,1662"metadata": {1663"id": "rcrUavhHJAke"1664},1665"outputs": [],1666"source": [1667"# Implement ChEES transition kernel.\n",1668"init_step_size = 1\n",1669"warmup_length = 250\n",1670"\n",1671"kernel = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",1672"kernel = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel, warmup_length)\n",1673"kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",1674" kernel, warmup_length, target_accept_prob = 0.75,\n",1675" reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n"1676]1677},1678{1679"cell_type": "code",1680"execution_count": null,1681"metadata": {1682"id": "FyoGwFy9Gsb7"1683},1684"outputs": [],1685"source": [1686"num_chains = 4\n",1687"\n",1688"# NOTE: It takes 29 seconds to run one iteration. So running 500 iterations\n",1689"# would take ~4 hours :(\n",1690"# QUESTION: why does JAX struggle so much to solve this type of problem??\n",1691"result = tfp.mcmc.sample_chain(\n",1692" num_results = 1, \n",1693" current_state = initial_state, \n",1694" kernel = kernel,\n",1695" seed = random.PRNGKey(1954))"1696]1697},1698{1699"cell_type": "markdown",1700"metadata": {1701"id": "Z5WsAlMA0ZBl"1702},1703"source": [1704"## Draft Code\n"1705]1706},1707{1708"cell_type": "code",1709"execution_count": null,1710"metadata": {1711"id": "cwZzVUChAus5"1712},1713"outputs": [],1714"source": [1715"R = 1.62\n",1716"1 / (R * R - 1)"1717]1718},1719{1720"cell_type": "code",1721"execution_count": null,1722"metadata": {1723"id": "CUk2qgx5GtLe"1724},1725"outputs": [],1726"source": [1727"a = np.array(range(4, 1024, 4))\n",1728"d = np.repeat(6., len(a))\n",1729"\n",1730"# Two optimization solutions, solving quadratic equations (+ / -)\n",1731"# Remark: + solution gives a negative upper-bound for delta_u\n",1732"alpha_1 = 2 * a + d / 2 - np.sqrt(np.square(2 * a + d / 2) - 2 * a)\n",1733"alpha_2 = a - alpha_1\n",1734"delta_u = (np.square(alpha_1 + d / 2) / (alpha_1 * alpha_2)) / 2"1735]1736},1737{1738"cell_type": "code",1739"execution_count": null,1740"metadata": {1741"id": "yi6HD2znsYIK"1742},1743"outputs": [],1744"source": [1745"eps = 0.01\n",1746"delta = np.square(1 + eps) - 1\n",1747"print(delta)"1748]1749},1750{1751"cell_type": "code",1752"execution_count": null,1753"metadata": {1754"id": "vyqywrFuq6hH"1755},1756"outputs": [],1757"source": [1758"semilogy(a / d, delta_u)\n",1759"hlines(delta, (a / d)[0], (a / d)[len(a) - 1], linestyles = '--',\n",1760" label = \"delta for 1.01 threshold\")\n",1761"xlabel(\"a / d\")"1762]1763},1764{1765"cell_type": "code",1766"execution_count": null,1767"metadata": {1768"id": "KOxZiFUGidI3"1769},1770"outputs": [],1771"source": [1772"semilogy(a / d, alpha_1 / a, label = \"alpha_1\")\n",1773"semilogy(a / d, alpha_2 / a, label = \"alpha_2\")\n",1774"legend(loc = 'best')\n",1775"xlabel(\"a / d\")\n",1776"ylabel(\"alpha\")"1777]1778},1779{1780"cell_type": "code",1781"execution_count": null,1782"metadata": {1783"id": "RxO-KhRocA13"1784},1785"outputs": [],1786"source": [1787"aindex_location = np.where(a / d == 100)\n",1788"print(index_location)\n",1789"print(delta_u[index_location])\n",1790"delta"1791]1792},1793{1794"cell_type": "code",1795"execution_count": null,1796"metadata": {1797"id": "-ItlDq-o0dBd"1798},1799"outputs": [],1800"source": [1801"pop_model = tfd.JointDistributionSequentialAutoBatched([\n",1802" tfd.LogNormal(loc = jnp.log(1.), scale = 0.5, name = \"k1_pop\"),\n",1803" tfd.LogNormal(loc = jnp.log(.5), scale = 0.25, name = \"k2_pop\"),\n",1804" tfd.Normal(loc = jnp.log(0.5), scale = 1., name = \"log_scale_k1\"),\n",1805" tfd.Normal(loc = jnp.log(0.5), scale = 1., name = \"log_scale_k2\"),\n",1806" tfd.HalfNormal(scale = 1., name = \"sigma\"),\n",1807"\n",1808" # non-centered parameterization for hierarchy\n",1809" tfd.Independent(tfd.Normal(loc = jnp.zeros(n_patients),\n",1810" scale = jnp.ones(n_patients),\n",1811" name = \"eta_k1\"),\n",1812" reinterpreted_batch_ndims = 1),\n",1813" \n",1814" tfd.Independent(tfd.Normal(loc = jnp.zeros(n_patients),\n",1815" scale = jnp.ones(n_patients),\n",1816" name = \"eta_k2\"),\n",1817" reinterpreted_batch_ndims = 1),\n",1818"\n",1819" lambda eta_k2, eta_k1, sigma, log_scale_k2, log_scale_k1,\n",1820" k2_pop, k1_pop: (\n",1821" tfd.Independent(tfd.LogNormal(\n",1822" loc = jnp.log(\n",1823" ode_map_event(theta = jnp.array(\n",1824" [jnp.exp(jnp.log(k1_pop[..., jnp.newaxis]) + eta_k1 * jnp.exp(log_scale_k1[..., jnp.newaxis])),\n",1825" jnp.exp(jnp.log(k2_pop[..., jnp.newaxis]) + eta_k2 * jnp.exp(log_scale_k2[..., jnp.newaxis]))]),\n",1826" use_second_axis = True)),\n",1827" scale = sigma[..., jnp.newaxis], name = \"y_obs\")))\n",1828"])\n"1829]1830},1831{1832"cell_type": "code",1833"execution_count": null,1834"metadata": {1835"id": "6zYQufGa2xtm"1836},1837"outputs": [],1838"source": [1839"num_hyper = 5\n",1840"num_dimensions = num_hyper + 2 * n_patients\n",1841"\n",1842"def pop_initialize(shape, key = random.PRNGKey(37272710)) :\n",1843" # init for k1_pop, k2_pop, and sigma\n",1844" hyper_prior_location = jnp.array([jnp.log(1.5), jnp.log(0.25), 0.])\n",1845" hyper_prior_scale = jnp.array([0.5, 0.1, 0.5])\n",1846" init_hyper_param = jnp.exp(hyper_prior_scale * random.normal(key, shape + \\\n",1847" (3, )) + hyper_prior_location)\n",1848"\n",1849" # init for log_scale_k1 and log_scale_k2\n",1850" scale_prior_location = jnp.array([-1., -1.])\n",1851" scale_prior_scale = jnp.array([0.25, 0.25])\n",1852" init_scale = scale_prior_scale * random.normal(key, shape + (2, )) +\\\n",1853" scale_prior_location\n",1854"\n",1855" # inits for the etas\n",1856" init_eta = random.normal(key, shape + (2 * n_patients, ))\n",1857" return jnp.append(jnp.append(init_hyper_param, init_scale, axis = 1), \n",1858" init_eta, axis = 1)\n",1859"\n",1860"initial_state = pop_initialize((4, ))"1861]1862},1863{1864"cell_type": "code",1865"execution_count": null,1866"metadata": {1867"id": "T3Zh9soc2t-M"1868},1869"outputs": [],1870"source": [1871"initial_list = [initial_state[:, 0], # k1_pop\n",1872" initial_state[:, 1], # k2_pop\n",1873" initial_state[:, 2], # log_scale_k1\n",1874" initial_state[:, 3], # log_scale_k2\n",1875" initial_state[:, 4], # sigma\n",1876" initial_state[:, 5:(5 + n_patients)], # eta_k1\n",1877" initial_state[:, (5 + n_patients):(5 + 2 * n_patients)] # eta_k2 \n",1878" ]"1879]1880}1881],1882"metadata": {1883"colab": {1884"collapsed_sections": [1885"_mjLAI5RvRNp",1886"eKflNcfnC4l1"1887],1888"last_runtime": {1889"build_target": "//learning/deepmind/dm_python:dm_notebook3_tpu",1890"kind": "private"1891},1892"name": "PKmodels.ipynb",1893"private_outputs": true,1894"provenance": [1895{1896"file_id": "1RrCAGXR3VT2pJEqMbTYWPNQSnyIWraV4",1897"timestamp": 16323442089821898}1899],1900"toc_visible": true1901},1902"kernelspec": {1903"display_name": "Python 3",1904"name": "python3"1905},1906"language_info": {1907"name": "python"1908}1909},1910"nbformat": 4,1911"nbformat_minor": 01912}1913