pdft 0.2.3__tar.gz → 0.2.4__tar.gz

pdft-0.2.4/PKG-INFO

@@ -0,0 +1,202 @@
+Metadata-Version: 2.5
+Name: pdft
+Version: 0.2.4
+Summary: Python port of ParametricDFT.jl: learning parametric quantum Fourier transforms via manifold optimization.
+Project-URL: Homepage, https://github.com/zazabap/pdft
+Project-URL: Issues, https://github.com/zazabap/pdft/issues
+Project-URL: Upstream, https://github.com/nzy1997/ParametricDFT.jl
+Author: zazabap
+License: MIT License
+        
+        Copyright (c) 2026 zazabap
+        
+        Portions of this project are a Python port of ParametricDFT.jl
+        (https://github.com/nzy1997/ParametricDFT.jl), Copyright (c) 2025 nzy1997,
+        originally released under the MIT License.
+        
+        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.
+License-File: LICENSE
+Keywords: dft,fourier,jax,manifold-optimization,quantum,variational
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.10.0
+Requires-Dist: typing-extensions>=4.5
+Provides-Extra: dev
+Requires-Dist: jaxtyping>=0.2; extra == 'dev'
+Requires-Dist: matplotlib>=3.8; extra == 'dev'
+Requires-Dist: numpy>=1.26; extra == 'dev'
+Requires-Dist: pytest-cov>=4; extra == 'dev'
+Requires-Dist: pytest>=7; extra == 'dev'
+Requires-Dist: ruff<0.17,>=0.6; extra == 'dev'
+Provides-Extra: gpu
+Requires-Dist: jax[cuda12]>=0.10.0; extra == 'gpu'
+Provides-Extra: plot
+Requires-Dist: matplotlib>=3.8; extra == 'plot'
+Description-Content-Type: text/markdown
+
+# pdft
+
+[![arXiv](https://img.shields.io/badge/arXiv-2608.00053-b31b1b.svg)](https://arxiv.org/abs/2608.00053)
+
+A Python port of [ParametricDFT.jl](https://github.com/nzy1997/ParametricDFT.jl):
+learning parametric quantum Fourier transforms via manifold optimization. The
+package implements a variational approach that approximates the Discrete
+Fourier Transform (DFT) with parameterized quantum circuits.
+
+This is the reference implementation accompanying the paper
+[*Fast Trainable Multilinear Bases for Image Compression*](https://arxiv.org/abs/2608.00053)
+(An, Ni, Zhou, Liu, 2026).
+
+> Status: feature-complete port. All bases (QFT, entangled QFT, TEBD, MERA,
+> Rich/RealRich, DCT-IV, blocked), both Riemannian optimizers (GD + Adam),
+> training, JSON/compression I/O, and visualization are implemented, with
+> parity against the Julia reference verified by committed goldens.
+
+## Installation
+
+From PyPI (Python 3.11+):
+
+```bash
+pip install "pdft>=0.2.3"
+```
+
+> **Note:** the older `pdft==0.2.2` wheel predates `DCT4Basis` and the
+> `parametrization="u4"` option of `TEBDBasis` / `MERABasis`, so it cannot run
+> the paper's DCT-IV, TEBD-U4, or MERA-U4 configurations. If
+> `pdft.__version__` reports `0.2.2`, upgrade with `pip install -U pdft`.
+
+From source:
+
+```bash
+git clone https://github.com/zazabap/pdft.git
+cd pdft
+pip install -e ".[dev]"
+```
+
+## Quick start
+
+Train a parametric QFT basis on a target image with Riemannian gradient
+descent:
+
+```python
+import jax
+import jax.numpy as jnp
+import pdft
+
+target = jax.random.normal(jax.random.PRNGKey(7), (4, 4)).astype(jnp.complex128)
+basis = pdft.QFTBasis(m=2, n=2)
+
+result = pdft.train_basis(
+    basis,
+    target=target,
+    loss=pdft.L1Norm(),
+    optimizer=pdft.RiemannianGD(lr=0.01),
+    steps=50,
+    seed=0,
+)
+print(result.loss_history[0], "->", result.loss_history[-1])
+```
+
+Runnable demos live in [`examples/`](examples/) (each finishes in under
+10 seconds):
+
+```bash
+python examples/basis_demo.py           # train a QFTBasis, plot the loss
+python examples/optimizer_benchmark.py  # GD vs Adam comparison
+python examples/mera_demo.py            # MERA basis training
+```
+
+## Coherence with the pixel basis
+
+Compression cares only how few coefficients a basis needs. Anything that
+recovers an image from a *subset of its pixels* — inpainting, completion,
+compressed sensing — is governed by a second quantity, and it is not sparsity:
+
+    mu(U) = N max_ij |U_ij|^2  in [1, N]
+
+`mu = 1` is maximal incoherence with the pixel basis, the best case for
+recovery from pointwise samples; `mu = N` is an atom living on one pixel,
+invisible to any sample set that misses it.
+
+Every basis here starts at `mu = 1`, and there is a structural reason it can
+stay there: if the only non-diagonal gates are one Hadamard per wire, then
+`|U_ij| = N^{-1/2}` for **every** parameter value, so `mu = 1` identically and
+`sqrt(N) U` is a complex Hadamard matrix. Training the controlled-phase gates
+arbitrarily hard, on any objective, cannot move it. Training the Hadamard /
+`U(4)` gates can and does — randomising them on a `(3, 3)` `QFTBasis` reaches
+`mu = 24.8` out of 64, and on `RichBasis`, which has no diagonal gates at all,
+`mu = 32.5`.
+
+`certify_flat_modulus` answers that before a run rather than measuring it
+after, using the same `frozen_indices` that `train_basis_batched` takes:
+
+```python
+from pdft.coherence import certify_flat_modulus, coherence, diagonal_tensor_indices
+
+basis = pdft.QFTBasis(m=3, n=3)
+coherence(basis)                      # 1.0
+
+cert = certify_flat_modulus(basis)    # nothing frozen
+print(cert.holds, cert.reason)        # False: the Hadamards are trainable
+
+frozen = cert.offending_indices       # exactly what must be held fixed
+assert certify_flat_modulus(basis, frozen_indices=frozen)
+result = pdft.train_basis_batched(basis, frozen_indices=frozen, ...)
+```
+
+Freezing gates is a real trade — it removes the freedom that `RichBasis` and
+`TEBDBasis` add for compression. The point is that the trade is now visible and
+checkable, so it can be made deliberately per task.
+
+## Background
+
+For the theory, see the paper:
+- [Fast Trainable Multilinear Bases for Image Compression](https://arxiv.org/abs/2608.00053) (arXiv:2608.00053)
+
+and the upstream notes:
+- [`note/stepbystep.pdf`](https://github.com/nzy1997/ParametricDFT.jl/blob/main/note/stepbystep.pdf)
+- [`note/main.pdf`](https://github.com/nzy1997/ParametricDFT.jl/blob/main/note/main.pdf)
+
+## Citation
+
+If you use this package in your research, please cite:
+
+```bibtex
+@misc{an2026fast,
+  title         = {Fast Trainable Multilinear Bases for Image Compression},
+  author        = {An, Shiwen and Ni, Zhongyi and Zhou, Huanhai and Liu, Jin-Guo},
+  year          = {2026},
+  eprint        = {2608.00053},
+  archivePrefix = {arXiv},
+  primaryClass  = {eess.IV},
+  url           = {https://arxiv.org/abs/2608.00053},
+}
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE). This project is a derivative port of
+ParametricDFT.jl (Copyright © 2025 nzy1997, MIT).

pdft-0.2.4/README.md

@@ -0,0 +1,143 @@
+# pdft
+
+[![arXiv](https://img.shields.io/badge/arXiv-2608.00053-b31b1b.svg)](https://arxiv.org/abs/2608.00053)
+
+A Python port of [ParametricDFT.jl](https://github.com/nzy1997/ParametricDFT.jl):
+learning parametric quantum Fourier transforms via manifold optimization. The
+package implements a variational approach that approximates the Discrete
+Fourier Transform (DFT) with parameterized quantum circuits.
+
+This is the reference implementation accompanying the paper
+[*Fast Trainable Multilinear Bases for Image Compression*](https://arxiv.org/abs/2608.00053)
+(An, Ni, Zhou, Liu, 2026).
+
+> Status: feature-complete port. All bases (QFT, entangled QFT, TEBD, MERA,
+> Rich/RealRich, DCT-IV, blocked), both Riemannian optimizers (GD + Adam),
+> training, JSON/compression I/O, and visualization are implemented, with
+> parity against the Julia reference verified by committed goldens.
+
+## Installation
+
+From PyPI (Python 3.11+):
+
+```bash
+pip install "pdft>=0.2.3"
+```
+
+> **Note:** the older `pdft==0.2.2` wheel predates `DCT4Basis` and the
+> `parametrization="u4"` option of `TEBDBasis` / `MERABasis`, so it cannot run
+> the paper's DCT-IV, TEBD-U4, or MERA-U4 configurations. If
+> `pdft.__version__` reports `0.2.2`, upgrade with `pip install -U pdft`.
+
+From source:
+
+```bash
+git clone https://github.com/zazabap/pdft.git
+cd pdft
+pip install -e ".[dev]"
+```
+
+## Quick start
+
+Train a parametric QFT basis on a target image with Riemannian gradient
+descent:
+
+```python
+import jax
+import jax.numpy as jnp
+import pdft
+
+target = jax.random.normal(jax.random.PRNGKey(7), (4, 4)).astype(jnp.complex128)
+basis = pdft.QFTBasis(m=2, n=2)
+
+result = pdft.train_basis(
+    basis,
+    target=target,
+    loss=pdft.L1Norm(),
+    optimizer=pdft.RiemannianGD(lr=0.01),
+    steps=50,
+    seed=0,
+)
+print(result.loss_history[0], "->", result.loss_history[-1])
+```
+
+Runnable demos live in [`examples/`](examples/) (each finishes in under
+10 seconds):
+
+```bash
+python examples/basis_demo.py           # train a QFTBasis, plot the loss
+python examples/optimizer_benchmark.py  # GD vs Adam comparison
+python examples/mera_demo.py            # MERA basis training
+```
+
+## Coherence with the pixel basis
+
+Compression cares only how few coefficients a basis needs. Anything that
+recovers an image from a *subset of its pixels* — inpainting, completion,
+compressed sensing — is governed by a second quantity, and it is not sparsity:
+
+    mu(U) = N max_ij |U_ij|^2  in [1, N]
+
+`mu = 1` is maximal incoherence with the pixel basis, the best case for
+recovery from pointwise samples; `mu = N` is an atom living on one pixel,
+invisible to any sample set that misses it.
+
+Every basis here starts at `mu = 1`, and there is a structural reason it can
+stay there: if the only non-diagonal gates are one Hadamard per wire, then
+`|U_ij| = N^{-1/2}` for **every** parameter value, so `mu = 1` identically and
+`sqrt(N) U` is a complex Hadamard matrix. Training the controlled-phase gates
+arbitrarily hard, on any objective, cannot move it. Training the Hadamard /
+`U(4)` gates can and does — randomising them on a `(3, 3)` `QFTBasis` reaches
+`mu = 24.8` out of 64, and on `RichBasis`, which has no diagonal gates at all,
+`mu = 32.5`.
+
+`certify_flat_modulus` answers that before a run rather than measuring it
+after, using the same `frozen_indices` that `train_basis_batched` takes:
+
+```python
+from pdft.coherence import certify_flat_modulus, coherence, diagonal_tensor_indices
+
+basis = pdft.QFTBasis(m=3, n=3)
+coherence(basis)                      # 1.0
+
+cert = certify_flat_modulus(basis)    # nothing frozen
+print(cert.holds, cert.reason)        # False: the Hadamards are trainable
+
+frozen = cert.offending_indices       # exactly what must be held fixed
+assert certify_flat_modulus(basis, frozen_indices=frozen)
+result = pdft.train_basis_batched(basis, frozen_indices=frozen, ...)
+```
+
+Freezing gates is a real trade — it removes the freedom that `RichBasis` and
+`TEBDBasis` add for compression. The point is that the trade is now visible and
+checkable, so it can be made deliberately per task.
+
+## Background
+
+For the theory, see the paper:
+- [Fast Trainable Multilinear Bases for Image Compression](https://arxiv.org/abs/2608.00053) (arXiv:2608.00053)
+
+and the upstream notes:
+- [`note/stepbystep.pdf`](https://github.com/nzy1997/ParametricDFT.jl/blob/main/note/stepbystep.pdf)
+- [`note/main.pdf`](https://github.com/nzy1997/ParametricDFT.jl/blob/main/note/main.pdf)
+
+## Citation
+
+If you use this package in your research, please cite:
+
+```bibtex
+@misc{an2026fast,
+  title         = {Fast Trainable Multilinear Bases for Image Compression},
+  author        = {An, Shiwen and Ni, Zhongyi and Zhou, Huanhai and Liu, Jin-Guo},
+  year          = {2026},
+  eprint        = {2608.00053},
+  archivePrefix = {arXiv},
+  primaryClass  = {eess.IV},
+  url           = {https://arxiv.org/abs/2608.00053},
+}
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE). This project is a derivative port of
+ParametricDFT.jl (Copyright © 2025 nzy1997, MIT).

{pdft-0.2.3 → pdft-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pdft"
-version = "0.2.3"
+version = "0.2.4"
 description = "Python port of ParametricDFT.jl: learning parametric quantum Fourier transforms via manifold optimization."
 readme = "README.md"
 requires-python = ">=3.11"
@@ -30,7 +30,7 @@ dependencies = [
 dev = [
     "pytest>=7",
     "pytest-cov>=4",
-    "ruff>=0.5",
+    "ruff>=0.6,<0.17",
     "numpy>=1.26",
     "jaxtyping>=0.2",
     "matplotlib>=3.8",
@@ -50,6 +50,29 @@ packages = ["src/pdft"]
 line-length = 100
 target-version = "py311"
 
+[tool.ruff.lint]
+# Declared deliberately, because ruff's *default* rule set widens between
+# releases. With no `select` here and an unpinned `ruff>=0.5` in [dev], lint
+# was a property of the release calendar rather than of the repository: CI
+# went red on `main` with no source change when a newer ruff turned on rules
+# 0.5 did not have (RUF100, I001, C408, B008, ...).
+#
+# This set is pyflakes + import order + pyupgrade + comprehensions: the rules
+# that catch real defects, all of which the codebase already satisfies.
+select = ["E4", "E7", "E9", "F", "I", "UP", "C4"]
+
+# The package configures jax x64 mode before importing anything that reads it,
+# which is a module-level import after a statement by construction. This is
+# what the scattered `# noqa: E402` comments were for; ignoring the rule once
+# here is the same intent without directives that later ruff versions then
+# flag as unused.
+ignore = ["E402"]
+
+# Not enabled, but each is a reasonable thing to opt into later:
+#   "B"   -> 24 findings, almost all `zip()` without `strict=`
+#   "RUF" -> 40 findings, mostly ambiguous-unicode in maths docstrings where
+#            the multiplication sign is intended, so it needs per-site review
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = "-ra --strict-markers"

{pdft-0.2.3 → pdft-0.2.4}/src/pdft/__init__.py

@@ -17,6 +17,7 @@ Public API: import from subpackages directly, e.g.
     from pdft.optimizers import RiemannianGD, RiemannianAdam
     from pdft.loss import L1Norm, MSELoss
     from pdft.io import save_basis, load_basis, compress
+    from pdft.coherence import coherence, certify_flat_modulus
 
 The names re-exported at the package root below are kept only for the
 small set most commonly used in interactive sessions and notebooks.
@@ -41,12 +42,12 @@ if not _os.environ.get("PDFT_DISABLE_COMPILE_CACHE"):
     except OSError:
         pass
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 __upstream_ref__ = "nzy1997/ParametricDFT.jl@a201a27e47df2f0f3ab460f83d49b6e5f5d1e9ef"
 
 # Slim public re-export hub: most-used names only. Anything else: import
 # from the relevant subpackage directly.
-from .bases import (  # noqa: E402
+from .bases import (
     AbstractSparseBasis,
     BlockedBasis,
     DCT4Basis,
@@ -59,9 +60,10 @@ from .bases import (  # noqa: E402
     bases_allclose,
     freeze_as_blocked,
 )
-from .loss import AbstractLoss, L1Norm, MSELoss, loss_function  # noqa: E402
-from .optimizers import RiemannianAdam, RiemannianGD, optimize  # noqa: E402
-from .training import TrainingResult, train_basis, train_basis_batched  # noqa: E402
+from .coherence import certify_flat_modulus, coherence
+from .loss import AbstractLoss, L1Norm, MSELoss, loss_function
+from .optimizers import RiemannianAdam, RiemannianGD, optimize
+from .training import TrainingResult, train_basis, train_basis_batched
 
 __all__ = [
     "AbstractLoss",
@@ -82,6 +84,8 @@ __all__ = [
     "__upstream_ref__",
     "__version__",
     "bases_allclose",
+    "certify_flat_modulus",
+    "coherence",
     "freeze_as_blocked",
     "loss_function",
     "optimize",

{pdft-0.2.3 → pdft-0.2.4}/src/pdft/bases/block/block.py

@@ -73,7 +73,7 @@ def _make_block_code(
     # axes (rows then cols) trailing. The trailing layout matches the inner
     # einsum's expected (2,)^(m_inner+n_inner) shape.
     perm = (
-        list(range(0, block_log_m))  # block row axes
+        list(range(block_log_m))  # block row axes
         + list(range(m_outer, m_outer + block_log_n))  # block col axes
         + list(range(block_log_m, m_outer))  # within row axes
         + list(range(m_outer + block_log_n, m_outer + n_outer))  # within col axes

pdft-0.2.4/src/pdft/coherence.py

@@ -0,0 +1,203 @@
+"""Mutual coherence of a learned basis with the pixel basis.
+
+For compression, the only thing that matters about a basis is how few
+coefficients it needs. For any task that recovers an image from a *subset of
+its pixels* --- inpainting, completion, compressed sensing --- a second
+quantity governs the outcome, and it is not sparsity:
+
+    mu(U) = N * max_ij |U_ij|^2   in [1, N]
+
+with N the operator dimension. mu = 1 is maximal incoherence with the pixel
+basis, the most favourable case for recovery from pointwise samples; mu = N is
+the worst, an atom supported on a single pixel, invisible to any sample set
+that misses it. Sample-complexity bounds for recovery from random pixel
+observations scale linearly in mu (Candes & Plan 2011, Krahmer & Ward 2014).
+
+A trained basis can drift in mu while its compression loss improves, so the
+quantity is worth being able to measure. This module measures it --- and, more
+usefully, certifies when it cannot drift at all.
+
+The guarantee
+-------------
+Let the circuit act on n wires, and suppose every gate is diagonal in the
+computational basis except for exactly one Hadamard per wire. Then
+
+    |U_ij| = N^{-1/2}   for every i, j and every parameter value,
+
+so mu(U) = 1 identically over the whole parameter space, and sqrt(N) U is a
+complex Hadamard matrix (Tadej & Zyczkowski 2006).
+
+*Proof.* Track the amplitude vector produced by an input basis vector.
+Diagonal gates multiply amplitudes by unit moduli and never move amplitude
+between the two values of any wire; a permutation permutes. Consider the moment
+the Hadamard on wire q is applied. Every earlier gate is diagonal or a Hadamard
+on a different wire, so none has moved amplitude across wire q, and the vector
+is supported on one value of it: one branch carries a, the other 0. Since
+H (a, 0)^T = (a, a)^T / sqrt(2), the gate replaces each modulus by two equal
+moduli |a|/sqrt(2) and never combines unequal ones. Each wire receives exactly
+one Hadamard, so every amplitude ends at 2^{-n/2}, and the remaining diagonal
+gates and permutations preserve moduli. QED
+
+The practical consequence is a classification of this package's bases by which
+gates are left free during training:
+
+  * free only the controlled-phase gates --- the compact-CP tensors --- and
+    mu is pinned at 1 no matter how hard the basis is trained, on any
+    objective, with no penalty, constraint or monitoring;
+  * free the Hadamard / U(4) gates as well and the guarantee is void. It is
+    void in practice, not only in principle: randomising them on an (m, n) =
+    (3, 3) QFTBasis reaches mu = 24.8 of a possible 64, and on RichBasis ---
+    which has no diagonal gates at all --- mu = 32.5.
+
+`certify_flat_modulus` performs that check for a given training configuration,
+taking the same `frozen_indices` that `train_basis_batched` accepts, so the
+question "does this training run preserve incoherence?" can be answered before
+the run rather than measured after it.
+
+Discovered while applying this package's basis family to image *completion*,
+where the reversal matters: the transform that compresses an image best is not
+the one that completes it best.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import jax
+import jax.numpy as jnp
+
+from .circuit.builder import is_compact_cp
+
+Array = jax.Array
+
+__all__ = [
+    "FlatModulusCertificate",
+    "certify_flat_modulus",
+    "coherence",
+    "dense_operator",
+    "diagonal_tensor_indices",
+    "is_flat_modulus",
+]
+
+
+def dense_operator(basis) -> Array:
+    """The basis as an explicit matrix acting on `vec(image)`.
+
+    Column j is the transform of the image that is 1 at pixel j and 0
+    elsewhere, so the result is (2^m 2^n) x (2^m 2^n).
+
+    This is a diagnostic, not a code path: it costs N^2 transforms of an
+    N-pixel image and is meant for small m, n. Nothing in training needs it ---
+    `certify_flat_modulus` gives the parameter-space guarantee without forming
+    it.
+    """
+    rows, cols = basis.image_size
+    dim = rows * cols
+    eye = jnp.eye(dim, dtype=jnp.complex128)
+    return jax.vmap(lambda e: basis.forward_transform(e.reshape(rows, cols)).reshape(-1))(eye).T
+
+
+def coherence(basis, operator: Array | None = None) -> float:
+    """mu(U) = N max_ij |U_ij|^2, in [1, N]. 1 is maximal incoherence.
+
+    Pass `operator` to reuse a matrix from `dense_operator`.
+    """
+    u = dense_operator(basis) if operator is None else operator
+    return float(u.shape[0] * jnp.max(jnp.abs(u) ** 2))
+
+
+def is_flat_modulus(basis, operator: Array | None = None, atol: float = 1e-8) -> bool:
+    """True if |U_ij| = N^{-1/2} everywhere, i.e. sqrt(N) U is complex Hadamard."""
+    u = dense_operator(basis) if operator is None else operator
+    return bool(jnp.allclose(jnp.abs(u), u.shape[0] ** -0.5, atol=atol))
+
+
+def diagonal_tensor_indices(basis) -> list[int]:
+    """Indices of the tensors that are diagonal in the computational basis.
+
+    These are the compact controlled-phase gates, `[[1, 1], [1, e^{i phi}]]`,
+    whose only freedom is a phase. Freeing exactly these is what preserves
+    mu == 1; see the module docstring.
+    """
+    return [i for i, t in enumerate(basis.tensors) if is_compact_cp(t)]
+
+
+@dataclass(frozen=True)
+class FlatModulusCertificate:
+    """Whether mu == 1 holds over the whole reachable parameter set.
+
+    `holds` is the answer; `reason` says why in one line. `mu` is the coherence
+    at the basis's current parameters, and `offending_indices` names the
+    trainable tensors that are not diagonal, which are exactly the ones that
+    would have to be frozen for the guarantee to apply.
+    """
+
+    holds: bool
+    reason: str
+    mu: float
+    offending_indices: list[int] = field(default_factory=list)
+
+    def __bool__(self) -> bool:
+        return self.holds
+
+
+def certify_flat_modulus(
+    basis,
+    frozen_indices: list[int] | None = None,
+    *,
+    atol: float = 1e-8,
+) -> FlatModulusCertificate:
+    """Certify that training cannot raise mu above 1.
+
+    Two conditions, both necessary. The basis must be flat-modulus *now*, and
+    every tensor left trainable must be diagonal --- then by the proposition in
+    the module docstring, mu == 1 for every value those parameters can take.
+
+    `frozen_indices` is the same argument `train_basis_batched` accepts, so:
+
+        cert = certify_flat_modulus(basis, frozen_indices=frozen)
+        if not cert:
+            raise ValueError(cert.reason)
+        result = train_basis_batched(basis, frozen_indices=frozen, ...)
+
+    Frozen non-diagonal gates are fine: a fixed Hadamard is what the
+    proposition assumes. It is *training* them that voids it.
+    """
+    operator = dense_operator(basis)
+    mu = coherence(basis, operator)
+
+    if not is_flat_modulus(basis, operator, atol=atol):
+        return FlatModulusCertificate(
+            holds=False,
+            reason=(
+                f"basis is not flat-modulus at its current parameters "
+                f"(mu = {mu:.6g}); the guarantee applies to circuits with one "
+                f"Hadamard per wire and diagonal gates elsewhere"
+            ),
+            mu=mu,
+        )
+
+    frozen = set(frozen_indices or [])
+    diagonal = set(diagonal_tensor_indices(basis))
+    offending = [i for i in range(len(basis.tensors)) if i not in frozen and i not in diagonal]
+
+    if offending:
+        return FlatModulusCertificate(
+            holds=False,
+            reason=(
+                f"mu = {mu:.6g} now, but {len(offending)} trainable tensor(s) "
+                f"are not diagonal, so training may raise it; freeze indices "
+                f"{offending} to obtain the guarantee"
+            ),
+            mu=mu,
+            offending_indices=offending,
+        )
+
+    return FlatModulusCertificate(
+        holds=True,
+        reason=(
+            f"mu == 1 for every reachable parameter value: all "
+            f"{len(basis.tensors) - len(frozen)} trainable tensor(s) are diagonal"
+        ),
+        mu=mu,
+    )

{pdft-0.2.3 → pdft-0.2.4}/src/pdft/io/serialize.py

@@ -71,8 +71,7 @@ def format_float_julia_like(x) -> str:
     mantissa, exp = s.split("e", 1)
 
     # Python always includes sign in exponent (e+20, e-07); Julia drops '+'.
-    if exp.startswith("+"):
-        exp = exp[1:]
+    exp = exp.removeprefix("+")
 
     # Strip leading zeros in the exponent magnitude, but keep at least one digit.
     if exp.startswith("-"):

{pdft-0.2.3 → pdft-0.2.4}/src/pdft/loss.py

@@ -19,14 +19,12 @@ Array = jax.Array
 class AbstractLoss(Protocol):
     """Marker protocol for loss types. No behavior; dispatch is functional."""
 
-    ...
 
 
 @dataclass(frozen=True)
 class L1Norm:
     """L1 norm loss: minimizes `sum(|T(x)|)` to encourage sparsity."""
 
-    pass
 
 
 @dataclass(frozen=True)

{pdft-0.2.3 → pdft-0.2.4}/src/pdft/optimizers/loop.py

@@ -17,7 +17,7 @@ Array = jax.Array
 
 
 def optimize(
-    opt: "RiemannianGD | RiemannianAdam",
+    opt: RiemannianGD | RiemannianAdam,
     tensors: list[Array],
     loss_fn: Callable[[list[Array]], Array],
     grad_fn: Callable[[list[Array]], list[Array]],