ENH(zuko): restore zuko training controls

CHANGELOG.md

@@ -8,11 +8,21 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 
+- Optional Zuko gradient clipping, cosine learning-rate annealing, early stopping
+  and progress display, available through flow adapters' `fit_kwargs`.
+- Zuko training-data checks before and after particle transforms, reporting
+  affected dimensions for non-finite values.
 - Software archive DOI and method references for flow preconditioning and
   masked autoregressive flows.
 - Documentation sidebar links to the main Sequince documentation and
   sequince-examples under Related Packages.
 
+### Changed
+
+- Zuko's `validation_fraction=0` skips validation and keeps final parameters,
+  with an empty validation-loss history. Combining it with `patience` raises
+  `ValueError`.
+
 ## [0.1.0b1] - 2026-10-27
 
 First public beta release.

docs/normalizing_flows.rst

@@ -34,6 +34,41 @@ Estimation*. Credit the architecture you select as well as the backend
 implementation: `Zuko <https://zuko.readthedocs.io/stable/>`_ or
 `FlowJAX <https://danielward27.github.io/flowjax/>`_.
 
+Torch training controls
+-----------------------
+
+``ZukoFlow.fit`` accepts these optional controls, all disabled by default:
+
+* ``clip_grad=None``: set a finite positive total gradient-norm limit.
+* ``lr_annealing=False``: enable cosine annealing over ``n_epochs``, starting at
+  ``lr`` and stepping after each epoch.
+* ``patience=None``: set a positive number of consecutive epochs without a
+  strictly lower validation loss before stopping. Ties count; improvement
+  resets the counter.
+* ``show_progress=False``: display epochs and training/validation losses.
+
+.. code-block:: python
+
+    flow.fit(
+        particles, n_epochs=100, clip_grad=5.0, lr_annealing=True,
+        patience=10, show_progress=True, rng=42,
+    )
+
+Pass the same options through ``fit_kwargs`` in ``FlowTransform`` or
+``FlowProposal``. See :py:meth:`~sequince_density.flows.zuko.ZukoFlow.fit`
+for all training defaults.
+
+``fit`` returns the adapter with losses for every completed epoch in
+``fit_history`` and restores the best validation parameters. With
+``validation_fraction=0``, validation is skipped, ``validation_loss`` is empty,
+and final parameters are retained; setting ``patience`` raises ``ValueError``.
+For nonzero fractions, an empty split after rounding uses all particles for
+both training and validation.
+
+Raw and transformed particles must be finite with nonempty shape
+``(n_samples, dims)``. Invalid data raises ``ValueError``; non-finite errors
+identify the stage and affected dimensions (zero-based).
+
 Particle transform restrictions
 -------------------------------
 

pyproject.toml

