v0.5.6: lowercase factories + auto-compiled linear_combination

Highlights

Three coupled API moves that simplify how user code talks to zflows.potential, all naturally backward-incompatible at the import line.

linear_combination lowercase factory. New linear_combination(potentials, coeffs=None) -> Linear_Combination returns a ready-to-use instance — the project-wide convention is uppercase = class, lowercase = instance, and call-sites that use the combination should look like instance-creation, not class-construction. The Linear_Combination class is still exported for isinstance(u, Linear_Combination) checks and subclassing. Every user-facing example in README.md, every test script, and every writeup under tests/*.md now uses the lowercase form.

potential_from(fn) returns an instance. The old potential_from returned a Potential subclass that the caller then instantiated; the companion potential_instance_from was a one-line wrapper around potential_from(fn)(). Both collapse into a single factory: potential_from(fn) now returns the instance directly, so the canonical line is u = potential_from(myforward).to(device). potential_instance_from is removed; callers using it must rename to potential_from.

Linear_Combination auto-inherits its children's compiled .grad / .eval fast paths. The combined _grad_fn / _eval_fn are linked at __init__ time as Python closures sum_k self.coeffs[k] * U_k.grad(x) (and same for .eval(x)). They read self.coeffs[k] fresh on every call, so:

• Enabling the children once is enough (u0.enable_grad(), u1.enable_grad()) — linear_combination([u0, u1], [c0, c1]).grad(x) then works immediately, with no lc.enable_grad() call needed.
• set_coeffs([c0_new, c1_new]) is a pure coefficient update: no compiled artifact is invalidated, no Dynamo recompile is triggered. Annealed-bridge workflows (e.g. tests/4D_Boltzmann_generator.py) now pay each child's vmap(grad(forward)) and compile(forward) cost exactly once for the entire $M$-rung anneal.
• The "needs .enable_grad() first" gate moves from a pre-flight check on Linear_Combination to a runtime check at the child level (an un-enabled child raises RuntimeError the first time the closure tries to call its .grad(x)).

Smaller additions

• Device-consistency check at Linear_Combination construction. Scans every parameter() and buffer() across the child potentials; raises RuntimeError if more than one device shows up, with the offending list spelled out. Pure-function children (e.g. potential_from(fn) instances, Uniform) contribute no device and are silently allowed.
• set_coeffs(coeffs) simplified. The previous enable_grad: bool = False / enable_eval: bool = False flags are gone — under the new closure-based design they were no-ops on hot instances and confusing on cold ones. Pure coefficient update only.
• Linear_Combination(potentials, coeffs=None) defaults to uniform 1/N if coeffs is omitted. coeffs is always stored as a plain list[float]; tensor inputs are detached and .tolist()-converted at construction, so there's no buffer-registration / requires_grad subtlety.
• Docstring caveats on loss_compile / loss_compile_beta. These helpers have only been tested against NSF / NCSF. Behaviour on CNF, OTFlow, RealNVP, or any other flow class is unknown — the docstrings now flag this and recommend verifying against the un-compiled reverse_KL(x, target, flow.t()) baseline before relying on the speedup.
• Docstring caveat on nesting Linear_Combination. Using a Linear_Combination instance as a child of another is untested and not recommended; the buffer-aliasing assumption (every child's compiled .grad(x) returning a static reduce-overhead buffer that's consumed before the next child runs) has not been exercised under nesting. List every constituent potential in one flat linear_combination([u_0, ..., u_N], [c_0, ..., c_N]) call instead.

Documentation

• README's "Auto-compiled gradient feature" section reordered: the leaf-Potential .enable_grad() story is introduced first, then linear_combination is shown inheriting it. The duplicate "Precompiled gradients on Potential" sub-header is gone.
• README user-facing examples scrubbed for the U = class / u = instance convention: the user-subclass example uses class MyPotential(Potential) + u = MyPotential().to(device); the factory example uses u = potential_from(myforward).to(device) (the function is myforward, not U_fn); the linear_combination example uses u0, u1, u.
• ESS narrative for the 4D Boltzmann generator refreshed against this release's run: $\sim 0.68$ at $k=1$, jumps straight to the $0.96$–$0.97$ band by $k=2$, holds there for the remaining ten rungs.
• Compiled-loss benchmark table refreshed from the new tests/compare_compiled_loss.csv: same RTX 5070 Ti hardware, $\sim$5–13 ms raw / 1.3–1.7 ms mode='default' / 0.5–1.7 ms mode='reduce-overhead', 4–9× / 7–12× speedups respectively.
• Test writeups (4D BG, 3D periodic, 2D reverse-KL, …) updated to describe the new factory pattern and the refreshed line-number pointers.

Tests

• Stateless Potential subclasses → potential_from(forward_fn). tests/2D_reverse_KL.py's class U1, tests/_verify_potential.py's class U0 / class U1, tests/_verify_utils.py's class Quartic / class Quartic1D all converted. Stateful subclasses with constructor args — U_target in tests/4D_Boltzmann_generator.py (r0, a, q2, eps), U_target in tests/3D_periodic.py (kappa, renamed from U1), Multi_Well in tests/multi_well_compare.py — kept as classes since potential_from only wraps stateless callables.
• u0 / u1 → u_source / u_target across every relevant test (2D forward-KL, 2D reverse-KL, 2D two-moons, 2D RealNVP, 3D periodic, _verify_potential). The bridge-instance naming now matches what 4D BG has been using since v0.5.5.
• Linear_Combination(...) → linear_combination(...) at every call-site in the test suite.
• New tests/_verify_potential.py section B.5 covers the new closure-based behaviour: combined _grad_fn / _eval_fn linked at __init__, runtime gate at the child level, set_coeffs preserves every id() and the closure picks up the new coefficients with err = 0.0 against autograd / forward, enable_grad / enable_eval propagate unconditionally, hot children survive via child-level idempotency, and a linear_combination of pre-enabled children works immediately without lc.enable_grad().

Verification

Full sweep via python -m tests._all: 11 passed, 0 failed in 677.5 s on the same RTX 5070 Ti. The committed tests/_all.out records the per-script wall time (62.8 s _verify_flow, 16.1 s _verify_potential, 7.2 s _verify_utils, 254.8 s compare_compiled_loss, 257.2 s multi_well_compare, etc.) and full stdout.

Migration

```python

v0.5.5

from zflows.potential import Linear_Combination, potential_from, potential_instance_from
U = potential_from(my_fn) # returned a class
u = U().to(device) # had to instantiate yourself
u_inst = potential_instance_from(my_fn).to(device) # one-liner alternative
lc = Linear_Combination([u0, u1], [c0, c1])
lc.enable_grad() # had to enable on the combination

v0.5.6

from zflows.potential import linear_combination, potential_from
u = potential_from(my_fn).to(device) # returns instance directly
u0.enable_grad(); u1.enable_grad() # enable on the children
lc = linear_combination([u0, u1], [c0, c1]) # lc.grad(x) works immediately
```

Dependencies

Unchanged: torch>=2.7, numpy>=2.0. No new runtime imports.

🤖 Generated with Claude Code