auto-fix

pymc-devs · Dec 22, 2024 · 185a76d · 185a76d
1 parent 5b88163
commit 185a76d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 18 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -3,7 +3,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: debug-statements
       - id: check-merge-conflict
@@ -16,7 +16,7 @@ repos:
       -   id: trailing-whitespace
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.6
+    rev: v0.8.3
     hooks:
       - id: ruff
         args: ["--fix", "--output-format=full"]

diff --git a/notebooks/pytensor_logp.ipynb b/notebooks/pytensor_logp.ipynb
@@ -34,15 +34,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pytensor\n",
-    "import pytensor.tensor as pt\n",
-    "import pymc as pm\n",
-    "import numpy as np\n",
-    "import nutpie\n",
     "import arviz\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
     "import pandas as pd\n",
+    "import pymc as pm\n",
     "import seaborn as sns\n",
-    "import matplotlib.pyplot as plt"
+    "\n",
+    "import nutpie"
    ]
   },
   {
@@ -404,7 +403,7 @@
     "    filename=\"radon_model.stan\",\n",
     "    coords=coords_stan,\n",
     "    dims=dims_stan,\n",
-    "    cache=False\n",
+    "    cache=False,\n",
     ")"
    ]
   },
@@ -484,8 +483,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import stan\n",
     "import nest_asyncio\n",
+    "import stan\n",
     "\n",
     "nest_asyncio.apply()"
    ]
@@ -522,7 +521,7 @@
    ],
    "source": [
     "%%time\n",
-    "with open(\"radon_model.stan\", \"r\") as file:\n",
+    "with open(\"radon_model.stan\") as file:\n",
     "    model = stan.build(file.read(), data=data_stan)"
    ]
   },
@@ -748,7 +747,10 @@
     }
    ],
    "source": [
-    "plt.plot((trace_pymc.warmup_sample_stats.n_steps).isel(draw=slice(0, 1000)).cumsum(\"draw\").T, np.log(trace_pymc.warmup_sample_stats.energy.isel(draw=slice(0, 1000)).T));\n",
+    "plt.plot(\n",
+    "    (trace_pymc.warmup_sample_stats.n_steps).isel(draw=slice(0, 1000)).cumsum(\"draw\").T,\n",
+    "    np.log(trace_pymc.warmup_sample_stats.energy.isel(draw=slice(0, 1000)).T),\n",
+    ")\n",
     "plt.xlim(0, 10000)\n",
     "plt.ylabel(\"log-energy\")\n",
     "plt.xlabel(\"gradient evaluations\");"
@@ -782,7 +784,13 @@
     }
    ],
    "source": [
-    "plt.plot((trace_cmdstan.warmup_sample_stats.n_steps).isel(draw=slice(0, 1000)).cumsum(\"draw\").T, np.log(trace_cmdstan.warmup_sample_stats.energy.isel(draw=slice(0, 1000)).T));\n",
+    "plt.plot(\n",
+    "    (trace_cmdstan.warmup_sample_stats.n_steps)\n",
+    "    .isel(draw=slice(0, 1000))\n",
+    "    .cumsum(\"draw\")\n",
+    "    .T,\n",
+    "    np.log(trace_cmdstan.warmup_sample_stats.energy.isel(draw=slice(0, 1000)).T),\n",
+    ")\n",
     "plt.xlim(0, 10000)\n",
     "plt.ylabel(\"log-energy\")\n",
     "plt.xlabel(\"gradient evaluations\");"
@@ -1607,7 +1615,12 @@
     }
    ],
    "source": [
-    "type({name: int(val) if isinstance(val, int) else list(val) for name, val in data_stan.items()}[\"county_idx\"][0])"
+    "type(\n",
+    "    {\n",
+    "        name: int(val) if isinstance(val, int) else list(val)\n",
+    "        for name, val in data_stan.items()\n",
+    "    }[\"county_idx\"][0]\n",
+    ")"
    ]
   },
   {
@@ -1622,13 +1635,13 @@
     "    if isinstance(val, int):\n",
     "        data_json[name] = int(val)\n",
     "        continue\n",
-    "    \n",
+    "\n",
     "    if val.dtype == np.int64:\n",
     "        data_json[name] = list(int(x) for x in val)\n",
     "        continue\n",
-    "    \n",
+    "\n",
     "    data_json[name] = list(val)\n",
-    "    \n",
+    "\n",
     "with open(\"radon.json\", \"w\") as file:\n",
     "    json.dump(data_json, file)"
    ]

diff --git a/python/nutpie/__init__.py b/python/nutpie/__init__.py
@@ -4,4 +4,4 @@
 from nutpie.sample import sample
 
 __version__: str = _lib.__version__
-__all__ = ["__version__", "sample", "compile_pymc_model", "compile_stan_model"]
+__all__ = ["__version__", "compile_pymc_model", "compile_stan_model", "sample"]