cuthbert 0.0.0__tar.gz → 0.0.1__tar.gz

cuthbert-0.0.1/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

cuthbert-0.0.1/PKG-INFO

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: cuthbert
+Version: 0.0.1
+Summary: State-space model inference with JAX
+Author-email: Sam Duffield <s@mduffield.com>, Sahel Iqbal <sahel13miqbal@proton.me>, Adrien Corenflos <adrien.corenflos.stats@gmail.com>
+License: Apache-2.0
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.35
+Requires-Dist: numba>=0.60.0
+Provides-Extra: tests
+Requires-Dist: chex; extra == "tests"
+Requires-Dist: pre-commit; extra == "tests"
+Requires-Dist: ruff; extra == "tests"
+Requires-Dist: pyright; extra == "tests"
+Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-xdist; extra == "tests"
+Requires-Dist: entangled-cli; python_version >= "3.12" and extra == "tests"
+Provides-Extra: docs
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocs-include-markdown-plugin; extra == "docs"
+Requires-Dist: mkdocstrings-python; extra == "docs"
+Provides-Extra: examples
+Requires-Dist: matplotlib; extra == "examples"
+Requires-Dist: pandas; extra == "examples"
+Requires-Dist: pandas-stubs; extra == "examples"
+Requires-Dist: yfinance; extra == "examples"
+Requires-Dist: dynamax; extra == "examples"
+Requires-Dist: ipython; extra == "examples"
+Requires-Dist: tfp-nightly[jax]; extra == "examples"
+Dynamic: license-file
+
+<!--intro-start-->
+<div align="center">
+<img src="docs/assets/cuthbert.png" alt="logo"></img>
+</div>
+
+A JAX library for state-space model inference
+(filtering, smoothing, static parameter estimation).
+
+> Disclaimer: The name `cuthbert` was chosen as a playful nod to the well-known
+> caterpillar cake rivalry between Aldi and M&S in the UK, as the classic state-space
+> model diagram looks vaguely like a caterpillar. However, this software project
+> has no formal connection to Aldi, M&S, or any food products (notwithstanding the coffee drunk during its writeup).
+> `cuthbert` is simply a fun name for this state-space model library and should not be interpreted as an
+> endorsement, association, or affiliation with any brand or animal themed baked goods.
+
+[![Discord](https://img.shields.io/badge/Discord-5865F2?logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/sBnr6JhT)
+[![GitHub](https://img.shields.io/badge/GitHub-181717?logo=github&logoColor=white&style=for-the-badge)](https://github.com/state-space-models/cuthbert)
+[![PyPI](https://img.shields.io/pypi/v/cuthbert?style=for-the-badge)](https://pypi.org/project/cuthbert/)
+[![Docs](https://img.shields.io/badge/Docs-b6d7a8?logo=materialformkdocs&logoColor=black&style=for-the-badge)](https://state-space-models.github.io/cuthbert/)
+<!--intro-end-->
+
+<!--goals-start-->
+### Goals
+- Simple, flexible and performant interface for state-space model inference.
+- Decoupling of model specification and inference. `cuthbert` is built to swap between
+different **inference** methods without be tied to a specific model specification.
+- Compose with the [JAX ecosystem](#ecosystem) for extensive external tools.
+- Functional API: The only classes in `cuthbert` are `NamedTuple`s and `Protocol`s.
+All functions are pure and work seamlessly with `jax.grad`, `jax.jit`, `jax.vmap` etc.
+- Methods for filtering: $p(x_t \mid y_{0:t}, \theta)$.
+- Methods for smoothing: $p(x_{0:T} \mid y_{0:T}, \theta)$ or $p(x_{t} \mid y_{0:T}, \theta)$.
+- Methods for static parameter estimation: $p(\theta \mid y_{0:T})$
+or $\text{argmax} p(y_{0:T} \mid \theta)$.
+- This includes support for forward-backward/Baum-Welch, particle filtering/sequential Monte Carlo,
+Kalman filtering (+ extended/unscented/ensemble), expectation-maximization and more!
+
+### Non-goals
+- Tools for defining models and distributions. `cuthbert` is not a probabilistic programming language (PPL).
+But can easily compose with [`dynamax`](https://github.com/probml/dynamax), [`distrax`](https://github.com/google-deepmind/distrax), [`numpyro`](https://github.com/pyro-ppl/numpyro) and [`pymc`](https://github.com/pymc-devs/pymc) in a similar way to how [`blackjax` does](https://blackjax-devs.github.io/blackjax/).
+- ["SMC Samplers"](https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf) which sample from a posterior
+distribution which is not (necessarily) a state-space model - [`blackjax` is great for this](https://github.com/blackjax-devs/blackjax/tree/main/blackjax/smc).
+<!--goals-end-->
+
+<!--codebase-structure-start-->
+### Codebase structure
+
+The codebase is structured as follows:
+
+- `cuthbert`: The main package with unified interface for filtering and smoothing.
+- `cuthbertlib`: A collection of atomic, smaller-scoped tools useful for state-space model inference,
+that represent the building blocks that power the main `cuthbert` package.
+<!--codebase-structure-end-->
+- `docs`: Source code for the documentation for `cuthbert` and `cuthbertlib`.
+- `tests`: Tests for the `cuthbert` and `cuthbertlib` packages.
+
+
+<!--installation-start-->
+## Installation
+
+`cuthbert` depends on JAX, so you'll need to [install JAX](https://docs.jax.dev/en/latest/installation.html) for the available hardware (CPU, GPU, or TPU).
+For example, on computers with NVIDIA GPUs:
+
+```bash
+pip install -U "jax[cuda13]"
+```
+
+Now install `cuthbert` from PyPI:
+
+```bash
+pip install -U cuthbert
+```
+
+Installing `cuthbert` will also install `cuthbertlib`.
+
+<!--installation-end-->
+
+<!--ecosystem-start-->
+## Ecosystem
+- `cuthbert` is built on top of [`jax`](https://github.com/google/jax) and composes
+easily with other JAX packages, e.g. [`optax`](https://github.com/google-deepmind/optax)
+for optimization, [`flax`](https://github.com/google/flax) for neural networks, and
+[`blackjax`](https://github.com/blackjax-devs/blackjax) for (SG)MCMC as well as the PPLs
+mentioned [above](#non-goals).
+- What about [`dynamax`](https://github.com/probml/dynamax)?
+    - `dynamax` is a great library for state-space model specification and inference with
+    discrete or Gaussian state-space models. `cuthbert` is focused on inference
+    with arbitrary state-space models via  e.g. SMC that is not supported in `dynamax`.
+    However as they are both built on [`jax`](https://github.com/google/jax)
+    they can be used together! A `dynamax`
+    model can be passed to `cuthbert` for inference.
+- And [`particles`](https://github.com/nchopin/particles)?
+    - [`particles`](https://github.com/nchopin/particles) and the accompanying book
+    [Sequential Monte Carlo Methods in Practice](https://link.springer.com/book/10.1007/978-3-030-47845-2)
+    are wonderful learning materials for state-space models and SMC.
+    `cuthbert` is more focused on performance and composability with the JAX ecosystem.
+- Much of the code in `cuthbert` is built on work from [`sqrt-parallel-smoothers`](https://github.com/EEA-sensors/sqrt-parallel-smoothers), [`mocat`](https://github.com/SamDuffield/mocat) and [`abile`](https://github.com/SamDuffield/abile).
+<!--ecosystem-end-->
+
+## Contributing
+
+We're always looking for contributions!
+Check out the [contributing guide](CONTRIBUTING.md) for more information.
+

cuthbert-0.0.1/README.md

@@ -0,0 +1,103 @@
+<!--intro-start-->
+<div align="center">
+<img src="docs/assets/cuthbert.png" alt="logo"></img>
+</div>
+
+A JAX library for state-space model inference
+(filtering, smoothing, static parameter estimation).
+
+> Disclaimer: The name `cuthbert` was chosen as a playful nod to the well-known
+> caterpillar cake rivalry between Aldi and M&S in the UK, as the classic state-space
+> model diagram looks vaguely like a caterpillar. However, this software project
+> has no formal connection to Aldi, M&S, or any food products (notwithstanding the coffee drunk during its writeup).
+> `cuthbert` is simply a fun name for this state-space model library and should not be interpreted as an
+> endorsement, association, or affiliation with any brand or animal themed baked goods.
+
+[![Discord](https://img.shields.io/badge/Discord-5865F2?logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/sBnr6JhT)
+[![GitHub](https://img.shields.io/badge/GitHub-181717?logo=github&logoColor=white&style=for-the-badge)](https://github.com/state-space-models/cuthbert)
+[![PyPI](https://img.shields.io/pypi/v/cuthbert?style=for-the-badge)](https://pypi.org/project/cuthbert/)
+[![Docs](https://img.shields.io/badge/Docs-b6d7a8?logo=materialformkdocs&logoColor=black&style=for-the-badge)](https://state-space-models.github.io/cuthbert/)
+<!--intro-end-->
+
+<!--goals-start-->
+### Goals
+- Simple, flexible and performant interface for state-space model inference.
+- Decoupling of model specification and inference. `cuthbert` is built to swap between
+different **inference** methods without be tied to a specific model specification.
+- Compose with the [JAX ecosystem](#ecosystem) for extensive external tools.
+- Functional API: The only classes in `cuthbert` are `NamedTuple`s and `Protocol`s.
+All functions are pure and work seamlessly with `jax.grad`, `jax.jit`, `jax.vmap` etc.
+- Methods for filtering: $p(x_t \mid y_{0:t}, \theta)$.
+- Methods for smoothing: $p(x_{0:T} \mid y_{0:T}, \theta)$ or $p(x_{t} \mid y_{0:T}, \theta)$.
+- Methods for static parameter estimation: $p(\theta \mid y_{0:T})$
+or $\text{argmax} p(y_{0:T} \mid \theta)$.
+- This includes support for forward-backward/Baum-Welch, particle filtering/sequential Monte Carlo,
+Kalman filtering (+ extended/unscented/ensemble), expectation-maximization and more!
+
+### Non-goals
+- Tools for defining models and distributions. `cuthbert` is not a probabilistic programming language (PPL).
+But can easily compose with [`dynamax`](https://github.com/probml/dynamax), [`distrax`](https://github.com/google-deepmind/distrax), [`numpyro`](https://github.com/pyro-ppl/numpyro) and [`pymc`](https://github.com/pymc-devs/pymc) in a similar way to how [`blackjax` does](https://blackjax-devs.github.io/blackjax/).
+- ["SMC Samplers"](https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf) which sample from a posterior
+distribution which is not (necessarily) a state-space model - [`blackjax` is great for this](https://github.com/blackjax-devs/blackjax/tree/main/blackjax/smc).
+<!--goals-end-->
+
+<!--codebase-structure-start-->
+### Codebase structure
+
+The codebase is structured as follows:
+
+- `cuthbert`: The main package with unified interface for filtering and smoothing.
+- `cuthbertlib`: A collection of atomic, smaller-scoped tools useful for state-space model inference,
+that represent the building blocks that power the main `cuthbert` package.
+<!--codebase-structure-end-->
+- `docs`: Source code for the documentation for `cuthbert` and `cuthbertlib`.
+- `tests`: Tests for the `cuthbert` and `cuthbertlib` packages.
+
+
+<!--installation-start-->
+## Installation
+
+`cuthbert` depends on JAX, so you'll need to [install JAX](https://docs.jax.dev/en/latest/installation.html) for the available hardware (CPU, GPU, or TPU).
+For example, on computers with NVIDIA GPUs:
+
+```bash
+pip install -U "jax[cuda13]"
+```
+
+Now install `cuthbert` from PyPI:
+
+```bash
+pip install -U cuthbert
+```
+
+Installing `cuthbert` will also install `cuthbertlib`.
+
+<!--installation-end-->
+
+<!--ecosystem-start-->
+## Ecosystem
+- `cuthbert` is built on top of [`jax`](https://github.com/google/jax) and composes
+easily with other JAX packages, e.g. [`optax`](https://github.com/google-deepmind/optax)
+for optimization, [`flax`](https://github.com/google/flax) for neural networks, and
+[`blackjax`](https://github.com/blackjax-devs/blackjax) for (SG)MCMC as well as the PPLs
+mentioned [above](#non-goals).
+- What about [`dynamax`](https://github.com/probml/dynamax)?
+    - `dynamax` is a great library for state-space model specification and inference with
+    discrete or Gaussian state-space models. `cuthbert` is focused on inference
+    with arbitrary state-space models via  e.g. SMC that is not supported in `dynamax`.
+    However as they are both built on [`jax`](https://github.com/google/jax)
+    they can be used together! A `dynamax`
+    model can be passed to `cuthbert` for inference.
+- And [`particles`](https://github.com/nchopin/particles)?
+    - [`particles`](https://github.com/nchopin/particles) and the accompanying book
+    [Sequential Monte Carlo Methods in Practice](https://link.springer.com/book/10.1007/978-3-030-47845-2)
+    are wonderful learning materials for state-space models and SMC.
+    `cuthbert` is more focused on performance and composability with the JAX ecosystem.
+- Much of the code in `cuthbert` is built on work from [`sqrt-parallel-smoothers`](https://github.com/EEA-sensors/sqrt-parallel-smoothers), [`mocat`](https://github.com/SamDuffield/mocat) and [`abile`](https://github.com/SamDuffield/abile).
+<!--ecosystem-end-->
+
+## Contributing
+
+We're always looking for contributions!
+Check out the [contributing guide](CONTRIBUTING.md) for more information.
+

cuthbert-0.0.1/cuthbert/__init__.py

@@ -0,0 +1,12 @@
+from cuthbert import discrete, gaussian, smc
+from cuthbert.filtering import filter
+from cuthbert.inference import (
+    Filter,
+    FilterCombine,
+    FilterPrepare,
+    InitPrepare,
+    Smoother,
+    SmootherCombine,
+    SmootherPrepare,
+)
+from cuthbert.smoothing import smoother

cuthbert-0.0.1/cuthbert/filtering.py

@@ -0,0 +1,81 @@
+"""Unified cuthbert filtering interface."""
+
+import warnings
+
+from jax import numpy as jnp
+from jax import random, tree, vmap
+from jax.lax import associative_scan, scan
+
+from cuthbert.inference import Filter
+from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
+
+
+def filter(
+    filter_obj: Filter,
+    model_inputs: ArrayTreeLike,
+    parallel: bool = False,
+    key: KeyArray | None = None,
+) -> ArrayTree:
+    """Applies offline filtering given a filter object and model inputs.
+
+    `model_inputs` should have leading temporal dimension of length T + 1,
+    where T is the number of time steps excluding the initial state.
+
+    Args:
+        filter_obj: The filter inference object.
+        model_inputs: The model inputs (with leading temporal dimension of length T + 1).
+        parallel: Whether to run the filter in parallel.
+            Requires `filter.associative_filter` to be `True`.
+        key: The key for the random number generator.
+
+    Returns:
+        The filtered states (NamedTuple with leading temporal dimension of length T + 1).
+    """
+    if parallel and not filter_obj.associative:
+        warnings.warn(
+            f"Parallel filtering attempted but filter.associative is False for {filter_obj}"
+        )
+
+    T = tree.leaves(model_inputs)[0].shape[0] - 1
+
+    if key is None:
+        # This will throw error if used as a key, which is desired behavior
+        # (albeit not a useful error, we could improve this)
+        prepare_keys = jnp.empty(T + 1)
+    else:
+        prepare_keys = random.split(key, T + 1)
+
+    init_model_input = tree.map(lambda x: x[0], model_inputs)
+    init_state = filter_obj.init_prepare(init_model_input, key=prepare_keys[0])
+
+    prep_model_inputs = tree.map(lambda x: x[1:], model_inputs)
+
+    if parallel:
+        other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(
+            prep_model_inputs, prepare_keys[1:]
+        )
+        prep_states = tree.map(
+            lambda x, y: jnp.concatenate([x[None], y]), init_state, other_prep_states
+        )
+        states = associative_scan(
+            vmap(filter_obj.filter_combine),
+            prep_states,
+        )
+    else:
+
+        def body(prev_state, prep_inp_and_k):
+            prep_inp, k = prep_inp_and_k
+            prep_state = filter_obj.filter_prepare(prep_inp, key=k)
+            state = filter_obj.filter_combine(prev_state, prep_state)
+            return state, state
+
+        _, states = scan(
+            body,
+            init_state,
+            (prep_model_inputs, prepare_keys[1:]),
+        )
+        states = tree.map(
+            lambda x, y: jnp.concatenate([x[None], y]), init_state, states
+        )
+
+    return states

cuthbert-0.0.1/cuthbert/inference.py

@@ -0,0 +1,220 @@
+"""Provides protocols and types for representing unified inference objects."""
+
+from typing import NamedTuple, Protocol
+
+from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
+
+
+class InitPrepare(Protocol):
+    """Protocol for preparing the initial state for the inference."""
+
+    def __call__(
+        self,
+        model_inputs: ArrayTreeLike,
+        key: KeyArray | None = None,
+    ) -> ArrayTree:
+        """Prepare the initial state for the inference.
+
+        The state at the first time point, prior to any observations.
+
+        Args:
+            model_inputs: The model inputs at the first time point.
+            key: The key for the random number generator.
+                Optional, as only used for stochastic inference methods
+
+        Returns:
+            The initial state, a NamedTuple with inference-specific fields.
+        """
+        ...
+
+
+class FilterPrepare(Protocol):
+    """Protocol for preparing the state for the filter at the next time point."""
+
+    def __call__(
+        self,
+        model_inputs: ArrayTreeLike,
+        key: KeyArray | None = None,
+    ) -> ArrayTree:
+        """Prepare the state for the filter at the next time point.
+
+        Converts the model inputs (and any stochasticity) into a unified state
+        object which can be combined with a state (of the same form) from the
+        previous time point with FilterCombine.
+
+        state = FilterCombine(prev_state, FilterPrepare(model_inputs, key))
+
+        Args:
+            model_inputs: The model inputs at the next time point.
+            key: The key for the random number generator.
+                Optional, as only used for stochastic inference methods
+
+        Returns:
+            The state prepared for FilterCombine,
+                a NamedTuple with inference-specific fields.
+        """
+        ...
+
+
+class FilterCombine(Protocol):
+    """Protocol for combining the previous state with the state from FilterPrepare."""
+
+    def __call__(
+        self,
+        state_1: ArrayTreeLike,
+        state_2: ArrayTreeLike,
+    ) -> ArrayTree:
+        """Combine state from previous time point with state from FilterPrepare.
+
+        ```python
+        state = FilterCombine(prev_state, FilterPrepare(model_inputs, key))
+        ```
+
+        Args:
+            state_1: The state from the previous time point.
+            state_2: The state from FilterPrepare for the current time point.
+
+        Returns:
+            The combined filter state, a NamedTuple with inference-specific fields.
+        """
+        ...
+
+
+class SmootherPrepare(Protocol):
+    """Protocol for preparing the state for the smoother."""
+
+    def __call__(
+        self,
+        filter_state: ArrayTreeLike,
+        model_inputs: ArrayTreeLike,
+        key: KeyArray | None = None,
+    ) -> ArrayTree:
+        """Prepare the state for the smoother at the next time point.
+
+        Converts `filter_state` with `model_inputs` (and any stochasticity) into a
+        unified state object which can be combined with a state (of the same form)
+        from the next time point with `SmootherCombine`.
+
+        Remember smoothing iterates backwards in time.
+
+        ```python
+        state = SmootherCombine(
+            SmootherPrepare(filter_state, model_inputs, key), next_smoother_state
+        )
+        ```
+
+        Note that the `model_inputs` here are different to `filter_state.model_inputs`.
+        The `model_inputs` required here are for the transition from t to t+1.
+        `filter_state.model_inputs` represents the transition from t-1 to t.
+
+        Args:
+            filter_state: The state from the filter at the previous time point.
+            model_inputs: Model inputs for the transition from t to t+1.
+            key: The key for the random number generator.
+                Optional, as only used for stochastic inference methods
+
+        Returns:
+            The state prepared for `SmootherCombine`,
+                a NamedTuple with inference-specific fields.
+        """
+        ...
+
+
+class SmootherCombine(Protocol):
+    """Protocol for combining the next smoother state with the state prepared with latest model inputs."""
+
+    def __call__(
+        self,
+        state_1: ArrayTreeLike,
+        state_2: ArrayTreeLike,
+    ) -> ArrayTree:
+        """Combine the state from the next time point with the state from `SmootherPrepare`.
+
+        Remember smoothing iterates backwards in time.
+
+        ```python
+        state = SmootherCombine(
+            SmootherPrepare(filter_state, model_inputs, key), next_smoother_state
+        )
+        ```
+
+        Args:
+            state_1: The state from `SmootherPrepare` for the current time point.
+            state_2: The state from the next time point.
+
+        Returns:
+            The combined smoother state, a NamedTuple with inference-specific fields.
+        """
+        ...
+
+
+class ConvertFilterToSmootherState(Protocol):
+    """Protocol for converting a filter state to a smoother state."""
+
+    def __call__(
+        self,
+        filter_state: ArrayTreeLike,
+        model_inputs: ArrayTreeLike | None = None,
+        key: KeyArray | None = None,
+    ) -> ArrayTree:
+        """Convert the filter state to a smoother state.
+
+        Useful for offline smoothing where the final filter state is statistically
+        equivalent to the final smoother state.
+        This function converts the filter state to the smoother state data structure.
+
+        Args:
+            filter_state: The filter state.
+            model_inputs: Only used to create an empty `model_inputs` tree
+                (the values are ignored).
+                Useful so that the final smoother state has the same structure as the rest.
+                By default, `filter_state.model_inputs` is used. So this
+                is only needed if the smoother `model_inputs` have a different tree
+                structure to `filter_state.model_inputs`.
+            key: The key for the random number generator.
+                Optional, as only used for stochastic inference methods
+
+        Returns:
+            The smoother state.
+        """
+        ...
+
+
+class Filter(NamedTuple):
+    """Filter object.
+
+    Typically passed to [cuthbert.filtering.filter][].
+
+    Attributes:
+        init_prepare: Function to prepare the initial state for the filter.
+        filter_prepare: Function to prepare intermediate states for the filter.
+        filter_combine: Function that combines two filter states to produce another.
+        associative: Whether `filter_combine` is an associative operator. Temporally
+            parallelized filters are guaranteed to produce correct results only if
+            `associative=True`.
+    """
+
+    init_prepare: InitPrepare
+    filter_prepare: FilterPrepare
+    filter_combine: FilterCombine
+    associative: bool = False
+
+
+class Smoother(NamedTuple):
+    """Smoother object.
+
+    Typically passed to [cuthbert.smoothing.smoother][].
+
+    Attributes:
+        convert_filter_to_smoother_state: Function to convert the final filter state to a smoother state.
+        smoother_prepare: Function to prepare intermediate states for the smoother.
+        smoother_combine: Function that combines two smoother states to produce another.
+        associative: Whether `smoother_combine` is an associative operator. Temporally
+            parallelized smoothers are guaranteed to produce correct results only if
+            `associative=True`.
+    """
+
+    convert_filter_to_smoother_state: ConvertFilterToSmootherState
+    smoother_prepare: SmootherPrepare
+    smoother_combine: SmootherCombine
+    associative: bool = False

cuthbert-0.0.1/cuthbert/smoothing.py

@@ -0,0 +1,128 @@
+"""Unified cuthbert smoothing interface."""
+
+import warnings
+
+from jax import numpy as jnp
+from jax import random, tree, vmap
+from jax.lax import associative_scan, scan
+
+from cuthbert.inference import Smoother
+from cuthbert.utils import dummy_tree_like
+from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
+
+
+def smoother(
+    smoother_obj: Smoother,
+    filter_states: ArrayTreeLike,
+    model_inputs: ArrayTreeLike | None = None,
+    parallel: bool = False,
+    key: KeyArray | None = None,
+) -> ArrayTree:
+    """Applies offline smoothing given a smoother object, output from filter, and model inputs.
+
+    `filter_states` should have leading temporal dimension of length T + 1, where
+    T is the number of time steps excluding the initial state.
+
+    Each element of `model_inputs` refers to the transition from t to t+1, except for the
+    first element which refers to the initial state. The initial state `model_inputs`
+    are not used for smoothing. Thus the `model_inputs` used here have length T.
+    By default, `filter_states.model_inputs[1:]` are used (i.e. the `model_inputs`
+    used for the initial state is ignored).
+
+    Args:
+        smoother_obj: The smoother inference object.
+        filter_states: The filtered states (with leading temporal dimension of length T + 1).
+        model_inputs: The model inputs (with leading temporal dimension of length T).
+            Optional, if None then `filter_states.model_inputs[1:]` are used.
+        parallel: Whether to run the smoother in parallel.
+            Requires `smoother_obj.associative_smoother` to be `True`.
+        key: The key for the random number generator.
+
+    Returns:
+        The smoothed states (NamedTuple with leading temporal dimension of length T + 1).
+    """
+    if parallel and not smoother_obj.associative:
+        warnings.warn(
+            "Parallel smoothing attempted but smoother.associative is False "
+            f"for {smoother}"
+        )
+
+    if model_inputs is None:
+        model_inputs = filter_states.model_inputs
+
+    T = tree.leaves(filter_states)[0].shape[0] - 1
+
+    # model_inputs for the dynamics distribution from t-1 to t is stored
+    # in model_inputs[t] thus we need model_inputs[1:]
+    # model_inputs[0] is only used for init_prepare and not for smoothing.
+    # Therefore, we allow model_inputs to be either of length T + 1 or T
+    # where if length is T + 1 then we simply discard model_inputs[0]
+    model_inputs_length = tree.leaves(model_inputs)[0].shape[0]
+    if model_inputs_length == T + 1:
+        model_inputs = tree.map(lambda x: x[1:], model_inputs)
+    elif model_inputs_length != T:
+        raise ValueError(
+            "model_inputs must have length T + 1 or T, got length "
+            f"{model_inputs_length}"
+        )
+
+    if key is None:
+        # This will throw error if used as a key, which is desired
+        # (albeit not a useful error, we could improve this)
+        prepare_keys = jnp.empty(T + 1)
+    else:
+        prepare_keys = random.split(key, T + 1)
+
+    final_filter_state = tree.map(lambda x: x[-1], filter_states)
+    other_filter_states = tree.map(lambda x: x[:-1], filter_states)
+
+    # Final smoother state doesn't need model inputs, so we create a dummy one
+    # with the same structure as model_inputs but with all values set to dummy values.
+    dummy_single_model_inputs = dummy_tree_like(tree.map(lambda x: x[0], model_inputs))
+
+    final_smoother_state = smoother_obj.convert_filter_to_smoother_state(
+        final_filter_state, model_inputs=dummy_single_model_inputs, key=prepare_keys[0]
+    )
+
+    if parallel:
+        prep_states = vmap(
+            lambda fs, inp, k: smoother_obj.smoother_prepare(
+                fs, model_inputs=inp, key=k
+            )
+        )(other_filter_states, model_inputs, prepare_keys[1:])
+        prep_states = tree.map(
+            lambda x, y: jnp.concatenate([x, y[None]]),
+            prep_states,
+            final_smoother_state,
+        )
+
+        states = associative_scan(
+            vmap(lambda current, next: smoother_obj.smoother_combine(next, current)),
+            # TODO: Maybe change cuthbertlib direction so that this lambda isn't needed
+            prep_states,
+            reverse=True,
+        )
+    else:
+
+        def body(next_state, filt_state_and_prep_inp_and_k):
+            filt_state, prep_inp, k = filt_state_and_prep_inp_and_k
+            prep_state = smoother_obj.smoother_prepare(
+                filt_state, model_inputs=prep_inp, key=k
+            )
+            state = smoother_obj.smoother_combine(prep_state, next_state)
+            return state, state
+
+        _, states = scan(
+            body,
+            final_smoother_state,
+            (other_filter_states, model_inputs, prepare_keys[1:]),
+            reverse=True,
+        )
+
+        states = tree.map(
+            lambda x, y: jnp.concatenate([x, y[None]]),
+            states,
+            final_smoother_state,
+        )
+
+    return states

cuthbert-0.0.1/cuthbert/utils.py

@@ -0,0 +1,32 @@
+"""Utility functions (filling dummy arrays and trees) for cuthbert."""
+
+import jax
+import jax.numpy as jnp
+
+from cuthbertlib.types import Array, ArrayLike, ArrayTree, ArrayTreeLike
+
+
+def _dummy_array(leaf: ArrayLike | jax.ShapeDtypeStruct) -> Array:
+    """Returns an array of the same shape and dtype filled with dummy values."""
+    if isinstance(leaf, jax.ShapeDtypeStruct):
+        leaf = jnp.empty_like(leaf)
+
+    leaf = jnp.asarray(leaf)
+    dtype = leaf.dtype
+    shape = leaf.shape
+
+    if jnp.issubdtype(dtype, jnp.integer):
+        min_val = jnp.iinfo(dtype).min
+    elif jnp.issubdtype(dtype, jnp.floating):
+        min_val = jnp.finfo(dtype).min
+    elif jnp.issubdtype(dtype, jnp.bool_):
+        min_val = False
+    else:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+    return jnp.full(shape, min_val, dtype=dtype)
+
+
+def dummy_tree_like(pytree: ArrayTreeLike) -> ArrayTree:
+    """Returns a pytree with the same structure filled with dummy values."""
+    return jax.tree.map(_dummy_array, pytree)

cuthbert-0.0.1/cuthbert.egg-info/PKG-INFO

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: cuthbert
+Version: 0.0.1
+Summary: State-space model inference with JAX
+Author-email: Sam Duffield <s@mduffield.com>, Sahel Iqbal <sahel13miqbal@proton.me>, Adrien Corenflos <adrien.corenflos.stats@gmail.com>
+License: Apache-2.0
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.35
+Requires-Dist: numba>=0.60.0
+Provides-Extra: tests
+Requires-Dist: chex; extra == "tests"
+Requires-Dist: pre-commit; extra == "tests"
+Requires-Dist: ruff; extra == "tests"
+Requires-Dist: pyright; extra == "tests"
+Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-xdist; extra == "tests"
+Requires-Dist: entangled-cli; python_version >= "3.12" and extra == "tests"
+Provides-Extra: docs
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocs-include-markdown-plugin; extra == "docs"
+Requires-Dist: mkdocstrings-python; extra == "docs"
+Provides-Extra: examples
+Requires-Dist: matplotlib; extra == "examples"
+Requires-Dist: pandas; extra == "examples"
+Requires-Dist: pandas-stubs; extra == "examples"
+Requires-Dist: yfinance; extra == "examples"
+Requires-Dist: dynamax; extra == "examples"
+Requires-Dist: ipython; extra == "examples"
+Requires-Dist: tfp-nightly[jax]; extra == "examples"
+Dynamic: license-file
+
+<!--intro-start-->
+<div align="center">
+<img src="docs/assets/cuthbert.png" alt="logo"></img>
+</div>
+
+A JAX library for state-space model inference
+(filtering, smoothing, static parameter estimation).
+
+> Disclaimer: The name `cuthbert` was chosen as a playful nod to the well-known
+> caterpillar cake rivalry between Aldi and M&S in the UK, as the classic state-space
+> model diagram looks vaguely like a caterpillar. However, this software project
+> has no formal connection to Aldi, M&S, or any food products (notwithstanding the coffee drunk during its writeup).
+> `cuthbert` is simply a fun name for this state-space model library and should not be interpreted as an
+> endorsement, association, or affiliation with any brand or animal themed baked goods.
+
+[![Discord](https://img.shields.io/badge/Discord-5865F2?logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/sBnr6JhT)
+[![GitHub](https://img.shields.io/badge/GitHub-181717?logo=github&logoColor=white&style=for-the-badge)](https://github.com/state-space-models/cuthbert)
+[![PyPI](https://img.shields.io/pypi/v/cuthbert?style=for-the-badge)](https://pypi.org/project/cuthbert/)
+[![Docs](https://img.shields.io/badge/Docs-b6d7a8?logo=materialformkdocs&logoColor=black&style=for-the-badge)](https://state-space-models.github.io/cuthbert/)
+<!--intro-end-->
+
+<!--goals-start-->
+### Goals
+- Simple, flexible and performant interface for state-space model inference.
+- Decoupling of model specification and inference. `cuthbert` is built to swap between
+different **inference** methods without be tied to a specific model specification.
+- Compose with the [JAX ecosystem](#ecosystem) for extensive external tools.
+- Functional API: The only classes in `cuthbert` are `NamedTuple`s and `Protocol`s.
+All functions are pure and work seamlessly with `jax.grad`, `jax.jit`, `jax.vmap` etc.
+- Methods for filtering: $p(x_t \mid y_{0:t}, \theta)$.
+- Methods for smoothing: $p(x_{0:T} \mid y_{0:T}, \theta)$ or $p(x_{t} \mid y_{0:T}, \theta)$.
+- Methods for static parameter estimation: $p(\theta \mid y_{0:T})$
+or $\text{argmax} p(y_{0:T} \mid \theta)$.
+- This includes support for forward-backward/Baum-Welch, particle filtering/sequential Monte Carlo,
+Kalman filtering (+ extended/unscented/ensemble), expectation-maximization and more!
+
+### Non-goals
+- Tools for defining models and distributions. `cuthbert` is not a probabilistic programming language (PPL).
+But can easily compose with [`dynamax`](https://github.com/probml/dynamax), [`distrax`](https://github.com/google-deepmind/distrax), [`numpyro`](https://github.com/pyro-ppl/numpyro) and [`pymc`](https://github.com/pymc-devs/pymc) in a similar way to how [`blackjax` does](https://blackjax-devs.github.io/blackjax/).
+- ["SMC Samplers"](https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf) which sample from a posterior
+distribution which is not (necessarily) a state-space model - [`blackjax` is great for this](https://github.com/blackjax-devs/blackjax/tree/main/blackjax/smc).
+<!--goals-end-->
+
+<!--codebase-structure-start-->
+### Codebase structure
+
+The codebase is structured as follows:
+
+- `cuthbert`: The main package with unified interface for filtering and smoothing.
+- `cuthbertlib`: A collection of atomic, smaller-scoped tools useful for state-space model inference,
+that represent the building blocks that power the main `cuthbert` package.
+<!--codebase-structure-end-->
+- `docs`: Source code for the documentation for `cuthbert` and `cuthbertlib`.
+- `tests`: Tests for the `cuthbert` and `cuthbertlib` packages.
+
+
+<!--installation-start-->
+## Installation
+
+`cuthbert` depends on JAX, so you'll need to [install JAX](https://docs.jax.dev/en/latest/installation.html) for the available hardware (CPU, GPU, or TPU).
+For example, on computers with NVIDIA GPUs:
+
+```bash
+pip install -U "jax[cuda13]"
+```
+
+Now install `cuthbert` from PyPI:
+
+```bash
+pip install -U cuthbert
+```
+
+Installing `cuthbert` will also install `cuthbertlib`.
+
+<!--installation-end-->
+
+<!--ecosystem-start-->
+## Ecosystem
+- `cuthbert` is built on top of [`jax`](https://github.com/google/jax) and composes
+easily with other JAX packages, e.g. [`optax`](https://github.com/google-deepmind/optax)
+for optimization, [`flax`](https://github.com/google/flax) for neural networks, and
+[`blackjax`](https://github.com/blackjax-devs/blackjax) for (SG)MCMC as well as the PPLs
+mentioned [above](#non-goals).
+- What about [`dynamax`](https://github.com/probml/dynamax)?
+    - `dynamax` is a great library for state-space model specification and inference with
+    discrete or Gaussian state-space models. `cuthbert` is focused on inference
+    with arbitrary state-space models via  e.g. SMC that is not supported in `dynamax`.
+    However as they are both built on [`jax`](https://github.com/google/jax)
+    they can be used together! A `dynamax`
+    model can be passed to `cuthbert` for inference.
+- And [`particles`](https://github.com/nchopin/particles)?
+    - [`particles`](https://github.com/nchopin/particles) and the accompanying book
+    [Sequential Monte Carlo Methods in Practice](https://link.springer.com/book/10.1007/978-3-030-47845-2)
+    are wonderful learning materials for state-space models and SMC.
+    `cuthbert` is more focused on performance and composability with the JAX ecosystem.
+- Much of the code in `cuthbert` is built on work from [`sqrt-parallel-smoothers`](https://github.com/EEA-sensors/sqrt-parallel-smoothers), [`mocat`](https://github.com/SamDuffield/mocat) and [`abile`](https://github.com/SamDuffield/abile).
+<!--ecosystem-end-->
+
+## Contributing
+
+We're always looking for contributions!
+Check out the [contributing guide](CONTRIBUTING.md) for more information.
+

cuthbert-0.0.1/cuthbert.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+cuthbert/__init__.py
+cuthbert/filtering.py
+cuthbert/inference.py
+cuthbert/smoothing.py
+cuthbert/utils.py
+cuthbert.egg-info/PKG-INFO
+cuthbert.egg-info/SOURCES.txt
+cuthbert.egg-info/dependency_links.txt
+cuthbert.egg-info/requires.txt
+cuthbert.egg-info/top_level.txt
+cuthbertlib/__init__.py
+cuthbertlib/types.py
+tests/test_examples_scripts.py

cuthbert-0.0.1/cuthbert.egg-info/requires.txt

@@ -0,0 +1,27 @@
+jax>=0.4.35
+numba>=0.60.0
+
+[docs]
+mkdocs-material
+mkdocs-include-markdown-plugin
+mkdocstrings-python
+
+[examples]
+matplotlib
+pandas
+pandas-stubs
+yfinance
+dynamax
+ipython
+tfp-nightly[jax]
+
+[tests]
+chex
+pre-commit
+ruff
+pyright
+pytest
+pytest-xdist
+
+[tests:python_version >= "3.12"]
+entangled-cli

cuthbert-0.0.1/cuthbert.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+cuthbert
+cuthbertlib

cuthbert-0.0.1/cuthbertlib/types.py

@@ -0,0 +1,20 @@
+"""Type aliases for common types used in the library."""
+
+from typing import Any, Callable, TypeAlias
+
+from jax import Array
+from jax.typing import ArrayLike
+
+KeyArray: TypeAlias = Array  # No native JAX type annotation for keys https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-16-sept-18-2023
+ArrayTree: TypeAlias = Any  # No native JAX type annotation for PyTrees https://github.com/google/jax/issues/3340
+ArrayTreeLike: TypeAlias = Any  # Tree with all leaves castable to jax.Array https://jax.readthedocs.io/en/latest/jax.typing.html#module-jax.typing
+ScalarArray: TypeAlias = (
+    Array  # jax.Array with just a single float element, i.e. shape ()
+)
+ScalarArrayLike: TypeAlias = ArrayLike  # Object that will be cast to a ScalarArray
+
+LogDensity: TypeAlias = Callable[[ArrayTreeLike], ScalarArray]
+LogConditionalDensity: TypeAlias = Callable[[ArrayTreeLike, ArrayTreeLike], ScalarArray]
+LogConditionalDensityAux: TypeAlias = Callable[
+    [ArrayTreeLike, ArrayTreeLike], tuple[ScalarArray, ArrayTree]
+]

cuthbert-0.0.1/pyproject.toml

@@ -0,0 +1,64 @@
+[build-system]
+requires = ["setuptools >= 61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cuthbert"
+version = "0.0.1"
+description = "State-space model inference with JAX"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+authors = [
+    { name = "Sam Duffield", email = "s@mduffield.com" },
+    { name = "Sahel Iqbal", email = "sahel13miqbal@proton.me" },
+    { name = "Adrien Corenflos", email = "adrien.corenflos.stats@gmail.com" },
+]
+
+dependencies = [
+    "jax>=0.4.35",
+    "numba>=0.60.0",
+]
+
+[tool.uv]
+# The tensorflow-probability override forces installation on undefined platforms only
+# effectively disabling it for all normal platforms
+override-dependencies = [
+    "tensorflow-probability ; sys_platform == 'undefined'",
+]
+
+[tool.setuptools.packages.find]
+include = ["cuthbert", "cuthbertlib"]
+
+[project.optional-dependencies]
+tests = [
+    "chex",  # test additional dependencies
+    "pre-commit", "ruff", "pyright",  # linting etc
+    "pytest", "pytest-xdist",  # testing
+    "entangled-cli; python_version >= '3.12'", # convert .md examples to .py, requires Python 3.12+
+]
+docs = ["mkdocs-material", "mkdocs-include-markdown-plugin", "mkdocstrings-python"]
+examples = [
+    "matplotlib", "pandas", "pandas-stubs", "yfinance",
+    "dynamax", "ipython", "tfp-nightly[jax]",
+]
+# ipython just required for dynamax, can be removed when they do a new release
+# tfp-nightly[jax] provides tensorflow_probability module that dynamax needs
+
+[tool.ruff]
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401", "F821", "E402", "D104"]
+"tests/**/*.py" = ["D"] # no docstring checking for tests
+"cuthbertlib/resampling/**/*.py" = ["D103"] # resampling uses decorator for docstrings
+[tool.ruff.lint]
+select = ["D"]
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.pytest]
+[tool.pytest.ini_options]
+markers = "examples: Run tangled example scripts as tests"
+
+[tool.entangled]
+version = "2.3.0"
+ignore_list = ["CONTRIBUTING.md"]

cuthbert-0.0.1/tests/test_examples_scripts.py

@@ -0,0 +1,13 @@
+import glob
+import runpy
+
+import pytest
+
+EXAMPLES_DIR = "examples_scripts"
+
+
+@pytest.mark.examples
+@pytest.mark.parametrize("script", glob.glob(f"{EXAMPLES_DIR}/*.py"))
+def test_example_scripts(script):
+    """Run each tangled example script to ensure it runs without error."""
+    runpy.run_path(script)

cuthbert-0.0.0/PKG-INFO

@@ -1,5 +0,0 @@
-Metadata-Version: 2.1
-Name: cuthbert
-Version: 0.0.0
-Summary: Coming soon
-Author-email: Sam Duffield <s@mduffield.com>

cuthbert-0.0.0/cuthbert.egg-info/PKG-INFO

@@ -1,5 +0,0 @@
-Metadata-Version: 2.1
-Name: cuthbert
-Version: 0.0.0
-Summary: Coming soon
-Author-email: Sam Duffield <s@mduffield.com>

cuthbert-0.0.0/cuthbert.egg-info/SOURCES.txt

@@ -1,6 +0,0 @@
-pyproject.toml
-cuthbert/__init__.py
-cuthbert.egg-info/PKG-INFO
-cuthbert.egg-info/SOURCES.txt
-cuthbert.egg-info/dependency_links.txt
-cuthbert.egg-info/top_level.txt

cuthbert-0.0.0/cuthbert.egg-info/top_level.txt

@@ -1 +0,0 @@
-cuthbert

cuthbert-0.0.0/pyproject.toml

@@ -1,7 +0,0 @@
-[project]
-name = "cuthbert"
-authors = [
-	{name = "Sam Duffield", email = "s@mduffield.com"},
-]
-version = "0.0.0"
-description = "Coming soon"