![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "# Initialize Code"
+ ],
+ "metadata": {
+ "id": "RRV9wE2JCZ-8"
+ }
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "ZrGngHdjR70X",
+ "outputId": "cc0489bc-704a-47fe-b2c9-6934bf6e5ffb"
+ },
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
+ "Collecting BayesLDM\n",
+ " Downloading BayesLDM-1.0.9.tar.gz (22 kB)\n",
+ " Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+ "Collecting textx\n",
+ " Downloading textX-3.1.1-py2.py3-none-any.whl (67 kB)\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.8/67.8 KB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[?25hCollecting numpyro\n",
+ " Downloading numpyro-0.11.0-py3-none-any.whl (300 kB)\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m300.2/300.2 KB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[?25hCollecting funsor\n",
+ " Downloading funsor-0.4.5-py3-none-any.whl (174 kB)\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m174.9/174.9 KB\u001b[0m \u001b[31m8.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[?25hRequirement already satisfied: numpy>=1.7 in /usr/local/lib/python3.9/dist-packages (from funsor->BayesLDM) (1.22.4)\n",
+ "Requirement already satisfied: multipledispatch in /usr/local/lib/python3.9/dist-packages (from funsor->BayesLDM) (0.6.0)\n",
+ "Collecting makefun\n",
+ " Downloading makefun-1.15.1-py2.py3-none-any.whl (22 kB)\n",
+ "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.9/dist-packages (from funsor->BayesLDM) (3.3.0)\n",
+ "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from funsor->BayesLDM) (4.5.0)\n",
+ "Requirement already satisfied: jaxlib>=0.4 in /usr/local/lib/python3.9/dist-packages (from numpyro->BayesLDM) (0.4.7+cuda11.cudnn86)\n",
+ "Requirement already satisfied: tqdm in /usr/local/lib/python3.9/dist-packages (from numpyro->BayesLDM) (4.65.0)\n",
+ "Requirement already satisfied: jax>=0.4 in /usr/local/lib/python3.9/dist-packages (from numpyro->BayesLDM) (0.4.7)\n",
+ "Collecting Arpeggio>=2.0.0\n",
+ " Downloading Arpeggio-2.0.0-py2.py3-none-any.whl (54 kB)\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m55.0/55.0 KB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[?25hRequirement already satisfied: scipy>=1.7 in /usr/local/lib/python3.9/dist-packages (from jax>=0.4->numpyro->BayesLDM) (1.10.1)\n",
+ "Requirement already satisfied: ml-dtypes>=0.0.3 in /usr/local/lib/python3.9/dist-packages (from jax>=0.4->numpyro->BayesLDM) (0.0.4)\n",
+ "Requirement already satisfied: six in /usr/local/lib/python3.9/dist-packages (from multipledispatch->funsor->BayesLDM) (1.16.0)\n",
+ "Building wheels for collected packages: BayesLDM\n",
+ " Building wheel for BayesLDM (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+ " Created wheel for BayesLDM: filename=BayesLDM-1.0.9-py3-none-any.whl size=22081 sha256=80633d3ef4b0cd6af93f2784550d74aaae30e6ff0162e604a254796f8282df9b\n",
+ " Stored in directory: /root/.cache/pip/wheels/55/65/41/3017a3b0760c4ce5a99a4c66431bbebcf3dca52c46f9f0508b\n",
+ "Successfully built BayesLDM\n",
+ "Installing collected packages: makefun, Arpeggio, textx, funsor, numpyro, BayesLDM\n",
+ "Successfully installed Arpeggio-2.0.0 BayesLDM-1.0.9 funsor-0.4.5 makefun-1.15.1 numpyro-0.11.0 textx-3.1.1\n"
+ ]
+ }
+ ],
+ "source": [
+ "!pip install BayesLDM"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {
+ "id": "xR3WiGXXIVuj"
+ },
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import matplotlib.pyplot as plt\n",
+ "import BayesLDM.BayesLDM as BayesLDM\n",
+ "import BayesLDM.utils as utils"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "EHZg_RW6aazl"
+ },
+ "source": [
+ "# Section 1. Introduction\n",
+ "\n",
+ "BayesLDM provides a domain specific probabilistic programming language that focuses on modeling multi-level data. Such data naturally arise in multiple health and social science domains where data sets contain time series data from multiple participants. BayesLDM is built on top of the NumPyro python package and enables the simplified specification of a wide range of probabilistic time series models. BayesLDM focuses on the application of Bayesian inference methods to infer posterior probability distributions over unknown model parameters. It can also automatically handle missing missing data using the same Bayesian inference approach and has the ability to simultaneously produce dustributions over missing data variables as well as over unknown model parameters. Relative to implementing models in Pyro itself, BayesLDM offers a number of optimizations to minimize inference run time, allowing users to focus on high-level modeling and not low-level implementation details. "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "nMqvZ9PTxkkk"
+ },
+ "source": [
+ "# Section 2. Defining Models with BayesLDM\n",
+ "\n",
+ "As in other probabilistic programming languages including the NumPyro language that BayesLDM is based on, BayesLDM models are specified using programs that describe a probabilistic generative process. A BayesLDM program consists of a header section specifying properties of the program, followed by the body of the program consisting of a sequence of program statements. A program is specified as a string that is passed to an interpreter. Below we describe the structure of the header and the program body."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "pWgL9JfqazYW"
+ },
+ "source": [
+ "##2.1 BayesLDM Program Header Structure\n",
+ "\n",
+ "The elements of the header section of a program are specified below.\n",
+ "\n",
+ "```\n",
+ "ProgramName: [name] \n",
+ "Indices: [index_name_1] [index_1_start] [index_1_end], ... , [index_name_N] [index_N_start] [index_N_end]\n",
+ "Inputs: [input_name_1] ... [input_name_J]\n",
+ "```\n",
+ "\n",
+ "* **ProgameName**: This property allows for specifying a name for the program. The argument [name] can be any sequence of alphanumeric characters. No white space is permitted.\n",
+ "\n",
+ "* **Indices**: In multi-level data sets, each level has a corresponding index variable. The indices property allows for declaring names and ranges of the index variables. Any number of index variables can be specified. Variables used as indices must be declared. \n",
+ "\n",
+ "* **Inputs**: The inputs property allows for declaring the names of fixed constants that can be used in a program. Any nuber of inputs can be specified. The values of input variables must be specified at the time that inference is run. If no inputs are required, this property can be ommitted."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "eahf0AYQa6Qk"
+ },
+ "source": [
+ "## 2.2 BayesLDM Program Body Structure\n",
+ "\n",
+ "The body of a BayesLDM program consists of a sequence of program statements. There are two types of program statements: assignemnt statements and sampling statements. Statements are constructed from numerical expressions and functions, which are in turn built from constants, variables and basic operators. To be valid, a BayesLDM program must contain at least one sampling statement."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Lt93WYqUZMRG"
+ },
+ "source": [
+ "### 2.2.1 Variables and Constants\n",
+ "\n",
+ "BayesLDM supports floating point and integer constants. Floating point constants are numbers with a decimal point. Integer constants are numbers without a decimal point.\n",
+ "\n",
+ "BayesLDM follows the same rules as Python for variable names. Variable names must start with an upper or lower case letter followed by a sequence of any alphanumeric characters. Variable names can contain underscores, but can not contain other punctuation characters or whitespace."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "BkXF-ba6dOXH"
+ },
+ "source": [
+ "### 2.2.2 Operators\n",
+ "\n",
+ "BayesLDM supports the basic arithmetic operations defined below. \n",
+ "\n",
+ "* Multiplication: *\n",
+ "* Division: /\n",
+ "* Addition: +\n",
+ "* Subtraction: -\n",
+ "\n",
+ "For example, the expression below scales and shifts the variable $y$ using an expression composed only of basic arithmetic operations. \n",
+ "\n",
+ "```\n",
+ "5*y+1\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "4p8xa95df8pa"
+ },
+ "source": [
+ "### 2.2.3 Numerical Functions\n",
+ "\n",
+ "BayesLDM can also leverage any numerical function imported from the JAX Numpy library. The JAX Numpy API is available here: https://jax.readthedocs.io/en/latest/jax.numpy.html. \n",
+ "\n",
+ "For example, assuming that the JAX Numpy library is imported as jnp, the statement below applies a log transform to the variable $y$. This expression will only be valid if $y$ is greater than $0$. \n",
+ "\n",
+ "```\n",
+ "jnp.log(y)\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "7M817jczasvO"
+ },
+ "source": [
+ "### 2.2.4 Numerical Expressions\n",
+ "\n",
+ "A numerical expression is a variable, a constant, or any valid composition of variables and constants defined using arithmetic operations and numerical functions that evaluates to a numerical value. \n",
+ "\n",
+ "For example, assuming the variable $y$ is defined and JAX Numpy is imported as jnp, the expression below applies and log transform to $y$ and then scales and shifts the result.\n",
+ "\n",
+ "```\n",
+ "5*jnp.log(y)+1\n",
+ "```\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "EG3By3Kga9rM"
+ },
+ "source": [
+ "\n",
+ "\n",
+ "### 2.2.5 Assignment Statements\n",
+ "\n",
+ "An assignment statement assigns a numerical value to a variable. The left hand side of the assignment statemnt must be a valid BayesLDM variable name. The right hand side must be a valid BayesLDM expression producing a numerical value. The basic structure of an assignment statement is shown below. The expression on the right hand side can reference any variables introduced on the left hand side of earlier lines in the program.\n",
+ "\n",
+ "```\n",
+ "[variable name] = [expression]\n",
+ "```\n",
+ "\n",
+ "Below, we show a basic example of assigning the constant 5 to the variable x.\n",
+ "\n",
+ "```\n",
+ "x = 5\n",
+ "```\n",
+ "\n",
+ "Assuming that the variable y is defined earlier in the program, the statement below assigns the value of the expression *5*y+1* to the variable $x$. \n",
+ "\n",
+ "```\n",
+ "x = 5*y+1\n",
+ "```\n",
+ "\n",
+ "Assuming that the JAX Numpy library is imported as jnp and the variable $y$ is previously defined, the statement below assigns the value of the log transform of $y$ to the variable $x$. This expression will only be valid if y is a positive variable.\n",
+ "\n",
+ "```\n",
+ "x = jnp.log(y)\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "jSZ3nf1da-dG"
+ },
+ "source": [
+ "### 2.2.6 Sampling Statements\n",
+ "\n",
+ "A sampling statement specifies a probability distribution for a random variable. The left hand side of the must be a valid BayesLDM variable name. The right hand side must be a valid BayesLDM probability distribution expression. A probability distribtuon expression consists of a distribution name and BayesLDM expressions for each of the distribution's parmeters. The basic structure of a sampling statement is shown below. The expressions on the right hand side can reference any variables introduced on the left hand side of earlier lines in the program.\n",
+ "\n",
+ "```\n",
+ "[variable name] ~ [Distribution Name]([expression1], ..., [expressionK])\n",
+ "```\n",
+ "\n",
+ "BayesLDM supports the following probability distributions:\n",
+ "\n",
+ "* Ber(p): The Bernoulli distribution on $\\{0,1\\}$ with $p\\in[0,1]$ as the probability of 1. \n",
+ "* Beta(a,b): The beta distribution on $[0,1]$ with parameters $a\\geq 0$, $b\\geq 0$.\n",
+ "* Exp(r): The exponential distribution on $[0,\\infty)$ with rate parameter $r>0$.\n",
+ "* Gamma(c, r): The gamma distribution on $[0,\\infty)$ with concentration parameter $c>0$ and rate parameter $r>0$.\n",
+ "* Poisson(r): The Poisson distribution on $0,1,2,...$ with rate parameter $r>=0$.\n",
+ "* N(m,s): The normal distribution on $(-\\infty, \\infty)$ with mean $m\\in\\mathbb{R}$ and standard deviation $s>0$. \n",
+ "* StudentT(d): The student's t distribution on $(-\\infty, \\infty)$ with degree of freedom parameter $d>0$. "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0rZUBWNZGTGf"
+ },
+ "source": [
+ "**Example: Standard Normal Distribution.** In this example, we construct a standard normal distribution for the variable $x$ with the mean fixed at $0$ and the standard deviation fixed at $1$. We call the BayesLDM.compile() on the program text to construct the model. We then call model.sample() to draw samples from the model. Sampling from the model produces simulations from the generative model specified by the program.\n",
+ "\n",
+ "In the call to model.sample(), a random seed is used to initialize the random number generator (by default n_seed=0. We can set n_seed=integer or chosen_rng_key=Jax random number). num_samples is the number of samples to draw. The output of running the program is the set of samples. Statistics on the samples are also shown. The default sampler used is the NumPyro \"No U-Turn Sampler\" (NUTS)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "51y0HS7L9aZz",
+ "outputId": "e104517c-74cf-4916-ff40-57af0409139b"
+ },
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stderr",
+ "text": [
+ "WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
+ ]
+ },
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "\n",
+ " mean std median 5.0% 95.0% n_eff r_hat\n",
+ " x 0.02 1.00 0.02 -1.63 1.75 376.59 1.01\n",
+ "\n",
+ "Number of divergences: 0\n"
+ ]
+ }
+ ],
+ "source": [
+ "program = \"\"\"\n",
+ "ProgramName: Normal\n",
+ "x ~ N(0,1)\n",
+ "\"\"\"\n",
+ "\n",
+ "model = BayesLDM.compile(program) \n",
+ "samples = model.sample(b_post_process=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "-TQovg5AIKfG"
+ },
+ "source": [
+ "The output of the call to model.sample() is a Python distionary with one key for each variable defined through a sampling statement. We can perform any desired analysis on the output samples. In the example below, we plot a histogram of the samples of x using MatplotLib."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 430
+ },
+ "id": "LP3t-FVKH7At",
+ "outputId": "9582506e-1ddf-459b-a188-1625f634b671"
+ },
+ "outputs": [
+ {
+ "output_type": "display_data",
+ "data": {
+ "text/plain": [
+ "| \n", + " | x | \n", + "
|---|---|
| n | \n", + "\n", + " |
| 0 | \n", + "5.626577 | \n", + "
| 1 | \n", + "4.268006 | \n", + "
| 2 | \n", + "5.877037 | \n", + "
| 3 | \n", + "3.538267 | \n", + "
| 4 | \n", + "4.953043 | \n", + "
| 5 | \n", + "3.667212 | \n", + "
| 6 | \n", + "6.642749 | \n", + "
| 7 | \n", + "5.303233 | \n", + "
| 8 | \n", + "5.022714 | \n", + "
| 9 | \n", + "4.661080 | \n", + "
| \n", + " | x1 | \n", + "
|---|---|
| n | \n", + "\n", + " |
| 0 | \n", + "-0.416758 | \n", + "
| 1 | \n", + "-0.056267 | \n", + "
| 2 | \n", + "-2.136196 | \n", + "
| 3 | \n", + "1.640271 | \n", + "
| 4 | \n", + "-1.793436 | \n", + "
| \n", + " | \n", + " | x2 | \n", + "y | \n", + "
|---|---|---|---|
| n | \n", + "t | \n", + "\n", + " | \n", + " |
| 0 | \n", + "0 | \n", + "-0.878108 | \n", + "0.644733 | \n", + "
| 1 | \n", + "-0.156434 | \n", + "-0.356260 | \n", + "|
| 2 | \n", + "0.256570 | \n", + "-1.462154 | \n", + "|
| 3 | \n", + "-0.988779 | \n", + "-1.531914 | \n", + "|
| 1 | \n", + "0 | \n", + "-0.338822 | \n", + "0.311810 | \n", + "
| \n", + " | x | \n", + "y | \n", + "
|---|---|---|
| n | \n", + "\n", + " | \n", + " |
| 0 | \n", + "0.221993 | \n", + "0.495405 | \n", + "
| 1 | \n", + "0.870732 | \n", + "1.122057 | \n", + "
| 2 | \n", + "0.206719 | \n", + "0.577291 | \n", + "
| 3 | \n", + "0.918611 | \n", + "1.271317 | \n", + "
| 4 | \n", + "0.488411 | \n", + "0.169269 | \n", + "