From 72e286990c0ef0d1015df6a91dfce11995bed73f Mon Sep 17 00:00:00 2001 From: Ariel Shurygin Date: Fri, 6 Sep 2024 19:40:09 +0000 Subject: [PATCH 1/2] adding debug mode to inferer --- mechanistic_model/mechanistic_inferer.py | 29 + poetry.lock | 1042 +++++++++++++++++++++- pyproject.toml | 1 + 3 files changed, 1071 insertions(+), 1 deletion(-) diff --git a/mechanistic_model/mechanistic_inferer.py b/mechanistic_model/mechanistic_inferer.py index 50c05804..59182df4 100644 --- a/mechanistic_model/mechanistic_inferer.py +++ b/mechanistic_model/mechanistic_inferer.py @@ -8,6 +8,7 @@ import warnings from typing import Union +import bayeux as bx import jax.numpy as jnp import jax.typing import numpy as np @@ -336,6 +337,34 @@ def infer(self, obs_metrics: jax.Array) -> MCMC: self.infer_complete = True return self.inference_algo + def _debug_likelihood(self, obs_metrics) -> bx.Model: + """uses Bayeux to recreate the self.likelihood function for purposes of basic sanity checking + + Parameters + ---------- + obs_metrics: jnp.array + observed metrics on which likelihood will be calculated on to tune parameters. + See `likelihood()` method for implemented definition of `obs_metrics` + + Returns + ------- + Bayeux.Model + model object used to debug + """ + bx_model = bx.Model.from_numpyro( + jax.tree_util.Partial( + self.likelihood, + tf=len(obs_metrics), + obs_metrics=obs_metrics, + ), + # this does not work for non-one/sampled self.INITIAL_INFECTIONS_SCALE + initial_state=self.INITIAL_STATE, + ) + bx_model.mcmc.numpyro_nuts.debug( + seed=PRNGKey(self.config.INFERENCE_PRNGKEY) + ) + return bx_model + def checkpoint( self, checkpoint_path: str, group_by_chain: bool = True ) -> None: diff --git a/poetry.lock b/poetry.lock index cd9d7255..45ffd48a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,5 +1,16 @@ # This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +[[package]] +name = "absl-py" +version = "2.1.0" +description = "Abseil Python Common Libraries, see https://github.com/abseil/abseil-py." +optional = false +python-versions = ">=3.7" +files = [ + {file = "absl-py-2.1.0.tar.gz", hash = "sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff"}, + {file = "absl_py-2.1.0-py3-none-any.whl", hash = "sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308"}, +] + [[package]] name = "adal" version = "1.2.7" @@ -50,6 +61,34 @@ files = [ {file = "appdirs-1.4.4.tar.gz", hash = "sha256:7d5d0167b2b1ba821647616af46a749d1c653740dd0d2415100fe26e27afdf41"}, ] +[[package]] +name = "arviz" +version = "0.19.0" +description = "Exploratory analysis of Bayesian models" +optional = false +python-versions = ">=3.10" +files = [ + {file = "arviz-0.19.0-py3-none-any.whl", hash = "sha256:652e739f1931b898922ebe9a9dd393840a0f3f7b803b606938edadc9c263a5bc"}, + {file = "arviz-0.19.0.tar.gz", hash = "sha256:04a369cd05293a17a24fcc1e133f9f350d202cf307f8df6c1dbd319e0f6a9f4f"}, +] + +[package.dependencies] +dm-tree = ">=0.1.8" +h5netcdf = ">=1.0.2" +matplotlib = ">=3.5" +numpy = ">=1.23.0" +packaging = "*" +pandas = ">=1.5.0" +scipy = ">=1.9.0" +setuptools = ">=60.0.0" +typing-extensions = ">=4.1.0" +xarray = ">=2022.6.0" +xarray-einstats = ">=0.3" + +[package.extras] +all = ["bokeh (>=3)", "contourpy", "dask[distributed]", "netcdf4", "numba", "ujson", "xarray-datatree", "zarr (>=2.5.0,<3)"] +preview = ["arviz-base[h5netcdf]", "arviz-plots", "arviz-stats[xarray]"] + [[package]] name = "asgiref" version = "3.8.1" @@ -259,6 +298,65 @@ typing-extensions = ">=4.3.0" [package.extras] aio = ["azure-core[aio] (>=1.28.0,<2.0.0)"] +[[package]] +name = "bayeux-ml" +version = "0.1.14" +description = "Stitching together probabilistic models and inference." +optional = false +python-versions = ">=3.9" +files = [ + {file = "bayeux_ml-0.1.14-py3-none-any.whl", hash = "sha256:c29040c7184d1b17607800c24badb2502b091283bad9d23c245f82874e6086df"}, + {file = "bayeux_ml-0.1.14.tar.gz", hash = "sha256:310b736ac7f3d9399b4fdab5129c85cb36d27cdef1059764e19d2c482fe57eae"}, +] + +[package.dependencies] +arviz = "*" +blackjax = "*" +flowmc = ">=0.3.0" +jax = ">=0.4.6" +jaxopt = "*" +numpyro = "*" +nutpie = "*" +optax = "*" +optimistix = "*" +oryx = ">=0.2.7" +pymc = "*" +tensorflow-probability = {version = ">=0.19.0", extras = ["jax"]} + +[package.extras] +dev = ["pyink", "pylint (>=2.6.0)", "pytest", "pytest-xdist"] +docs = ["mkdocs (==1.5.3)", "mkdocs-material (==9.5.11)", "mkdocstrings (==0.24.1)", "mknotebooks (==0.8.0)", "pymdown-extensions (==10.7)"] + +[[package]] +name = "blackjax" +version = "1.2.3" +description = "Flexible and fast sampling in Python" +optional = false +python-versions = ">=3.9" +files = [ + {file = "blackjax-1.2.3-py3-none-any.whl", hash = "sha256:57f0336a23c9075ca57bcefacb12d1b6de68d8a77075bbfe2107fc427b43626c"}, + {file = "blackjax-1.2.3.tar.gz", hash = "sha256:c1e0ad412faa229eb75966e1f7b881844455295d33b58841d683e8a99c22c117"}, +] + +[package.dependencies] +fastprogress = ">=1.0.0" +jax = ">=0.4.16" +jaxlib = ">=0.4.16" +jaxopt = ">=0.8" +optax = ">=0.1.7" +typing-extensions = ">=4.4.0" + +[[package]] +name = "cachetools" +version = "5.5.0" +description = "Extensible memoizing collections and decorators" +optional = false +python-versions = ">=3.7" +files = [ + {file = "cachetools-5.5.0-py3-none-any.whl", hash = "sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292"}, + {file = "cachetools-5.5.0.tar.gz", hash = "sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a"}, +] + [[package]] name = "certifi" version = "2024.7.4" @@ -487,6 +585,26 @@ files = [ {file = "charset_normalizer-3.3.2-py3-none-any.whl", hash = "sha256:3e4d1f6587322d2788836a99c69062fbb091331ec940e02d12d179c1d53e25fc"}, ] +[[package]] +name = "chex" +version = "0.1.86" +description = "Chex: Testing made fun, in JAX!" +optional = false +python-versions = ">=3.9" +files = [ + {file = "chex-0.1.86-py3-none-any.whl", hash = "sha256:251c20821092323a3d9c28e1cf80e4a58180978bec368f531949bd9847eee568"}, + {file = "chex-0.1.86.tar.gz", hash = "sha256:e8b0f96330eba4144659e1617c0f7a57b161e8cbb021e55c6d5056c7378091d1"}, +] + +[package.dependencies] +absl-py = ">=0.9.0" +jax = ">=0.4.16" +jaxlib = ">=0.1.37" +numpy = ">=1.24.1" +setuptools = {version = "*", markers = "python_version >= \"3.12\""} +toolz = ">=0.9.0" +typing-extensions = ">=4.2.0" + [[package]] name = "click" version = "8.1.7" @@ -501,6 +619,17 @@ files = [ [package.dependencies] colorama = {version = "*", markers = "platform_system == \"Windows\""} +[[package]] +name = "cloudpickle" +version = "3.0.0" +description = "Pickler class to extend the standard pickle.Pickler functionality" +optional = false +python-versions = ">=3.8" +files = [ + {file = "cloudpickle-3.0.0-py3-none-any.whl", hash = "sha256:246ee7d0c295602a036e86369c77fecda4ab17b506496730f2f576d9016fd9c7"}, + {file = "cloudpickle-3.0.0.tar.gz", hash = "sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882"}, +] + [[package]] name = "colorama" version = "0.4.6" @@ -512,6 +641,19 @@ files = [ {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] +[[package]] +name = "cons" +version = "0.4.6" +description = "An implementation of Lisp/Scheme-like cons in Python." +optional = false +python-versions = ">=3.6" +files = [ + {file = "cons-0.4.6.tar.gz", hash = "sha256:669fe9d5ee916d5e42b9cac6acc911df803d04f2e945c1604982a04d27a29b47"}, +] + +[package.dependencies] +logical-unification = ">=0.4.0" + [[package]] name = "contourpy" version = "1.2.1" @@ -575,6 +717,25 @@ mypy = ["contourpy[bokeh,docs]", "docutils-stubs", "mypy (==1.8.0)", "types-Pill test = ["Pillow", "contourpy[test-no-images]", "matplotlib"] test-no-images = ["pytest", "pytest-cov", "pytest-xdist", "wurlitzer"] +[[package]] +name = "corner" +version = "2.2.2" +description = "Make some beautiful corner plots" +optional = false +python-versions = ">=3.9" +files = [ + {file = "corner-2.2.2-py3-none-any.whl", hash = "sha256:e7577cdb59cfa304effa243b0c7ac0e3777030d3dc2f2e217a387e87a47074bb"}, + {file = "corner-2.2.2.tar.gz", hash = "sha256:4bc79f3b6778c270103f0926e64ef2606c48c3b6f92daf5382fc4babf5d608d1"}, +] + +[package.dependencies] +matplotlib = ">=2.1" + +[package.extras] +arviz = ["arviz (>=0.9)"] +docs = ["arviz (>=0.9)", "ipython", "myst-nb", "pandoc", "sphinx (>=1.7.5)", "sphinx-book-theme"] +test = ["arviz (>=0.9)", "pytest", "scipy"] + [[package]] name = "cryptography" version = "43.0.0" @@ -639,6 +800,17 @@ files = [ docs = ["ipython", "matplotlib", "numpydoc", "sphinx"] tests = ["pytest", "pytest-cov", "pytest-xdist"] +[[package]] +name = "decorator" +version = "5.1.1" +description = "Decorators for Humans" +optional = false +python-versions = ">=3.5" +files = [ + {file = "decorator-5.1.1-py3-none-any.whl", hash = "sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186"}, + {file = "decorator-5.1.1.tar.gz", hash = "sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330"}, +] + [[package]] name = "diffrax" version = "0.5.0" @@ -669,6 +841,61 @@ files = [ {file = "distlib-0.3.8.tar.gz", hash = "sha256:1530ea13e350031b6312d8580ddb6b27a104275a31106523b8f123787f494f64"}, ] +[[package]] +name = "dm-tree" +version = "0.1.8" +description = "Tree is a library for working with nested data structures." +optional = false +python-versions = "*" +files = [ + {file = "dm-tree-0.1.8.tar.gz", hash = "sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430"}, + {file = "dm_tree-0.1.8-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60"}, + {file = "dm_tree-0.1.8-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f"}, + {file = "dm_tree-0.1.8-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef"}, + {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436"}, + {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410"}, + {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca"}, + {file = "dm_tree-0.1.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144"}, + {file = "dm_tree-0.1.8-cp310-cp310-win_amd64.whl", hash = "sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee"}, + {file = "dm_tree-0.1.8-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7"}, + {file = "dm_tree-0.1.8-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b"}, + {file = "dm_tree-0.1.8-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5"}, + {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de"}, + {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e"}, + {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d"}, + {file = "dm_tree-0.1.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393"}, + {file = "dm_tree-0.1.8-cp311-cp311-win_amd64.whl", hash = "sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80"}, + {file = "dm_tree-0.1.8-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:ea9e59e0451e7d29aece402d9f908f2e2a80922bcde2ebfd5dcb07750fcbfee8"}, + {file = "dm_tree-0.1.8-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:94d3f0826311f45ee19b75f5b48c99466e4218a0489e81c0f0167bda50cacf22"}, + {file = "dm_tree-0.1.8-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:435227cf3c5dc63f4de054cf3d00183790bd9ead4c3623138c74dde7f67f521b"}, + {file = "dm_tree-0.1.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:09964470f76a5201aff2e8f9b26842976de7889300676f927930f6285e256760"}, + {file = "dm_tree-0.1.8-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:75c5d528bb992981c20793b6b453e91560784215dffb8a5440ba999753c14ceb"}, + {file = "dm_tree-0.1.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0a94aba18a35457a1b5cd716fd7b46c5dafdc4cf7869b4bae665b91c4682a8e"}, + {file = "dm_tree-0.1.8-cp312-cp312-win_amd64.whl", hash = "sha256:96a548a406a6fb15fe58f6a30a57ff2f2aafbf25f05afab00c8f5e5977b6c715"}, + {file = "dm_tree-0.1.8-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571"}, + {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d"}, + {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb"}, + {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6"}, + {file = "dm_tree-0.1.8-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1"}, + {file = "dm_tree-0.1.8-cp37-cp37m-win_amd64.whl", hash = "sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6"}, + {file = "dm_tree-0.1.8-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf"}, + {file = "dm_tree-0.1.8-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a"}, + {file = "dm_tree-0.1.8-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d"}, + {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c"}, + {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8"}, + {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68"}, + {file = "dm_tree-0.1.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134"}, + {file = "dm_tree-0.1.8-cp38-cp38-win_amd64.whl", hash = "sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5"}, + {file = "dm_tree-0.1.8-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f"}, + {file = "dm_tree-0.1.8-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf"}, + {file = "dm_tree-0.1.8-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7"}, + {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb"}, + {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913"}, + {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426"}, + {file = "dm_tree-0.1.8-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317"}, + {file = "dm_tree-0.1.8-cp39-cp39-win_amd64.whl", hash = "sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368"}, +] + [[package]] name = "docker" version = "7.1.0" @@ -691,6 +918,17 @@ docs = ["myst-parser (==0.18.0)", "sphinx (==5.1.1)"] ssh = ["paramiko (>=2.4.3)"] websockets = ["websocket-client (>=1.3.0)"] +[[package]] +name = "dotmap" +version = "1.3.30" +description = "ordered, dynamically-expandable dot-access dictionary" +optional = false +python-versions = "*" +files = [ + {file = "dotmap-1.3.30-py3-none-any.whl", hash = "sha256:bd9fa15286ea2ad899a4d1dc2445ed85a1ae884a42effb87c89a6ecce71243c6"}, + {file = "dotmap-1.3.30.tar.gz", hash = "sha256:5821a7933f075fb47563417c0e92e0b7c031158b4c9a6a7e56163479b658b368"}, +] + [[package]] name = "epiweeks" version = "2.3.0" @@ -718,6 +956,78 @@ jax = ">=0.4.13" jaxtyping = ">=0.2.20" typing-extensions = ">=4.5.0" +[[package]] +name = "etils" +version = "1.9.4" +description = "Collection of common python utils" +optional = false +python-versions = ">=3.10" +files = [ + {file = "etils-1.9.4-py3-none-any.whl", hash = "sha256:4387e7a4911a3b5cc4b92b99a9211386d176b43bae1dac8e2fe345fc2cb95e4b"}, + {file = "etils-1.9.4.tar.gz", hash = "sha256:fad950414f0a1ca58c70c70915b0014f9953dd9bcf8aa951a0f75ff9becbeb24"}, +] + +[package.dependencies] +fsspec = {version = "*", optional = true, markers = "extra == \"epath\""} +importlib_resources = {version = "*", optional = true, markers = "extra == \"epath\""} +typing_extensions = {version = "*", optional = true, markers = "extra == \"epy\""} +zipp = {version = "*", optional = true, markers = "extra == \"epath\""} + +[package.extras] +all = ["etils[array-types]", "etils[eapp]", "etils[ecolab]", "etils[edc]", "etils[enp]", "etils[epath-gcs]", "etils[epath-s3]", "etils[epath]", "etils[epy]", "etils[etqdm]", "etils[etree-dm]", "etils[etree-jax]", "etils[etree-tf]", "etils[etree]"] +array-types = ["etils[enp]"] +dev = ["chex", "optree", "pydantic", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-subtests", "pytest-xdist", "tensorflow_datasets", "torch"] +docs = ["etils[all,dev]", "sphinx-apitree[ext]"] +eapp = ["absl-py", "etils[epy]", "simple_parsing"] +ecolab = ["etils[enp]", "etils[epy]", "etils[etree]", "jupyter", "mediapy", "numpy", "packaging", "protobuf"] +edc = ["etils[epy]"] +enp = ["etils[epy]", "numpy"] +epath = ["etils[epy]", "fsspec", "importlib_resources", "typing_extensions", "zipp"] +epath-gcs = ["etils[epath]", "gcsfs"] +epath-s3 = ["etils[epath]", "s3fs"] +epy = ["typing_extensions"] +etqdm = ["absl-py", "etils[epy]", "tqdm"] +etree = ["etils[array-types]", "etils[enp]", "etils[epy]", "etils[etqdm]"] +etree-dm = ["dm-tree", "etils[etree]"] +etree-jax = ["etils[etree]", "jax[cpu]"] +etree-tf = ["etils[etree]", "tensorflow"] +lazy-imports = ["etils[ecolab]"] + +[[package]] +name = "etuples" +version = "0.3.9" +description = "Python S-expression emulation using tuple-like objects." +optional = false +python-versions = ">=3.8" +files = [ + {file = "etuples-0.3.9.tar.gz", hash = "sha256:a474e586683d8ba8d842ba29305005ceed1c08371a4b4b0e0e232527137e5ea3"}, +] + +[package.dependencies] +cons = "*" +multipledispatch = "*" + +[[package]] +name = "evosax" +version = "0.1.6" +description = "JAX-Based Evolution Strategies" +optional = false +python-versions = ">=3.7" +files = [ + {file = "evosax-0.1.6-py3-none-any.whl", hash = "sha256:195145d68c67f5f1de9c580b6f248f624ac9d089de628e67665f7b3616d638e4"}, + {file = "evosax-0.1.6.tar.gz", hash = "sha256:bc3dd032191f2f878b08127e0c88d270d08fc95e80b379509b14dc265f29484b"}, +] + +[package.dependencies] +chex = "*" +dotmap = "*" +flax = "*" +jax = "*" +jaxlib = ">=0.1.37" +matplotlib = "*" +numpy = "*" +pyyaml = "*" + [[package]] name = "exceptiongroup" version = "1.2.2" @@ -732,6 +1042,17 @@ files = [ [package.extras] test = ["pytest (>=6)"] +[[package]] +name = "fastprogress" +version = "1.0.3" +description = "A nested progress with plotting options for fastai" +optional = false +python-versions = ">=3.6" +files = [ + {file = "fastprogress-1.0.3-py3-none-any.whl", hash = "sha256:6dfea88f7a4717b0a8d6ee2048beae5dbed369f932a368c5dd9caff34796f7c5"}, + {file = "fastprogress-1.0.3.tar.gz", hash = "sha256:7a17d2b438890f838c048eefce32c4ded47197ecc8ea042cecc33d3deb8022f5"}, +] + [[package]] name = "filelock" version = "3.15.4" @@ -748,6 +1069,56 @@ docs = ["furo (>=2023.9.10)", "sphinx (>=7.2.6)", "sphinx-autodoc-typehints (>=1 testing = ["covdefaults (>=2.3)", "coverage (>=7.3.2)", "diff-cover (>=8.0.1)", "pytest (>=7.4.3)", "pytest-asyncio (>=0.21)", "pytest-cov (>=4.1)", "pytest-mock (>=3.12)", "pytest-timeout (>=2.2)", "virtualenv (>=20.26.2)"] typing = ["typing-extensions (>=4.8)"] +[[package]] +name = "flax" +version = "0.8.4" +description = "Flax: A neural network library for JAX designed for flexibility" +optional = false +python-versions = ">=3.9" +files = [ + {file = "flax-0.8.4-py3-none-any.whl", hash = "sha256:785707e3a48f782a1bec17aa665697b7618c113a357d5f975791dcb090d818d8"}, + {file = "flax-0.8.4.tar.gz", hash = "sha256:968683f850198e1aa5eb2d9d1e20bead880ef7423c14f042db9d60848cb1c90b"}, +] + +[package.dependencies] +jax = ">=0.4.19" +msgpack = "*" +numpy = [ + {version = ">=1.22", markers = "python_version < \"3.11\""}, + {version = ">=1.23.2", markers = "python_version >= \"3.11\" and python_version < \"3.12\""}, + {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, +] +optax = "*" +orbax-checkpoint = "*" +PyYAML = ">=5.4.1" +rich = ">=11.1" +tensorstore = "*" +typing-extensions = ">=4.2" + +[package.extras] +all = ["matplotlib"] +testing = ["black[jupyter] (==23.7.0)", "clu", "clu (<=0.0.9)", "einops", "gymnasium[accept-rom-license,atari]", "jaxlib", "jraph (>=0.0.6dev0)", "ml-collections", "mypy", "nbstripout", "opencv-python", "penzai", "pytest", "pytest-cov", "pytest-custom-exit-code", "pytest-xdist", "pytype", "sentencepiece", "tensorflow", "tensorflow-datasets", "tensorflow-text (>=2.11.0)", "torch"] + +[[package]] +name = "flowmc" +version = "0.3.4" +description = "Normalizing flow exhanced sampler in jax" +optional = false +python-versions = ">=3.10" +files = [ + {file = "flowMC-0.3.4-py3-none-any.whl", hash = "sha256:b65c2cc892c07056ea51f1056a4ef4df8d8075b924a76f242efb5b37654716bf"}, + {file = "flowmc-0.3.4.tar.gz", hash = "sha256:ad64011efdb58ff624b176bc0853aee549fdeb37e5a7a7081580eb5009b2945f"}, +] + +[package.dependencies] +corner = "*" +equinox = ">=0.10.6" +evosax = ">=0.1.4" +jax = ">=0.4.12" +jaxlib = ">=0.4.12" +optax = ">=0.1.5" +tqdm = "*" + [[package]] name = "fonttools" version = "4.53.1" @@ -813,6 +1184,56 @@ ufo = ["fs (>=2.2.0,<3)"] unicode = ["unicodedata2 (>=15.1.0)"] woff = ["brotli (>=1.0.1)", "brotlicffi (>=0.8.0)", "zopfli (>=0.1.4)"] +[[package]] +name = "fsspec" +version = "2024.9.0" +description = "File-system specification" +optional = false +python-versions = ">=3.8" +files = [ + {file = "fsspec-2024.9.0-py3-none-any.whl", hash = "sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b"}, + {file = "fsspec-2024.9.0.tar.gz", hash = "sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8"}, +] + +[package.extras] +abfs = ["adlfs"] +adl = ["adlfs"] +arrow = ["pyarrow (>=1)"] +dask = ["dask", "distributed"] +dev = ["pre-commit", "ruff"] +doc = ["numpydoc", "sphinx", "sphinx-design", "sphinx-rtd-theme", "yarl"] +dropbox = ["dropbox", "dropboxdrivefs", "requests"] +full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "dask", "distributed", "dropbox", "dropboxdrivefs", "fusepy", "gcsfs", "libarchive-c", "ocifs", "panel", "paramiko", "pyarrow (>=1)", "pygit2", "requests", "s3fs", "smbprotocol", "tqdm"] +fuse = ["fusepy"] +gcs = ["gcsfs"] +git = ["pygit2"] +github = ["requests"] +gs = ["gcsfs"] +gui = ["panel"] +hdfs = ["pyarrow (>=1)"] +http = ["aiohttp (!=4.0.0a0,!=4.0.0a1)"] +libarchive = ["libarchive-c"] +oci = ["ocifs"] +s3 = ["s3fs"] +sftp = ["paramiko"] +smb = ["smbprotocol"] +ssh = ["paramiko"] +test = ["aiohttp (!=4.0.0a0,!=4.0.0a1)", "numpy", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "requests"] +test-downstream = ["aiobotocore (>=2.5.4,<3.0.0)", "dask-expr", "dask[dataframe,test]", "moto[server] (>4,<5)", "pytest-timeout", "xarray"] +test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas", "panel", "paramiko", "pyarrow", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "smbprotocol", "tqdm", "urllib3", "zarr", "zstandard"] +tqdm = ["tqdm"] + +[[package]] +name = "gast" +version = "0.6.0" +description = "Python AST that abstracts the underlying Python version" +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,>=2.7" +files = [ + {file = "gast-0.6.0-py3-none-any.whl", hash = "sha256:52b182313f7330389f72b069ba00f174cfe2a06411099547288839c6cbafbd54"}, + {file = "gast-0.6.0.tar.gz", hash = "sha256:88fc5300d32c7ac6ca7b515310862f71e6fdf2c029bbec7c66c0f5dd47b6b1fb"}, +] + [[package]] name = "gitdb" version = "4.0.11" @@ -856,6 +1277,57 @@ files = [ {file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"}, ] +[[package]] +name = "h5netcdf" +version = "1.3.0" +description = "netCDF4 via h5py" +optional = false +python-versions = ">=3.9" +files = [ + {file = "h5netcdf-1.3.0-py3-none-any.whl", hash = "sha256:f2df69dcd3665dc9c4d43eb6529dedd113b2508090d12ac973573305a8406465"}, + {file = "h5netcdf-1.3.0.tar.gz", hash = "sha256:a171c027daeb34b24c24a3b6304195b8eabbb6f10c748256ed3cfe19806383cf"}, +] + +[package.dependencies] +h5py = "*" +packaging = "*" + +[package.extras] +test = ["netCDF4", "pytest"] + +[[package]] +name = "h5py" +version = "3.11.0" +description = "Read and write HDF5 files from Python" +optional = false +python-versions = ">=3.8" +files = [ + {file = "h5py-3.11.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:1625fd24ad6cfc9c1ccd44a66dac2396e7ee74940776792772819fc69f3a3731"}, + {file = "h5py-3.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c072655ad1d5fe9ef462445d3e77a8166cbfa5e599045f8aa3c19b75315f10e5"}, + {file = "h5py-3.11.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:77b19a40788e3e362b54af4dcf9e6fde59ca016db2c61360aa30b47c7b7cef00"}, + {file = "h5py-3.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:ef4e2f338fc763f50a8113890f455e1a70acd42a4d083370ceb80c463d803972"}, + {file = "h5py-3.11.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:bbd732a08187a9e2a6ecf9e8af713f1d68256ee0f7c8b652a32795670fb481ba"}, + {file = "h5py-3.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:75bd7b3d93fbeee40860fd70cdc88df4464e06b70a5ad9ce1446f5f32eb84007"}, + {file = "h5py-3.11.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:52c416f8eb0daae39dabe71415cb531f95dce2d81e1f61a74537a50c63b28ab3"}, + {file = "h5py-3.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:083e0329ae534a264940d6513f47f5ada617da536d8dccbafc3026aefc33c90e"}, + {file = "h5py-3.11.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:a76cae64080210389a571c7d13c94a1a6cf8cb75153044fd1f822a962c97aeab"}, + {file = "h5py-3.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f3736fe21da2b7d8a13fe8fe415f1272d2a1ccdeff4849c1421d2fb30fd533bc"}, + {file = "h5py-3.11.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aa6ae84a14103e8dc19266ef4c3e5d7c00b68f21d07f2966f0ca7bdb6c2761fb"}, + {file = "h5py-3.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:21dbdc5343f53b2e25404673c4f00a3335aef25521bd5fa8c707ec3833934892"}, + {file = "h5py-3.11.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:754c0c2e373d13d6309f408325343b642eb0f40f1a6ad21779cfa9502209e150"}, + {file = "h5py-3.11.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:731839240c59ba219d4cb3bc5880d438248533366f102402cfa0621b71796b62"}, + {file = "h5py-3.11.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8ec9df3dd2018904c4cc06331951e274f3f3fd091e6d6cc350aaa90fa9b42a76"}, + {file = "h5py-3.11.0-cp38-cp38-win_amd64.whl", hash = "sha256:55106b04e2c83dfb73dc8732e9abad69d83a436b5b82b773481d95d17b9685e1"}, + {file = "h5py-3.11.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:f4e025e852754ca833401777c25888acb96889ee2c27e7e629a19aee288833f0"}, + {file = "h5py-3.11.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:6c4b760082626120031d7902cd983d8c1f424cdba2809f1067511ef283629d4b"}, + {file = "h5py-3.11.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:67462d0669f8f5459529de179f7771bd697389fcb3faab54d63bf788599a48ea"}, + {file = "h5py-3.11.0-cp39-cp39-win_amd64.whl", hash = "sha256:d9c944d364688f827dc889cf83f1fca311caf4fa50b19f009d1f2b525edd33a3"}, + {file = "h5py-3.11.0.tar.gz", hash = "sha256:7b7e8f78072a2edec87c9836f25f34203fd492a4475709a18b417a33cfb21fa9"}, +] + +[package.dependencies] +numpy = ">=1.17.3" + [[package]] name = "htmltools" version = "0.5.3" @@ -875,6 +1347,20 @@ typing-extensions = ">=3.10.0.0" dev = ["Flake8-pyproject", "black (>=24.2.0)", "build", "flake8 (>=6.0.0)", "isort (>=5.11.2)", "pre-commit (>=2.15.0)", "pyright (>=1.1.348)", "wheel"] test = ["pytest (>=6.2.4)", "syrupy (>=4.6.0)"] +[[package]] +name = "humanize" +version = "4.10.0" +description = "Python humanize utilities" +optional = false +python-versions = ">=3.8" +files = [ + {file = "humanize-4.10.0-py3-none-any.whl", hash = "sha256:39e7ccb96923e732b5c2e27aeaa3b10a8dfeeba3eb965ba7b74a3eb0e30040a6"}, + {file = "humanize-4.10.0.tar.gz", hash = "sha256:06b6eb0293e4b85e8d385397c5868926820db32b9b654b932f57fa41c23c9978"}, +] + +[package.extras] +tests = ["freezegun", "pytest", "pytest-cov"] + [[package]] name = "identify" version = "2.6.0" @@ -900,6 +1386,25 @@ files = [ {file = "idna-3.7.tar.gz", hash = "sha256:028ff3aadf0609c1fd278d8ea3089299412a7a8b9bd005dd08b9f8285bcb5cfc"}, ] +[[package]] +name = "importlib-resources" +version = "6.4.4" +description = "Read resources from Python packages" +optional = false +python-versions = ">=3.8" +files = [ + {file = "importlib_resources-6.4.4-py3-none-any.whl", hash = "sha256:dda242603d1c9cd836c3368b1174ed74cb4049ecd209e7a1a0104620c18c5c11"}, + {file = "importlib_resources-6.4.4.tar.gz", hash = "sha256:20600c8b7361938dc0bb2d5ec0297802e575df486f5a544fa414da65e13721f7"}, +] + +[package.extras] +check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1)"] +cover = ["pytest-cov"] +doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] +enabler = ["pytest-enabler (>=2.2)"] +test = ["jaraco.test (>=5.4)", "pytest (>=6,!=8.1.*)", "zipp (>=3.17)"] +type = ["pytest-mypy"] + [[package]] name = "iniconfig" version = "2.0.0" @@ -1001,6 +1506,23 @@ scipy = [ [package.extras] cuda12-pip = ["nvidia-cublas-cu12 (>=12.1.3.1)", "nvidia-cuda-cupti-cu12 (>=12.1.105)", "nvidia-cuda-nvcc-cu12 (>=12.1.105)", "nvidia-cuda-runtime-cu12 (>=12.1.105)", "nvidia-cudnn-cu12 (>=8.9.2.26,<9.0)", "nvidia-cufft-cu12 (>=11.0.2.54)", "nvidia-cusolver-cu12 (>=11.4.5.107)", "nvidia-cusparse-cu12 (>=12.1.0.106)", "nvidia-nccl-cu12 (>=2.18.1)", "nvidia-nvjitlink-cu12 (>=12.1.105)"] +[[package]] +name = "jaxopt" +version = "0.8.3" +description = "Hardware accelerated, batchable and differentiable optimizers in JAX." +optional = false +python-versions = "*" +files = [ + {file = "jaxopt-0.8.3-py3-none-any.whl", hash = "sha256:4be2f82798393682529c9ca5046e5397ac6c8657b8acb6bf275e773b28df15b6"}, + {file = "jaxopt-0.8.3.tar.gz", hash = "sha256:4b06dfa6f915a4f3291699606245af6069371a48dc5c92d4c507840d62990646"}, +] + +[package.dependencies] +jax = ">=0.2.18" +jaxlib = ">=0.1.69" +numpy = ">=1.18.4" +scipy = ">=1.0.0" + [[package]] name = "jaxtyping" version = "0.2.33" @@ -1165,6 +1687,20 @@ dev = ["black", "flake8", "isort", "pre-commit", "pyproject-flake8"] doc = ["myst-parser", "sphinx", "sphinx-book-theme"] test = ["coverage", "pytest", "pytest-cov"] +[[package]] +name = "logical-unification" +version = "0.4.6" +description = "Logical unification in Python" +optional = false +python-versions = ">=3.6" +files = [ + {file = "logical-unification-0.4.6.tar.gz", hash = "sha256:908435123f8a106fa4dcf9bf1b75c7beb309fa2bbecf277868af8f1c212650a0"}, +] + +[package.dependencies] +multipledispatch = "*" +toolz = "*" + [[package]] name = "markdown-it-py" version = "3.0.0" @@ -1271,6 +1807,23 @@ files = [ {file = "mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"}, ] +[[package]] +name = "minikanren" +version = "1.0.3" +description = "Relational programming in Python" +optional = false +python-versions = ">=3.6" +files = [ + {file = "miniKanren-1.0.3.tar.gz", hash = "sha256:1ec8bdb01144ad5e8752c7c297fb8a122db920f859276d25a72d164e998d7f6e"}, +] + +[package.dependencies] +cons = ">=0.4.0" +etuples = ">=0.3.1" +logical-unification = ">=0.4.1" +multipledispatch = "*" +toolz = "*" + [[package]] name = "ml-dtypes" version = "0.4.0" @@ -1341,6 +1894,71 @@ files = [ msal = ">=1.29,<2" portalocker = ">=1.4,<3" +[[package]] +name = "msgpack" +version = "1.0.8" +description = "MessagePack serializer" +optional = false +python-versions = ">=3.8" +files = [ + {file = "msgpack-1.0.8-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:505fe3d03856ac7d215dbe005414bc28505d26f0c128906037e66d98c4e95868"}, + {file = "msgpack-1.0.8-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6b7842518a63a9f17107eb176320960ec095a8ee3b4420b5f688e24bf50c53c"}, + {file = "msgpack-1.0.8-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:376081f471a2ef24828b83a641a02c575d6103a3ad7fd7dade5486cad10ea659"}, + {file = "msgpack-1.0.8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5e390971d082dba073c05dbd56322427d3280b7cc8b53484c9377adfbae67dc2"}, + {file = "msgpack-1.0.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00e073efcba9ea99db5acef3959efa45b52bc67b61b00823d2a1a6944bf45982"}, + {file = "msgpack-1.0.8-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:82d92c773fbc6942a7a8b520d22c11cfc8fd83bba86116bfcf962c2f5c2ecdaa"}, + {file = "msgpack-1.0.8-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9ee32dcb8e531adae1f1ca568822e9b3a738369b3b686d1477cbc643c4a9c128"}, + {file = "msgpack-1.0.8-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:e3aa7e51d738e0ec0afbed661261513b38b3014754c9459508399baf14ae0c9d"}, + {file = "msgpack-1.0.8-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:69284049d07fce531c17404fcba2bb1df472bc2dcdac642ae71a2d079d950653"}, + {file = "msgpack-1.0.8-cp310-cp310-win32.whl", hash = "sha256:13577ec9e247f8741c84d06b9ece5f654920d8365a4b636ce0e44f15e07ec693"}, + {file = "msgpack-1.0.8-cp310-cp310-win_amd64.whl", hash = "sha256:e532dbd6ddfe13946de050d7474e3f5fb6ec774fbb1a188aaf469b08cf04189a"}, + {file = "msgpack-1.0.8-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:9517004e21664f2b5a5fd6333b0731b9cf0817403a941b393d89a2f1dc2bd836"}, + {file = "msgpack-1.0.8-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d16a786905034e7e34098634b184a7d81f91d4c3d246edc6bd7aefb2fd8ea6ad"}, + {file = "msgpack-1.0.8-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e2872993e209f7ed04d963e4b4fbae72d034844ec66bc4ca403329db2074377b"}, + {file = "msgpack-1.0.8-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5c330eace3dd100bdb54b5653b966de7f51c26ec4a7d4e87132d9b4f738220ba"}, + {file = "msgpack-1.0.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83b5c044f3eff2a6534768ccfd50425939e7a8b5cf9a7261c385de1e20dcfc85"}, + {file = "msgpack-1.0.8-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1876b0b653a808fcd50123b953af170c535027bf1d053b59790eebb0aeb38950"}, + {file = "msgpack-1.0.8-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:dfe1f0f0ed5785c187144c46a292b8c34c1295c01da12e10ccddfc16def4448a"}, + {file = "msgpack-1.0.8-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:3528807cbbb7f315bb81959d5961855e7ba52aa60a3097151cb21956fbc7502b"}, + {file = "msgpack-1.0.8-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e2f879ab92ce502a1e65fce390eab619774dda6a6ff719718069ac94084098ce"}, + {file = "msgpack-1.0.8-cp311-cp311-win32.whl", hash = "sha256:26ee97a8261e6e35885c2ecd2fd4a6d38252246f94a2aec23665a4e66d066305"}, + {file = "msgpack-1.0.8-cp311-cp311-win_amd64.whl", hash = "sha256:eadb9f826c138e6cf3c49d6f8de88225a3c0ab181a9b4ba792e006e5292d150e"}, + {file = "msgpack-1.0.8-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:114be227f5213ef8b215c22dde19532f5da9652e56e8ce969bf0a26d7c419fee"}, + {file = "msgpack-1.0.8-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:d661dc4785affa9d0edfdd1e59ec056a58b3dbb9f196fa43587f3ddac654ac7b"}, + {file = "msgpack-1.0.8-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d56fd9f1f1cdc8227d7b7918f55091349741904d9520c65f0139a9755952c9e8"}, + {file = "msgpack-1.0.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0726c282d188e204281ebd8de31724b7d749adebc086873a59efb8cf7ae27df3"}, + {file = "msgpack-1.0.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8db8e423192303ed77cff4dce3a4b88dbfaf43979d280181558af5e2c3c71afc"}, + {file = "msgpack-1.0.8-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:99881222f4a8c2f641f25703963a5cefb076adffd959e0558dc9f803a52d6a58"}, + {file = "msgpack-1.0.8-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b5505774ea2a73a86ea176e8a9a4a7c8bf5d521050f0f6f8426afe798689243f"}, + {file = "msgpack-1.0.8-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:ef254a06bcea461e65ff0373d8a0dd1ed3aa004af48839f002a0c994a6f72d04"}, + {file = "msgpack-1.0.8-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:e1dd7839443592d00e96db831eddb4111a2a81a46b028f0facd60a09ebbdd543"}, + {file = "msgpack-1.0.8-cp312-cp312-win32.whl", hash = "sha256:64d0fcd436c5683fdd7c907eeae5e2cbb5eb872fafbc03a43609d7941840995c"}, + {file = "msgpack-1.0.8-cp312-cp312-win_amd64.whl", hash = "sha256:74398a4cf19de42e1498368c36eed45d9528f5fd0155241e82c4082b7e16cffd"}, + {file = "msgpack-1.0.8-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0ceea77719d45c839fd73abcb190b8390412a890df2f83fb8cf49b2a4b5c2f40"}, + {file = "msgpack-1.0.8-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1ab0bbcd4d1f7b6991ee7c753655b481c50084294218de69365f8f1970d4c151"}, + {file = "msgpack-1.0.8-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:1cce488457370ffd1f953846f82323cb6b2ad2190987cd4d70b2713e17268d24"}, + {file = "msgpack-1.0.8-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3923a1778f7e5ef31865893fdca12a8d7dc03a44b33e2a5f3295416314c09f5d"}, + {file = "msgpack-1.0.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a22e47578b30a3e199ab067a4d43d790249b3c0587d9a771921f86250c8435db"}, + {file = "msgpack-1.0.8-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bd739c9251d01e0279ce729e37b39d49a08c0420d3fee7f2a4968c0576678f77"}, + {file = "msgpack-1.0.8-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:d3420522057ebab1728b21ad473aa950026d07cb09da41103f8e597dfbfaeb13"}, + {file = "msgpack-1.0.8-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:5845fdf5e5d5b78a49b826fcdc0eb2e2aa7191980e3d2cfd2a30303a74f212e2"}, + {file = "msgpack-1.0.8-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:6a0e76621f6e1f908ae52860bdcb58e1ca85231a9b0545e64509c931dd34275a"}, + {file = "msgpack-1.0.8-cp38-cp38-win32.whl", hash = "sha256:374a8e88ddab84b9ada695d255679fb99c53513c0a51778796fcf0944d6c789c"}, + {file = "msgpack-1.0.8-cp38-cp38-win_amd64.whl", hash = "sha256:f3709997b228685fe53e8c433e2df9f0cdb5f4542bd5114ed17ac3c0129b0480"}, + {file = "msgpack-1.0.8-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:f51bab98d52739c50c56658cc303f190785f9a2cd97b823357e7aeae54c8f68a"}, + {file = "msgpack-1.0.8-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:73ee792784d48aa338bba28063e19a27e8d989344f34aad14ea6e1b9bd83f596"}, + {file = "msgpack-1.0.8-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f9904e24646570539a8950400602d66d2b2c492b9010ea7e965025cb71d0c86d"}, + {file = "msgpack-1.0.8-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e75753aeda0ddc4c28dce4c32ba2f6ec30b1b02f6c0b14e547841ba5b24f753f"}, + {file = "msgpack-1.0.8-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5dbf059fb4b7c240c873c1245ee112505be27497e90f7c6591261c7d3c3a8228"}, + {file = "msgpack-1.0.8-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4916727e31c28be8beaf11cf117d6f6f188dcc36daae4e851fee88646f5b6b18"}, + {file = "msgpack-1.0.8-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:7938111ed1358f536daf311be244f34df7bf3cdedb3ed883787aca97778b28d8"}, + {file = "msgpack-1.0.8-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:493c5c5e44b06d6c9268ce21b302c9ca055c1fd3484c25ba41d34476c76ee746"}, + {file = "msgpack-1.0.8-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5fbb160554e319f7b22ecf530a80a3ff496d38e8e07ae763b9e82fadfe96f273"}, + {file = "msgpack-1.0.8-cp39-cp39-win32.whl", hash = "sha256:f9af38a89b6a5c04b7d18c492c8ccf2aee7048aff1ce8437c4683bb5a1df893d"}, + {file = "msgpack-1.0.8-cp39-cp39-win_amd64.whl", hash = "sha256:ed59dd52075f8fc91da6053b12e8c89e37aa043f8986efd89e61fae69dc1b011"}, + {file = "msgpack-1.0.8.tar.gz", hash = "sha256:95c02b0e27e706e48d0e5426d1710ca78e0f0628d6e89d5b5a5b91a5f12274f3"}, +] + [[package]] name = "msrest" version = "0.7.1" @@ -1447,6 +2065,17 @@ files = [ {file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"}, ] +[[package]] +name = "nest-asyncio" +version = "1.6.0" +description = "Patch asyncio to allow nested event loops" +optional = false +python-versions = ">=3.5" +files = [ + {file = "nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c"}, + {file = "nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe"}, +] + [[package]] name = "nodeenv" version = "1.9.1" @@ -1539,6 +2168,45 @@ examples = ["arviz", "jupyter", "matplotlib", "pandas", "scikit-learn", "seaborn test = ["importlib-metadata (<5.0)", "pyro-api (>=0.1.1)", "pytest (>=4.1)", "ruff (>=0.1.8)", "scikit-learn", "scipy (>=1.9)"] tpu = ["jax[tpu] (>=0.4.25)"] +[[package]] +name = "nutpie" +version = "0.13.2" +description = "Sample Stan or PyMC models" +optional = false +python-versions = ">=3.10" +files = [ + {file = "nutpie-0.13.2-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:4c731b6b32f51407ca973aefdcb0241c6dadfebcf47e781557344d28d346c0fa"}, + {file = "nutpie-0.13.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:b69e62c4d25e62e670ef31244e65556ed562650dfbc56a068972e177c5e5e291"}, + {file = "nutpie-0.13.2-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:a7cfe73f29769f7185e677587755ba63818e9334d161a69216c8d6cefd9d66b7"}, + {file = "nutpie-0.13.2-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:0202a5b2352b065a269dd1467cacd4b9ef4020665373e4d12eede232425eaea8"}, + {file = "nutpie-0.13.2-cp310-none-win_amd64.whl", hash = "sha256:fa2f5f46fad31d9cdac486510a656a7e85df470662ffcd6c3c84534eb7d24c28"}, + {file = "nutpie-0.13.2-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:024fb04ddcaa2ce8a2cf6864bebe68acfb68518f6199c6d3de0c6b9b49d1ac75"}, + {file = "nutpie-0.13.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:225f17a15e33f731db43c55f821b988df2781568e2dc6f22ae9798e259386009"}, + {file = "nutpie-0.13.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:1a7a5e7012976327485349b581ae762cd6e60bb1805f9d323e0eed2d945c73a3"}, + {file = "nutpie-0.13.2-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:be1635cdd6ec19cc541e212ee95e11288dda7a234a2ae7f70c2c91fdaa677fe0"}, + {file = "nutpie-0.13.2-cp311-none-win_amd64.whl", hash = "sha256:d7d297a975737ca997890cae284adca74e429567503596cbf66a37640faf4f10"}, + {file = "nutpie-0.13.2-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:1656a4e45981db30d9ca850e889c10ac69c3e327a994607924c2db1dcefb49c7"}, + {file = "nutpie-0.13.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:57b6f6640996d88b290285acdcf7978bf9f6257c2a80d38eb5d1903e11bb0301"}, + {file = "nutpie-0.13.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e1419e53a5ce3bfba39157cb1381eb18f1835bd1b73312d485e1f543f9ce3748"}, + {file = "nutpie-0.13.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:6d29babf3773544692153799b3579f9de1e084a06fd2dcc851e97bef4c92768b"}, + {file = "nutpie-0.13.2-cp312-none-win_amd64.whl", hash = "sha256:5b6f45e2e475eee1519f18b6cbcd56ef225dbcaeb6f35e248d829467097ab385"}, + {file = "nutpie-0.13.2-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:db240a317b1ded7eddf2ca8e2b4bcfcdbd4624256655aac61625c8f7d5ca39d0"}, + {file = "nutpie-0.13.2-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:2100024275ec6ba6de899188a3a2111f4b68aee7bfdbd4e4eb02ed4c922a9f22"}, + {file = "nutpie-0.13.2.tar.gz", hash = "sha256:f14282e2ac045c67a9b262a865b02a243178c55b541b236b21dfcb0c3678bcea"}, +] + +[package.dependencies] +arviz = ">=0.15.0" +pandas = ">=2.0" +pyarrow = ">=12.0.0" +xarray = ">=2023.6.0" + +[package.extras] +all = ["bridgestan (>=2.4.1)", "jax (>=0.4.27)", "numba (>=0.59.1)", "pymc (>=5.15.0)"] +pymc = ["numba (>=0.59.1)", "pymc (>=5.15.0)"] +pymc-jax = ["jax (>=0.4.27)", "pymc (>=5.15.0)"] +stan = ["bridgestan (>=2.4.1)"] + [[package]] name = "oauthlib" version = "3.2.2" @@ -1573,6 +2241,30 @@ numpy = ">=1.7" docs = ["numpydoc", "sphinx (==1.2.3)", "sphinx-rtd-theme", "sphinxcontrib-napoleon"] tests = ["pytest", "pytest-cov", "pytest-pep8"] +[[package]] +name = "optax" +version = "0.2.2" +description = "A gradient processing and optimisation library in JAX." +optional = false +python-versions = ">=3.9" +files = [ + {file = "optax-0.2.2-py3-none-any.whl", hash = "sha256:411c414a76aae259f4191a60b712663968741a5163ca92fc250b5d5c7d36fb57"}, + {file = "optax-0.2.2.tar.gz", hash = "sha256:f09bf790ef4b09fb9c35f79a07594c6196a719919985f542dc84b0bf97812e0e"}, +] + +[package.dependencies] +absl-py = ">=0.7.1" +chex = ">=0.1.86" +jax = ">=0.1.55" +jaxlib = ">=0.1.37" +numpy = ">=1.18.0" + +[package.extras] +docs = ["flax", "ipython (>=8.8.0)", "matplotlib (>=3.5.0)", "myst-nb (>=1.0.0)", "sphinx (>=6.0.0)", "sphinx-autodoc-typehints", "sphinx-book-theme (>=1.0.1)", "sphinx-collections (>=0.0.1)", "sphinx-gallery (>=0.14.0)", "sphinx_contributors", "sphinxcontrib-katex", "tensorflow (>=2.4.0)", "tensorflow-datasets (>=4.2.0)"] +dp-accounting = ["absl-py (>=1.0.0)", "attrs (>=21.4.0)", "mpmath (>=1.2.1)", "numpy (>=1.21.4)", "scipy (>=1.7.1)"] +examples = ["dp_accounting (>=0.4)", "flax", "tensorflow (>=2.4.0)", "tensorflow-datasets (>=4.2.0)"] +test = ["dm-tree (>=0.1.7)", "flax (>=0.5.3)"] + [[package]] name = "optimistix" version = "0.0.7" @@ -1591,6 +2283,49 @@ jaxtyping = ">=0.2.23" lineax = ">=0.0.4" typing-extensions = ">=4.5.0" +[[package]] +name = "orbax-checkpoint" +version = "0.6.1" +description = "Orbax Checkpoint" +optional = false +python-versions = ">=3.9" +files = [ + {file = "orbax_checkpoint-0.6.1-py3-none-any.whl", hash = "sha256:f7fcb1ef528cee294ea244e769eaee17de2379c68a00d6df4c3a463e5cf716a1"}, + {file = "orbax_checkpoint-0.6.1.tar.gz", hash = "sha256:0d41fa2c38e2fa7e84ebefa974f87f6fed152f3130af188e7911b5a42d2dc067"}, +] + +[package.dependencies] +absl-py = "*" +etils = {version = "*", extras = ["epath", "epy"]} +humanize = "*" +jax = ">=0.4.26" +jaxlib = "*" +msgpack = "*" +nest_asyncio = "*" +numpy = "*" +protobuf = "*" +pyyaml = "*" +tensorstore = ">=0.1.60" +typing_extensions = "*" + +[package.extras] +testing = ["flax", "google-cloud-logging", "mock", "pytest", "pytest-xdist"] + +[[package]] +name = "oryx" +version = "0.2.7" +description = "Probabilistic programming and deep learning in JAX" +optional = false +python-versions = "<4.0,>=3.9" +files = [ + {file = "oryx-0.2.7-py3-none-any.whl", hash = "sha256:f8d3816849fedd5dc208021699223d964caab1f316fe95c8c0e53af68c6bbfbc"}, + {file = "oryx-0.2.7.tar.gz", hash = "sha256:6fa1fea1beae126a3d289cdda54c3852934e0587541cc60cbc53a366c6c9c953"}, +] + +[package.dependencies] +jax = ">=0.4.6,<0.5.0" +tensorflow-probability = ">=0.19.0" + [[package]] name = "packaging" version = "24.1" @@ -1876,6 +2611,77 @@ files = [ [package.dependencies] wcwidth = "*" +[[package]] +name = "protobuf" +version = "5.28.0" +description = "" +optional = false +python-versions = ">=3.8" +files = [ + {file = "protobuf-5.28.0-cp310-abi3-win32.whl", hash = "sha256:66c3edeedb774a3508ae70d87b3a19786445fe9a068dd3585e0cefa8a77b83d0"}, + {file = "protobuf-5.28.0-cp310-abi3-win_amd64.whl", hash = "sha256:6d7cc9e60f976cf3e873acb9a40fed04afb5d224608ed5c1a105db4a3f09c5b6"}, + {file = "protobuf-5.28.0-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:532627e8fdd825cf8767a2d2b94d77e874d5ddb0adefb04b237f7cc296748681"}, + {file = "protobuf-5.28.0-cp38-abi3-manylinux2014_aarch64.whl", hash = "sha256:018db9056b9d75eb93d12a9d35120f97a84d9a919bcab11ed56ad2d399d6e8dd"}, + {file = "protobuf-5.28.0-cp38-abi3-manylinux2014_x86_64.whl", hash = "sha256:6206afcb2d90181ae8722798dcb56dc76675ab67458ac24c0dd7d75d632ac9bd"}, + {file = "protobuf-5.28.0-cp38-cp38-win32.whl", hash = "sha256:eef7a8a2f4318e2cb2dee8666d26e58eaf437c14788f3a2911d0c3da40405ae8"}, + {file = "protobuf-5.28.0-cp38-cp38-win_amd64.whl", hash = "sha256:d001a73c8bc2bf5b5c1360d59dd7573744e163b3607fa92788b7f3d5fefbd9a5"}, + {file = "protobuf-5.28.0-cp39-cp39-win32.whl", hash = "sha256:dde9fcaa24e7a9654f4baf2a55250b13a5ea701493d904c54069776b99a8216b"}, + {file = "protobuf-5.28.0-cp39-cp39-win_amd64.whl", hash = "sha256:853db610214e77ee817ecf0514e0d1d052dff7f63a0c157aa6eabae98db8a8de"}, + {file = "protobuf-5.28.0-py3-none-any.whl", hash = "sha256:510ed78cd0980f6d3218099e874714cdf0d8a95582e7b059b06cabad855ed0a0"}, + {file = "protobuf-5.28.0.tar.gz", hash = "sha256:dde74af0fa774fa98892209992295adbfb91da3fa98c8f67a88afe8f5a349add"}, +] + +[[package]] +name = "pyarrow" +version = "17.0.0" +description = "Python library for Apache Arrow" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pyarrow-17.0.0-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:a5c8b238d47e48812ee577ee20c9a2779e6a5904f1708ae240f53ecbee7c9f07"}, + {file = "pyarrow-17.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:db023dc4c6cae1015de9e198d41250688383c3f9af8f565370ab2b4cb5f62655"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da1e060b3876faa11cee287839f9cc7cdc00649f475714b8680a05fd9071d545"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:75c06d4624c0ad6674364bb46ef38c3132768139ddec1c56582dbac54f2663e2"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:fa3c246cc58cb5a4a5cb407a18f193354ea47dd0648194e6265bd24177982fe8"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:f7ae2de664e0b158d1607699a16a488de3d008ba99b3a7aa5de1cbc13574d047"}, + {file = "pyarrow-17.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:5984f416552eea15fd9cee03da53542bf4cddaef5afecefb9aa8d1010c335087"}, + {file = "pyarrow-17.0.0-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:1c8856e2ef09eb87ecf937104aacfa0708f22dfeb039c363ec99735190ffb977"}, + {file = "pyarrow-17.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2e19f569567efcbbd42084e87f948778eb371d308e137a0f97afe19bb860ccb3"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b244dc8e08a23b3e352899a006a26ae7b4d0da7bb636872fa8f5884e70acf15"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b72e87fe3e1db343995562f7fff8aee354b55ee83d13afba65400c178ab2597"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:dc5c31c37409dfbc5d014047817cb4ccd8c1ea25d19576acf1a001fe07f5b420"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:e3343cb1e88bc2ea605986d4b94948716edc7a8d14afd4e2c097232f729758b4"}, + {file = "pyarrow-17.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:a27532c38f3de9eb3e90ecab63dfda948a8ca859a66e3a47f5f42d1e403c4d03"}, + {file = "pyarrow-17.0.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:9b8a823cea605221e61f34859dcc03207e52e409ccf6354634143e23af7c8d22"}, + {file = "pyarrow-17.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f1e70de6cb5790a50b01d2b686d54aaf73da01266850b05e3af2a1bc89e16053"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0071ce35788c6f9077ff9ecba4858108eebe2ea5a3f7cf2cf55ebc1dbc6ee24a"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:757074882f844411fcca735e39aae74248a1531367a7c80799b4266390ae51cc"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9ba11c4f16976e89146781a83833df7f82077cdab7dc6232c897789343f7891a"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b0c6ac301093b42d34410b187bba560b17c0330f64907bfa4f7f7f2444b0cf9b"}, + {file = "pyarrow-17.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:392bc9feabc647338e6c89267635e111d71edad5fcffba204425a7c8d13610d7"}, + {file = "pyarrow-17.0.0-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:af5ff82a04b2171415f1410cff7ebb79861afc5dae50be73ce06d6e870615204"}, + {file = "pyarrow-17.0.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:edca18eaca89cd6382dfbcff3dd2d87633433043650c07375d095cd3517561d8"}, + {file = "pyarrow-17.0.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7c7916bff914ac5d4a8fe25b7a25e432ff921e72f6f2b7547d1e325c1ad9d155"}, + {file = "pyarrow-17.0.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f553ca691b9e94b202ff741bdd40f6ccb70cdd5fbf65c187af132f1317de6145"}, + {file = "pyarrow-17.0.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:0cdb0e627c86c373205a2f94a510ac4376fdc523f8bb36beab2e7f204416163c"}, + {file = "pyarrow-17.0.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:d7d192305d9d8bc9082d10f361fc70a73590a4c65cf31c3e6926cd72b76bc35c"}, + {file = "pyarrow-17.0.0-cp38-cp38-win_amd64.whl", hash = "sha256:02dae06ce212d8b3244dd3e7d12d9c4d3046945a5933d28026598e9dbbda1fca"}, + {file = "pyarrow-17.0.0-cp39-cp39-macosx_10_15_x86_64.whl", hash = "sha256:13d7a460b412f31e4c0efa1148e1d29bdf18ad1411eb6757d38f8fbdcc8645fb"}, + {file = "pyarrow-17.0.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9b564a51fbccfab5a04a80453e5ac6c9954a9c5ef2890d1bcf63741909c3f8df"}, + {file = "pyarrow-17.0.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:32503827abbc5aadedfa235f5ece8c4f8f8b0a3cf01066bc8d29de7539532687"}, + {file = "pyarrow-17.0.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a155acc7f154b9ffcc85497509bcd0d43efb80d6f733b0dc3bb14e281f131c8b"}, + {file = "pyarrow-17.0.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:dec8d129254d0188a49f8a1fc99e0560dc1b85f60af729f47de4046015f9b0a5"}, + {file = "pyarrow-17.0.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:a48ddf5c3c6a6c505904545c25a4ae13646ae1f8ba703c4df4a1bfe4f4006bda"}, + {file = "pyarrow-17.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:42bf93249a083aca230ba7e2786c5f673507fa97bbd9725a1e2754715151a204"}, + {file = "pyarrow-17.0.0.tar.gz", hash = "sha256:4beca9521ed2c0921c1023e68d097d0299b62c362639ea315572a58f3f50fd28"}, +] + +[package.dependencies] +numpy = ">=1.16.6" + +[package.extras] +test = ["cffi", "hypothesis", "pandas", "pytest", "pytz"] + [[package]] name = "pycparser" version = "2.22" @@ -1887,6 +2693,20 @@ files = [ {file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"}, ] +[[package]] +name = "pygments" +version = "2.18.0" +description = "Pygments is a syntax highlighting package written in Python." +optional = false +python-versions = ">=3.8" +files = [ + {file = "pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a"}, + {file = "pygments-2.18.0.tar.gz", hash = "sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199"}, +] + +[package.extras] +windows-terminal = ["colorama (>=0.4.6)"] + [[package]] name = "pyjwt" version = "2.9.0" @@ -1907,6 +2727,28 @@ dev = ["coverage[toml] (==5.0.4)", "cryptography (>=3.4.0)", "pre-commit", "pyte docs = ["sphinx", "sphinx-rtd-theme", "zope.interface"] tests = ["coverage[toml] (==5.0.4)", "pytest (>=6.0.0,<7.0.0)"] +[[package]] +name = "pymc" +version = "5.8.2" +description = "Probabilistic Programming in Python: Bayesian Modeling and Probabilistic Machine Learning with PyTensor" +optional = false +python-versions = ">=3.9" +files = [ + {file = "pymc-5.8.2-py3-none-any.whl", hash = "sha256:7bac7734427bf5a48e55cc72244569195d3d02bf985065672cce5fbacf9ca058"}, + {file = "pymc-5.8.2.tar.gz", hash = "sha256:53720dae99cfdc9752bd0420f1c63e4ee4d83e00808d5bcee0beaab213f7385b"}, +] + +[package.dependencies] +arviz = ">=0.13.0" +cachetools = ">=4.2.1" +cloudpickle = "*" +fastprogress = ">=0.2.0" +numpy = ">=1.15.0" +pandas = ">=0.24.0" +pytensor = ">=2.16.1,<2.17" +scipy = ">=1.4.1" +typing-extensions = ">=3.7.4" + [[package]] name = "pyparsing" version = "3.1.2" @@ -1921,6 +2763,50 @@ files = [ [package.extras] diagrams = ["jinja2", "railroad-diagrams"] +[[package]] +name = "pytensor" +version = "2.16.1" +description = "Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs." +optional = false +python-versions = ">=3.9" +files = [ + {file = "pytensor-2.16.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8d521e9fa556ee085c837098310ca6561772e6012cb310ff66eed1306d6c68d4"}, + {file = "pytensor-2.16.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2b9629ee883cd7881c0e06e17b93da3a672028d4fb0198dbf3cbd959b55e79b8"}, + {file = "pytensor-2.16.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:e6b2d78b84e9b22ac1c817409aeb629de49886c01c1f0e9607cf968416c8956a"}, + {file = "pytensor-2.16.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:0fa2956f5219711673b71825f6316cfbb8dcbf4d398b13d9d3e2e73d7578c157"}, + {file = "pytensor-2.16.1-cp310-cp310-win_amd64.whl", hash = "sha256:67f25d3c794a25430efad499f356dbe389dd6700387d245fa293fa474adcbcaf"}, + {file = "pytensor-2.16.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5ff70f49c66059a894b408d340f9d83c6cc5850bb096a035f1fbf8c609aedfab"}, + {file = "pytensor-2.16.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a8eb0dd68fa383d635e5fd0b89d22d3632c7cf8e90df3ccc37ea459bbf3e6c3f"}, + {file = "pytensor-2.16.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:5f89154c13d51e1d3b4f29d027d60f1dbce21e8bcf5a71d6fc9cd6c650b72353"}, + {file = "pytensor-2.16.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:37b787e5f864d8ecf6a285a3cfbf179df422e642f616a208da2f24f574d839a3"}, + {file = "pytensor-2.16.1-cp311-cp311-win_amd64.whl", hash = "sha256:78998d5f4a0501d30e9873232c6209165305ed236a50bee250c3743c258491c1"}, + {file = "pytensor-2.16.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:208b2fa548465f323ef5fd1779e37c7d4a14980c6263650e4a9f4befbe7c002e"}, + {file = "pytensor-2.16.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:096c4cbd4c22c50d01f3e3fe156d8182337011accc634cfec56043f8e93f39db"}, + {file = "pytensor-2.16.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:fa29fe4f15785e326d4af91b8882442820f94db8ce4293d97b4a83bd3b0e895a"}, + {file = "pytensor-2.16.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:47eca1d1fc0acd86e5b5867d40be2e361b399f6c4137510b158642cd66c76d75"}, + {file = "pytensor-2.16.1-cp39-cp39-win_amd64.whl", hash = "sha256:3191208f2e9816d205995acd94269f97db5a050f9e16f836fb7e0d8cbe824b67"}, + {file = "pytensor-2.16.1.tar.gz", hash = "sha256:dc586d5aea09e29dcc6927c3b99b13af86ea07736d1555be03e2d11282da821a"}, +] + +[package.dependencies] +cons = "*" +etuples = "*" +filelock = "*" +logical-unification = "*" +miniKanren = "*" +numpy = ">=1.17.0" +scipy = ">=0.14" +setuptools = ">=48.0.0" +typing-extensions = "*" + +[package.extras] +complete = ["pytensor[jax]", "pytensor[numba]"] +development = ["pytensor[complete]", "pytensor[rtd]", "pytensor[tests]"] +jax = ["jax", "jaxlib"] +numba = ["numba (>=0.55)", "numba-scipy (>=0.3.0)"] +rtd = ["pydot", "pydot-ng", "pydot2", "pygments", "sphinx (>=5.1.0,<6)"] +tests = ["coverage (>=5.1)", "pre-commit", "pytest", "pytest-benchmark", "pytest-cov (>=2.6.1)", "pytest-mock"] + [[package]] name = "pytest" version = "8.3.2" @@ -2120,6 +3006,24 @@ requests = ">=2.0.0" [package.extras] rsa = ["oauthlib[signedtoken] (>=3.0.0)"] +[[package]] +name = "rich" +version = "13.8.0" +description = "Render rich text, tables, progress bars, syntax highlighting, markdown and more to the terminal" +optional = false +python-versions = ">=3.7.0" +files = [ + {file = "rich-13.8.0-py3-none-any.whl", hash = "sha256:2e85306a063b9492dffc86278197a60cbece75bcb766022f3436f567cae11bdc"}, + {file = "rich-13.8.0.tar.gz", hash = "sha256:a5ac1f1cd448ade0d59cc3356f7db7a7ccda2c8cbae9c7a90c28ff463d3e91f4"}, +] + +[package.dependencies] +markdown-it-py = ">=2.2.0" +pygments = ">=2.13.0,<3.0.0" + +[package.extras] +jupyter = ["ipywidgets (>=7.5.1,<9)"] + [[package]] name = "scipy" version = "1.14.0" @@ -2288,6 +3192,66 @@ anyio = ">=3.4.0,<5" [package.extras] full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart (>=0.0.7)", "pyyaml"] +[[package]] +name = "tensorflow-probability" +version = "0.24.0" +description = "Probabilistic modeling and statistical inference in TensorFlow" +optional = false +python-versions = ">=3.9" +files = [ + {file = "tensorflow_probability-0.24.0-py2.py3-none-any.whl", hash = "sha256:8c1774683e38359dbcaf3697e79b7e6a4e69b9c7b3679e78ee18f43e59e5759b"}, +] + +[package.dependencies] +absl-py = "*" +cloudpickle = ">=1.3" +decorator = "*" +dm-tree = "*" +gast = ">=0.3.2" +jax = {version = "*", optional = true, markers = "extra == \"jax\""} +jaxlib = {version = "*", optional = true, markers = "extra == \"jax\""} +numpy = ">=1.13.3" +six = ">=1.10.0" + +[package.extras] +jax = ["jax", "jaxlib"] +tf = ["tensorflow (>=2.16)", "tf-keras (>=2.16)"] +tfds = ["tensorflow-datasets (>=2.2.0)"] + +[[package]] +name = "tensorstore" +version = "0.1.65" +description = "Read and write large, multi-dimensional arrays" +optional = false +python-versions = ">=3.9" +files = [ + {file = "tensorstore-0.1.65-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:07b6f515dfb5c5f0c5d9f4a858cd4c484f8b61bb4eeed45960bf0678c947245f"}, + {file = "tensorstore-0.1.65-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:dafb95c6624cc7f19e687833512eab45dcac4051eb795cd9e484f96120d426ba"}, + {file = "tensorstore-0.1.65-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:39c7e0dd502b33fd786b4e9ca4cb890d83ff308a3dac2bed14ee0277701a840e"}, + {file = "tensorstore-0.1.65-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a01ebc1636339910115a96da8d0b8ba40cd6c20f420a38b8b4b651142b33b639"}, + {file = "tensorstore-0.1.65-cp310-cp310-win_amd64.whl", hash = "sha256:4140e5f5a678da2b16d12b28d3186ac0ce758ab772b99b9a62d3a69626019777"}, + {file = "tensorstore-0.1.65-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:b2e183f7f0a8fc9f5792bc917bcee3aecebba970db33ba5a6d8917925bb29538"}, + {file = "tensorstore-0.1.65-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d314550d20c4f6ea73db253b9b3a5add8ea5e8ce4558e3003ca5810bc798dc4a"}, + {file = "tensorstore-0.1.65-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:85e805795de6d23a634d7354eae7bfde419ed42cbe3df7a1e060def7a76f2fba"}, + {file = "tensorstore-0.1.65-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90611e701bbb6f74c51942118ab23dab65b79e2a62430b21ff5381f7646a0c1e"}, + {file = "tensorstore-0.1.65-cp311-cp311-win_amd64.whl", hash = "sha256:015a9c09d0f93c1cf72707b1589a414438e15e275cac29884e8f25c9d52ef391"}, + {file = "tensorstore-0.1.65-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:33072ea4c19f2c102c136320cce583540fead4c4904562735a783f8acd066feb"}, + {file = "tensorstore-0.1.65-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8b64e62ce9766870f98b2cc5e98a636c728598118fc798b13d8c8d28a41f44d3"}, + {file = "tensorstore-0.1.65-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3b4555b52a9211de1daa399a5cc4d1bdffc171cf5847bcf36f3f4dffc1957179"}, + {file = "tensorstore-0.1.65-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:570f2c4f991a852376d0142cdcc3987ac49466a9380aebdf651d73e6700eee0f"}, + {file = "tensorstore-0.1.65-cp312-cp312-win_amd64.whl", hash = "sha256:5c4003abcbccbb7f4340f18ed64ebe0626a7f5cde5bf18c48fb052374049f37b"}, + {file = "tensorstore-0.1.65-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:a7ff77ad23366bb0355f37b86ae6ab5194935b3113cb9cc88eb22be17ced5aee"}, + {file = "tensorstore-0.1.65-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:aac453e66d43910122fe767e39d295cfde3a5706b021c4d217845c07c65f411e"}, + {file = "tensorstore-0.1.65-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b780725b24295ddae82abb5d5bce8ba3692b480e30d611752ed41b5f90e0d68b"}, + {file = "tensorstore-0.1.65-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2316325d741d2303021763a0085ccb8350b2cf498edb8b76729d73f2dce913fb"}, + {file = "tensorstore-0.1.65-cp39-cp39-win_amd64.whl", hash = "sha256:380f5946f0df9503bab8895b3d57f3f3c6b0a0bfaba1e265b67d3f8f26bf48da"}, + {file = "tensorstore-0.1.65.tar.gz", hash = "sha256:65cbe5a600c32569bb0b9f597ea318cc298a13b42d5fc98168c97bb11f320eae"}, +] + +[package.dependencies] +ml-dtypes = ">=0.3.1" +numpy = ">=1.22.0" + [[package]] name = "toml" version = "0.10.2" @@ -2310,6 +3274,17 @@ files = [ {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, ] +[[package]] +name = "toolz" +version = "0.12.1" +description = "List processing tools and functional utilities" +optional = false +python-versions = ">=3.7" +files = [ + {file = "toolz-0.12.1-py3-none-any.whl", hash = "sha256:d22731364c07d72eea0a0ad45bafb2c2937ab6fd38a3507bf55eae8744aa7d85"}, + {file = "toolz-0.12.1.tar.gz", hash = "sha256:ecca342664893f177a13dac0e6b41cbd8ac25a358e5f215316d43e2100224f4d"}, +] + [[package]] name = "tqdm" version = "4.66.5" @@ -2628,7 +3603,72 @@ files = [ {file = "websockets-12.0.tar.gz", hash = "sha256:81df9cbcbb6c260de1e007e58c011bfebe2dafc8435107b0537f393dd38c8b1b"}, ] +[[package]] +name = "xarray" +version = "2024.7.0" +description = "N-D labeled arrays and datasets in Python" +optional = false +python-versions = ">=3.9" +files = [ + {file = "xarray-2024.7.0-py3-none-any.whl", hash = "sha256:1b0fd51ec408474aa1f4a355d75c00cc1c02bd425d97b2c2e551fd21810e7f64"}, + {file = "xarray-2024.7.0.tar.gz", hash = "sha256:4cae512d121a8522d41e66d942fb06c526bc1fd32c2c181d5fe62fe65b671638"}, +] + +[package.dependencies] +numpy = ">=1.23" +packaging = ">=23.1" +pandas = ">=2.0" + +[package.extras] +accel = ["bottleneck", "flox", "numbagg", "opt-einsum", "scipy"] +complete = ["xarray[accel,dev,io,parallel,viz]"] +dev = ["hypothesis", "mypy", "pre-commit", "pytest", "pytest-cov", "pytest-env", "pytest-timeout", "pytest-xdist", "ruff", "xarray[complete]"] +io = ["cftime", "fsspec", "h5netcdf", "netCDF4", "pooch", "pydap", "scipy", "zarr"] +parallel = ["dask[complete]"] +viz = ["matplotlib", "nc-time-axis", "seaborn"] + +[[package]] +name = "xarray-einstats" +version = "0.7.0" +description = "Stats, linear algebra and einops for xarray" +optional = false +python-versions = ">=3.9" +files = [ + {file = "xarray_einstats-0.7.0-py3-none-any.whl", hash = "sha256:f39403341ebf5b634ab1f1bd0e1bb2dc51046e0df31aa908dfbe2fa6a493712e"}, + {file = "xarray_einstats-0.7.0.tar.gz", hash = "sha256:2d7b571b3bbad3cf2fd10c6c75fd949d247d14c29574184c8489d9d607278d38"}, +] + +[package.dependencies] +numpy = ">=1.22" +scipy = ">=1.8" +xarray = ">=2022.09.0" + +[package.extras] +doc = ["furo", "jupyter-sphinx", "matplotlib", "myst-nb", "myst-parser[linkify]", "numpydoc", "sphinx (>=5)", "sphinx-copybutton", "sphinx-design", "sphinx-togglebutton", "watermark"] +einops = ["einops"] +numba = ["numba (>=0.55)"] +test = ["hypothesis", "packaging", "pytest", "pytest-cov"] + +[[package]] +name = "zipp" +version = "3.20.1" +description = "Backport of pathlib-compatible object wrapper for zip files" +optional = false +python-versions = ">=3.8" +files = [ + {file = "zipp-3.20.1-py3-none-any.whl", hash = "sha256:9960cd8967c8f85a56f920d5d507274e74f9ff813a0ab8889a5b5be2daf44064"}, + {file = "zipp-3.20.1.tar.gz", hash = "sha256:c22b14cc4763c5a5b04134207736c107db42e9d3ef2d9779d465f5f1bcba572b"}, +] + +[package.extras] +check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1)"] +cover = ["pytest-cov"] +doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] +enabler = ["pytest-enabler (>=2.2)"] +test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more-itertools", "pytest (>=6,!=8.1.*)", "pytest-ignore-flaky"] +type = ["pytest-mypy"] + [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "90842421bd5b896092991666760e8bd1ac9dd09af55e29edc8cc5e86c7b64905" +content-hash = "fb1bcac78684d833ced5a924aca0e97194bbc4a54aaa93d4e49b68f2032808b8" diff --git a/pyproject.toml b/pyproject.toml index 488823c9..4e147de0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ pre-commit = "^3.7.1" mypy = "^1.10.0" requests = "^2.32.3" docker = "^7.1.0" +bayeux-ml = "^0.1.14" [build-system] From e31cd3f857c34efee0297534e501861308a06f73 Mon Sep 17 00:00:00 2001 From: Ariel Shurygin Date: Thu, 26 Sep 2024 17:06:16 +0000 Subject: [PATCH 2/2] kwargs passing to _debug_inferer and a test --- src/resp_ode/mechanistic_inferer.py | 15 ++++----------- tests/test_inferer.py | 7 +++++++ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/resp_ode/mechanistic_inferer.py b/src/resp_ode/mechanistic_inferer.py index 831edcee..f935dc4b 100644 --- a/src/resp_ode/mechanistic_inferer.py +++ b/src/resp_ode/mechanistic_inferer.py @@ -336,14 +336,11 @@ def infer(self, obs_metrics: jax.Array) -> MCMC: self.infer_complete = True return self.inference_algo - def _debug_likelihood(self, obs_metrics) -> bx.Model: + def _debug_likelihood(self, **kwargs) -> bx.Model: """uses Bayeux to recreate the self.likelihood function for purposes of basic sanity checking - Parameters - ---------- - obs_metrics: jnp.array - observed metrics on which likelihood will be calculated on to tune parameters. - See `likelihood()` method for implemented definition of `obs_metrics` + passes all parameters given to it to `self.likelihood`, initializes with `self.INITIAL_STATE` + and passes `self.config.INFERENCE_PRNGKEY` as seed for randomness. Returns ------- @@ -351,11 +348,7 @@ def _debug_likelihood(self, obs_metrics) -> bx.Model: model object used to debug """ bx_model = bx.Model.from_numpyro( - jax.tree_util.Partial( - self.likelihood, - tf=len(obs_metrics), - obs_metrics=obs_metrics, - ), + jax.tree_util.Partial(self.likelihood, **kwargs), # this does not work for non-one/sampled self.INITIAL_INFECTIONS_SCALE initial_state=self.INITIAL_STATE, ) diff --git a/tests/test_inferer.py b/tests/test_inferer.py index bbe75ce6..03673bd8 100644 --- a/tests/test_inferer.py +++ b/tests/test_inferer.py @@ -176,3 +176,10 @@ def test_random_sampling_across_chains_and_particles(): "Unable to run all tests within test_random_sampling_across_chains_and_particles " "since you have only one chain! check test_config_inferer.json" ) + + +def test_debug_inferer(): + """A simple test to make sure the _debug_likelihood function does not explode and correctly passes kwargs""" + inferer._debug_likelihood( + tf=len(synthetic_hosp_obs), obs_metrics=synthetic_hosp_obs + )