nbqa NB

Author: twiecki

docs/source/notebooks/GLM-hierarchical-jax.ipynb

@@ -54,7 +54,8 @@
     "import pymc3 as pm\n",
     "import theano\n",
     "import pymc3.sampling_jax\n",
-    "print('Running on PyMC3 v{}'.format(pm.__version__))"
+    "\n",
+    "print(\"Running on PyMC3 v{}\".format(pm.__version__))"
    ]
   },
   {
@@ -64,7 +65,7 @@
    "outputs": [],
    "source": [
     "%config InlineBackend.figure_format = 'retina'\n",
-    "az.style.use('arviz-darkgrid')"
+    "az.style.use(\"arviz-darkgrid\")"
    ]
   },
   {
@@ -73,10 +74,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data = pd.read_csv(pm.get_data('radon.csv'))\n",
-    "data['log_radon'] = data['log_radon'].astype(theano.config.floatX)\n",
+    "data = pd.read_csv(pm.get_data(\"radon.csv\"))\n",
+    "data[\"log_radon\"] = data[\"log_radon\"].astype(theano.config.floatX)\n",
     "county_names = data.county.unique()\n",
-    "county_idx = data.county_code.values.astype('int32')\n",
+    "county_idx = data.county_code.values.astype(\"int32\")\n",
     "\n",
     "n_counties = len(data.county.unique())"
    ]
@@ -96,27 +97,26 @@
    "source": [
     "with pm.Model() as hierarchical_model:\n",
     "    # Hyperpriors for group nodes\n",
-    "    mu_a = pm.Normal('mu_a', mu=0., sigma=100.)\n",
-    "    sigma_a = pm.HalfNormal('sigma_a', 5.)\n",
-    "    mu_b = pm.Normal('mu_b', mu=0., sigma=100.)\n",
-    "    sigma_b = pm.HalfNormal('sigma_b', 5.)\n",
+    "    mu_a = pm.Normal(\"mu_a\", mu=0.0, sigma=100.0)\n",
+    "    sigma_a = pm.HalfNormal(\"sigma_a\", 5.0)\n",
+    "    mu_b = pm.Normal(\"mu_b\", mu=0.0, sigma=100.0)\n",
+    "    sigma_b = pm.HalfNormal(\"sigma_b\", 5.0)\n",
     "\n",
     "    # Intercept for each county, distributed around group mean mu_a\n",
     "    # Above we just set mu and sd to a fixed value while here we\n",
     "    # plug in a common group distribution for all a and b (which are\n",
     "    # vectors of length n_counties).\n",
-    "    a = pm.Normal('a', mu=mu_a, sigma=sigma_a, shape=n_counties)\n",
+    "    a = pm.Normal(\"a\", mu=mu_a, sigma=sigma_a, shape=n_counties)\n",
     "    # Intercept for each county, distributed around group mean mu_a\n",
-    "    b = pm.Normal('b', mu=mu_b, sigma=sigma_b, shape=n_counties)\n",
+    "    b = pm.Normal(\"b\", mu=mu_b, sigma=sigma_b, shape=n_counties)\n",
     "\n",
     "    # Model error\n",
-    "    eps = pm.HalfCauchy('eps', 5.)\n",
+    "    eps = pm.HalfCauchy(\"eps\", 5.0)\n",
     "\n",
-    "    radon_est = a[county_idx] + b[county_idx]*data.floor.values\n",
+    "    radon_est = a[county_idx] + b[county_idx] * data.floor.values\n",
     "\n",
     "    # Data likelihood\n",
-    "    radon_like = pm.Normal('radon_like', mu=radon_est,\n",
-    "                           sigma=eps, observed=data.log_radon)"
+    "    radon_like = pm.Normal(\"radon_like\", mu=radon_est, sigma=eps, observed=data.log_radon)"
    ]
   },
   {
@@ -193,8 +193,9 @@
    "source": [
     "%%time\n",
     "with hierarchical_model:\n",
-    "    hierarchical_trace = pm.sample(2000, tune=2000, target_accept=.9, \n",
-    "                                   compute_convergence_checks=False)"
+    "    hierarchical_trace = pm.sample(\n",
+    "        2000, tune=2000, target_accept=0.9, compute_convergence_checks=False\n",
+    "    )"
    ]
   },
   {
@@ -238,8 +239,7 @@
     "%%time\n",
     "# Inference button (TM)!\n",
     "with hierarchical_model:\n",
-    "    hierarchical_trace_jax = pm.sampling_jax.sample_numpyro_nuts(\n",
-    "        2000, tune=2000, target_accept=.9)"
+    "    hierarchical_trace_jax = pm.sampling_jax.sample_numpyro_nuts(2000, tune=2000, target_accept=0.9)"
    ]
   },
   {
@@ -281,10 +281,10 @@
     }
    ],
    "source": [
-    "pm.traceplot(hierarchical_trace_jax, \n",
-    "             var_names=['mu_a', 'mu_b',\n",
-    "                        'sigma_a_log__', 'sigma_b_log__',\n",
-    "                        'eps_log__']);"
+    "pm.traceplot(\n",
+    "    hierarchical_trace_jax,\n",
+    "    var_names=[\"mu_a\", \"mu_b\", \"sigma_a_log__\", \"sigma_b_log__\", \"eps_log__\"],\n",
+    ");"
    ]
   },
   {
@@ -309,8 +309,7 @@
     }
    ],
    "source": [
-    "pm.traceplot(hierarchical_trace_jax, \n",
-    "             var_names=['a'], coords={'a_dim_0': range(5)});"
+    "pm.traceplot(hierarchical_trace_jax, var_names=[\"a\"], coords={\"a_dim_0\": range(5)});"
    ]
   },
   {