jflows-md 0.5.4

Mixed-domain molecular potentials, flows, and samplers for jflows.

jflows_md

jflows_md is the mixed-domain molecular companion to jflows. It supplies bundle-backed molecular potentials, flows on R^p x T^q, molecular KLX/KLXX training, sampling kernels, linear regularization sharpening, and complete-stage resume.

The numerical core is intentionally a raw-performance layer. Constructors and compiled kernels assume valid shapes, dtypes, schedules, chunk counts, and physical parameters. Bundle construction and structural validation are offline concerns; persistence is separate from computation.

Reduced potential and inverse temperature

Molecular_Potential owns the inverse temperature. For internal coordinates $q$ and their canonical Cartesian representative $x(q)$, it evaluates

$$ \begin{aligned} U(q)&=\beta E(x(q))-\log J(q), \ \beta&=\frac{1}{k_{\mathrm B}T}. \end{aligned} $$

The trainers and Boltzmann controllers consume reduced potentials and do not apply another temperature factor. A regularized potential modifies the Cartesian energy first and then applies the same beta; the coordinate Jacobian is never regularized.

import jax

from jflows_md import Mixed_NSF, Molecular_Potential

target = Molecular_Potential.from_bundle("adp_ff96_obc1")
source = target.source()
samples = source.samples(jax.random.key(0), 32)

flow = Mixed_NSF(
    jax.random.key(1),
    target.domain,
    bins=32,
    transforms=6,
    hidden_features=(256, 256),
).zeros()

soft = target.regularized((50.0, 0.10))  # (e [kJ/mol], r [nm])

For rg_param=(e,r), $r$ floors only the regular and exception Amber Coulomb/Lennard-Jones pair distances. If $E_r$ is that floor-aware energy and $E_{\mathrm{ref},r}$ its value at the bundle reference, the excess $d=E_r-E_{\mathrm{ref},r}$ is mapped by

$$ R_e(d)= \begin{cases} d, & d\le e, \ e\left[1+\log(d/e)\right], & d>e. \end{cases} $$

The regularized reduced potential is

$$ U_{\mathrm{rg}}(q) =\beta\left[E_{\mathrm{ref},r}+R_e(d)\right]-\log J(q). $$

The pair $(e,r)$ is the complete regularization state; no optimizer clipping parameter participates in this definition.

Layout

jflows_md/
├── artifacts.py              generic template-based artifacts
├── boltzmann/
│   ├── __init__.py           pure adaptive-staging BG computation
│   ├── write.py              atomic complete-stage writer
│   └── load.py               load, inspect, fork, and resume
├── bundle_build/             optional OpenMM-side construction
├── core/                     BAT, Amber/OBC, domain, and spline kernels
├── flow.py                   Mixed_Identity and Mixed_NSF
├── openmm/                   native OpenMM potential and samplers
├── potential.py              physical and regularized potentials
├── source.py                 Gaussian x uniform-torus source
├── system.py                 minimal runtime-bundle loading
├── train.py                  compiled molecular KLX/KLXX trainers
└── utils/
    ├── anneal.py             SMC and AIS
    ├── quench.py             quench and temper
    └── rejuvenation.py       wrapped mixed-domain MALA

Eager Python controllers split work into chunks; compiled kernels operate on fixed array shapes. Filesystem paths, manifests, and resume policy never enter the computation functions.

The complete low-/medium-/high-level interface manual is in doc/. Narrow executable contracts are organized under smoke/.

Native OpenMM

The same bundle can instantiate an independent Cartesian OpenMM potential. It evaluates beta E(x) without JAX or the internal-coordinate Jacobian, and its (e,r) regularization matches Molecular_Potential.regularized. Native Langevin and replica-exchange runs accept either the physical or regularized potential.

from jflows_md.openmm import OpenMM_Potential, langevin, parallel_tempering

target = OpenMM_Potential.from_bundle("adp_ff96_obc1")
soft = target.regularized((50.0, 0.10))

trajectory, energy = langevin(soft, steps=10000, sample_interval=100)
replicas, energy, swap_acceptance = parallel_tempering(
    target,
    (300.0, 360.0, 432.0, 518.4),
    rounds=1000,
    steps_per_round=100,
)

The samplers use fresh OpenMM systems and contexts; platform="Reference", "CPU", "CUDA", or "OpenCL" can be selected explicitly.

Direct training

train_forward_KLX_G and train_forward_KLXX_G train the molecular G direction: the flow maps target-like samples toward the source and flow.inv generates target-like proposals. Both functions return only (trained_flow, batch_ess_history).

from jflows_md.train import train_forward_KLX_G

trained, batch_ess = train_forward_KLX_G(
    target_samples,
    source_samples,
    source,
    target,
    flow,
    batch_size=128,
    train_steps=1000,
    lr=1e-3,
)

Direct trainers default to initialize_from_identity=False. Adaptive-staging trained Boltzmann functions default to True, so each accepted-stage attempt starts from an identity parameterization unless warm-starting is explicitly selected. boltzmann_identity performs no flow training at all.

