jflows 0.5.3__tar.gz

jflows-0.5.3/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Xuda
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jflows-0.5.3/PKG-INFO

@@ -0,0 +1,254 @@
+Metadata-Version: 2.4
+Name: jflows
+Version: 0.5.3
+Summary: JAX normalizing flows for unconditional energy-based sampling and Boltzmann generators.
+Author-email: Xuda <abneryepku@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/xuda-ye-math/jflows
+Keywords: normalizing-flows,boltzmann-generator,jax,equinox,sampling,smc,neural-spline-flows,ot-flow
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.6.0
+Requires-Dist: equinox>=0.13.0
+Requires-Dist: numpy>=2.0.0
+Provides-Extra: examples
+Requires-Dist: matplotlib>=3.8; extra == "examples"
+Dynamic: license-file
+
+<p align="center"><img src="jflows.png" alt="jflows banner" width="800px"></p>
+<p align="center"><sub><em>Banner designed by ChatGPT: the character Jax from "The Amazing Digital Circus", a nod to jflows being built on Google JAX.</em></sub></p>
+
+# jflows
+
+`jflows` is a JAX/Equinox package for unconditional normalizing flows,
+energy-based sampling, and staged Boltzmann generators. It provides a small
+public API for building flows, defining unnormalized target energies, training
+one transport stage, and connecting accepted stages into a complete source-to-
+target sampler.
+
+> **Status: experimental.** The package is developed and tested on Linux with
+> NVIDIA CUDA JAX. JAX GPU is not supported on Windows, including WSL.
+>
+> This project was developed with Codex.
+
+## Core features
+
+- **First-class spline flows.** `NSF` models bounded Euclidean boxes and
+  `NCSF` models periodic boxes or tori. `CNF`, `OTFlow`, and `RealNVP` provide
+  complementary continuous and affine-coupling architectures.
+- **One immutable flow interface.** Every flow provides `flow(x)`,
+  `flow.call_and_ladj(x)`, `flow.inv(y)`, and `flow.inv_and_ladj(y)` and works
+  as an Equinox pytree under JIT, vectorization, and differentiation.
+- **Unified energy interface.** A `Potential` maps `[N,d]` samples to `[N]`
+  energies for densities proportional to `exp(-U)`. Built-in uniform, Gaussian,
+  and Gaussian-mixture potentials support explicit-key sampling, while
+  `potential_from` wraps a plain batched energy function.
+- **Potential algebra.** Addition, subtraction, scaling, and
+  `linear_combination` construct bridge energies without introducing a second
+  target abstraction.
+- **Explicit randomness.** Every random constructor and sampler takes a JAX
+  PRNG key. There is no package-global random state.
+- **Sampling and diagnostics.** Importance weights, normalized ESS, coverage,
+  resampling, Langevin/MALA, stochastic Heun, HMC, SMC, flow-proposal AIS,
+  L-BFGS, AdamW, and quench-and-temper are available from `jflows.utils`.
+- **Runtime backend report.** `jflows.backend()` lists CPU and installed JAX
+  accelerator plugins backed by visible hardware, identifies the configured
+  runtime, and reports JAX/Equinox versions and accelerator models without
+  creating a JAX device client or performing device computation.
+- **Progressive training objectives.** Reverse KL, forward KL, KLX, and KLXX
+  share one validation convention while adding progressively richer
+  target-batch and density-ratio information.
+- **Complete-stage persistence.** Boltzmann computation is separate from
+  optional stage writers/loaders, so an interrupted stage can be recomputed
+  from the last published population.
+
+Flows are immutable. Identity initialization returns a new object:
+
+```python
+flow = flow.zeros()
+```
+
+`OTFlow` is the training exception: use `flow.near_identity()` so its
+quadratic factor retains a nonzero gradient. Keep `OTFlow.zeros()` for an exact
+identity map.
+
+## Three-level interface
+
+The public API is deliberately arranged from basic numerical components to a
+complete staged generator.
+
+### LOW — building blocks
+
+The low level contains flows, potentials, per-sample losses, metrics, MCMC,
+SMC/AIS, optimization, and quench-and-temper:
+
+```python
+from jflows.flow import NSF, NCSF, CNF, OTFlow, RealNVP
+from jflows.potential import Potential, Nlog_Gaussian, potential_from
+from jflows.loss import reverse_KL_F, forward_KL_G, forward_KLX_G, forward_X_G
+from jflows import backend
+from jflows.utils import compute_ESS_log, importance_weights_log, langevin, smc
+```
+
+`F` denotes the source-to-target map and `G = F^{-1}` the target-to-source
+map. An `_F` or `_G` suffix fixes the native training direction.
+
+### MEDIUM — one-stage training
+
+The medium level trains one map between a source and target potential:
+
+```python
+from jflows.train import (
+    Monitor,
+    train_reverse_KL_F,
+    train_forward_KL_G,
+    train_forward_KLX_G,
+    train_forward_KLXX_G,
+)
+```
+
+Every trainer returns a new flow and its optimizer-batch ESS history. Simple
+flow, sample, and history serialization lives in `jflows.artifacts`.
+
+### HIGH — staged Boltzmann generation
+
+The high level connects accepted bridge stages and compares every trained map
+with an exact identity fallback using full-validation incremental ESS:
+
+```python
+from jflows.boltzmann import (
+    boltzmann_identity,
+    boltzmann_reverse_KL_F,
+    boltzmann_forward_KL_G,
+    boltzmann_forward_KLX_G,
+    boltzmann_forward_KLXX_G,
+)
+```
+
+- identity runs the adaptive-staging controller for stage selection, ESS
+  gating, resampling, and rejuvenation with no flow or optimizer;
+- reverse KL is the simplest source-sampled F baseline;
+- forward KL manufactures target-side batches and trains G;
+- KLX adds target-measure log-weight variation control; and
+- KLXX adds a quench-and-temper/proposal mixture for wider mode and leakage
+  regularization.
+
+The four trained generators compare their candidate with identity. The
+standalone `boltzmann_identity` function instead advances every accepted
+stage by exact identity reweighting, so it requires no flow, batch size,
+training steps, or learning rate. Complete-stage persistence is exposed
+separately through `jflows.boltzmann.write` and `jflows.boltzmann.load`.
+
+## Minimal setup
+
+```bash
+pip install "jax[cuda13]" equinox
+pip install -e .
+```
+
+```python
+import jflows
+
+jflows.backend()
+```
+
+A minimal training setup is:
+
+```python
+import jax
+
+from jflows.flow import NSF
+from jflows.potential import Nlog_Gaussian, potential_from
+from jflows.train import train_forward_KLX_G
+from jflows.utils import compute_ESS_log, importance_weights_log
+
+source = Nlog_Gaussian([0.0, 0.0], [1.0, 1.0])
+target = potential_from(
+    lambda x: 10.0 * (jax.numpy.linalg.norm(x, axis=-1) - 2.0) ** 2
+)
+
+flow = NSF(
+    jax.random.key(0),
+    [-4.0, -4.0],
+    [4.0, 4.0],
+    bins=16,
+    transforms=4,
+).zeros()
+x_valid = source.samples(jax.random.key(1), 20000)
+
+flow, batch_ess = train_forward_KLX_G(
+    x_valid,
+    source,
+    target,
+    flow,
+    batch_size=500,
+    train_steps=200,
+    lr=1e-3,
+    ladder=1,
+    mc_dt=1e-3,
+    mc_steps=50,
+    coeff_lambda=1.0,
+)
+
+log_weights = importance_weights_log(
+    x_valid, source, target, flow, type="G"
+)
+validation_ess = compute_ESS_log(log_weights)
+y_valid = flow.inv(x_valid)
+```
+
+## Repository map
+
+```text
+jflows/
+├── doc/
+│   ├── README.md                 # detailed documentation map
+│   ├── 01-low-level.md
+│   ├── 02-medium-level.md
+│   ├── 03-high-level.md
+│   ├── 04-smoke-tests.md
+│   └── 05-examples.md
+├── jflows/
+│   ├── flow.py                   # LOW
+│   ├── potential.py              # LOW
+│   ├── loss.py                   # LOW
+│   ├── utils/                    # LOW
+│   ├── train.py                  # MEDIUM
+│   ├── artifacts.py              # MEDIUM
+│   ├── boltzmann/                # HIGH
+│   └── core/                     # private implementation
+├── smoke/                        # narrow executable contracts
+├── example/                      # complete workflows and results
+└── README.md                     # this concise project introduction
+```
+
+## Read next
+
+- **Detailed interface documentation:** [`doc/`](doc/README.md)
+- **Executable API contracts:** [`smoke/`](smoke/)
+- **Complete scripts, figures, and numerical results:**
+  [`example/`](example/) and [`example/results.md`](example/results.md)
+
+The detailed documentation gives exact signatures, direction conventions,
+return values, memory controls, stage records, and persistence behavior. Smoke
+tests establish narrow software contracts; examples demonstrate end-to-end
+scientific workflows.
+
+## Acknowledgements
+
+`jflows` is strongly inspired by [zuko](https://github.com/probabilists/zuko).
+The private flow, transform, and masked-MLP machinery in `jflows.core` is a
+stripped-down port of zuko's clean and composable transform design.

jflows-0.5.3/README.md

@@ -0,0 +1,224 @@
+<p align="center"><img src="jflows.png" alt="jflows banner" width="800px"></p>
+<p align="center"><sub><em>Banner designed by ChatGPT: the character Jax from "The Amazing Digital Circus", a nod to jflows being built on Google JAX.</em></sub></p>
+
+# jflows
+
+`jflows` is a JAX/Equinox package for unconditional normalizing flows,
+energy-based sampling, and staged Boltzmann generators. It provides a small
+public API for building flows, defining unnormalized target energies, training
+one transport stage, and connecting accepted stages into a complete source-to-
+target sampler.
+
+> **Status: experimental.** The package is developed and tested on Linux with
+> NVIDIA CUDA JAX. JAX GPU is not supported on Windows, including WSL.
+>
+> This project was developed with Codex.
+
+## Core features
+
+- **First-class spline flows.** `NSF` models bounded Euclidean boxes and
+  `NCSF` models periodic boxes or tori. `CNF`, `OTFlow`, and `RealNVP` provide
+  complementary continuous and affine-coupling architectures.
+- **One immutable flow interface.** Every flow provides `flow(x)`,
+  `flow.call_and_ladj(x)`, `flow.inv(y)`, and `flow.inv_and_ladj(y)` and works
+  as an Equinox pytree under JIT, vectorization, and differentiation.
+- **Unified energy interface.** A `Potential` maps `[N,d]` samples to `[N]`
+  energies for densities proportional to `exp(-U)`. Built-in uniform, Gaussian,
+  and Gaussian-mixture potentials support explicit-key sampling, while
+  `potential_from` wraps a plain batched energy function.
+- **Potential algebra.** Addition, subtraction, scaling, and
+  `linear_combination` construct bridge energies without introducing a second
+  target abstraction.
+- **Explicit randomness.** Every random constructor and sampler takes a JAX
+  PRNG key. There is no package-global random state.
+- **Sampling and diagnostics.** Importance weights, normalized ESS, coverage,
+  resampling, Langevin/MALA, stochastic Heun, HMC, SMC, flow-proposal AIS,
+  L-BFGS, AdamW, and quench-and-temper are available from `jflows.utils`.
+- **Runtime backend report.** `jflows.backend()` lists CPU and installed JAX
+  accelerator plugins backed by visible hardware, identifies the configured
+  runtime, and reports JAX/Equinox versions and accelerator models without
+  creating a JAX device client or performing device computation.
+- **Progressive training objectives.** Reverse KL, forward KL, KLX, and KLXX
+  share one validation convention while adding progressively richer
+  target-batch and density-ratio information.
+- **Complete-stage persistence.** Boltzmann computation is separate from
+  optional stage writers/loaders, so an interrupted stage can be recomputed
+  from the last published population.
+
+Flows are immutable. Identity initialization returns a new object:
+
+```python
+flow = flow.zeros()
+```
+
+`OTFlow` is the training exception: use `flow.near_identity()` so its
+quadratic factor retains a nonzero gradient. Keep `OTFlow.zeros()` for an exact
+identity map.
+
+## Three-level interface
+
+The public API is deliberately arranged from basic numerical components to a
+complete staged generator.
+
+### LOW — building blocks
+
+The low level contains flows, potentials, per-sample losses, metrics, MCMC,
+SMC/AIS, optimization, and quench-and-temper:
+
+```python
+from jflows.flow import NSF, NCSF, CNF, OTFlow, RealNVP
+from jflows.potential import Potential, Nlog_Gaussian, potential_from
+from jflows.loss import reverse_KL_F, forward_KL_G, forward_KLX_G, forward_X_G
+from jflows import backend
+from jflows.utils import compute_ESS_log, importance_weights_log, langevin, smc
+```
+
+`F` denotes the source-to-target map and `G = F^{-1}` the target-to-source
+map. An `_F` or `_G` suffix fixes the native training direction.
+
+### MEDIUM — one-stage training
+
+The medium level trains one map between a source and target potential:
+
+```python
+from jflows.train import (
+    Monitor,
+    train_reverse_KL_F,
+    train_forward_KL_G,
+    train_forward_KLX_G,
+    train_forward_KLXX_G,
+)
+```
+
+Every trainer returns a new flow and its optimizer-batch ESS history. Simple
+flow, sample, and history serialization lives in `jflows.artifacts`.
+
+### HIGH — staged Boltzmann generation
+
+The high level connects accepted bridge stages and compares every trained map
+with an exact identity fallback using full-validation incremental ESS:
+
+```python
+from jflows.boltzmann import (
+    boltzmann_identity,
+    boltzmann_reverse_KL_F,
+    boltzmann_forward_KL_G,
+    boltzmann_forward_KLX_G,
+    boltzmann_forward_KLXX_G,
+)
+```
+
+- identity runs the adaptive-staging controller for stage selection, ESS
+  gating, resampling, and rejuvenation with no flow or optimizer;
+- reverse KL is the simplest source-sampled F baseline;
+- forward KL manufactures target-side batches and trains G;
+- KLX adds target-measure log-weight variation control; and
+- KLXX adds a quench-and-temper/proposal mixture for wider mode and leakage
+  regularization.
+
+The four trained generators compare their candidate with identity. The
+standalone `boltzmann_identity` function instead advances every accepted
+stage by exact identity reweighting, so it requires no flow, batch size,
+training steps, or learning rate. Complete-stage persistence is exposed
+separately through `jflows.boltzmann.write` and `jflows.boltzmann.load`.
+
+## Minimal setup
+
+```bash
+pip install "jax[cuda13]" equinox
+pip install -e .
+```
+
+```python
+import jflows
+
+jflows.backend()
+```
+
+A minimal training setup is:
+
+```python
+import jax
+
+from jflows.flow import NSF
+from jflows.potential import Nlog_Gaussian, potential_from
+from jflows.train import train_forward_KLX_G
+from jflows.utils import compute_ESS_log, importance_weights_log
+
+source = Nlog_Gaussian([0.0, 0.0], [1.0, 1.0])
+target = potential_from(
+    lambda x: 10.0 * (jax.numpy.linalg.norm(x, axis=-1) - 2.0) ** 2
+)
+
+flow = NSF(
+    jax.random.key(0),
+    [-4.0, -4.0],
+    [4.0, 4.0],
+    bins=16,
+    transforms=4,
+).zeros()
+x_valid = source.samples(jax.random.key(1), 20000)
+
+flow, batch_ess = train_forward_KLX_G(
+    x_valid,
+    source,
+    target,
+    flow,
+    batch_size=500,
+    train_steps=200,
+    lr=1e-3,
+    ladder=1,
+    mc_dt=1e-3,
+    mc_steps=50,
+    coeff_lambda=1.0,
+)
+
+log_weights = importance_weights_log(
+    x_valid, source, target, flow, type="G"
+)
+validation_ess = compute_ESS_log(log_weights)
+y_valid = flow.inv(x_valid)
+```
+
+## Repository map
+
+```text
+jflows/
+├── doc/
+│   ├── README.md                 # detailed documentation map
+│   ├── 01-low-level.md
+│   ├── 02-medium-level.md
+│   ├── 03-high-level.md
+│   ├── 04-smoke-tests.md
+│   └── 05-examples.md
+├── jflows/
+│   ├── flow.py                   # LOW
+│   ├── potential.py              # LOW
+│   ├── loss.py                   # LOW
+│   ├── utils/                    # LOW
+│   ├── train.py                  # MEDIUM
+│   ├── artifacts.py              # MEDIUM
+│   ├── boltzmann/                # HIGH
+│   └── core/                     # private implementation
+├── smoke/                        # narrow executable contracts
+├── example/                      # complete workflows and results
+└── README.md                     # this concise project introduction
+```
+
+## Read next
+
+- **Detailed interface documentation:** [`doc/`](doc/README.md)
+- **Executable API contracts:** [`smoke/`](smoke/)
+- **Complete scripts, figures, and numerical results:**
+  [`example/`](example/) and [`example/results.md`](example/results.md)
+
+The detailed documentation gives exact signatures, direction conventions,
+return values, memory controls, stage records, and persistence behavior. Smoke
+tests establish narrow software contracts; examples demonstrate end-to-end
+scientific workflows.
+
+## Acknowledgements
+
+`jflows` is strongly inspired by [zuko](https://github.com/probabilists/zuko).
+The private flow, transform, and masked-MLP machinery in `jflows.core` is a
+stripped-down port of zuko's clean and composable transform design.

jflows-0.5.3/jflows/__init__.py

@@ -0,0 +1,93 @@
+"""jflows — JAX normalizing flows for unconditional energy-based sampling.
+
+Built on JAX and equinox.
+
+Public surface:
+
+    jflows.flow      : Flow, NSF, NCSF, CNF, OTFlow, RealNVP, ComposedTransform
+    jflows.potential : Potential, potential_from, Nlog_Uniform, Nlog_Gaussian,
+                       Nlog_Gaussian_Mixture, linear_combination (+ the operator
+                       algebra c*U, U+V, U-V, -U, U/c, sum([...]))
+    jflows.loss      : reverse_KL_F, forward_KL_G, forward_KLX_G,
+                       forward_X_G
+    jflows.train     : Monitor and all train_* stage drivers
+    jflows.boltzmann : adaptive-staging and fixed-schedule boltzmann_* generators
+    jflows.artifacts : save and load medium-level training outputs
+    jflows.backend   : selected/available JAX backends and accelerator model
+    jflows.utils     : metrics / optimization / rejuvenation / anneal / quench
+
+Internals (`jflows.core.*`) adapt zuko's clean flow and transform design.
+The public computation API uses explicit PRNG keys and takes `Flow` objects
+directly for losses, importance weights, AIS, and training.
+"""
+
+from importlib.metadata import entry_points
+from importlib.util import find_spec
+import os
+from shutil import which
+import subprocess
+
+import equinox
+import jax
+
+from . import artifacts, boltzmann, flow, loss, potential, train, utils
+from .version import __version__
+
+
+def backend():
+    """Print the available JAX backends and selected accelerator."""
+    plugins = tuple(item.name.lower() for item in entry_points(group="jax_plugins"))
+    available = ["CPU"]
+    details = ["- CPU — cpu"]
+
+    cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") not in ("", "-1")
+    if cuda_visible and any("cuda" in name for name in plugins) and which("nvidia-smi"):
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+        )
+        models = tuple(
+            dict.fromkeys(line for line in result.stdout.splitlines() if line)
+        )
+        if result.returncode == 0 and models:
+            available.append("CUDA")
+            details.append(f"- CUDA — {', '.join(models)}")
+
+    rocm_visible = os.environ.get("ROCR_VISIBLE_DEVICES") not in ("", "-1")
+    if rocm_visible and any("rocm" in name for name in plugins) and which("rocm-smi"):
+        available.append("ROCm")
+        details.append("- ROCm")
+
+    if find_spec("libtpu"):
+        available.append("TPU")
+        details.append("- TPU")
+
+    configured = os.environ.get("JAX_PLATFORMS", "").split(",")[0].lower()
+    selected = {"cpu": "CPU", "cuda": "CUDA", "rocm": "ROCm", "tpu": "TPU"}.get(
+        configured
+    )
+    if selected not in available:
+        selected = next(
+            name for name in ("CUDA", "ROCm", "TPU", "CPU") if name in available
+        )
+
+    print(f"JAX {jax.__version__}")
+    print(f"Equinox {equinox.__version__}")
+    print(f"Selected backend: {selected}")
+    print(f"Available backends: {', '.join(available)}")
+    for detail in details:
+        print(detail)
+
+
+__all__ = [
+    "__version__",
+    "artifacts",
+    "backend",
+    "boltzmann",
+    "flow",
+    "loss",
+    "potential",
+    "train",
+    "utils",
+]

jflows-0.5.3/jflows/artifacts.py

@@ -0,0 +1,46 @@
+"""Save and load medium-level training outputs."""
+
+from pathlib import Path
+
+import equinox as eqx
+import numpy as np
+
+__all__ = [
+    "load_flow",
+    "load_history",
+    "load_samples",
+    "save_flow",
+    "save_history",
+    "save_samples",
+]
+
+
+def save_flow(path, flow) -> None:
+    """Save one flow's array leaves."""
+    eqx.tree_serialise_leaves(Path(path), flow)
+
+
+def load_flow(path, template):
+    """Load flow leaves into a matching template."""
+    return eqx.tree_deserialise_leaves(Path(path), template)
+
+
+def save_samples(path, samples) -> None:
+    """Save one sample array."""
+    np.save(Path(path), np.asarray(samples), allow_pickle=False)
+
+
+def load_samples(path):
+    """Load one sample array."""
+    return np.load(Path(path), allow_pickle=False)
+
+
+def save_history(path, **history) -> None:
+    """Save named training-history arrays."""
+    np.savez(Path(path), **history)
+
+
+def load_history(path) -> dict:
+    """Load named training-history arrays."""
+    with np.load(Path(path), allow_pickle=False) as history:
+        return {name: history[name].copy() for name in history.files}