Skip to content

WIP: Build on the new MeasureBase - #581

Draft
oschulz wants to merge 18 commits into
mainfrom
new-mb
Draft

oschulz wants to merge 18 commits into
mainfrom
new-mb

Conversation

@oschulz

@oschulz oschulz commented Sep 19, 2026 •

Copy link
Copy Markdown
Member

Description written by Claude Code (model Claude Fable 5.1, standard effort) on behalf of the maintainer; the commits carry the trailer Assisted-By: generative AI.

BAT 6 builds on the rebuilt MeasureBase (JuliaMath/MeasureBase.jl#182): BAT's own measure layer, distribution wrappers and transforms are gone, hierarchical priors are mbinds, walker densities are evaluated in batches, and sample storage is structured (StructArrays laid out like MeasureBase's rand, constants as Fill columns). The branch is a sequence of 18 commits over main; redesign.md (first commit, to be dropped before merging) records the design and open points.

The release notes below will go into the registrator comment of v6.0.0.


BAT.jl v6.0.0

Breaking changes

BAT's measures are now MeasureBase measures. BAT used to carry its own
measure layer (BATMeasure and friends), which duplicated what MeasureBase
now provides, and its own distribution-transform implementation. Both are
gone, BAT builds on MeasureBase 0.15 instead.

  • BAT.BATMeasure is gone, BAT's measure types subtype
    MeasureBase.AbstractMeasure directly.

  • The wrapper types BAT.BATDistMeasure, BAT.BATPwrMeasure,
    BAT.BATPushFwdMeasure, BAT.BATWeightedMeasure and
    BAT.BATSuperpositionMeasure are gone. Use MeasureBase.asmeasure,
    powermeasure (^), pushfwd, weightedmeasure and superpose.

  • BAT.StandardUvUniform, BAT.StandardMvUniform, BAT.StandardUvNormal
    and BAT.StandardMvNormal are gone. Use MeasureBase.StdUniform(),
    MeasureBase.StdNormal() and their powers, like StdNormal()^n.

  • BAT.DistributionTransform is gone. Use MeasureBase.transport_to.
    bat_transform and the transform intents (NormalBased, UniformBased,
    ToRealVector, DoNotTransform) are unchanged, but the transformation
    functions they return are MeasureBase.TransportFunctions and transformed
    measures are MeasureBase.PushforwardMeasures.

  • HierarchicalDistribution is gone. Use MeasureBase.mbind.

  • distprod now returns a MeasureBase.ProductMeasure instead of a
    ValueShapes.NamedTupleDist resp. a Distributions.Product. Marginals
    are specified as before (distributions, intervals, arrays of those, or
    constant values), constant marginals become MeasureBase.Dirac.

  • lbqintegral and distbind are deprecated in favour of
    MeasureBase.mintegrate/mintegrate_exp and MeasureBase.mbind.

  • DensitySampleVector stores the variates in the layout that rand(μ^n)
    produces for their measure μ, instead of a ShapedAsNTArray over flat
    storage: numbers in vectors, arrays as nested arrays (ArraysOfArrays),
    named tuples as StructArrays of their components and constant
    components as Fill columns. varshape(samples) is derived from this
    layout. Vectors of named tuples passed to the constructor are converted
    to it. The new constructor
    DensitySampleVector(X, μ; logd, weight, info, aux) takes a batch of
    variates of μ or a matrix of their flat free components. Constant
    components must match the point masses of μ, and all variates must
    share one layout.

  • Statistics of samples have the structure of the variates: mean, var,
    std, quantile, median, minimum, maximum and mode return named
    tuples resp. arrays instead of ShapedAsNT values, and so do effective
    sample sizes and credible intervals. minimum and maximum are
    elementwise, they used to return the lexicographically smallest resp.
    largest variate. cov and cor work on the flat free components of any
    samples.

  • unshaped.(samples) and bat_transform(Vector, samples) return samples
    with the flat free components of the variates. vcat and merge of
    samples allocate and promote the types of the weights and log-densities,
    so samples with integer and real weights can be combined (merge used to
    keep the types of its first argument). push!, append! and merge!
    require growable storage, like the flat storage of MCMC output, and throw
    an ArgumentError otherwise.

  • Broadcasting a value shape over samples, vs.(samples), and
    bat_transform with a value shape on samples are deprecated. Use
    DensitySampleVector(X, μ) to get samples in the layout of a measure.

  • BAT no longer defines mean, var, std and cov of measures.
    MeasureBase defines them for measures with exactly known moments and
    throws an ArgumentError for others. BAT keeps its median and mode
    methods.

  • MeasureBase renamed insupport and params to insupportof and
    paramsof. BAT implements insupportof for its measures and
    likelihoods, custom measures and likelihoods used with BAT must add
    methods to the new names as well.

  • batmeasure is BAT's canonicalization: it turns distributions into
    measures, named tuples of distributions and NamedTupleDists into product
    measures, and MeasureBase density measures into PosteriorMeasures.

  • Densities of measures are -Inf outside their support instead of
    throwing, and transports are NaN there; the boundary of the support can
    transport to infinite variates. Accordingly,
    BAT.checked_logdensityof no longer throws for a NaN density at a
    non-finite variate, it returns a zero density there.

  • GridSampler places its grid points at the cell midpoints of the unit
    hypercube instead of including its boundary.

  • DistributionsAD and ForwardDiffPullbacks are no longer dependencies.

Further behavior changes

  • MeasureBase.basemeasure of a posterior measure is its prior and
    MeasureBase.logdensity_def is the log-density of its likelihood, so a
    posterior behaves like the mintegrate_exp(likelihood, prior) it
    represents.

  • MeasureBase.weightedmeasure no longer special-cases posterior measures,
    it yields a WeightedMeasure for them like for any other measure. BAT's
    own renormalization (e.g. bat_sample on an auto-renormalized target)
    still rescales the likelihood and keeps the posterior structure.

  • ValueShapes.varshape and ValueShapes.unshaped are defined for
    MeasureBase measures and derived from the measure structure.

  • truncate_batmeasure supports measures with scalar variates, via
    MeasureBase's truncate. Products of univariate distributions and
    posterior measures (via their prior) are truncated as before. It still
    doesn't renormalize.

  • bat_write marks the layout of each object in HDF5 files with an
    attribute bat_layout (format version 6, see BATHDF5IO), so bat_read
    restores nested named tuples, array-valued and constant parameters
    exactly. Files of BAT 5.x are read as before, with the layout inferred
    from the values, now including matrix-valued parameters.

  • Samples without parameter names, like flat samples, can be plotted.
    They are selected by position (among the free parameters, in the order
    of unshaped.(samples)) and labeled v_1, v_2 and so on.
    BAT.bat_marginalize keeps the sample weights.

  • Sample reports work for samples without free parameters, they list the
    constant parameters only.

  • EllipticalSliceMCMCSampling requires a posterior measure whose prior
    transforms to a standard multivariate normal, and says so instead of
    failing later.

  • Algorithms that record the degrees of freedom of a measure (e.g. the
    importance samplers) record nothing for measures that declare none,
    instead of throwing.

  • The MCMC samplers evaluate the target density for all walkers of a step in
    one batched call (MeasureBase.logdensities) instead of walker by walker.
    Measures with fused batched kernels may return results that differ from the
    per-walker ones in the last ulp, so MCMC runs with a fixed RNG seed can take
    a different path than before.

  • Support for GPUs is provisional. Sample storage, its flat conversion and
    concatenation work with device arrays. Batches of variates are drawn in
    one call, rand(rng, μ^n), so on a GPU they come from a single device
    RNG instead of per-chain partitions of BAT's counter-based RNG. Reports,
    plots and HDF5 output require samples on the host.

https://claude.ai/code/session_01L5yPndh2K6u14z3CwF9YXf

Working notes on the migration of BAT to MeasureBase 0.15: what goes and
what stays, variate shapes and ValueShapes, the conventions adopted from
MeasureBase, the accelerator plan and the open questions. They accompany
the migration and will be removed before it is merged.

Assisted-By: generative AI
GelmanRubinConvergence and BrooksGelmanConvergence passed the value of
max(R²) as the name parameter of their ValueAndThreshold results instead
of the symbol `:max_Rsqr`. Both return a `ValueAndThreshold{:max_Rsqr}`
now, which is tested.

Assisted-By: generative AI
`maximum(maximum(logds))` took the maximum of the lexicographically
largest walker's log-densities, not the maximum over all walkers. Take
the maximum over the log-densities of all walkers instead.

Assisted-By: generative AI
The docstrings of `InitFromSamples` and `InitFromIID` were swapped, and
the `DensitySample` docstring named `DensitySampleVector` as its
constructor. The `MCMCIterator` docstring named functions and types that
don't exist (or no longer do), such as `getproposal` for chain states,
`MCMCIteratorInfo`, `isvalidchain` and a `DensitySampleVector`
constructor for chains.

Assisted-By: generative AI
The histograms had a hard-coded 40 bins between the minimum and maximum,
ignoring `nbins`, and the maximum fell outside the last left-closed bin.
StatsBase now chooses edges that cover all values for the given number
of bins.

Assisted-By: generative AI
`calculate_levels` left a level undefined when a single bin took the
cumulative probability past an interval, which raised an UndefRefError.
Each level is now the weight of the bin at which the cumulative
probability of the bins, in descending order, reaches the interval.

Assisted-By: generative AI
A ternary `cond ? subplot := 3 :` took the next attribute line as its
else branch, so marginal plots dropped a marker attribute and other
plots set it only by accident. Also accept a Dict for `std`, as for the
other estimators.

Assisted-By: generative AI
BAT will store batches of variates in the layout of `rand(μ^n)`: vectors
of numbers, nested arrays and struct vectors of named tuples, with
`Fill` columns for constant components. The internal batch layer
provides the variate shape of a batch (`batch_varshape`), conversion to
and from the flat standard-variate layout (`flatbatch`, `unflatbatch`),
leafwise mapping, reshaping and concatenation of batches, and
`_own_batch`, which brings `rand` and transport outputs into this layout
and checks constant components against their point masses. All of it
works with device arrays. Nothing uses it yet, its tests need
MeasureBase 0.15 and come with the structured sample storage.

Assisted-By: generative AI
MeasureBase 0.15 provides everything BAT's own measure layer used to
duplicate, so BAT's measure types subtype `MeasureBase.AbstractMeasure`
directly, BAT requires MeasureBase 0.15, HeterogeneousComputing 0.3 and
ValueShapes 0.11.11, and DistributionsAD and ForwardDiffPullbacks are no
longer used. Removed are the `BATMeasure` supertype, the wrappers
`BATDistMeasure`, `BATPwrMeasure`, `BATPushFwdMeasure`,
`BATWeightedMeasure` and `BATSuperpositionMeasure` (use `asmeasure`,
`powermeasure`, `pushfwd`, `weightedmeasure` and `superpose`), BAT's
standard uniform and normal distributions (use `StdUniform()`,
`StdNormal()` and their powers) and the distribution transforms
(replaced by `transport_to`). Posteriors are MeasureBase measures with
the prior as base measure and the likelihood as density, `batmeasure`
canonicalizes inputs to measures, and `UniformBased` and `NormalBased`
transport to `StdUniform()^n` resp. `StdNormal()^n`. Samples of
named-tuple measures are still stored as `ShapedAsNTArray`s, and the
tests of the removed types go.

Assisted-By: generative AI
`MeasureBase.mbind` is the general form of what
`HierarchicalDistribution` provided: a primary measure, a transition
kernel and a combination function, with densities, draws and transports
derived structurally. The distribution and its `UnshapedHDist` flat view
go away. `distbind` and `lbqintegral` are stable API, so they stay as
deprecated wrappers for `mbind` resp. `mintegrate`/`mintegrate_exp`.

Assisted-By: generative AI
The MCMC samplers evaluated the target density walker by walker, now
`BAT.checked_logdensities(target, X)` hands the batch of walker
positions to `MeasureBase.logdensities` and applies BAT's NaN/+Inf
policy to the results. The zero-prior short-circuit of posterior
measures moves from an override of `logdensityof`, which bypassed
MeasureBase's checks and its batched path, to MeasureBase's extension
points `logdensityof_impl` and `batched_logdensityof_impl`, so
posteriors evaluate the prior for the whole batch and the likelihood
only where the prior is nonzero. Proposals are built as
`VectorOfSimilarVectors` and evaluated from the chain state's buffer, so
that MeasureBase sees flat storage. Fused batched kernels can differ
from per-variate results in the last ulp, so MCMC runs with a fixed seed
may take a different path.

Assisted-By: generative AI
`DensitySampleVector` stores its variates in the layout of `rand(μ^n)`,
with `Fill` columns for constants, instead of `ShapedAsNTArray`s over
flat storage, and everything that produces or consumes samples follows.
The new `DensitySampleVector(X, μ)` takes a batch or a flat matrix of
variates of a measure and checks constant components against their point
masses, `varshape(samples)` is derived from the batch, and statistics,
effective sample sizes, credible intervals and sample reports have the
structure of the variates. `transform_samples` transports whole batches
via `transports_with_ladj`, with log-densities corrected by batched
Jacobians, and the importance samplers, initial values and the HDF5,
MGVI, NestedSamplers and UltraNest extensions follow the new layout.
Broadcasting a value shape over samples and `bat_transform` with a value
shape on samples are deprecated.

Assisted-By: generative AI
`reshape_variate`, `reshape_variates` and their shape-compatibility
checks have no callers anymore. The samplers produce the structured
layout directly, and initial values are reshaped to the variate shape of
the target via the flat layout.

Assisted-By: generative AI
Each object written by `bat_write` carries a `bat_layout` attribute
(samples, namedtuple, scalar, array, const or group), plus the field
order of named tuples, the length of constant columns, the element type
of sample provenance batches and the format version of sample groups.
The reader rebuilds the exact structured batch. Files without layout
attributes, written by BAT 5.x, go through the value-based heuristic,
which now also reads matrix-valued parameters.

Assisted-By: generative AI
The recipes and `bat_marginalize` resolve parameter selections against
the variate shape derived from the batch, as positions of the free
parameters in the standard-variate layout, and read values via
`flatbatch`, with one resolver for names, positions, ranges and mixed
selections. `bat_marginalize` keeps the sample weights, `MarginalDist`
stores the variate shape of the selection, marginals of a `MarginalDist`
sum bin probabilities, custom bins apply to the selected parameters and
KDE panels of MCMC diagnostics use the sample weights. The unused
plotting leftovers `Marginalization`, `chainresults` and
`partition_tree` go, and the recipes are tested via RecipesBase.

Assisted-By: generative AI
Describe parameter selection by name, path and position (labeled `v_1`,
`v_2`, ... for samples without names) and the MCMC diagnostics plots,
and drop references to the removed `bat_stats`. Remove the dev-internal
plotting examples, which were not part of the docs and used the BAT 1.x
result of `bat_sample`.

Assisted-By: generative AI
In the tutorial, statistics of samples have the structure of the
parameters, the covariance is computed on the samples directly and HDF5
files keep the structure of the variates. The internal API docs list the
batch layer and the batched density checks, since the docs build stops
on documented internals missing from the manual.

Assisted-By: generative AI
BAT 6 builds on MeasureBase 0.15 and stores samples in the layout of
their measures. See the pull request for the release notes.

Assisted-By: generative AI

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant