flowMC 0.4.0__tar.gz → 0.4.2__tar.gz

flowmc-0.4.2/.all-contributorsrc

@@ -0,0 +1,113 @@
+{
+  "projectName": "flowMC",
+  "projectOwner": "kazewong",
+  "files": [
+    "README.md"
+  ],
+  "commitType": "docs",
+  "commitConvention": "angular",
+  "contributorsPerLine": 7,
+  "contributors": [
+    {
+      "login": "HajimeKawahara",
+      "name": "Hajime Kawahara",
+      "avatar_url": "https://avatars.githubusercontent.com/u/15956904?v=4",
+      "profile": "http://secondearths.sakura.ne.jp/en/index.html",
+      "contributions": [
+        "bug"
+      ]
+    },
+    {
+      "login": "daniel-dodd",
+      "name": "Daniel Dodd",
+      "avatar_url": "https://avatars.githubusercontent.com/u/68821880?v=4",
+      "profile": "https://github.com/daniel-dodd",
+      "contributions": [
+        "doc",
+        "review",
+        "test",
+        "bug"
+      ]
+    },
+    {
+      "login": "matt-graham",
+      "name": "Matt Graham",
+      "avatar_url": "https://avatars.githubusercontent.com/u/6746980?v=4",
+      "profile": "http://matt-graham.github.io",
+      "contributions": [
+        "bug",
+        "test",
+        "review",
+        "doc"
+      ]
+    },
+    {
+      "login": "kazewong",
+      "name": "Kaze Wong",
+      "avatar_url": "https://avatars.githubusercontent.com/u/8803931?v=4",
+      "profile": "https://www.kaze-wong.com/",
+      "contributions": [
+        "bug",
+        "blog",
+        "code",
+        "content",
+        "doc",
+        "example",
+        "infra",
+        "maintenance",
+        "research",
+        "review",
+        "test",
+        "tutorial"
+      ]
+    },
+    {
+      "login": "marylou-gabrie",
+      "name": "Marylou Gabrié",
+      "avatar_url": "https://avatars.githubusercontent.com/u/11092071?v=4",
+      "profile": "https://marylou-gabrie.github.io/",
+      "contributions": [
+        "bug",
+        "code",
+        "content",
+        "doc",
+        "example",
+        "maintenance",
+        "research",
+        "test",
+        "tutorial"
+      ]
+    },
+    {
+      "login": "Qazalbash",
+      "name": "Meesum Qazalbash",
+      "avatar_url": "https://avatars.githubusercontent.com/u/62182585?v=4",
+      "profile": "https://github.com/Qazalbash",
+      "contributions": [
+        "code",
+        "maintenance"
+      ]
+    },
+    {
+      "login": "thomasckng",
+      "name": "Thomas Ng",
+      "avatar_url": "https://avatars.githubusercontent.com/u/97585527?v=4",
+      "profile": "https://github.com/thomasckng",
+      "contributions": [
+        "code",
+        "maintenance"
+      ]
+    },
+    {
+      "login": "tedwards2412",
+      "name": "Thomas Edwards",
+      "avatar_url": "https://avatars.githubusercontent.com/u/6105841?v=4",
+      "profile": "https://github.com/tedwards2412",
+      "contributions": [
+        "bug",
+        "code"
+      ]
+    }
+  ],
+  "repoType": "github"
+}

{flowmc-0.4.0 → flowmc-0.4.2}/.gitignore

@@ -13,3 +13,6 @@ README.md
 *settings.json
 *build*
 test_docs/test_quickstart.ipynb
+node_modules
+package.json
+yarn.lock

