add self.stats to selector

Author: martinjankowiak

bvas/bvas_selector.py

@@ -7,6 +7,21 @@
 from bvas.util import namespace_to_numpy
 
 
+def populate_alpha_beta_stats(container, stats):
+    for s in ['h_alpha', 'h_beta', 'h']:
+        if hasattr(container, s):
+            stats['Mean ' + s] = getattr(container, s)
+
+
+def populate_weight_stats(selector, stats, weights, quantiles=[5.0, 10.0, 20.0, 50.0, 90.0, 95.0]):
+    q5, q10, q20, q50, q90, q95 = np.percentile(weights, quantiles).tolist()
+    s = "5/10/20/50/90/95:  {:.2e}  {:.2e}  {:.2e}  {:.2e}  {:.2e}  {:.2e}"
+    stats['Weight quantiles'] = s.format(q5, q10, q20, q50, q90, q95)
+    s = "mean/std/min/max:  {:.2e}  {:.2e}  {:.2e}  {:.2e}"
+    stats['Weight moments'] = s.format(weights.mean().item(), weights.std().item(),
+                                       weights.min().item(), weights.max().item())
+
+
 class BVASSelector(object):
     r"""
     Main analysis class for Bayesian Viral Allele Selection (BVAS).
@@ -89,6 +104,7 @@ def run(self, T, T_burnin, seed=None):
             pango = pd.Series(self.variant_names, name="Variant Name")
             self.growth_rates = pd.concat([growth_rate, growth_rate_std, pango], axis=1)
 
-        if hasattr(self.container, 'h'):
-            if self.container.h.size == 1:
-                print('[h_ratio] {:.4f}'.format(self.container.h.item()))
+        self.stats = {}
+        self.weights = np.array(self.container._weights)
+        populate_alpha_beta_stats(self.container, self.stats)
+        populate_weight_stats(self, self.stats, self.weights)

notebooks/basic_demo.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "6657ec71",
+   "id": "e84d968e",
    "metadata": {},
    "source": [
     "# Basic BVAS demo using simulated data"
@@ -11,7 +11,7 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "a1456578",
+   "id": "27b49fb2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -24,7 +24,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d9224796",
+   "id": "fe5e4e14",
    "metadata": {},
    "source": [
     "### Simulate data"
@@ -33,7 +33,7 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "0ec74519",
+   "id": "750d0541",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -51,7 +51,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "935c020e",
+   "id": "9f8274ff",
    "metadata": {},
    "outputs": [
     {
@@ -77,7 +77,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "dfc65b90",
+   "id": "d995ccbc",
    "metadata": {},
    "source": [
     "### Instantiate BVASSelector object"
@@ -86,7 +86,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "79fd0337",
+   "id": "ad539a7c",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -103,7 +103,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3ad5765d",
+   "id": "5da317dd",
    "metadata": {},
    "source": [
     "### Run BVAS MCMC-based inference"
@@ -112,13 +112,13 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "51be4150",
+   "id": "a11e7ef9",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "345f12e6501141989ec6d910ecb9c4c2",
+       "model_id": "838905aacba843549b78082c99f55479",
        "version_major": 2,
        "version_minor": 0
       },
@@ -136,7 +136,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0c60cc7f",
+   "id": "e5cbb4df",
    "metadata": {},
    "source": [
     "### Inspect results\n",
@@ -153,7 +153,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "d1d58d65",
+   "id": "5c26864a",
    "metadata": {},
    "outputs": [
     {
@@ -186,7 +186,7 @@
   {
    "cell_type": "code",
    "execution_count": 7,
-   "id": "67d6406f",
+   "id": "76c1bf93",
    "metadata": {},
    "outputs": [
     {
@@ -215,7 +215,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "id": "1814216c",
+   "id": "d41b59d1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -225,7 +225,39 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d0939917",
+   "id": "5a91ef89",
+   "metadata": {},
+   "source": [
+    "### MCMC stats\n",
+    "\n",
+    "Additional MCMC stats are in the `selector.stats` dictionary."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "4902488d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'Weight quantiles': '5/10/20/50/90/95:  1.84e-16  1.84e-16  2.71e+00  5.11e+00  1.17e+01  1.17e+01',\n",
+       " 'Weight moments': 'mean/std/min/max:  6.19e+00  4.16e+00  1.84e-16  1.17e+01'}"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "selector.stats"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5ed607dd",
    "metadata": {},
    "source": [
     "# Compare to MAP inference\n",
@@ -235,8 +267,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
-   "id": "d3d22182",
+   "execution_count": 10,
+   "id": "e2c3a9c3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -245,8 +277,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
-   "id": "ff352512",
+   "execution_count": 11,
+   "id": "6759a828",
    "metadata": {},
    "outputs": [
     {
@@ -373,7 +405,7 @@
        "Spurious70  0.011479    15"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -385,16 +417,16 @@
   },
   {
    "cell_type": "markdown",
-   "id": "40e3bd65",
+   "id": "c1349f1a",
    "metadata": {},
    "source": [
     "# Compare to MAP inference with a Sparsity-inducing Laplace prior"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
-   "id": "b6b6cecc",
+   "execution_count": 12,
+   "id": "06e2e0fb",
    "metadata": {},
    "outputs": [
     {
@@ -431,8 +463,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
-   "id": "63d92a7b",
+   "execution_count": 13,
+   "id": "178836fa",
    "metadata": {},
    "outputs": [
     {
@@ -559,7 +591,7 @@
        "Spurious88  0.002844    15"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }