Skip to content

DOC: clean up readme #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mj-will merged 1 commit into from
Oct 14, 2026
Merged

DOC: clean up readme #5

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
214 changes: 97 additions & 117 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,168 +1,148 @@
# sequince-density

Fitted proposal distributions and optional normalizing-flow backends for
`sequince`.
`sequince-density` provides density estimators for use with the `sequince` SMC
framework, including fitted proposal distributions and normalizing flows for
preconditioned sampling.

## Installation

`sequince-density` can be installed directly from PyPI:

```bash
pip install sequince-density
```

## Gaussian-mixture proposal
However, we recommend also installing the relevant dependencies for the features
you plan to use. See below for more details.

## Features

- Normalizing-flow adapters for Zuko (PyTorch) and FlowJAX (JAX).
- Gaussian-mixture proposals fitted with scikit-learn.
- Integration with Sequince's `FlowTransform` for mutation in a learned latent
space and `FlowProposal` for fitted proposal distributions.

## Quick start

### Normalizing flows

Install the scikit-learn integration with:
Install the optional dependencies for your chosen backend:

```bash
pip install "sequince-density[gmm]"
pip install "sequince-density[zuko]"
# Alternatively, for the JAX backend:
pip install "sequince-density[flowjax]"
```

Fit a Zuko flow to a two-dimensional sample population, then draw samples and
evaluate their log densities:

```python
from orng import RandomGenerator
from sequince_density import GaussianMixtureProposal
import numpy as np

proposal = GaussianMixtureProposal(n_components=4).fit(
particles,
rng=RandomGenerator(seed=122),
)
samples, log_q = proposal.sample_and_log_prob(
256,
rng=RandomGenerator(seed=123),
)
```
from sequince_density import ZukoFlow

Fitting uses scikit-learn's `GaussianMixture`. Inputs may use NumPy, PyTorch,
JAX, or CuPy arrays; samples and densities are restored to the input backend
and device unless an explicit `xp` is requested when sampling. Full covariance
matrices, regularization, and checkpoint state are supported. Weighted fitting
is not currently supported; uniform post-resampling populations are the
intended input.
rng = np.random.default_rng(42)
particles = rng.normal(size=(2000, 2))

flow = ZukoFlow(dims=2, flow_class="MAF", hidden_features=[32, 32])
flow.fit(particles, n_epochs=100, batch_size=500, rng=42)
samples, log_q = flow.sample_and_log_prob(256, rng=43)
```

## Sequince rejuvenation
`ZukoFlow` returns PyTorch tensors; `FlowJax` returns JAX arrays.

`sequince-density` depends on `sequince`; Sequince does not depend on this
package. Inject the proposal into Sequince's backend-agnostic rejuvenation
step and place it explicitly after resampling:
For preconditioned SMC, pass an unfitted flow to `FlowTransform` and use the
transform as the `particle_transform` argument to
`make_standard_bayesian_smc`. The mutation step fits it to the current particle
population:

```python
from sequince.steps import ProposalRejuvenation
from sequince_density import GaussianMixtureProposal
from sequince.transforms import FlowTransform

