aspire-inference 0.1.0a9__tar.gz → 0.1.0a11__tar.gz

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a9
+Version: 0.1.0a11
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -26,7 +26,8 @@ Requires-Dist: torch; extra == "torch"
 Requires-Dist: zuko; extra == "torch"
 Requires-Dist: tqdm; extra == "torch"
 Provides-Extra: minipcn
-Requires-Dist: minipcn; extra == "minipcn"
+Requires-Dist: minipcn[array-api]>=0.2.0a3; extra == "minipcn"
+Requires-Dist: orng; extra == "minipcn"
 Provides-Extra: emcee
 Requires-Dist: emcee; extra == "emcee"
 Provides-Extra: blackjax
@@ -49,10 +50,28 @@ aspire is a framework for reusing existing posterior samples to obtain new resul
 
 ## Installation
 
-aspire can be installed from PyPI using `pip`
+aspire can be installed from PyPI using `pip`. By default, you need to install
+one of the backends for the normalizing flows, either `torch` or `jax`.
+We also recommend installing `minipcn` if using the `smc` sampler:
+
+
+**Torch**
+
+We recommend installing `torch` manually to ensure correct CPU/CUDA versions are
+ installed. See the [PyTorch installation instructions](https://pytorch.org/)
+ for more details.
+
+```
+pip install aspire-inference[torch,minipcn]
+```
+
+**Jax**:
+
+We recommend install `jax` manually to ensure the correct GPU/CUDA versions
+are installed. See the [jax documentation for details](https://docs.jax.dev/en/latest/installation.html)
 
 ```
-pip install aspire-inference
+pip install aspire-inference[jax,minipcn]
 ```
 
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/README.md

@@ -10,10 +10,28 @@ aspire is a framework for reusing existing posterior samples to obtain new resul
 
 ## Installation
 
-aspire can be installed from PyPI using `pip`
+aspire can be installed from PyPI using `pip`. By default, you need to install
+one of the backends for the normalizing flows, either `torch` or `jax`.
+We also recommend installing `minipcn` if using the `smc` sampler:
+
+
+**Torch**
+
+We recommend installing `torch` manually to ensure correct CPU/CUDA versions are
+ installed. See the [PyTorch installation instructions](https://pytorch.org/)
+ for more details.
+
+```
+pip install aspire-inference[torch,minipcn]
+```
+
+**Jax**:
+
+We recommend install `jax` manually to ensure the correct GPU/CUDA versions
+are installed. See the [jax documentation for details](https://docs.jax.dev/en/latest/installation.html)
 
 ```
-pip install aspire-inference
+pip install aspire-inference[jax,minipcn]
 ```
 
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/aspire_inference.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a9
+Version: 0.1.0a11
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -26,7 +26,8 @@ Requires-Dist: torch; extra == "torch"
 Requires-Dist: zuko; extra == "torch"
 Requires-Dist: tqdm; extra == "torch"
 Provides-Extra: minipcn
-Requires-Dist: minipcn; extra == "minipcn"
+Requires-Dist: minipcn[array-api]>=0.2.0a3; extra == "minipcn"
+Requires-Dist: orng; extra == "minipcn"
 Provides-Extra: emcee
 Requires-Dist: emcee; extra == "emcee"
 Provides-Extra: blackjax
@@ -49,10 +50,28 @@ aspire is a framework for reusing existing posterior samples to obtain new resul
 
 ## Installation
 
-aspire can be installed from PyPI using `pip`
+aspire can be installed from PyPI using `pip`. By default, you need to install
+one of the backends for the normalizing flows, either `torch` or `jax`.
+We also recommend installing `minipcn` if using the `smc` sampler:
+
+
+**Torch**
+
+We recommend installing `torch` manually to ensure correct CPU/CUDA versions are
+ installed. See the [PyTorch installation instructions](https://pytorch.org/)
+ for more details.
+
+```
+pip install aspire-inference[torch,minipcn]
+```
+
+**Jax**:
+
+We recommend install `jax` manually to ensure the correct GPU/CUDA versions
+are installed. See the [jax documentation for details](https://docs.jax.dev/en/latest/installation.html)
 
 ```
-pip install aspire-inference
+pip install aspire-inference[jax,minipcn]
 ```
 
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/aspire_inference.egg-info/SOURCES.txt

@@ -13,6 +13,7 @@ aspire_inference.egg-info/dependency_links.txt
 aspire_inference.egg-info/requires.txt
 aspire_inference.egg-info/top_level.txt
 docs/Makefile
+docs/checkpointing.rst
 docs/conf.py
 docs/entry_points.rst
 docs/examples.rst
@@ -53,6 +54,7 @@ tests/test_samples.py
 tests/test_transforms.py
 tests/test_utils.py
 tests/integration_tests/conftest.py
+tests/integration_tests/test_checkpointing.py
 tests/integration_tests/test_integration.py
 tests/test_flows/test_flows_core.py
 tests/test_flows/test_jax_flows/test_flowjax_flows.py

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/aspire_inference.egg-info/requires.txt

@@ -16,7 +16,8 @@ jaxlib
 flowjax
 
 [minipcn]
-minipcn
+minipcn[array-api]>=0.2.0a3
+orng
 
 [scipy]
 scipy

aspire_inference-0.1.0a11/docs/checkpointing.rst

@@ -0,0 +1,88 @@
+Checkpointing and Resuming
+==========================
+
+Aspire provides a few simple patterns to resume long runs.
+
+Saving checkpoints while sampling
+---------------------------------
+
+- Pass ``checkpoint_path`` (an HDF5 file) to :py:meth:`aspire.Aspire.sample_posterior`
+  to write checkpoints as the sampler runs. Use ``checkpoint_every`` to control
+  frequency and ``checkpoint_save_config/flow`` to control what metadata is saved.
+- For a convenience wrapper, wrap your sampling in
+  ``with aspire.auto_checkpoint("run.h5", every=1): ...``. Inside the context,
+  ``sample_posterior`` will default to checkpointing to that file, and the config/flow
+  will be updated as needed.
+
+What gets saved
+^^^^^^^^^^^^^^^
+
+- The sampler stores checkpoints under ``/checkpoint/state`` in the HDF5 file.
+- Aspire writes ``/aspire_config`` (with ``sampler_type`` and ``sampler_config``) and
+  ``/flow``. If these already exist, they are overwritten when saving.
+
+Resuming from a file
+--------------------
+
+- Use :py:meth:`aspire.Aspire.resume_from_file` to rebuild an Aspire instance and flow
+  from a checkpoint file:
+
+  .. code-block:: python
+
+      aspire = Aspire.resume_from_file(
+          "run.h5",
+          log_likelihood=log_likelihood,
+          log_prior=log_prior,
+      )
+      # Optionally continue checkpointing to the same file
+      with aspire.auto_checkpoint("run.h5", every=1):
+          samples = aspire.sample_posterior()
+
+- ``resume_from_file`` loads config, flow, and the last checkpoint (if present), and
+  primes the instance to resume sampling; you can still override sampler kwargs when
+  calling ``sample_posterior``.
+
+Manual resume via ``sample_posterior`` args
+-------------------------------------------
+
+- If you have a checkpoint blob (bytes or dict) already, you can pass it directly:
+
+  .. code-block:: python
+
+      samples = aspire.sample_posterior(
+          n_samples=...,
+          sampler="smc",
+          resume_from=checkpoint_bytes_or_dict,
+          checkpoint_path="run.h5",  # optional: keep writing checkpoints
+      )
+
+- To resume from a file without using ``resume_from_file``, load the checkpoint bytes
+  and flow yourself, then call ``sample_posterior``:
+
+  .. code-block:: python
+
+      from aspire.utils import AspireFile
+
+      aspire = Aspire(..., flow_backend="zuko")
+      with AspireFile("run.h5", "r") as f:
+          aspire.load_flow(f, path="flow")
+          # Standard layout is /checkpoint/state; adjust if you used a different path
+          checkpoint_bytes = f["checkpoint"]["state"][...].tobytes()
+      samples = aspire.sample_posterior(
+          n_samples=..., sampler="smc", resume_from="run.h5"
+          # or resume_from=checkpoint_bytes if you prefer to pass bytes directly
+      )
+
+Notes and tips
+--------------
+
+- Checkpoint files must be HDF5 (``.h5``/``.hdf5``).
+- If a checkpoint is missing in the file (e.g., sampling never wrote one), the flow
+  and config are still loaded; you can simply start sampling again and checkpointing
+  will continue to the same file.
+- For manual control, you can always call ``save_config`` / ``save_flow`` yourself
+  on an :class:`aspire.utils.AspireFile`.
+- SMC samplers also accept a custom ``checkpoint_callback`` and ``checkpoint_every`` if
+  you want full control over how checkpoints are persisted or inspected. Provide a
+  callable that accepts the checkpoint state dict; from there you can, for example,
+  serialize to another format or push to remote storage.

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/docs/conf.py

@@ -46,7 +46,7 @@ exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 autoapi_type = "python"
 autoapi_dirs = ["../src/aspire/"]
-autoapi_add_toctree_entry = True
+autoapi_add_toctree_entry = False
 autoapi_options = [
     "members",
     "imported-members",

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/docs/index.rst

@@ -57,7 +57,16 @@ examples, and the complete API reference.
 
     installation
     user_guide
+    checkpointing
     recipes
     multiprocessing
     examples
     entry_points
+    API Reference </autoapi/aspire/index>
+
+
+.. toctree::
+    :maxdepth: 2
+    :caption: Related Packages
+
+    aspire-bilby <https://aspire.readthedocs.io/projects/aspire-bilby/en/latest/>

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/docs/installation.rst

@@ -14,7 +14,11 @@ Install the library from PyPI (note the published name):
 
    $ python -m pip install aspire-inference
 
-The installed distribution exposes the ``aspire`` import namespace.
+The installed distribution exposes the ``aspire`` import namespace. By default,
+this doesn't include any optional dependencies beyond the core ones listed above.
+We recommend installing with at least one backend for normalizing flows, e.g.
+``torch`` (PyTorch + ``zuko``) or ``jax`` (JAX + ``flowjax``).
+and optionally the ``minipcn`` SMC kernel.
 
 Optional extras
 ---------------

{aspire_inference-0.1.0a9 → aspire_inference-0.1.0a11}/pyproject.toml

@@ -39,7 +39,8 @@ torch = [
     "tqdm",
 ]
 minipcn = [
-    "minipcn",
+    "minipcn[array-api]>=0.2.0a3",
+    "orng",
 ]
 emcee = [
     "emcee",