martingales.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bf2daabd-1463-4363-a862-5f4f21c327ef",
   "metadata": {},
   "source": [
    "# Supermartingale-based Tests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54df817a-e17c-4369-971c-8da3523f6797",
   "metadata": {},
   "source": [
    "## Martingales, supermartingales, submartingales\n",
    "\n",
    "Let $(X_j)_{j \\in \\mathbb{N}} = X_1, X_2, \\ldots$ and \n",
    "$(Y_j)_{j \\in \\mathbb{N}} = Y_1, Y_2, \\ldots$ be sequences of random variables (stochastic processes).\n",
    "Let $Y^{i} := (Y_1, \\ldots, Y_i)$ be the first $i$ elements of the sequence $(Y_j)_{j \\in \\mathbb{N}}$.\n",
    "Suppose $\\mathbb{E}|X_j| < \\infty$ for all $j \\in \\mathbb{N}$.\n",
    "\n",
    "Then $(X_j)_{j \\in \\mathbb{N}}$ is a _martingale_ with respect to $(Y_j)_{j \\in \\mathbb{N}}$ if \n",
    "\n",
    "\\begin{equation} \\mathbb{E}(X_j | Y^{j-1}) = X_{j-1}. \\end{equation}\n",
    "\n",
    "It is a _supermartingale_ if \n",
    "\\begin{equation} \\mathbb{E}(X_j | Y^{j-1}) \\le  X_{j-1}, \\end{equation}\n",
    "and a _submartingale_ if\n",
    "\\begin{equation} \\mathbb{E}(X_j | Y^{j-1}) \\ge  X_{j-1}. \\end{equation}\n",
    "\n",
    "A martingale, supermartingale, or submartingale is _nonnegative_ if $\\mathbb{P} \\{X_j \\ge 0 \\} = 1$ for all \n",
    "$j \\in \\mathbb{N}$.\n",
    "\n",
    "Terminology: A sequence $(\\lambda_i)_{i \\in \\mathbb{N}}$ is _predictable_ with respect to $(Y_j)_{j \\in \\mathbb{N}}$\n",
    "if $\\lambda_j$ does not depend on $(Y_i)_{i>j}$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974c833e-0c58-4967-89de-ef0204d9b62f",
   "metadata": {},
   "source": [
    "## Examples\n",
    "\n",
    "### Product of independent random variables with expected value 1.\n",
    "\n",
    "Suppose $(Y_j)_{j \\in \\mathbb{N}}$ are independent and have expected value 1.\n",
    "Define $X_j := \\prod_{i=1}^j Y_i$, $j \\in \\mathbb{N}$.\n",
    "Then\n",
    "\\begin{equation}\n",
    "    \\mathbb{E} (X_j | Y^{j-1}) = \\mathbb{E} (X_{j-1}Y_j | Y^{j-1}) =\n",
    "    X_{j-1} \\mathbb{E} (Y_j | Y^{j-1}) =  X_{j-1} \\mathbb{E} (Y_j) =  X_{j-1} \\cdot 1 = X_{j-1}.\n",
    "\\end{equation}\n",
    "\n",
    "Note that if $(Z_i)$ are independent random variables, then $\\prod_{i=1}^j Z_i/\\mathbb{E}Z_i$ is a martingale (if $\\mathbb{E} Z_i \\ne 0$).\n",
    "\n",
    "### Sum of independent random variables with expected value 0.\n",
    "\n",
    "Suppose $(Y_j)_{j \\in \\mathbb{N}}$ are independent and have expected value 0.\n",
    "Define $X_j := \\sum_{i=1}^j Y_j$, $j \\in \\mathbb{N}$.\n",
    "Then\n",
    "\\begin{equation}\n",
    "    \\mathbb{E} (X_j | Y^{j-1}) = \\mathbb{E} (X_{j-1}+Y_j | Y^{j-1}) =\n",
    "    X_{j-1} + \\mathbb{E} (Y_j | Y^{j-1}) =  X_{j-1} + \\mathbb{E} (Y_j ) = \n",
    "     X_{j-1} + 0 = X_{j-1}.\n",
    "\\end{equation}\n",
    "\n",
    "### Gambler's fortune in repeated bets on independent fair coin tosses.\n",
    "\n",
    "Suppose we bet on the outcomes of a sequence of independent fair coin tosses $Y_1, Y_2, \\ldots$, i.e., $\\{Y_j\\}_{j \\in \\mathbb{N}}$ are IID $\\mbox{Bernoulli}(1/2)$.\n",
    "Our initial fortune $X_0$; our fortune after the $j$th wager is $X_j$.\n",
    "We are not allowed to wager more than our current fortune on the next bet.\n",
    "If $X_i$ is zero then $X_k = 0$ for all $k > i$: we can't bet again if we go broke.\n",
    "The wager on the $j$ toss is $b_j \\in [0, X_{j-1}]$; $b_j$ may depend on $Y^{j-1}$.\n",
    "If $Y_j=1$, we win the $j$th bet and our fortune\n",
    "increases by $b_j$;\n",
    "if $Y_j=0$, we lose the $j$th bet and our fortune decreases by $b_j$.\n",
    "Then $X_j = X_{j-1}+b_jY_j - b_j(1-Y_j)$.\n",
    "\n",
    "\\begin{equation}\n",
    " \\mathbb{E} (X_j | Y^{j-1}) = X_{j-1}+ b_j \\mathbb{E} (2Y_j-1) = X_{j-1}.\n",
    "\\end{equation}\n",
    "Thus the fortune $(X_j)_{j \\in \\mathbb{N}}$ is a nonnegative martingale with respect to the Bernoulli \n",
    "coin tosses $(Y_j)_{j \\in \\mathbb{N}}$.\n",
    "\n",
    "### Compensated process\n",
    "\n",
    "Let $\\{Y_i\\}_{i \\in \\mathbb{N}}$ be IID with equal chance of being equal to $\\pm 1$, so $\\mathbb{E}Y_i = 0$ and\n",
    "$\\mathbb{E}Y_i^2 = 1$.\n",
    "Let\n",
    "$X_j := \\sum_{i \\le j} Y_i$, and $Z_j := X_j^2-j$.\n",
    "That is, $Z_j$ is the square of a gambler's fortune in a series of 1-unit bets on a series of independent, fair coin tosses, minus the number of tosses.\n",
    "\\begin{eqnarray}\n",
    " \\mathbb{E} (Z_j | Y^{j-1}) &=& \\mathbb{E} ((X_{j-1}+Y_j)^2 - j  | Y^{j-1}) \\\\\n",
    " &=& \\mathbb{E} (X_{j-1}^2+2X_{j-1}Y_j + Y_j^2 - j  | Y^{j-1}) \\\\\n",
    " &=& X_{j-1}^2 - j + \\mathbb{E}Y_j^2 \\\\\n",
    " &=& X_{j-1}^2 - (j-1) \\\\\n",
    " &=& Z_{j-1}.\n",
    "\\end{eqnarray}\n",
    "\n",
    "\n",
    "### Likelihood ratios\n",
    "\n",
    "We have a family $\\mathcal{P} := \\{ \\mathbb{P}_\\theta : \\theta \\in \\Theta \\}$\n",
    "of probability distributions on a measurable space $\\mathcal{X}$, dominated by a\n",
    "common $\\sigma$-finite measure $\\mu$. \n",
    "Let $p_\\theta(x)$, $x \\in \\mathcal{X}$ denote the density of $P_\\theta$ with respect to $\\mu$. \n",
    "For fixed $x \\in \\mathcal{X}$, the function\n",
    "\\begin{eqnarray}\n",
    "  \\mathcal{L} &=& \\mathcal{L}_x : \\Theta \\mapsto \\Re^+  \\\\\n",
    "  \\theta &\\rightarrow& p_\\theta(x)\n",
    "\\end{eqnarray}\n",
    "is the _likelihood_.\n",
    "If we observe $X=x$, then for $\\eta \\in \\Theta$, \n",
    "$\\mathcal{L}(\\eta) = \\mathcal{L}(\\eta | X=x)$\n",
    "is the _likelihood of $\\eta$ given_ $X=x$.\n",
    "The likelihood ratio of $\\eta$ to $\\nu$ given $X=x$ is $\\mathcal{L}(\\eta | X=x)/\\mathcal{L}(\\nu | X=x)$.\n",
    "\n",
    "Suppose $\\mathbb{P}_1$ and $\\mathbb{P}_2$ are distributions on some outcome space and that they have densities\n",
    "with respect to some dominating measure $\\mu$. \n",
    "Let $f$ and $g$ denote those densities. \n",
    "Suppose $(Y_j)_{j \\in \\mathbb{N}}$ are IID with density $f$. \n",
    "Define the likelihood ratio process\n",
    "\\begin{equation}\n",
    "  L_j := \\prod_{i=1}^j g(Y_i)/f(Y_i), \\;\\; j \\in \\mathbb{N}.\n",
    "\\end{equation}\n",
    "Now \n",
    "\\begin{equation}\n",
    "\\mathbb{E} (L_j | Y^{j-1}) = L_{j-1} \\mathbb{E}(g(Y_j)/f(Y_j)) \n",
    "    = L_{j-1} \\int_\\Omega \\frac{g(\\omega)}{f(\\omega)} f(\\omega) d\\mu(\\omega) \n",
    "    = L_{j-1} \\int_\\Omega g(\\omega) d\\mu(\\omega) = L_{j-1} \n",
    "\\end{equation}\n",
    "since $g$ is a probability density with respect to $\\mu$.\n",
    "\n",
    "### Prior-posterior ratio (PPR) martingale\n",
    "\n",
    "See [Waudby-Smith and Ramdas (2021)](https://arxiv.org/pdf/2006.04347.pdf).\n",
    "\n",
    "Again, we have a family of distributions $\\{\\mathbb{P}_\\theta : \\theta \\in \\Theta\\}$.\n",
    "Let $\\Theta$ be countable and suppose\n",
    "$\\{\\mathbb{P}_\\theta : \\theta \\in \\Theta\\}$ are dominated by a common sigma-finite measure $\\mu$,\n",
    "where $p_\\theta$ is the density of $\\mathbb{P}_\\theta$ with respect to $\\mu$.\n",
    "Let $\\pi$ (the _prior_) be a probability distribution on $\\Theta$.\n",
    "We observe $(X_i)_{i \\in \\mathbb{N}}$ sequentially, with $X_1 \\sim p_\\theta$ and $X_{j+1} \\sim p_\\theta(x | X^j)$,\n",
    "for all $j \\in \\mathbb{N}$, for some fixed $\\theta \\in \\Theta$.\n",
    "The _predictive distribution of $x^j = (x_1, \\ldots, x_j)$_ is \n",
    "\\begin{equation}\n",
    "    m(x^j) := \\int_\\eta p_\\eta(x^j) \\pi(d\\eta).\n",
    "\\end{equation}\n",
    "The _posterior_ distribution of $\\theta$ given $X^j=x^j$ is\n",
    "\\begin{equation}\n",
    "   \\pi_j(\\theta) = \\pi(d\\theta | X^j) := \\frac{p_\\theta(x) \\pi(\\theta)}{m(x^j)}.\n",
    "\\end{equation}\n",
    "\n",
    "Consider the sequence\n",
    "\\begin{equation}\n",
    "    Y_j(\\theta) := \\frac{\\pi(\\theta)}{\\pi_j(\\theta)}, \\;\\; j \\in \\mathbb{N}.\n",
    "\\end{equation}\n",
    "Then $Y_j$ is nonnegative.\n",
    "Suppose the true value of $\\theta$ is $\\eta$ and that $\\pi$ assigns positive mass to $\\eta$.\n",
    "Then $(Y_j(\\mu))$ is a nonnegative martingale with expected value 1.\n",
    "\n",
    "### Sampling with or without replacement from a finite population  \n",
    "\n",
    "Suppose we have a population $\\{x_i\\}_{i=1}^N$.\n",
    "Let $\\theta = \\bar{x} := \\frac{1}{N} \\sum_{i=1}^N x_i$ be the population mean.\n",
    "Let $X_k$ be the value of $x_i$ selected on the $k$th draw. \n",
    "Let $X^j := (X_1, \\ldots, X_j)$.\n",
    "Let $\\theta_j := \\mathbb{E}(X_j | X^{j-1})$. \n",
    "For sampling with replacement, $\\theta_j = \\theta$.\n",
    "Define $S_j := \\sum_{i=1}^j X_j$.\n",
    "For sampling without replacement, $\\theta_j = (N\\theta - S_{j-1})/(N-j)$: the population total was originally\n",
    "$N\\theta$, from which items totalling $S_{j-1}$ have been removed, leaving $(N-j)$ items in the population.\n",
    "\n",
    "Then $\\sum_{i=1}^j (X_i-\\theta_i)$ and $\\prod_{i=1}^j X_i/\\theta_i$ (if $\\theta_i \\ne 0$) are martingales.\n",
    "\n",
    "### Betting martingales\n",
    "\n",
    "[Waudby-Smith and Ramdas (2021)](https://arxiv.org/pdf/2010.09686.pdf)\n",
    "\n",
    "\\begin{equation}\n",
    "    Y_j(\\eta) := \\prod_{i=1}^j (1 + \\lambda_i(\\eta) ) (X_i-\\eta_i)\n",
    "\\end{equation}\n",
    "with $\\eta_i := \\mathbb{E}_\\eta (X_i | X^{i-1})$, $Y_0 := 1$, $\\lambda_i \\in [-1/(1-\\eta_i), 1/\\eta_i]$ predictable.\n",
    "\n",
    "Initial fortune of 1.\n",
    "Bet on the value of $\\theta_i = \\mathbb{E}X_i$.\n",
    "Positive $\\lambda_i$ pays if $X_i > \\eta_i$; negative $\\lambda_i$ pays if $X_i < \\eta_i$.\n",
    "The constraint on $\\lambda$ ensures the gambler can't go into debt:\n",
    "the fortune is always nonnegative.\n",
    "The value of $\\lambda_i$ is related to the fraction of the current fortune the\n",
    "gambler wagers on the next draw.\n",
    "\n",
    "**Terminology.**\n",
    "\n",
    "A stochastic process $(Y_j)_{j \\in \\mathbb{N}}$ is a _test (super)martingale_ for the distribution\n",
    "$\\mathbb{P}$ if it is a nonnegative (super)martingale with respect to $(X_j) \\sim \\mathbb{P}$\n",
    "with $Y_0 = 1$.\n",
    "If $(Y_j)_{j \\in \\mathbb{N}}$ is a test (super)martingale for all $\\mathbb{P} \\in \\mathcal{P}$, it\n",
    "is a test (super)martingale for $\\mathcal{P}$.\n",
    "\n",
    "Waudby-Smith and Ramdas Proposition 3:  \n",
    "For any arbitrary set of (possibly unbounded) distributions $\\mathcal{P}$, every\n",
    "test (super)martingale is necessarily of the form \n",
    "\\begin{equation}\n",
    "  Y_j = \\prod_{i=1}^j (1 + \\lambda_i Z_i)\n",
    "\\end{equation}\n",
    "for some predictable $\\lambda_i \\le 1$ and \n",
    "some $Z_i \\ge -1$, $\\mathbb{P}$-almost surely, \n",
    "and $\\mathbb{E}_{\\theta}(Z_i | X^{i-1}) = (\\le) 0$ for every $\\mathbb{P} \\in \\mathcal{P}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e6375a5-1048-4ada-b7c5-7113ff6b2c8f",
   "metadata": {},
   "source": [
    "## Properties of martingales, supermartingales, and submartingales\n",
    "\n",
    "+ If $(X_j)_{j \\in \\mathbb{N}}$ is a martingale, then $\\mathbb{E} X_j = \\mathbb{E} X_1$.\n",
    "+ If $(X_j)_{j \\in \\mathbb{N}}$ is a supermartingale, then $\\mathbb{E} X_j \\le \\mathbb{E} X_1$.\n",
    "+ If $(X_j)_{j \\in \\mathbb{N}}$ is a submartingale, then $\\mathbb{E} X_j \\ge \\mathbb{E} X_1$.\n",
    "\n",
    "__Proof.__\n",
    "Suppose $(X_j)_{j \\in \\mathbb{N}}$ is a martingale w.r.t. $(Y_j)_{j \\in \\mathbb{N}}$.\n",
    "Then, by the [law of total expectation](./math-inequalities.ipynb#law_of_total_expectation), \n",
    "\\begin{equation} \\mathbb{E}X_j = \\mathbb{E}\\mathbb{E}(X_j | Y^{j-1}) = \\mathbb{E} X_{j-1}. \\end{equation}\n",
    "\\begin{equation} \\mathbb{E}X_{j-1} = \\mathbb{E}\\mathbb{E}(X_{j-1} | Y^{j-2}) = \\mathbb{E} X_{j-2}. \\end{equation}\n",
    "Continuing, we end up with $\\mathbb{E}X_1$.\n",
    "\n",
    "For supermartingales and submartingales, $=$ is replaced with $\\le$ or $\\ge$, respectively.\n",
    "$\\Box$.\n",
    "\n",
    "\n",
    "A concave (convex) function of a martingale is a super- (sub-)submartingale. (Proof by Jensen's inequality.)\n",
    "\n",
    "The sum of a countable number of (super|sub)martingales (with respect to the same stochastic process or filtration) \n",
    "is also a (super|sub)martingale with respect to\n",
    "the same stochastic process, by the linearity of conditional expectation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06f7ac61-d4dc-4152-85ae-50621ade217c",
   "metadata": {},
   "source": [
    "## Stopping times\n",
    "\n",
    "Let $\\tau$ be a measurable mapping from $X_1, X_2, \\ldots$ to $\\mathbb{N} \\cup \\infty$\n",
    "and suppose that for $j=1, 2, \\ldots$, the event $\\tau = j$ depends only on $X^j = (X_1, \\ldots, X_j)$,\n",
    "and not on $X_k$ for $k>j$. (The usual definition of stopping times involves _filtrations_, \n",
    "which are beyond the scope of these materials.)\n",
    "Then $\\tau$ is a _stopping time_ for $X$.\n",
    "\n",
    "In the example where $X_j$ represents a sequence of wagers, some stopping rules are:\n",
    "+ stop after 10 bets\n",
    "+ stop if you go broke\n",
    "+ stop if your fortune reaches $100\n",
    "+ stop if you double your initial stake\n",
    "+ stop if you win three bets in a row\n",
    "\n",
    "These are _not_ stopping rules:\n",
    "+ stop when your fortune is as high as it will ever get\n",
    "+ stop if you would lose the next bet\n",
    "+ stop if you would lose three of the next five bets\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b04c4721-ebf5-48dd-933a-ae3b06798c80",
   "metadata": {},
   "source": [
    "## Doob's Optional Stopping Theorem (1953)\n",
    "\n",
    "Suppose $\\tau$ is a stopping time for the stochastic process $(X_j)_{j \\in \\mathbb{N}}$\n",
    "and that any of the following conditions holds:\n",
    "\n",
    "+ $\\mathbb{P} \\{\\tau \\le c \\} = 1$ for some $c \\in \\mathbb{N}$\n",
    "\n",
    "+ $\\mathbb{E} \\tau < \\infty$ and $ | X_j - X_{j-1}| \\le c$ for all $j$ and some $c$\n",
    "\n",
    "+ $|X_{j}| \\le c $ for all $j \\in \\mathbb{N}$ for some $c$ and $\\mathbb{P} \\{ \\tau < \\infty \\}=1$.\n",
    "\n",
    "Then if $(X_j)_{j \\in \\mathbb{N}}$ is a martingale, $\\mathbb{E} X_\\tau = \\mathbb{E} X_1$;\n",
    "if it is a supermartingale, $\\mathbb{E} X_\\tau \\le \\mathbb{E} X_1$.\n",
    "\n",
    "## Doob's Martingale Convergence Theorem\n",
    "\n",
    "Suppose $(X_j)_{j \\in \\mathbb{N}}$ is a martingale and $\\sup_{j \\in \\mathbb{N}} \\mathbb{E} (-X_j \\wedge 0) < \\infty$.\n",
    "Then $(X_j)$ converges almost surely to a random variable $X_\\infty$ with $\\mathbb{E} X_\\infty < \\infty$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81e057ab-5d25-40de-82db-c8743e3f0203",
   "metadata": {},
   "source": [
    "## Ville's Inequality (1939)\n",
    "\n",
    "Let $(X_j)_{j \\in \\mathbb{N}}$ be a nonnegative (super)martingale with respect to $(Y_j)_{j \\in \\mathbb{N}}$.\n",
    "Then\n",
    "\\begin{equation}\n",
    "\\mathbb{P} \\{ \\exists j : X_j \\ge c \\mathbb{E} X_1 \\} \\le 1/c.\n",
    "\\end{equation}\n",
    "\n",
    "__Proof.__ (informal: can be made rigorous using Doob's martingale convergence theorem)\n",
    "\n",
    "Define the stopping time $\\tau(k)$ to be the first $j \\le k $ such that $X_j \\ge x > 0$, or $k$ otherwise.\n",
    "For each $k$, $X_{\\tau(k)}$ is a stopping time, so $\\mathbb{E} X_{\\tau(k)} = \\mathbb{E} X_1$.\n",
    "Apply Markov's inequality to $X_{\\tau(k)}$ and let $k \\rightarrow \\infty$.\n",
    "\n",
    "\n",
    "Ville's inequality is the foundation of (super)martingale based tests. \n",
    "It also provides deep connections between gambling and probability."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a70d723-84c4-4ee7-b3f7-7eac2453ac16",
   "metadata": {},
   "source": [
    "## The Azuma-Hoeffding inequality\n",
    "\n",
    "Let $(X_j)_{j \\in \\mathbb{N}}$ be a supermartingale and suppose that there are finite constants \n",
    "$\\{c_k\\}_{k \\in \\mathbb{N}}$ and predictable values $(A_k)$ and $(B_k)$ such that for all $k$,\n",
    "\\begin{equation}\n",
    "\\mathbb{P} \\{A_k \\le X_k-X_{k-1} \\le B_k \\} = 1\n",
    "\\end{equation}\n",
    "with $B_k-A_k = c_k$,\n",
    "then for all $k \\in \\mathbb{N}$ and all $\\epsilon > 0$,\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbb{P} \\{ X_k-X_0 \\ge \\epsilon \\} \\le \\exp \\left( - \\frac{2\\epsilon^2}{\\sum _{j=1}^k c_j^2 }\\right).\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c469e31-76cd-421d-96b0-5488964320d7",
   "metadata": {},
   "source": [
    "## Supermartingale tests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "670f6087-967b-42af-9ea8-d06e7204f63a",
   "metadata": {},
   "source": [
    "## Wald's SPRT\n",
    "\n",
    "See [SPRT](./sprt.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be41e2a5-bd13-4d0a-8ddb-c9cfd1be8077",
   "metadata": {},
   "source": [
    "## ALPHA martingales\n",
    "\n",
    "Analogous to the SPRT for the Bernoulli $p$.\n",
    "\n",
    "Bernoulli SPRT:  Observe $\\{X_j \\}_{j \\in \\mathbb{N}}$ IID Bernoulli($p$). The null hypothesis is that $p = \\mu$; alternative hypothesis is that $p = \\eta$.\n",
    "\n",
    "\n",
    "\\begin{equation}\n",
    "T_0 := 1\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "T_j := T_{j-1} \\left (X_j \\frac{\\eta}{\\mu} + (1 - X_j)  \\frac{1-\\eta}{1-\\mu} \\right ).\n",
    "\\end{equation}\n",
    "If $p = \\mu$, this is a nonnegative martingale starting at 1.\n",
    "\n",
    "Here's a stochastic process of the same form, for which we will show a more general result:\n",
    "Assume $X_i \\in [0, u_i]$ for some known $u_i>0$.\n",
    "The joint distribution of $\\{X_i\\}$ depends on a parameter $\\theta$.\n",
    "Let $\\mu_j := \\mathbb{E}(X_j | X^{j-1})$ computed under the null hypothesis $\\theta = \\mu$. \n",
    "Let $\\eta_j = \\eta_j(X^{j-1}) \\in (\\mu_j, u_j]$, $j=1, \\ldots$, be a _predictable sequence_ in the sense that \n",
    "$\\eta_j$ may depend on $X^{j-1}$, but not on $X_k$ for $k \\ge j$.\n",
    "\n",
    "Define $T_0 := 1$ and\n",
    "\\begin{equation}\n",
    "    T_j := T_{j-1} u_j^{-1}\\left ( X_j\\frac{\\eta_j}{\\mu_j} + (u_j-X_j) \\frac{u_j-\\eta_j}{u_j-\\mu_j} \\right ), \\;\\; j=1, \\ldots . \\end{equation}\n",
    "For $u_j=1$, this is the same as the stochastic process above. \n",
    "This definition can be rearranged to yield\n",
    "\\begin{equation}\n",
    "    T_j := T_{j-1} \\left ( \\frac{X_j}{\\mu_j} \\cdot \\frac{\\eta_j-\\mu_j}{u_j-\\mu_j} + \\frac{u_j-\\eta_j}{u_j-\\mu_j} \\right ). \n",
    "\\end{equation}\n",
    "Equivalently,\n",
    "\\begin{equation}\n",
    "    T_j := \\prod_{i=1}^j \\left ( \\frac{X_i}{\\mu_i} \\cdot \\frac{\\eta_i-\\mu_i}{u_j-\\mu_i} + \\frac{u_j-\\eta_i}{u_j-\\mu_i} \\right ), \\;\\; j \\ge 1. \n",
    "\\end{equation}\n",
    "\n",
    "Under the null hypothesis that $\\theta = \\mu$, $T_j$ is nonnegative since $X_j$, $\\mu_j$, and $\\eta_j$\n",
    "are all in $[0, u]$.\n",
    "Also,\n",
    "\\begin{eqnarray}\n",
    "    \\mathbb{E} (T_j | X^{j-1} ) &=& T_{j-1} \\left ( \\frac{\\mu_j}{\\mu_j} \\cdot \\frac{\\eta_j-\\mu_j}{u_j-\\mu_j} + \\frac{u_j-\\eta_j}{u_j-\\mu_j} \\right ) \\nonumber \\\\\n",
    "    &=&  T_{j-1} \\left ( \\frac{\\eta_j-\\mu_j}{u_j-\\mu_j} + \\frac{u_j-\\eta_j}{u_j-\\mu_j} \\right ) \\nonumber \\\\\n",
    "    &=& T_{j-1}.\n",
    "\\end{eqnarray}\n",
    "Thus if $\\theta = \\mu$, $(T_j)_{j \\in \\mathbb{N}}$ is a nonnegative martingale with respect to $(X_j)_{j \\in \\mathbb{N}}$, starting at $1$.\n",
    "\n",
    "If $\\theta < \\mu$, then $\\mathbb{E} (X_j | X^{j-1}) < \\mu_j$ and $r_j = \\frac{\\mathbb{E} (X_j | X^{j-1})}{\\mu_j} < 1$, so\n",
    "\\begin{equation} \\label{eq:supermartingale}\n",
    "    \\mathbb{E} (T_j | X^{j-1} ) = T_{j-1} \\left ( r_j \\cdot \\frac{\\eta_j-\\mu_j}{u_j-\\mu_j} + \\frac{u_j-\\eta_j}{u-\\mu_j} \\right ) < T_{j-1}.\n",
    "\\end{equation}\n",
    "Thus $(T_j)$ is a nonnegative supermartingale starting at 1 if $\\theta \\le \\mu$.\n",
    "It follows from Ville's inequality that if $\\theta \\le \\mu$,\n",
    "\\begin{equation} \\label{eq:p-value-adapt}\n",
    "\\mathbb{P} \\{ \\exists j :  T_j \\ge \\alpha^{-1} \\} \\le \\alpha.\n",
    "\\end{equation}\n",
    "That is, $\\min(1, 1/T_j)$ is an \"anytime $P$-value\" for the composite null hypothesis $\\theta \\le \\mu$.\n",
    "\n",
    "Note that the derivation did not use any information about $\\{x_i\\}$ other than $x_i \\in [0, u]$:\n",
    "it applies to populations $\\{x_i\\}$ that are nonnegative and bounded, not merely binary populations.\n",
    "\n",
    "### Sampling without replacement\n",
    "\n",
    "To use ALPHA with a sample drawn without replacement, we need $\\mathbb{E}(X_j | X^{j-1})$ computed on the assumption that\n",
    "$\\theta :=  \\frac{1}{N} \\sum_{i=1}^N x_i = \\mu$.\n",
    "For sampling without replacement from a population with mean $\\mu$, after draw $j-1$, the mean of the remaining numbers is $(N\\mu - S_{j-1})/(N-j+1)$, where $S_j := \\sum_{i=1}^j X_j$ is the sample sum as of the $j$th draw.\n",
    "\n",
    "Thus the conditional expectation of $X_j$ given $X^{j-1}$ under the null is $(N\\mu - \\sum_{k=1}^{j-1}X_k)/(N-j+1)$.\n",
    "If $N\\mu - \\sum_{k=1}^{j-1}X_k < 0$ for any $k$, the null hypothesis $\\theta = \\mu$ is certainly false.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dab1d9c-c2f9-4970-a5b6-40224bd53882",
   "metadata": {},
   "source": [
    "## General betting martingales\n",
    "\n",
    "See Waudby-Smith and Ramdas (2022).\n",
    "\n",
    "Set up is the same as for ALPHA.\n",
    "\n",
    "Define $T_0 :=1$ and\n",
    "\\begin{equation}\n",
    "T_j := T_{j-1} (1 + \\lambda_j (X_j - \\mu_j)) = \\prod_{i=1}^j (1 + \\lambda_i (X_i - \\mu_i)),\n",
    "\\end{equation}\n",
    "where $\\lambda_j \\in  (−u/(u − \\mu_j), 1/\\mu_j)$ and $(\\lambda_j)_{j \\in \\mathbb{N}}$ is predictable.\n",
    "\n",
    "Claim: if $\\theta = \\mu$, this is a non-negative martingale starting at 1.\n",
    "\n",
    "Proof: since $\\mathbb{E}_\\mu (X_j | X^{j-1}) = \\mu_j$, $\\mathbb{E}_\\mu (T_j | X^{j-1}) = T_{j-1}\\times 1$.\n",
    "Since $\\lambda_j \\in (−u/(u − \\mu_j), 1/\\mu_j)$, $(1 + \\lambda_j (X_j - \\mu_j)) > 0$.\n",
    "\n",
    "Can think of this as the fortune of a gambler who starts with a stake of 1, bets on whether the next\n",
    "term is greater than or less than its (conditional) expected value, and is not allowed to bet more than\n",
    "the current stake (i.e., can't go into debt).\n",
    "\n",
    "Different choices of $\\lambda_j$ give power against different alternatives. Different betting strategies\n",
    "translate to different tests.\n",
    "\n",
    "### The ALPHA martingales are betting martingales\n",
    "\n",
    "The ALPHA martingales are products of terms of the form\n",
    "\\begin{eqnarray}\n",
    "\\frac{X_j}{\\mu_j} \\cdot \\frac{\\eta_j-\\mu_j}{u_j-\\mu_j} + \\frac{u_j-\\eta_j}{u_j-\\mu_j} &=&\n",
    "\\frac{ X_j(\\eta_j/\\mu_j - 1) + u_j - \\eta_j}{u_j-\\mu_j} \\\\\n",
    "&=& \\frac{ X_j(\\eta_j/\\mu_j - 1) + u_j - \\eta_j + \\mu_j - \\mu_j}{u_j-\\mu_j} \\\\\n",
    "&=& 1 + \\frac{ X_j(\\eta_j/\\mu_j - 1) - \\eta_j + \\mu_j}{u_j-\\mu_j} \\\\\n",
    "&=& 1 + \\frac{\\eta_j/\\mu_j - 1}{u_j-\\mu_j} (X_j - \\mu_j),\n",
    "\\end{eqnarray}\n",
    "which is of the form $(1 + \\lambda_j (X_j - \\mu_j))$ with $\\lambda_j = \\frac{\\eta_j/\\mu_j - 1}{u_j-\\mu_j}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54309fd6-a007-46d6-8340-4c9cb46a7db4",
   "metadata": {},
   "source": [
    "## Lower confidence bounds on the mean for sampling without replacement: Ville's inequality\n",
    "\n",
    "This is an alternative derivation of the Kaplan-Wald method, also from [Harold Kaplan's website](http://web.archive.org/web/20131209044835/http://printmacroj.com/martMean.htm), involving the fact that a suitably constructed sequence is a nonnegative martingale if the null hypothesis is true, then applying Ville's inequality.\n",
    "\n",
    "Suppose we are sampling without replacement from a finite population of $N$ \n",
    "non-negative items,  $\\{x_1, \\ldots, x_N\\}$, with $x_j \\ge 0$ $\\forall j$.\n",
    "The population mean is $\\theta = \\frac{1}{N} \\sum_{j=1}^N x_j \\ge 0$ and the population total\n",
    "is $N\\theta \\ge 0$.\n",
    "We draw $\\{X_1, \\ldots, X_j \\}$ sequentially without replacement.\n",
    "If $\\theta = \\mu$, $\\mathbb{E}X_1 = \\mu$, so $\\mathbb{E}(X_1/\\mu) = 1$.\n",
    "Conditional on $X_1, \\ldots, X_n$, the total of the remaining $N-n$ items is\n",
    "$N\\mu - \\sum_{j=1}^n X_j$, so the mean of the remaining items is\n",
    "\\begin{equation*} \n",
    "    \\frac{N\\mu-\\sum_{j=1}^n X_j}{N-n} = \\frac{\\mu - \\frac{1}{N} \\sum_{j=1}^n X_j}{1-n/N}.\n",
    "\\end{equation*}\n",
    "Thus, if $\\theta = \\mu$, the expected value of $X_{n+1}$ given $X_1, \\ldots, X_n$ is \n",
    "$\\frac{\\mu - \\frac{1}{N} \\sum_{j=1}^n X_j}{1-n/N}$.\n",
    "\n",
    "Define \n",
    "\\begin{equation*} \n",
    "Y_1(\\mu) :=\n",
    "\\begin{cases}\n",
    "X_1/\\mu,& N\\mu > 0, \\\\\n",
    "1,&   N\\mu = 0, \\\\\n",
    "\\end{cases}\n",
    "\\end{equation*}\n",
    "and for $1 \\le j \\le N-1$,\n",
    "\\begin{equation*} \n",
    "Y_{j+1}(\\mu) :=\n",
    "\\begin{cases}\n",
    "X_{j+1} \n",
    "\\frac\n",
    "{1 - \\frac{n}{N}}\n",
    "{\\mu - \\frac{1}{N} \\sum_{i=1}^j X_i},& \\sum_{i=1}^j X_i < N\\mu, \\\\\n",
    "1,& \\sum_{i=1}^j X_i \\ge N\\mu. \\\\\n",
    "\\end{cases}\n",
    "\\end{equation*} \n",
    "Then $\\mathbb{E}(Y_{j+1}(\\mu) | Y^j) = 1$.\n",
    "Let $Z_j(\\mu) := \\prod_{i=1}^j Y_j(\\mu)$.\n",
    "Note that $Y_k(\\mu)$ can be recovered from $\\{Z_j(\\mu), j \\le k\\}$, \n",
    "since $Y_k(\\mu) = Z_k(\\mu)/Z_{k-1}(\\mu)$.\n",
    "Now $\\mathbb{E}|Z_k| \\le \\max_j x_j < \\infty$ and\n",
    "\\begin{equation*} \n",
    "   \\mathbb{E}\\left ( Z_j(\\mu) | Z^{j-1}(\\mu) \\right ) = \n",
    "   \\mathbb{E} \\left (Y_j(\\mu)Z_{j-1}(\\mu) | Z^{j-1}(\\mu) \\right ) = Z_{j-1}(\\mu).\n",
    "\\end{equation*}\n",
    "Thus \n",
    "\\begin{equation*} \n",
    "    ( Z_j(\\mu))_{1 \\le j \\le N}\n",
    "\\end{equation*}\n",
    "is a non-negative martingale.\n",
    "\n",
    "By Ville's inequality, for any $p > 0$ and any $k \\in \\{1, \\ldots, N \\}$,\n",
    "\\begin{equation*} \n",
    "     \\mathbb{P} \\left ( \\max_{1 \\le j \\le k} Z_j(\\mu) > 1/p \\right ) \\le p \\; \\mathbb{E}|Z_k|.\n",
    "\\end{equation*}\n",
    "Since $(Z_j)$ is a non-negative martingale, $\\mathbb{E}|Z_k| = \\mathbb{E}Z_k = \\mathbb{E}Z_1 = 1$.\n",
    "\n",
    "Thus a $p$-value for the hypothesis $\\mu = t$ based on data $X_1, \\ldots X_k$ is \n",
    "$\\left (\\max_{1 \\le j \\le k} Z_j(\\mu) \\right )^{-1} \\wedge 1$.\n",
    "If $X_j = 0$ for some $j$, then $Z_k = 0$ for all $k \\ge j$.\n",
    "\n",
    "To avoid that problem, we can shift everything to the right: pick $\\gamma > 0$,\n",
    "find a lower confidence bound for $\\delta = \\theta+\\gamma > \\theta > 0$ from data $\\{X_j+\\gamma\\}$, then subtract $\\gamma$ from the lower \n",
    "confidence bound to get a lower confidence bound for $\\theta$.\n",
    "There are tradeoffs involved in picking $\\gamma$: if many $X_j$ turn out to be \n",
    "small, especially for small $j$, it helps to have $\\gamma$ large, and vice versa.\n",
    "\n",
    "Unpacking the math yields the $p$ value\n",
    "\n",
    "\\begin{equation*} \n",
    "p_{\\mathrm{KK}} := 1 \\wedge \\left ( \\max_{1 \\le j \\le J} \\prod_{k=1}^j (X_{k}+\\gamma) \\frac{1-(k-1)/N}{\\mu - \\frac{1}{N} \\sum_{\\ell=1}^{k-1} (X_\\ell+\\gamma)} \\right )^{-1}\n",
    "\\end{equation*}\n",
    "for the hypothesis that $\\theta \\le \\mu - \\gamma$.\n",
    "\n",
    "The corresponding $1-\\alpha$ lower confidence bound is\n",
    "\n",
    "\\begin{equation*} \n",
    "   \\sup \\left \\{\\mu \\ge 0: \\max_{1 \\le j \\le k}  \\prod_{i=1}^j (X_k +\\gamma) \\frac{1-(k-1)/N}{\\mu - \\frac{1}{N} \\sum_{\\ell=1}^{k-1} (X_\\ell+\\gamma)} \\le 1/\\alpha \\right \\} - \\gamma.\n",
    "\\end{equation*}\n",
    "\n",
    "Note that this is one of infinitely many possible martingales we might construct from the sequential sample $(X_j)$.\n",
    "In particular, instead of dividing the next draw by its expected value, we might subtract its expected value."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45e56fe3-0e0f-4aa1-afee-cb6909a9c62a",
   "metadata": {},
   "source": [
    "## Other test martingales\n",
    "\n",
    "See Waudby-Smith and Ramdas (2022). They summarize the performance of a number of different test martingales for _confidence sequences_ for the mean of a $[0, 1]$-valued population: a sequence of confidence intervals that are sequentially valid, \n",
    "in the sense that the chance that any interval (of potentially infinitely many) does not contain the mean is at most $1-\\alpha$.\n",
    "Each interval is derived by inverting tests.\n",
    "Some of the tests are designed so that inverting them produces an interval, rather than (possibly) a disconnected set.\n",
    "They find that betting martingales, with suitably constructed bets, dominate the other test martingales (the tests are more powerful and the resulting confidence sets are narrower).\n",
    "\n",
    "They also discuss strategies for constructing good test martingales\n",
    "\n",
    "We will now implement one of their approaches, but in the form of a sequentially valid one-sided test rather than a confidence sequence.\n",
    "The test martingale is the betting martingale\n",
    "\\begin{equation}\n",
    "  M_j := \\prod_{i=1}^j (1 + \\lambda_i (X_i - \\mu_i)),\n",
    "\\end{equation}\n",
    "where $\\mu_i := \\mathbb{E}_0 (X_i | \\mathcal{F}_{i-1})$\n",
    "The issue is how to choose $(\\lambda_i)$, which must be predictable ($\\lambda_i$ can depend on $\\mathcal{F}_{i-1}$),\n",
    "practical to calculate, and tend to make $\\mathbb{Q} (\\ln(1+\\lambda_i(X_i-\\mu_i)) | \\mathcal{F}_{i-1})$ large.\n",
    "\n",
    "Their preferred choice is:\n",
    "+ $\\lambda_i := \\sqrt{\\frac{2\\ln (2/\\alpha)}{\\hat{\\sigma}_{i-1}^2 i \\ln(i+1)}}$\n",
    "+ $\\hat{\\sigma}_i^2 := \\frac{ \\frac{1}{4} + \\sum_{j=1}^i (X_j-\\hat{\\mu}_j)^2}{i+1}$\n",
    "+ $\\hat{\\mu}_i := \\frac{\\frac{1}{2} + \\sum_{j=1}^i X_i}{i+1}$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36043c4f-04d1-4fa5-949e-42984ff9d621",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import scipy as sp\n",
    "import numpy as np\n",
    "from numpy.testing import assert_allclose\n",
    "import json\n",
    "np.random.seed(123456789)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a26042-7069-40a3-98c1-70f13efa86e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sprt_mart(x : np.array, N : int, mu : float=1/2, eta: float=1-np.finfo(float).eps, \\\n",
    "              u: float=1, random_order = True):\n",
    "    '''\n",
    "    Finds the p value for the hypothesis that the population \n",
    "    mean is less than or equal to mu against the alternative that it is eta,\n",
    "    for a population of size N of values in the interval [0, u].\n",
    "    \n",
    "    Generalizes Wald's SPRT for the Bernoulli to sampling without replacement and to bounded\n",
    "    values rather than binary values.\n",
    "\n",
    "    If N is finite, assumes the sample is drawn without replacement\n",
    "    If N is infinite, assumes the sample is with replacement\n",
    "    \n",
    "    Data are assumed to be in random order. \n",
    "\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    x : binary list, one element per draw. A list element is 1 if the \n",
    "        the corresponding trial was a success\n",
    "    N : int\n",
    "        population size for sampling without replacement, or np.infinity for \n",
    "        sampling with replacement\n",
    "    theta : float in (0,u)\n",
    "        hypothesized population mean\n",
    "    eta : float in (0,u)\n",
    "        alternative hypothesized population mean\n",
    "    random_order : Boolean\n",
    "        if the data are in random order, setting this to True can improve the power.\n",
    "        If the data are not in random order, set to False\n",
    "    '''\n",
    "    if any((xx < 0 or xx > u) for xx in x):\n",
    "        raise ValueError(f'Data out of range [0,{u}]')\n",
    "    if np.isfinite(N):\n",
    "        if not random_order:\n",
    "            raise ValueError(\"data must be in random order for samples without replacement\")\n",
    "        S = np.insert(np.cumsum(x),0,0)[0:-1]  # 0, x_1, x_1+x_2, ...,  \n",
    "        j = np.arange(1,len(x)+1)              # 1, 2, 3, ..., len(x)\n",
    "        m = (N*mu-S)/(N-j+1)                   # mean of population after (j-1)st draw, if null is true\n",
    "    else:\n",
    "        m = mu\n",
    "    with np.errstate(divide='ignore',invalid='ignore'): \n",
    "        terms = np.cumprod((x*eta/m + (u-x)*(u-eta)/(u-m))/u) # generalization of Bernoulli SPRT\n",
    "    terms[m<0] = np.inf                        # the null is surely false\n",
    "    return terms\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84ab2533-4463-482d-8459-0016653be17a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def shrink_trunc(x: np.array, N: int, mu: float=1/2, nu: float=1-np.finfo(float).eps, u: float=1, c: float=1/2, \n",
    "                 d: float=100) -> np.array: \n",
    "    '''\n",
    "    apply the shrinkage and truncation estimator to an array\n",
    "    \n",
    "    sample mean is shrunk towards nu, with relative weight d compared to a single observation.\n",
    "    estimate is truncated above at u-u*eps and below at mu_j+e_j(c,j)\n",
    "    \n",
    "    S_1 = 0\n",
    "    S_j = \\sum_{i=1}^{j-1} x_i, j > 1\n",
    "    m_j = (N*mu-S_j)/(N-j+1) if np.isfinite(N) else mu\n",
    "    e_j = c/sqrt(d+j-1)\n",
    "    eta_j =  ( (d*nu + S_j)/(d+j-1) \\vee (m_j+e_j) ) \\wedge u*(1-eps)\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x : np.array\n",
    "        input data       \n",
    "    mu : float in (0, 1)\n",
    "        hypothesized population mean\n",
    "    eta : float in (t, 1)\n",
    "        initial alternative hypothethesized value for the population mean\n",
    "    c : positive float\n",
    "        scale factor for allowing the estimated mean to approach t from above\n",
    "    d : positive float\n",
    "        relative weight of nu compared to an observation, in updating the alternative for each term\n",
    "    '''\n",
    "    S = np.insert(np.cumsum(x),0,0)[0:-1]  # 0, x_1, x_1+x_2, ...,  \n",
    "    j = np.arange(1,len(x)+1)              # 1, 2, 3, ..., len(x)\n",
    "    m = (N*mu-S)/(N-j+1) if np.isfinite(N) else mu   # mean of population after (j-1)st draw, if null is true \n",
    "    return np.minimum(u*(1-np.finfo(float).eps), np.maximum((d*nu+S)/(d+j-1),m+c/np.sqrt(d+j-1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23eb9a9c-e73e-41c8-b13c-9c45e9db3bf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_shrink_trunc():\n",
    "    epsj = lambda c, d, j: c/math.sqrt(d+j-1)\n",
    "    Sj = lambda x, j: 0 if j==1 else np.sum(x[0:j-1])\n",
    "    muj = lambda N, mu, x, j: (N*mu - Sj(x, j))/(N-j+1) if np.isfinite(N) else mu\n",
    "    nus = [.51, .55, .6]\n",
    "    mu = 1/2\n",
    "    u = 1\n",
    "    d = 10\n",
    "    vrand =  sp.stats.bernoulli.rvs(1/2, size=20)\n",
    "    v = [\n",
    "        np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]),\n",
    "        np.array([1, 1, 1, 1, 1, 1, 0, 0, 0, 0]),\n",
    "        vrand\n",
    "    ]\n",
    "    for nu in nus:\n",
    "        c = (nu-mu)/2\n",
    "        for x in v:\n",
    "            N = len(x)\n",
    "            xinf = shrink_trunc(x, np.inf, mu, nu, c=c, d=d)\n",
    "            xfin = shrink_trunc(x, len(x), mu, nu, c=c, d=d)\n",
    "            yinf = np.zeros(len(x))\n",
    "            yfin = np.zeros(len(x))\n",
    "            for j in range(1,len(x)+1):\n",
    "                est = (d*nu + Sj(x,j))/(d+j-1)\n",
    "                most = u*(1-np.finfo(float).eps)\n",
    "                yinf[j-1] = np.minimum(np.maximum(mu+epsj(c,d,j), est), most)\n",
    "                yfin[j-1] = np.minimum(np.maximum(muj(N,mu,x,j)+epsj(c,d,j), est), most)\n",
    "            np.testing.assert_allclose(xinf, yinf)    \n",
    "            np.testing.assert_allclose(xfin, yfin)    \n",
    "    \n",
    "test_shrink_trunc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67e14af3-0dec-4254-ad12-19956d033605",
   "metadata": {},
   "outputs": [],
   "source": [
    "def alpha_mart(x: np.array, N: int, mu: float=1/2, eta: float=1-np.finfo(float).eps, u: float=1, \\\n",
    "               estim: callable=shrink_trunc) -> np.array :\n",
    "    '''\n",
    "    Finds the ALPHA martingale for the hypothesis that the population \n",
    "    mean is less than or equal to t using a martingale method,\n",
    "    for a population of size N, based on a series of draws x.\n",
    "    \n",
    "    The draws must be in random order, or the sequence is not a martingale under the null\n",
    "    \n",
    "    If N is finite, assumes the sample is drawn without replacement\n",
    "    If N is infinite, assumes the sample is with replacement\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : list corresponding to the data\n",
    "    N : int\n",
    "        population size for sampling without replacement, or np.infinity for sampling with replacement\n",
    "    mu : float in (0,1)\n",
    "        hypothesized fraction of ones in the population\n",
    "    eta : float in (t,1) \n",
    "        alternative hypothesized population mean\n",
    "    estim : callable\n",
    "        estim(x, N, mu, eta, u) -> np.array of length len(x), the sequence of values of eta_j for ALPHA\n",
    "               \n",
    "    Returns\n",
    "    -------   \n",
    "    terms : array\n",
    "        sequence of terms that would be a nonnegative martingale under the null\n",
    "    '''\n",
    "    S = np.insert(np.cumsum(x),0,0)[0:-1]  # 0, x_1, x_1+x_2, ...,  \n",
    "    j = np.arange(1,len(x)+1)              # 1, 2, 3, ..., len(x)\n",
    "    m = (N*mu-S)/(N-j+1) if np.isfinite(N) else mu   # mean of population after (j-1)st draw, if null is true \n",
    "    etaj = estim(x, N, mu, eta, u) \n",
    "    with np.errstate(divide='ignore',invalid='ignore'):\n",
    "        terms = np.cumprod((x*etaj/m + (1-x)*(u-etaj)/(u-m))/u)\n",
    "    terms[m<0] = np.inf\n",
    "    return terms"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}