Methane Boltzmann generator example

example/methane_9d_raw/ is the repository's one minimal molecular Boltzmann generator example. It contains a complete 9D CH4 driver, its fixed 300 K GAFF2/AM1-BCC and OBC1/ACE bundle, and the saved result reports for identity, forward KL, and KLXX. The example is self-contained and loads its bundle directly from example/methane_9d_raw/bundle/.

The checked-in reports were copied from the completed KLXX experiment; they were not regenerated for this documentation update. To launch a new run from the example directory when desired, use python train.py --method id, python train.py --method kl, or python train.py --method klxx.

Linear sharpening

All three adaptive-staging entry points require two regularization states. The identity-only form is:

from jflows_md.boltzmann import boltzmann_identity

particles, stages = boltzmann_identity(
    x_valid,
    source,
    target,
    ladder=16,
    mc_dt=1e-4,
    mc_steps=4,
    rg_param_0=(20.0, 0.15),
    rg_param_1=(1000.0, 0.0),
)

It omits every flow, training, optimizer, pool, and checkpoint argument. Identity removes only the learned proposal: SMC stage selection, identity ESS over the complete validation set, resampling, pre-sharpen MALA, exact regularization sharpening, its ESS gate, and final MALA remain active.

The trained form is:

from jflows_md.boltzmann import boltzmann_forward_KLX_G

particles, stages = boltzmann_forward_KLX_G(
    x_valid,
    source,
    target,
    flow,
    pool_size=4096,
    batch_size=256,
    train_steps=1000,
    lr=1e-3,
    ladder=16,
    mc_dt=1e-4,
    mc_steps=4,
    rg_param_0=(20.0, 0.15),
    rg_param_1=(1000.0, 0.0),
)

The _rg helper computes

$$ \mathrm{rg}(t)

\mathtt{rg_param_0} +t\left(\mathtt{rg_param_1}-\mathtt{rg_param_0}\right). $$

Here $t$ is a dimensionless stage-interpolation parameter. It does not change the physical temperature_kelvin or inverse temperature $\beta$ (beta), which remain properties of the molecular target.

Let $U_{\mathrm{source}}$ denote source, and let $U_{\mathrm{rg}(t)}$ denote target.regularized(rg(t)). For an accepted stage interval $a\to b$, the code variables source_bridge and target_soft are

$$ \begin{aligned} B_a &=(1-a)U_{\mathrm{source}}+aU_{\mathrm{rg}(a)}, \ B_b^{-} &=(1-b)U_{\mathrm{source}}+bU_{\mathrm{rg}(a)}. \end{aligned} $$

After proposal selection and resampling, target_sharp is

$$ B_b^{+}=(1-b)U_{\mathrm{source}}+bU_{\mathrm{rg}(b)}, $$

and sharpen_log_weight stores

$$ \log w_{\mathrm{sharpen}}=B_b^{-}-B_b^{+}. $$

The controller then resamples and applies MALA at target_sharp. Thus the emitted samples at $t=1$ target target.regularized(rg_param_1).

The flow and sharpening transitions form one accepted stage. Both the selected-flow ESS and the sharpening ESS must reach bg_param["tau_ess"] (default 0.6). If either gate fails, the controller applies the configured shrink_factor to $b-a$ and retries the complete stage. A stage is emitted and made available to persistence only after both gates pass.

Each stage record deliberately distinguishes:

• flow_rg and flow_endpoint="pre_sharpen": the law used to train the proposal flow;
• population_rg: the law of the emitted post-sharpen particles;
• sharpen_ess_hist: attempt-aligned sharpening ESS (NaN when the flow gate rejected before sharpening was evaluated);
• sharpen_ess and sharpen_mala_acceptance: accepted sharpening diagnostics.

The per-stage flows are proposal maps for this particle algorithm. They do not compose into a deterministic source-to-final generator because sharpening is a stochastic transition between stages.

Complete-stage resume

Persistence mirrors jflows 0.5 and remains outside the compute API. jflows_md.boltzmann.write atomically writes the post-sharpen population, histories, and stage metadata before publishing the stage in run.json. Trained runs additionally write both flow artifacts; identity runs write no .eqx files. jflows_md.boltzmann.load.run resumes from the last published stage; an incomplete unpublished directory is ignored and recomputed.

from jflows_md.boltzmann import iterate_boltzmann
from jflows_md.boltzmann.load import run

def iterate(samples, continuation, accepted_t, start_stage):
    return iterate_boltzmann(
        samples,
        source,
        target,
        continuation,
        objective="forward_klx",
        accepted_t=accepted_t,
        start_stage=start_stage,
        # supply the same numerical controls used for a direct run
        **controls,
    )

config = {
    "target": "adp_ff96_obc1",
    "seed": 0,
    "rg_param_0": (20.0, 0.15),
    "rg_param_1": (1000.0, 0.0),
}

particles, stages = run(
    "runs/adp",
    "adp-regularization",
    config,
    x_valid,
    flow,
    iterate,
)