@@ -26,6 +26,7 @@ gmm = [
 zuko = [
     "torch>=2",
     "zuko>=1.6",
+    "tqdm>=4.66",
 ]
 flowjax = [
     "jax>0.5.3",
@@ -58,6 +59,7 @@ test = [
     "scikit-learn>=1.4",
     "torch>=2",
     "zuko>=1.6",
+    "tqdm>=4.66",
 ]
 dev = [
     "pre-commit",

src/sequince_density/flows/zuko.py

@@ -8,8 +8,9 @@
 import array_api_compat.torch as torch_api
 import torch
 import zuko
+from tqdm.auto import tqdm
 
-from ..utils import torch_random_context
+from ..utils import torch_random_context, validate_training_data
 from .base import Flow
 
 logger = logging.getLogger(__name__)
@@ -73,9 +74,54 @@ def fit(
         lr: float = 1e-3,
         batch_size: int = 500,
         validation_fraction: float = 0.2,
+        clip_grad: float | None = None,
+        lr_annealing: bool = False,
+        patience: int | None = None,
+        show_progress: bool = False,
         *,
         rng=None,
     ) -> "ZukoFlow":
+        """Fit by maximum likelihood, restoring the best epoch when validating.
+
+        Parameters
+        ----------
+        samples : array_like, shape (n_samples, dims)
+            Nonempty, finite training particles. Transformed particles must
+            also be finite and have the same shape.
+        context : object, optional
+            Context passed to the particle transform.
+        n_epochs : int, default=100
+            Maximum number of training epochs.
+        lr : float, default=1e-3
+            Initial Adam learning rate.
+        batch_size : int, default=500
+            Maximum number of particles per training or validation batch.
+        validation_fraction : float, default=0.2
+            Fraction held out for validation. Zero skips validation and retains
+            the final parameters. For nonzero fractions, an empty split after
+            rounding uses all particles for both training and validation.
+        clip_grad : float or None, default=None
+            Positive maximum total gradient norm, applied before each optimizer
+            step. None disables clipping.
+        lr_annealing : bool, default=False
+            Apply cosine learning-rate annealing from lr to zero over n_epochs,
+            stepping the scheduler after each completed epoch.
+        patience : int or None, default=None
+            Positive number of consecutive epochs without a strictly lower
+            validation loss before stopping. Improvement resets the counter;
+            ties count as unsuccessful epochs. Requires validation; None
+            disables early stopping.
+        show_progress : bool, default=False
+            Show epoch progress and training and validation losses.
+        rng : RandomGenerator or int or None, optional
+            RNG for fitting, isolated from global Torch state when supplied.
+
+        Returns
+        -------
+        ZukoFlow
+            This fitted adapter. fit_history contains losses for every
+            completed epoch; validation_loss is empty when validation is off.
+        """
         with torch_random_context(rng, device=self.device):
             return self._fit(
                 samples,
@@ -84,6 +130,10 @@ def fit(
                 lr=lr,
                 batch_size=batch_size,
                 validation_fraction=validation_fraction,
+                clip_grad=clip_grad,
+                lr_annealing=lr_annealing,
+                patience=patience,
+                show_progress=show_progress,
                 rng=rng,
             )
 
@@ -95,15 +145,36 @@ def _fit(
         lr: float = 1e-3,
         batch_size: int = 500,
         validation_fraction: float = 0.2,
+        clip_grad: float | None = None,
+        lr_annealing: bool = False,
+        patience: int | None = None,
+        show_progress: bool = False,
         rng=None,
     ) -> "ZukoFlow":
+        if clip_grad is not None and (not 0 < clip_grad < float("inf")):
+            raise ValueError("clip_grad must be finite and positive or None.")
+        if patience is not None and (
+            isinstance(patience, bool)
+            or not isinstance(patience, int)
+            or patience < 1
+        ):
+            raise ValueError("patience must be a positive integer or None.")
+        if validation_fraction == 0 and patience is not None:
+            raise ValueError("patience requires validation_fraction > 0.")
         x = self._as_torch(samples)
+        validate_training_data(x, self.dims, "Raw")
         x_prime = self.fit_particle_transform(x, context=context, rng=rng)
+        validate_training_data(x_prime, self.dims, "Transformed")
+        if x_prime.shape != x.shape:
+            raise ValueError("Transformed training data must preserve shape.")
         indices = torch.randperm(x_prime.shape[0], device=x_prime.device)
         x_prime = x_prime[indices]
 
         n_val = int(validation_fraction * x_prime.shape[0])
-        if n_val <= 0 or n_val >= x_prime.shape[0]:
+        if validation_fraction == 0:
+            x_train = x_prime
+            x_val = None
+        elif n_val <= 0 or n_val >= x_prime.shape[0]:
             x_train = x_prime
             x_val = x_prime
         else:
@@ -115,47 +186,83 @@ def _fit(
             shuffle=True,
             batch_size=min(batch_size, x_train.shape[0]),
         )
-        val_dataset = torch.utils.data.DataLoader(
-            torch.utils.data.TensorDataset(x_val),
-            shuffle=False,
-            batch_size=min(batch_size, x_val.shape[0]),
+        val_dataset = (
+            torch.utils.data.DataLoader(
+                torch.utils.data.TensorDataset(x_val),
+                shuffle=False,
+                batch_size=min(batch_size, x_val.shape[0]),
+            )
+            if x_val is not None
+            else None
         )
 
         optimizer = torch.optim.Adam(self.flow.parameters(), lr=lr)
+        scheduler = (
+            torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs)
+            if lr_annealing
+            else None
+        )
+        unsuccessful_epochs = 0
         best_state = None
         best_val_loss = float("inf")
         training_loss = []
         validation_loss = []
 
-        for _ in range(n_epochs):
-            self.flow.train()
-            train_loss = 0.0
-            n_train_batches = 0
-            for (x_batch,) in dataset:
-                loss = self.loss_fn(x_batch)
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-                train_loss += loss.item()
-                n_train_batches += 1
-
-            self.flow.eval()
-            val_loss = 0.0
-            n_batches = 0
-            for (x_batch,) in val_dataset:
-                with torch.no_grad():
-                    val_loss += self.loss_fn(x_batch).item()
-                n_batches += 1
-
-            avg_val_loss = val_loss / max(n_batches, 1)
-            training_loss.append(train_loss / max(n_train_batches, 1))
-            validation_loss.append(avg_val_loss)
-            if avg_val_loss < best_val_loss:
-                best_val_loss = avg_val_loss
-                best_state = {
-                    key: value.detach().clone()
-                    for key, value in self.flow.state_dict().items()
-                }
+        with tqdm(
+            range(n_epochs), desc="Epochs", disable=not show_progress
+        ) as progress:
+            for _ in progress:
+                self.flow.train()
+                train_loss = 0.0
+                n_train_batches = 0
+                for (x_batch,) in dataset:
+                    loss = self.loss_fn(x_batch)
+                    optimizer.zero_grad()
+                    loss.backward()
+                    if clip_grad is not None:
+                        torch.nn.utils.clip_grad_norm_(
+                            self.flow.parameters(), clip_grad
+                        )
+                    optimizer.step()
+                    train_loss += loss.item()
+                    n_train_batches += 1
+
+                if scheduler is not None:
+                    scheduler.step()
+                training_loss.append(train_loss / max(n_train_batches, 1))
+                if val_dataset is None:
+                    if show_progress:
+                        progress.set_postfix(
+                            train_loss=f"{training_loss[-1]:.4f}"
+                        )
+                    continue
+
+                self.flow.eval()
+                val_loss = 0.0
+                n_batches = 0
+                for (x_batch,) in val_dataset:
+                    with torch.no_grad():
+                        val_loss += self.loss_fn(x_batch).item()
+                    n_batches += 1
+
+                avg_val_loss = val_loss / max(n_batches, 1)
+                validation_loss.append(avg_val_loss)
+                if avg_val_loss < best_val_loss:
+                    unsuccessful_epochs = 0
+                    best_val_loss = avg_val_loss
+                    best_state = {
+                        key: value.detach().clone()
+                        for key, value in self.flow.state_dict().items()
+                    }
+                else:
+                    unsuccessful_epochs += 1
+                if show_progress:
+                    progress.set_postfix(
+                        train_loss=f"{training_loss[-1]:.4f}",
+                        val_loss=f"{avg_val_loss:.4f}",
+                    )
+                if patience is not None and unsuccessful_epochs >= patience:
+                    break
 
         if best_state is not None:
             self.flow.load_state_dict(best_state)

src/sequince_density/utils.py

@@ -7,7 +7,49 @@
 from typing import Any
 
 from orng import RandomGenerator
-from sequince.array_api import to_numpy
+from sequince.array_api import namespace, to_numpy
+from sequince.typing import Array
+
+
+def validate_training_data(
+    samples: Array, dims: int, stage: str = "Raw"
+) -> None:
+    """Validate training particles using their native array backend.
+
+    Parameters
+    ----------
+    samples : array, shape (n_samples, dims)
+        Nonempty, finite particle coordinates. Validation does not modify them.
+    dims : int
+        Expected number of dimensions.
+    stage : str, default="Raw"
+        Label for the data in error messages, for example "Transformed".
+
+    Raises
+    ------
+    ValueError
+        If the shape is invalid or any coordinates are non-finite. Non-finite
+        errors report the affected dimension indices (zero-based).
+    """
+    if samples.ndim != 2 or samples.shape[0] == 0 or samples.shape[1] != dims:
+        raise ValueError(
+            f"{stage} training data must have nonempty shape "
+            f"(n_samples, {dims})."
+        )
+    xp = namespace(samples)
+    invalid = ~xp.isfinite(samples)
+    if xp.any(invalid):
+        dimensions = [
+            i
+            for i, affected in enumerate(
+                to_numpy(xp.any(invalid, axis=0)).tolist()
+            )
+            if affected
+        ]
+        raise ValueError(
+            f"{stage} training data contains non-finite values "
+            f"in dimensions: {dimensions}"
+        )
 
 
 def draw_seed(rng: RandomGenerator | int | None) -> int | None:
@@ -47,4 +89,4 @@ def torch_random_context(
         yield
 
 
-__all__ = ["draw_seed", "torch_random_context"]
+__all__ = ["draw_seed", "torch_random_context", "validate_training_data"]

tests/test_training_data.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from sequince_density.utils import validate_training_data
+
+
+@pytest.mark.parametrize("backend", ["numpy", "torch", "jax.numpy"])
+def test_validate_training_data_across_backends(backend: str) -> None:
+    xp = pytest.importorskip(backend)
+    values = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]])
+    particles = xp.asarray(values)
+    validate_training_data(particles, dims=3)
+    np.testing.assert_array_equal(np.asarray(particles), values)
+
+    invalid = values.copy()
+    invalid[0, 0] = np.nan
+    invalid[1, 2] = np.inf
+    with pytest.raises(
+        ValueError, match=r"Transformed.*non-finite.*dimensions: \[0, 2\]"
+    ):
+        validate_training_data(
+            xp.asarray(invalid), dims=3, stage="Transformed"
+        )
+
+    with pytest.raises(ValueError, match="Raw training data.*shape"):
+        validate_training_data(particles, dims=2)