{flowmc-0.4.0 → flowmc-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flowMC
-Version: 0.4.0
+Version: 0.4.2
 Summary: Normalizing flow exhanced sampler in jax
 Project-URL: Documentation, https://github.com/kazewong/flowMC
 Author: Marylou Gabrié, Dan Foreman-Mackey
@@ -66,6 +66,9 @@ Description-Content-Type: text/markdown
 </a>
 <a href='https://coveralls.io/github/kazewong/flowMC?branch=main'><img src='https://img.shields.io/coverallsCoverage/github/kazewong/flowMC?style=flat-square' alt='Coverage Status' /></a>
 
+> [!WARNING]
+> Note that `flowMC` has not reached v1.0.0, meaning the API could subject to changes. In general, the higher level the API, the less likely it is going to change. However, intermediate level API such as the resource strategy interface could subject to major revision for performance concerns.
+
 ![flowMC_logo](./docs/logo_0810.png)
 
 flowMC is a Jax-based python package for normalizing-flow enhanced Markov chain Monte Carlo (MCMC) sampling.
@@ -163,4 +166,31 @@ This will help `flowMC` getting more recognition, and the main benefit *for you*
 > *Adaptive Monte Carlo augmented with normalizing flows.*
 Gabrié M., Rotskoff G. M., Vanden-Eijnden E. - PNAS 2022 - [doi](https://www.pnas.org/doi/10.1073/pnas.2109420119), [arxiv](https://arxiv.org/abs/2105.12603)
 
+## Contributors
+
+<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
+<!-- prettier-ignore-start -->
+<!-- markdownlint-disable -->
+<table>
+  <tbody>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="http://secondearths.sakura.ne.jp/en/index.html"><img src="https://avatars.githubusercontent.com/u/15956904?v=4?s=100" width="100px;" alt="Hajime Kawahara"/><br /><sub><b>Hajime Kawahara</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3AHajimeKawahara" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/daniel-dodd"><img src="https://avatars.githubusercontent.com/u/68821880?v=4?s=100" width="100px;" alt="Daniel Dodd"/><br /><sub><b>Daniel Dodd</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/commits?author=daniel-dodd" title="Documentation">📖</a> <a href="https://github.com/kazewong/flowMC/pulls?q=is%3Apr+reviewed-by%3Adaniel-dodd" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/kazewong/flowMC/commits?author=daniel-dodd" title="Tests">⚠️</a> <a href="https://github.com/kazewong/flowMC/issues?q=author%3Adaniel-dodd" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="http://matt-graham.github.io"><img src="https://avatars.githubusercontent.com/u/6746980?v=4?s=100" width="100px;" alt="Matt Graham"/><br /><sub><b>Matt Graham</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Amatt-graham" title="Bug reports">🐛</a> <a href="https://github.com/kazewong/flowMC/commits?author=matt-graham" title="Tests">⚠️</a> <a href="https://github.com/kazewong/flowMC/pulls?q=is%3Apr+reviewed-by%3Amatt-graham" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/kazewong/flowMC/commits?author=matt-graham" title="Documentation">📖</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://www.kaze-wong.com/"><img src="https://avatars.githubusercontent.com/u/8803931?v=4?s=100" width="100px;" alt="Kaze Wong"/><br /><sub><b>Kaze Wong</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Akazewong" title="Bug reports">🐛</a> <a href="#blog-kazewong" title="Blogposts">📝</a> <a href="https://github.com/kazewong/flowMC/commits?author=kazewong" title="Code">💻</a> <a href="#content-kazewong" title="Content">🖋</a> <a href="https://github.com/kazewong/flowMC/commits?author=kazewong" title="Documentation">📖</a> <a href="#example-kazewong" title="Examples">💡</a> <a href="#infra-kazewong" title="Infrastructure (Hosting, Build-Tools, etc)">🚇</a> <a href="#maintenance-kazewong" title="Maintenance">🚧</a> <a href="#research-kazewong" title="Research">🔬</a> <a href="https://github.com/kazewong/flowMC/pulls?q=is%3Apr+reviewed-by%3Akazewong" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/kazewong/flowMC/commits?author=kazewong" title="Tests">⚠️</a> <a href="#tutorial-kazewong" title="Tutorials">✅</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://marylou-gabrie.github.io/"><img src="https://avatars.githubusercontent.com/u/11092071?v=4?s=100" width="100px;" alt="Marylou Gabrié"/><br /><sub><b>Marylou Gabrié</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Amarylou-gabrie" title="Bug reports">🐛</a> <a href="https://github.com/kazewong/flowMC/commits?author=marylou-gabrie" title="Code">💻</a> <a href="#content-marylou-gabrie" title="Content">🖋</a> <a href="https://github.com/kazewong/flowMC/commits?author=marylou-gabrie" title="Documentation">📖</a> <a href="#example-marylou-gabrie" title="Examples">💡</a> <a href="#maintenance-marylou-gabrie" title="Maintenance">🚧</a> <a href="#research-marylou-gabrie" title="Research">🔬</a> <a href="https://github.com/kazewong/flowMC/commits?author=marylou-gabrie" title="Tests">⚠️</a> <a href="#tutorial-marylou-gabrie" title="Tutorials">✅</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/Qazalbash"><img src="https://avatars.githubusercontent.com/u/62182585?v=4?s=100" width="100px;" alt="Meesum Qazalbash"/><br /><sub><b>Meesum Qazalbash</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/commits?author=Qazalbash" title="Code">💻</a> <a href="#maintenance-Qazalbash" title="Maintenance">🚧</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/thomasckng"><img src="https://avatars.githubusercontent.com/u/97585527?v=4?s=100" width="100px;" alt="Thomas Ng"/><br /><sub><b>Thomas Ng</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/commits?author=thomasckng" title="Code">💻</a> <a href="#maintenance-thomasckng" title="Maintenance">🚧</a></td>
+    </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/tedwards2412"><img src="https://avatars.githubusercontent.com/u/6105841?v=4?s=100" width="100px;" alt="Thomas Edwards"/><br /><sub><b>Thomas Edwards</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Atedwards2412" title="Bug reports">🐛</a> <a href="https://github.com/kazewong/flowMC/commits?author=tedwards2412" title="Code">💻</a></td>
+    </tr>
+  </tbody>
+</table>
+
+<!-- markdownlint-restore -->
+<!-- prettier-ignore-end -->
+
+<!-- ALL-CONTRIBUTORS-LIST:END -->
+
  

flowmc-0.4.2/README.md

@@ -0,0 +1,140 @@
+# flowMC
+
+**Normalizing-flow enhanced sampling package for probabilistic inference**
+
+<a href="https://flowmc.readthedocs.io/en/main/">
+<img src="https://readthedocs.org/projects/flowmc/badge/?version=main&style=flat-square" alt="doc"/>
+</a>
+<a href="https://github.com/kazewong/flowMC/blob/Packaging/LICENSE">
+<img src="https://img.shields.io/badge/License-MIT-blue?style=flat-square" alt="doc"/>
+</a>
+<a href='https://coveralls.io/github/kazewong/flowMC?branch=main'><img src='https://img.shields.io/coverallsCoverage/github/kazewong/flowMC?style=flat-square' alt='Coverage Status' /></a>
+
+> [!WARNING]
+> Note that `flowMC` has not reached v1.0.0, meaning the API could subject to changes. In general, the higher level the API, the less likely it is going to change. However, intermediate level API such as the resource strategy interface could subject to major revision for performance concerns.
+
+![flowMC_logo](./docs/logo_0810.png)
+
+flowMC is a Jax-based python package for normalizing-flow enhanced Markov chain Monte Carlo (MCMC) sampling.
+The code is open source under MIT license, and it is under active development.
+
+- Just-in-time compilation is supported.
+- Native support for GPU acceleration.
+- Suit for problems with multi-modality.
+- Minimal tuning.
+
+# Installation 
+
+The simplest way to install the package is to do it through pip
+
+```
+pip install flowMC
+```
+
+This will install the latest stable release and its dependencies.
+flowMC is based on [Jax](https://github.com/google/jax) and [Equinox](https://github.com/patrick-kidger/equinox).
+By default, installing flowMC will automatically install Jax and Equinox available on [PyPI](https://pypi.org).
+Jax does not install GPU support by default.
+If you want to use GPU with Jax, you need to install Jax with GPU support according to their document.
+At the time of writing this documentation page, this is the command to install Jax with GPU support:
+
+```
+pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+
+If you want to install the latest version of flowMC, you can clone this repo and install it locally:
+
+```
+git clone https://github.com/kazewong/flowMC.git
+cd flowMC
+pip install -e .
+```
+
+## Requirements
+
+Here is a list of packages we use in the main library
+
+    * Python 3.9+
+    * Jax
+    * Jaxlib
+    * equinox
+
+To visualize the inference results in the examples, we requrie the following packages in addtion to the above:
+
+    * matplotlib
+    * corner
+    * arviz
+
+The test suite is based on pytest. To run the tests, one needs to install `pytest` and run `pytest` at the root directory of this repo.
+
+# Attribution
+
+If you used `flowMC` in your research, we would really appericiate it if you could at least cite the following papers:
+
+```
+@article{Wong:2022xvh,
+    author = "Wong, Kaze W. k. and Gabri\'e, Marylou and Foreman-Mackey, Daniel",
+    title = "{flowMC: Normalizing flow enhanced sampling package for probabilistic inference in JAX}",
+    eprint = "2211.06397",
+    archivePrefix = "arXiv",
+    primaryClass = "astro-ph.IM",
+    doi = "10.21105/joss.05021",
+    journal = "J. Open Source Softw.",
+    volume = "8",
+    number = "83",
+    pages = "5021",
+    year = "2023"
+}
+
+@article{Gabrie:2021tlu,
+    author = "Gabri\'e, Marylou and Rotskoff, Grant M. and Vanden-Eijnden, Eric",
+    title = "{Adaptive Monte Carlo augmented with normalizing flows}",
+    eprint = "2105.12603",
+    archivePrefix = "arXiv",
+    primaryClass = "physics.data-an",
+    doi = "10.1073/pnas.2109420119",
+    journal = "Proc. Nat. Acad. Sci.",
+    volume = "119",
+    number = "10",
+    pages = "e2109420119",
+    year = "2022"
+}
+```
+
+This will help `flowMC` getting more recognition, and the main benefit *for you* is this means the `flowMC` community will grow and it will be continuously improved. If you believe in the magic of open-source software, please support us by attributing our software in your work.
+
+
+`flowMC` is a Jax implementation of methods described in: 
+> *Efficient Bayesian Sampling Using Normalizing Flows to Assist Markov Chain Monte Carlo Methods* Gabrié M., Rotskoff G. M., Vanden-Eijnden E. - ICML INNF+ workshop 2021 - [pdf](https://openreview.net/pdf?id=mvtooHbjOwx)
+
+> *Adaptive Monte Carlo augmented with normalizing flows.*
+Gabrié M., Rotskoff G. M., Vanden-Eijnden E. - PNAS 2022 - [doi](https://www.pnas.org/doi/10.1073/pnas.2109420119), [arxiv](https://arxiv.org/abs/2105.12603)
+
+## Contributors
+
+<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
+<!-- prettier-ignore-start -->
+<!-- markdownlint-disable -->
+<table>
+  <tbody>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="http://secondearths.sakura.ne.jp/en/index.html"><img src="https://avatars.githubusercontent.com/u/15956904?v=4?s=100" width="100px;" alt="Hajime Kawahara"/><br /><sub><b>Hajime Kawahara</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3AHajimeKawahara" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/daniel-dodd"><img src="https://avatars.githubusercontent.com/u/68821880?v=4?s=100" width="100px;" alt="Daniel Dodd"/><br /><sub><b>Daniel Dodd</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/commits?author=daniel-dodd" title="Documentation">📖</a> <a href="https://github.com/kazewong/flowMC/pulls?q=is%3Apr+reviewed-by%3Adaniel-dodd" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/kazewong/flowMC/commits?author=daniel-dodd" title="Tests">⚠️</a> <a href="https://github.com/kazewong/flowMC/issues?q=author%3Adaniel-dodd" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="http://matt-graham.github.io"><img src="https://avatars.githubusercontent.com/u/6746980?v=4?s=100" width="100px;" alt="Matt Graham"/><br /><sub><b>Matt Graham</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Amatt-graham" title="Bug reports">🐛</a> <a href="https://github.com/kazewong/flowMC/commits?author=matt-graham" title="Tests">⚠️</a> <a href="https://github.com/kazewong/flowMC/pulls?q=is%3Apr+reviewed-by%3Amatt-graham" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/kazewong/flowMC/commits?author=matt-graham" title="Documentation">📖</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://www.kaze-wong.com/"><img src="https://avatars.githubusercontent.com/u/8803931?v=4?s=100" width="100px;" alt="Kaze Wong"/><br /><sub><b>Kaze Wong</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Akazewong" title="Bug reports">🐛</a> <a href="#blog-kazewong" title="Blogposts">📝</a> <a href="https://github.com/kazewong/flowMC/commits?author=kazewong" title="Code">💻</a> <a href="#content-kazewong" title="Content">🖋</a> <a href="https://github.com/kazewong/flowMC/commits?author=kazewong" title="Documentation">📖</a> <a href="#example-kazewong" title="Examples">💡</a> <a href="#infra-kazewong" title="Infrastructure (Hosting, Build-Tools, etc)">🚇</a> <a href="#maintenance-kazewong" title="Maintenance">🚧</a> <a href="#research-kazewong" title="Research">🔬</a> <a href="https://github.com/kazewong/flowMC/pulls?q=is%3Apr+reviewed-by%3Akazewong" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/kazewong/flowMC/commits?author=kazewong" title="Tests">⚠️</a> <a href="#tutorial-kazewong" title="Tutorials">✅</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://marylou-gabrie.github.io/"><img src="https://avatars.githubusercontent.com/u/11092071?v=4?s=100" width="100px;" alt="Marylou Gabrié"/><br /><sub><b>Marylou Gabrié</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Amarylou-gabrie" title="Bug reports">🐛</a> <a href="https://github.com/kazewong/flowMC/commits?author=marylou-gabrie" title="Code">💻</a> <a href="#content-marylou-gabrie" title="Content">🖋</a> <a href="https://github.com/kazewong/flowMC/commits?author=marylou-gabrie" title="Documentation">📖</a> <a href="#example-marylou-gabrie" title="Examples">💡</a> <a href="#maintenance-marylou-gabrie" title="Maintenance">🚧</a> <a href="#research-marylou-gabrie" title="Research">🔬</a> <a href="https://github.com/kazewong/flowMC/commits?author=marylou-gabrie" title="Tests">⚠️</a> <a href="#tutorial-marylou-gabrie" title="Tutorials">✅</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/Qazalbash"><img src="https://avatars.githubusercontent.com/u/62182585?v=4?s=100" width="100px;" alt="Meesum Qazalbash"/><br /><sub><b>Meesum Qazalbash</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/commits?author=Qazalbash" title="Code">💻</a> <a href="#maintenance-Qazalbash" title="Maintenance">🚧</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/thomasckng"><img src="https://avatars.githubusercontent.com/u/97585527?v=4?s=100" width="100px;" alt="Thomas Ng"/><br /><sub><b>Thomas Ng</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/commits?author=thomasckng" title="Code">💻</a> <a href="#maintenance-thomasckng" title="Maintenance">🚧</a></td>
+    </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/tedwards2412"><img src="https://avatars.githubusercontent.com/u/6105841?v=4?s=100" width="100px;" alt="Thomas Edwards"/><br /><sub><b>Thomas Edwards</b></sub></a><br /><a href="https://github.com/kazewong/flowMC/issues?q=author%3Atedwards2412" title="Bug reports">🐛</a> <a href="https://github.com/kazewong/flowMC/commits?author=tedwards2412" title="Code">💻</a></td>
+    </tr>
+  </tbody>
+</table>
+
+<!-- markdownlint-restore -->
+<!-- prettier-ignore-end -->
+
+<!-- ALL-CONTRIBUTORS-LIST:END -->
+
+ 

{flowmc-0.4.0 → flowmc-0.4.2}/docs/tutorials/custom_strategy.ipynb

@@ -31,6 +31,7 @@
     "import jax.numpy as jnp\n",
     "from flowMC.resource.local_kernel.Gaussian_random_walk import GaussianRandomWalk\n",
     "from flowMC.resource.buffers import Buffer\n",
+    "from flowMC.resource.logPDF import LogPDF\n",
     "from flowMC.strategy.optimization import AdamOptimization\n",
     "from flowMC.strategy.take_steps import TakeSerialSteps\n",
     "from flowMC.Sampler import Sampler\n",
@@ -82,16 +83,19 @@
     "    \"log_prob\": log_prob,\n",
     "    \"acceptance\": acceptance,\n",
     "    \"RWMCMC\": RWMCMC_sampler,\n",
+    "    \"logpdf\": LogPDF(target_log_prob, n_dims=n_dims),\n",
     "}\n",
     "\n",
     "# Defining strategy\n",
     "\n",
-    "strategy = TakeSerialSteps(\n",
-    "    logpdf=target_log_prob,\n",
-    "    kernel_name=\"RWMCMC\",\n",
-    "    buffer_names=[\"positions\", \"log_prob\", \"acceptance\"],\n",
-    "    n_steps=n_steps,\n",
-    ")\n",
+    "strategy = {\n",
+    "    \"take_serial_step\": TakeSerialSteps(\n",
+    "        logpdf_name=\"logpdf\",\n",
+    "        kernel_name=\"RWMCMC\",\n",
+    "        buffer_names=[\"positions\", \"log_prob\", \"acceptance\"],\n",
+    "        n_steps=n_steps,\n",
+    "    )\n",
+    "}\n",
     "\n",
     "# Initializing sampler\n",
     "\n",
@@ -100,7 +104,8 @@
     "    n_chains=n_chains,\n",
     "    rng_key=rng_key,\n",
     "    resources=resource,\n",
-    "    strategies=[strategy],\n",
+    "    strategies=strategy,\n",
+    "    strategy_order=[\"take_serial_step\"],\n",
     ")"
    ]
   },
@@ -119,7 +124,7 @@
     {
      "data": {
       "text/plain": [
-       "<matplotlib.colorbar.Colorbar at 0x7fc714262cf0>"
+       "<matplotlib.colorbar.Colorbar at 0x7f96b5f2f230>"
       ]
      },
      "execution_count": 3,
@@ -181,7 +186,7 @@
     {
      "data": {
       "text/plain": [
-       "<matplotlib.collections.PathCollection at 0x7fc6f5f6bc50>"
+       "<matplotlib.collections.PathCollection at 0x7f96b55efed0>"
       ]
      },
      "execution_count": 4,
@@ -220,10 +225,11 @@
     "    \"log_prob\": log_prob,\n",
     "    \"acceptance\": acceptance,\n",
     "    \"RWMCMC\": RWMCMC_sampler,\n",
+    "    \"logpdf\": LogPDF(target_log_prob, n_dims=n_dims),\n",
     "}\n",
     "\n",
     "strategy = TakeSerialSteps(\n",
-    "    logpdf=target_log_prob,\n",
+    "    logpdf_name=\"logpdf\",\n",
     "    kernel_name=\"RWMCMC\",\n",
     "    buffer_names=[\"positions\", \"log_prob\", \"acceptance\"],\n",
     "    n_steps=n_steps,\n",
@@ -236,7 +242,11 @@
     "    n_chains=n_chains,\n",
     "    rng_key=rng_key,\n",
     "    resources=resource,\n",
-    "    strategies=[adam_optimizer, strategy],\n",
+    "    strategies={\n",
+    "        \"take_serial_step\": strategy,\n",
+    "        \"optimize\": adam_optimizer,\n",
+    "    },\n",
+    "    strategy_order=[\"optimize\", \"take_serial_step\"],\n",
     ")\n",
     "\n",
     "nf_sampler.sample(initial_position, {})\n",
@@ -275,7 +285,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.2"
+   "version": "3.11.2"
   }
  },
  "nbformat": 4,