# After interruption, `flow` is the same architecture template.
particles, stages = run(
    "runs/adp",
    "adp-regularization",
    config,
    None,
    flow,
    iterate,
    resume=True,
)

Resume requires the same problem identifier, complete numerical configuration, and exact static flow template; mismatches are rejected. The saved population, not flow-only replay, is the continuation state. Use manifest, validate, load, load_stage_flow, load_validation_samples, load_training_history, and fork for inspection and run management.

Bundles and installation

Runtime bundles contain exactly:

coordinates.json
manifest.json
reference.pdb
system.json
system.xml
validation.json

A source checkout includes these seven audited named bundles:

CLI preset	Bundle	Model	Domain	Fixed centers
adp	adp_ff96_obc1	ff96; OBC1/ACE	R^42 x T^18	1 (L-Ala)
glycerol	glycerol_gaff2_am1bcc_obc1	GAFF2/AM1-BCC; OBC1/ACE	R^25 x T^11	0
diethanolamine	diethanolamine_gaff2_am1bcc_obc1	GAFF2/AM1-BCC; OBC1/ACE	R^33 x T^15	0
nma	nma_ff96_obc1	ff96; OBC1/ACE	R^21 x T^9	0
s_2_butanol	s_2_butanol_gaff2_am1bcc_obc1	GAFF2/AM1-BCC; OBC1/ACE	R^28 x T^11	1 (S)
rr_2_3_butanediol	rr_2_3_butanediol_gaff2_am1bcc_obc1	GAFF2/AM1-BCC; OBC1/ACE	R^31 x T^11	2 (R,R)
cyclohexane	cyclohexane_gaff2_am1bcc_obc1	GAFF2/AM1-BCC; OBC1/ACE	R^33 x T^15	0

NMA keeps its amide cis/trans coordinate periodic and unrestricted; the force field assigns its energetic preference. Cyclohexane likewise has no stereochemical support restriction, so ring conformations and chair inversion are not split into artificial components. The two alcohol presets restrict the chart to their named absolute configurations. Candidate manifests pin AmberTools 26.0.0 and their structure/parameter provenance. The two alcohol parmchk2 files add no terms; cyclohexane records the zero-penalty transfer of its c6 terms from the GAFF2 c3 types.

Wheels contain Python code only; bundle data must be downloaded or prepared separately. An installed user can search an external directory and select a bundle by name:

from jflows_md import Molecular_Bundle, available_bundles

bundle_root = "downloaded-bundles"
names = available_bundles(bundle_root)
bundle = Molecular_Bundle.load(names[0], root=bundle_root)

The same root= keyword is accepted by Molecular_Potential.from_bundle and OpenMM_Potential.from_bundle. available_bundles() returns an empty tuple when the requested directory is absent or contains no bundles.

Install matching JAX and OpenMM accelerator wheels separately. Choose one CUDA toolkit generation:

pip install "jax[cuda13]" "openmm[cuda13]"
# or: pip install "jax[cuda12]" "openmm[cuda12]"

Installing a JAX CUDA extra does not install the OpenMM CUDA plugin, and the OpenMM extra does not install the JAX plugin. These extras select the independently packaged JAX and OpenMM CUDA wheels; they must use the same CUDA generation. Then install the local jflows checkout and this project:

python -m pip install -e "${JFLOWS_ROOT:?set JFLOWS_ROOT}"
python -m pip install -e .

Inspect the installed runtimes without creating a JAX client or OpenMM context:

import jflows_md

jflows_md.backend()

Base OpenMM-side construction is optional and CPU-capable; select a CUDA extra above when native GPU execution is required:

python -m pip install -e ".[openmm]"
python -m pip install -e ".[bundles]"
python -m jflows_md.bundle_build \
  glycerol molecule.prmtop molecule.rst7 generated/glycerol

From a source checkout, the equivalent wrapper is python -m bundles.build_molecular_bundles ....

The command-line presets above carry explicit coordinate configurations. Programmatic jflows_md.bundle_build.write_bundle(...) calls may provide zmatrix, fixed_stereocenters, and signed_volume_diagnostics for any additional explicitly prepared molecule, but that low-level route is not a claim of curated target support. Multiple fixed tetrahedral centers use multiple half-chart torsions; no center is inferred from the molecule name. For named chiral presets, caller-supplied CIP-priority atom orders make bundle construction reject a reference with the wrong requested R/S configuration. Legacy schema-v2 achiral and single-center bundles remain runtime compatible. For source compatibility, a target-only builder call with all new coordinate options at their defaults still reproduces the historical schema-v2 defaults for adp, glycerol, and diethanolamine; any non-default coordinate option uses the new schema-v3 route. See doc/01-low-level.md for the configuration contract.

Verification

Run the complete local suite from the repository root only when the accelerator is available:

XLA_PYTHON_CLIENT_PREALLOCATE=false \
python smoke/run_all.py

smoke/benchmark_compile.py is an opt-in compilation benchmark and is not part of run_all.py.