rejuvenation = ProposalRejuvenation(
posterior,
GaussianMixtureProposal(n_components=8),
transform = FlowTransform(
ZukoFlow(dims=2, flow_class="MAF", hidden_features=[32, 32]),
fit_kwargs={"n_epochs": 100, "batch_size": 500},
)
```

This direct redraw is approximate and deliberately has no Metropolis-Hastings
correction. Sequince fits it after resampling in the standard step order.
### Gaussian mixture models

## Available flow backends
GMM support is provided via scikit-learn and can be installed with:

Install the Zuko integration with `pip install "sequince-density[zuko]"` or
the FlowJAX integration with `pip install "sequince-density[flowjax]"`.
```bash
pip install "sequince-density[gmm]"
```

```python
from sequince_density import FlowJax, ZukoFlow
import numpy as np

from sequince_density import GaussianMixtureProposal

particles = np.random.default_rng(42).normal(size=(2000, 2))
proposal = GaussianMixtureProposal(n_components=4).fit(particles, rng=42)
samples, log_q = proposal.sample_and_log_prob(256, rng=43)
```

`ZukoFlow` wraps flows from `zuko.flows`, while `FlowJax` wraps constructors
from `flowjax.flows`. Both satisfy the `sequince.flows.NormalizingFlow`
protocol and return arrays in their native backend; `FlowProposal` handles any
conversion to the SMC particle namespace.
Fitting expects particles shaped `(n_particles, n_dimensions)`. Nonuniform
weights are not supported; use an equal-weight population, such as particles
after resampling.

## Flow as a mutation transform
## Proposal distributions

Use `FlowTransform` when you want Sequince mutation steps to operate in a
learned latent space.
Both Gaussian mixtures and normalizing flows can provide a fitted starting
distribution for SMC. `GaussianMixtureProposal` implements Sequince's proposal
interface directly. Wrap a flow in `FlowProposal` to use the same interface,
including conversion to the SMC particle backend:

```python
import numpy as np
from sequince.proposals import FlowProposal

from sequince import ParticleState, make_standard_bayesian_smc
from sequince.transforms import AffineTransform, FlowTransform
from sequince_density import ZukoFlow
proposal = FlowProposal(
ZukoFlow(dims=2, flow_class="MAF", hidden_features=[32, 32]),
fit_kwargs={"n_epochs": 100, "batch_size": 500},
).fit(particles, rng=42)
```

rng = np.random.default_rng(123)
dims = 2
n_particles = 256
For either fitted proposal type, pass `proposal` as `log_proposal` to
`make_standard_bayesian_smc` and draw the initial particles from that same
proposal. SMC then bridges from the proposal to the target posterior.

initial_particles = rng.normal(size=(n_particles, dims))
initial_state = ParticleState(particles=initial_particles)
In an ASPIRE-style analysis, a flow is fitted to samples from an initial
analysis and used as the starting proposal for a new SMC run with an updated
likelihood.

flow = ZukoFlow(
dims=dims,
flow_class="MAF",
hidden_features=[64, 64],
particle_transform=AffineTransform(),
)
Full proposal examples will be linked here once published in
[`sequince-examples`](https://github.com/sequince-dev/sequince-examples).

smc, posterior, metadata = make_standard_bayesian_smc(
log_prior=lambda x, context=None: -0.5 * np.sum(x**2, axis=1),
log_likelihood=lambda x, context=None: np.zeros(x.shape[0]),
particle_transform=FlowTransform(
flow,
fit_kwargs={
"n_epochs": 100,
"lr": 1e-3,
"batch_size": 500,
},
),
mutation_kwargs={"transform_refit": "per-context"},
)
## Documentation

result = smc.run(initial_state, rng=rng, metadata=metadata)
```
See the [sequince density documentation](https://sequince.readthedocs.io/projects/sequince-density) for the SMC
framework and its proposal and transform interfaces.

A `FlowTransform` is fitted on the current particle cloud when the mutation
step prepares its transform, so you normally pass an unfitted flow adapter.
The mutation transform fit policy controls its lifetime:
## Examples

- `always` fits before every mutation call.
- `per-context` fits once per outer SMC iteration and reuses the flow across
repeated or alternating mutation steps.
- `once` fits once for the lifetime of the mutation step.
The [`sequince-examples`](https://github.com/sequince-dev/sequince-examples)
repository and corresponding documentation contain examples highlighting
different features in `sequince`.

## Flow as a proposal distribution
## Contributing

Use `FlowProposal` when you want a trained flow to provide `log q(x)` for a
bridged target, and optionally to draw proposal samples.
Bug reports, feature requests, and pull requests are welcome via the GitHub
repository. Please see the [contributing guidelines](https://github.com/sequince-dev/sequince/blob/main/CONTRIBUTING.md)
and [AI policy](https://github.com/sequince-dev/sequince/blob/main/AI_POLICY.md).

```python
import numpy as np

from sequince import TemperedPosterior
from sequince.proposals import FlowProposal
from sequince_density import ZukoFlow
## Related repositories

training_particles = np.random.default_rng(123).normal(size=(2000, 2))
The following repositories are part of the `sequince` ecosystem:

proposal = FlowProposal(
ZukoFlow(
dims=2,
flow_class="MAF",
hidden_features=[64, 64],
),
fit_kwargs={"n_epochs": 100, "lr": 1e-3, "batch_size": 500},
).fit(training_particles, rng=123)

posterior = TemperedPosterior(
log_prior=lambda x, context=None: -0.5 * np.sum(x**2, axis=1),
log_likelihood=lambda x, context=None: np.zeros(x.shape[0]),
log_proposal=proposal,
)
- [`sequince`](https://github.com/sequince-dev/sequince): the main `sequince` repository
- [`sequince-bilby`](https://github.com/sequince-dev/sequince-bilby): provides an interface for using prebuilt `sequince` samplers with `bilby`
- [`orng`](https://github.com/sequince-dev/orng): provides cross-backend random number generation
- [`acai`](https://github.com/sequince-dev/acai): provides array-API compatible control flow and iteration

samples, log_q = proposal.sample_and_log_prob(
128,
rng=123,
xp=np,
)
values = posterior.log_prob(samples, context=None)
```

`FlowProposal.sample_and_log_prob()` uses the flow backend's native
implementation when available and otherwise falls back to `sample()` plus
`log_prob()`. Both `GaussianMixtureProposal` and a flow wrapped by
`FlowProposal` satisfy Sequince's `TrainableProposal` protocol, with the same
`fit`, `sample`, `log_prob`, and `sample_and_log_prob` surface. Stochastic
fitting and sampling both accept an explicit `rng`.
## License

Licensed under the [MIT License](LICENSE).
Loading