pdft 0.2.1__py3-none-any.whl

pdft/bases/block/block.py

@@ -0,0 +1,261 @@
+"""BlockedBasis: apply any inner parametric basis independently to each block.
+
+Motivation: the paper's central finding is that **block size dominates the
+within-block basis** — BlockDCT 8x8 beats every full-image transform tested.
+The full-image learned circuits (QFT/EntangledQFT/TEBD/MERA) are competitive
+among full-image bases but lose to BlockDCT 8x8 by ~3 dB.
+
+`BlockedBasis` directly responds to that finding. It wraps an inner basis at
+smaller m_inner, n_inner and applies it independently to each
+(2^m_inner, 2^n_inner) block of a larger (2^m_outer, 2^n_outer) image.
+Parameters are SHARED across blocks (one basis tiled over many blocks).
+
+Concretely: with m_outer = m_inner + block_log_m and n_outer = n_inner +
+block_log_n, the image is conceptually decomposed into a grid of
+2^block_log_m * 2^block_log_n independent blocks, each transformed by the
+inner basis.
+
+Implementation strategy: BlockedBasis exposes ``m, n, tensors, code, inv_code``
+just like any other basis, so the existing training pipeline
+(``train_basis_batched``, ``loss_function``, ``_build_jit_adam_step``) works
+unchanged. The trick is in ``code``/``inv_code``: they are closures that
+permute axes, vmap the inner code over the block-index dims, and permute back.
+
+Yao little-endian convention preserved: block-index qubits are the
+HIGHER-numbered qubits per dimension (qubits m_inner+1..m_outer for rows),
+which correspond to LOW-INDEXED axes [0..block_log_m) in the (2,)^(m+n)
+tensor layout. See CLAUDE.md §2.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from jax import tree_util
+
+Array = jax.Array
+
+
+# ---------------------------------------------------------------------------
+# Block-aware einsum closure
+# ---------------------------------------------------------------------------
+
+
+def _make_block_code(
+    inner_code: Callable[..., Array],
+    *,
+    m_inner: int,
+    n_inner: int,
+    block_log_m: int,
+    block_log_n: int,
+) -> Callable[..., Array]:
+    """Return a callable with the same signature as ``inner_code`` that operates
+    on a (2,)^(m_outer+n_outer) image by vmap-ing ``inner_code`` over the
+    block-index axes.
+
+    Outer axis layout (Yao little-endian, see _circuit.build_circuit_einsum):
+        [0..block_log_m)              -- block-index ROW qubits (msbs)
+        [block_log_m..m_outer)        -- within-block ROW qubits (lsbs, m_inner of them)
+        [m_outer..m_outer+block_log_n) -- block-index COL qubits (msbs)
+        [m_outer+block_log_n..)        -- within-block COL qubits (lsbs, n_inner of them)
+    """
+    m_outer = m_inner + block_log_m
+    n_outer = n_inner + block_log_n
+    n_blocks = 2 ** (block_log_m + block_log_n)
+    inner_shape = (2,) * (m_inner + n_inner)
+    outer_shape = (2,) * (m_outer + n_outer)
+
+    # Permutation: block-index axes (rows then cols) to the front, within-block
+    # 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(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
+    )
+    inv_perm = [0] * len(perm)
+    for i, p in enumerate(perm):
+        inv_perm[p] = i
+
+    leading_block_shape = (2,) * (block_log_m + block_log_n)
+
+    def block_code(*args: Any) -> Array:
+        *tensors, image = args
+        if image.shape != outer_shape:
+            raise ValueError(f"BlockedBasis expected image shape {outer_shape}, got {image.shape}")
+        x = jnp.transpose(image, perm)
+        x_flat = x.reshape((n_blocks,) + inner_shape)
+
+        def apply_one(img_one: Array) -> Array:
+            return inner_code(*tensors, img_one)
+
+        out_flat = jax.vmap(apply_one)(x_flat)
+        out = out_flat.reshape(leading_block_shape + inner_shape)
+        return jnp.transpose(out, inv_perm)
+
+    return block_code
+
+
+# ---------------------------------------------------------------------------
+# BlockedBasis
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class BlockedBasis:
+    """Wraps an inner parametric basis as a within-block transform.
+
+    Parameters
+    ----------
+    inner : any pdft basis (QFTBasis, EntangledQFTBasis, TEBDBasis, MERABasis)
+        Within-block parametric circuit at smaller m_inner = inner.m,
+        n_inner = inner.n.
+    block_log_m, block_log_n : int
+        Number of block-index qubits per dimension. The image is
+        (2^(inner.m + block_log_m), 2^(inner.n + block_log_n)).
+        block_log_m=0 (and =0) reduces to the inner basis.
+
+    Notes
+    -----
+    All learnable parameters live in ``inner.tensors``; ``BlockedBasis`` is
+    a pure structural wrapper. Block parameters are SHARED across blocks
+    (one inner basis tiled).
+    """
+
+    inner: Any
+    block_log_m: int
+    block_log_n: int
+    code: object = field(compare=False, repr=False)
+    inv_code: object = field(compare=False, repr=False)
+
+    def __init__(
+        self,
+        inner: Any,
+        block_log_m: int,
+        block_log_n: int,
+        code: object | None = None,
+        inv_code: object | None = None,
+    ):
+        if not hasattr(inner, "m") or not hasattr(inner, "n"):
+            raise TypeError(
+                f"inner must be a pdft basis with m, n attributes; got {type(inner).__name__}"
+            )
+        if block_log_m < 0 or block_log_n < 0:
+            raise ValueError(
+                f"block_log_m and block_log_n must be >= 0; got "
+                f"block_log_m={block_log_m}, block_log_n={block_log_n}"
+            )
+        self.inner = inner
+        self.block_log_m = block_log_m
+        self.block_log_n = block_log_n
+        if code is None or inv_code is None:
+            built_code = _make_block_code(
+                inner.code,
+                m_inner=inner.m,
+                n_inner=inner.n,
+                block_log_m=block_log_m,
+                block_log_n=block_log_n,
+            )
+            built_inv = _make_block_code(
+                inner.inv_code,
+                m_inner=inner.m,
+                n_inner=inner.n,
+                block_log_m=block_log_m,
+                block_log_n=block_log_n,
+            )
+            self.code = code if code is not None else built_code
+            self.inv_code = inv_code if inv_code is not None else built_inv
+        else:
+            self.code = code
+            self.inv_code = inv_code
+
+    # ---- AbstractSparseBasis interface (matches QFTBasis) ----
+
+    @property
+    def m(self) -> int:
+        return self.inner.m + self.block_log_m
+
+    @property
+    def n(self) -> int:
+        return self.inner.n + self.block_log_n
+
+    @property
+    def tensors(self) -> list[Array]:
+        """Forward to inner — pytree leaves order is `inner.tensors`."""
+        return self.inner.tensors
+
+    @property
+    def inv_tensors(self) -> list[Array]:
+        return self.inner.tensors
+
+    @property
+    def image_size(self) -> tuple[int, int]:
+        return (2**self.m, 2**self.n)
+
+    @property
+    def num_parameters(self) -> int:
+        return self.inner.num_parameters
+
+    @property
+    def num_blocks(self) -> int:
+        return 2 ** (self.block_log_m + self.block_log_n)
+
+    @property
+    def block_shape(self) -> tuple[int, int]:
+        return (2**self.inner.m, 2**self.inner.n)
+
+    def forward_transform(self, pic: Array) -> Array:
+        from ...loss import _apply_circuit
+
+        return _apply_circuit(self.tensors, self.code, self.m, self.n, pic)
+
+    def inverse_transform(self, pic: Array) -> Array:
+        from ...loss import _apply_circuit
+
+        return _apply_circuit(
+            [jnp.conj(t) for t in self.tensors],
+            self.inv_code,
+            self.m,
+            self.n,
+            pic,
+        )
+
+
+# ---------------------------------------------------------------------------
+# JAX pytree registration
+# ---------------------------------------------------------------------------
+
+
+def _blockedbasis_flatten(b: BlockedBasis):
+    """Flatten by recursing into inner via JAX's pytree machinery.
+
+    This delegates inner reconstruction to JAX's tree_unflatten, which
+    works for ANY pytree-registered inner — including nested wrappers
+    like StackedBasis whose constructor doesn't take (m, n, tensors, ...).
+    """
+    inner_leaves, inner_treedef = jax.tree_util.tree_flatten(b.inner)
+    aux = (inner_treedef, b.block_log_m, b.block_log_n, b.code, b.inv_code)
+    return tuple(inner_leaves), aux
+
+
+def _blockedbasis_unflatten(aux, leaves) -> BlockedBasis:
+    inner_treedef, block_log_m, block_log_n, code, inv_code = aux
+    inner = jax.tree_util.tree_unflatten(inner_treedef, list(leaves))
+    return BlockedBasis(
+        inner=inner,
+        block_log_m=block_log_m,
+        block_log_n=block_log_n,
+        code=code,
+        inv_code=inv_code,
+    )
+
+
+tree_util.register_pytree_node(BlockedBasis, _blockedbasis_flatten, _blockedbasis_unflatten)
+
+
+__all__ = ["BlockedBasis"]

pdft/bases/block/real_rich.py

@@ -0,0 +1,164 @@
+"""RealRichBasis (Approach A): real-orthogonal restriction of RichBasis.
+
+Motivation: RichBasis (54 free real params per dim, complex U(4) gates) is
+a strict 54-dim submanifold of SU(8) — does not contain DCT, can only
+approximate. But fully complex U(4) is also wasteful for the natural-image
+problem: DCT is real-valued (lives in O(8)), so the IMAGINARY parts of
+RichBasis's parameters are doing no useful work for natural images.
+
+RealRichBasis keeps the same QFT topology and gate count but constrains
+each tensor to be REAL-valued:
+  - 3 H gates per dim → 2×2 real-orthogonal matrices (init: Hadamard).
+    Free params per gate: 1 (rotation angle of the connected component
+    of Hadamard in O(2)).
+  - 3 "U(4)" gates per dim → real 4×4 orthogonal matrices, 6 free real
+    params each (init: identity).
+
+Total per dim: 3·1 + 3·6 = 21 free real params (BELOW dim O(8) = 28).
+Strict submanifold of O(8); whether DCT is in this family is empirical.
+
+Storage: tensors are stored as complex128 with all-zero imaginary parts.
+Cayley retraction on the real-image gradient preserves real-ness
+automatically — no manifold change needed; the existing UnitaryManifold
+trains the orthogonal subset correctly when initialised with real values
+on a real-valued objective.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import tree_util
+
+from ...circuit.builder import HADAMARD, Gate, compile_circuit
+
+Array = jax.Array
+
+
+def _real_eye_u4() -> Array:
+    """4x4 identity reshaped to (2, 2, 2, 2), as real complex128."""
+    return jnp.asarray(np.eye(4).reshape(2, 2, 2, 2), dtype=jnp.complex128)
+
+
+def _real_rich_qft_gates_1d(n_qubits: int, offset: int) -> list[Gate]:
+    """QFT topology with H + REAL-orthogonal 2-qubit gates.
+
+    H slots are the canonical Hadamard. 2-qubit slots are the 4×4 identity
+    (a real-orthogonal matrix). Both are within the connected component of
+    O(d) reachable via Cayley retraction with real updates.
+    """
+    eye_u4 = _real_eye_u4()
+    gates: list[Gate] = []
+    for j in range(1, n_qubits + 1):
+        q = offset + j
+        gates.append(Gate(kind="H", qubits=(q,), tensor=HADAMARD, phase=0.0))
+        for target in range(j + 1, n_qubits + 1):
+            t = offset + target
+            gates.append(
+                Gate(
+                    kind="U4",
+                    qubits=(t, q),
+                    tensor=eye_u4,
+                    phase=0.0,
+                )
+            )
+    return gates
+
+
+def _real_rich_code(m: int, n: int, *, inverse: bool):
+    if m < 1 or n < 1:
+        raise ValueError(f"m and n must be >= 1, got m={m}, n={n}")
+    gates = _real_rich_qft_gates_1d(m, offset=0) + _real_rich_qft_gates_1d(n, offset=m)
+    return compile_circuit(gates, m, n, inverse=inverse)
+
+
+@dataclass
+class RealRichBasis:
+    """QFT topology with H + real-orthogonal 2-qubit gates.
+
+    The U(4) slots are *initialised* to the 4×4 identity (real-orthogonal,
+    not the complex controlled-phase) so the basis is NOT bit-identical
+    to QFTBasis at training step 0 — the forward circuit at init is
+    H ⊗ H ⊗ H per dim followed by identity 2-qubit ops, i.e. just the
+    Walsh-Hadamard transform. This is the appropriate starting point for
+    a real-valued search; the Walsh-Hadamard is the simplest real-orthogonal
+    basis and a natural baseline for natural-image transforms.
+
+    Pytree contract:
+        leaves   = tensors                                (one list)
+        aux data = (m, n, len(tensors), code, inv_code)
+    """
+
+    m: int
+    n: int
+    tensors: list[Array]
+    code: object = field(compare=False, repr=False)
+    inv_code: object = field(compare=False, repr=False)
+
+    def __init__(
+        self,
+        m: int,
+        n: int,
+        tensors: Sequence[Array] | None = None,
+        code: object | None = None,
+        inv_code: object | None = None,
+    ):
+        if m < 1 or n < 1:
+            raise ValueError(f"m and n must be >= 1, got m={m}, n={n}")
+        self.m = m
+        self.n = n
+        _code, init_tensors = _real_rich_code(m, n, inverse=False)
+        _inv_code, _ = _real_rich_code(m, n, inverse=True)
+        self.tensors = list(tensors) if tensors is not None else init_tensors
+        self.code = code if code is not None else _code
+        self.inv_code = inv_code if inv_code is not None else _inv_code
+
+    @property
+    def inv_tensors(self) -> list[Array]:
+        return self.tensors
+
+    @property
+    def image_size(self) -> tuple[int, int]:
+        return (2**self.m, 2**self.n)
+
+    @property
+    def num_parameters(self) -> int:
+        return sum(int(t.size) for t in self.tensors)
+
+    def forward_transform(self, pic: Array) -> Array:
+        from ...loss import _apply_circuit
+
+        return _apply_circuit(self.tensors, self.code, self.m, self.n, pic)
+
+    def inverse_transform(self, pic: Array) -> Array:
+        from ...loss import _apply_circuit
+
+        return _apply_circuit(
+            [jnp.conj(t) for t in self.tensors],
+            self.inv_code,
+            self.m,
+            self.n,
+            pic,
+        )
+
+
+def _realrichbasis_flatten(b: RealRichBasis):
+    leaves = tuple(b.tensors)
+    aux = (b.m, b.n, len(b.tensors), b.code, b.inv_code)
+    return leaves, aux
+
+
+def _realrichbasis_unflatten(aux, leaves) -> RealRichBasis:
+    m, n, n_fwd, code, inv_code = aux
+    assert len(leaves) == n_fwd
+    return RealRichBasis(m=m, n=n, tensors=list(leaves), code=code, inv_code=inv_code)
+
+
+tree_util.register_pytree_node(RealRichBasis, _realrichbasis_flatten, _realrichbasis_unflatten)
+
+
+__all__ = ["RealRichBasis"]

pdft/bases/block/rich.py

@@ -0,0 +1,265 @@
+"""RichBasis: single-layer parametric circuit with full 2-qubit unitary gates.
+
+Motivation: at small block sizes (m=n=3 = 8x8) the existing H+CP gate
+family hits an expressivity ceiling — all topologies converge ~1.75 dB
+below 8x8 DCT. The cause is that diagonal CP gates have only 1 free
+parameter each. A general 2-qubit unitary (U(4)) has 15 free parameters,
+and the H + U(4) gate family is provably universal for SU(2^n) at any
+qubit count >= 2, so it CONTAINS DCT as a special case.
+
+RichBasis emits the same QFT topology gate sequence (H per qubit + 2-qubit
+gates between qubit pairs) but each 2-qubit gate is a learnable U(4)
+instead of a 1-parameter CP. Initialised so the circuit is BIT-IDENTICAL
+to QFT at training step 0 (each U(4) gate equals the 4×4 controlled-phase
+diag(1, 1, 1, exp(iφ)) at its standard QFT phase). This gives Adam a
+gentle starting point: the optimiser begins exactly where plain QFT does
+and can only improve.
+
+Parameter count at m=n=3 (8x8 block):
+  - 6 H gates (3 per dim) at 4 real params each = 24
+  - 6 U(4) gates (3 per dim) at 15 real params each = 90
+  - total: 114 real params per dim
+  vs SU(8) dimension = 63 free real params
+  → strictly more parameters than needed for any 8x8 unitary.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+import jax
+import jax.numpy as jnp
+from jax import tree_util
+
+from ...circuit.builder import HADAMARD, Gate, compile_circuit, u4_from_phase
+
+Array = jax.Array
+
+
+def _rich_qft_gates_1d(n_qubits: int, offset: int) -> list[Gate]:
+    """Same QFT topology as qft._qft_gates_1d, but with U(4) gates instead of CP.
+
+    Each U(4) gate is initialised to the 4x4 unitary equivalent of the
+    standard QFT phase (so the basis is bit-identical to QFTBasis at init).
+    """
+    gates: list[Gate] = []
+    for j in range(1, n_qubits + 1):
+        q = offset + j
+        gates.append(Gate(kind="H", qubits=(q,), tensor=HADAMARD, phase=0.0))
+        for target in range(j + 1, n_qubits + 1):
+            k = target - j + 1
+            t = offset + target
+            phi = float(2 * jnp.pi / (2**k))
+            gates.append(
+                Gate(
+                    kind="U4",
+                    qubits=(t, q),  # control, target
+                    tensor=u4_from_phase(phi),
+                    phase=phi,
+                )
+            )
+    return gates
+
+
+def _rich_code(m: int, n: int, *, inverse: bool):
+    """Build the rich (H + U(4)) circuit einsum + initial tensor list."""
+    if m < 1 or n < 1:
+        raise ValueError(f"m and n must be >= 1, got m={m}, n={n}")
+    gates = _rich_qft_gates_1d(m, offset=0) + _rich_qft_gates_1d(n, offset=m)
+    return compile_circuit(gates, m, n, inverse=inverse)
+
+
+@dataclass
+class RichBasis:
+    """QFT topology with H + learnable U(4) gates instead of H + CP.
+
+    Parameter count at m=n=3 (within an 8x8 block):
+      - 3 H gates per dim × 3 free real params (SU(2)) = 9
+      - 3 U(4) gates per dim × 15 free real params (SU(4)) = 45
+      - total per dim: 54 (BELOW the 63-dim of SU(8) — meaningful structure)
+
+    Initialised so the circuit is BIT-IDENTICAL to QFTBasis at training step 0
+    (each U(4) starts at the 4×4 controlled-phase diag(1, 1, 1, exp(iφ)) of
+    its corresponding QFT slot). Adam can then deform the U(4) gates into any
+    4×4 unitary, but the family is a strict 54-dim submanifold of SU(8) and
+    does NOT contain 8×8 DCT exactly (empirically: fit_to_dct plateaus at
+    Frobenius² ≈ 63.7).
+
+    Pytree contract:
+        leaves   = tensors                                (one list)
+        aux data = (m, n, len(tensors), code, inv_code)
+    """
+
+    m: int
+    n: int
+    tensors: list[Array]
+    code: object = field(compare=False, repr=False)
+    inv_code: object = field(compare=False, repr=False)
+
+    def __init__(
+        self,
+        m: int,
+        n: int,
+        tensors: Sequence[Array] | None = None,
+        code: object | None = None,
+        inv_code: object | None = None,
+    ):
+        if m < 1 or n < 1:
+            raise ValueError(f"m and n must be >= 1, got m={m}, n={n}")
+        self.m = m
+        self.n = n
+        _code, init_tensors = _rich_code(m, n, inverse=False)
+        _inv_code, _ = _rich_code(m, n, inverse=True)
+        self.tensors = list(tensors) if tensors is not None else init_tensors
+        self.code = code if code is not None else _code
+        self.inv_code = inv_code if inv_code is not None else _inv_code
+
+    @property
+    def inv_tensors(self) -> list[Array]:
+        return self.tensors
+
+    @property
+    def image_size(self) -> tuple[int, int]:
+        return (2**self.m, 2**self.n)
+
+    @property
+    def num_parameters(self) -> int:
+        return sum(int(t.size) for t in self.tensors)
+
+    def forward_transform(self, pic: Array) -> Array:
+        from ...loss import _apply_circuit
+
+        return _apply_circuit(self.tensors, self.code, self.m, self.n, pic)
+
+    def inverse_transform(self, pic: Array) -> Array:
+        from ...loss import _apply_circuit
+
+        return _apply_circuit(
+            [jnp.conj(t) for t in self.tensors],
+            self.inv_code,
+            self.m,
+            self.n,
+            pic,
+        )
+
+
+def _richbasis_flatten(b: RichBasis):
+    leaves = tuple(b.tensors)
+    aux = (b.m, b.n, len(b.tensors), b.code, b.inv_code)
+    return leaves, aux
+
+
+def _richbasis_unflatten(aux, leaves) -> RichBasis:
+    m, n, n_fwd, code, inv_code = aux
+    assert len(leaves) == n_fwd
+    return RichBasis(m=m, n=n, tensors=list(leaves), code=code, inv_code=inv_code)
+
+
+tree_util.register_pytree_node(RichBasis, _richbasis_flatten, _richbasis_unflatten)
+
+
+def _dct_matrix(n: int) -> Array:
+    """Orthonormal 1D DCT-II matrix of size n × n."""
+    import numpy as _np
+
+    k = _np.arange(n).reshape(-1, 1)
+    j = _np.arange(n)
+    M = _np.cos(_np.pi * (2 * j + 1) * k / (2 * n))
+    M[0, :] *= 1.0 / _np.sqrt(n)
+    M[1:, :] *= _np.sqrt(2.0 / n)
+    return jnp.asarray(M, dtype=jnp.complex128)
+
+
+def fit_to_dct(
+    basis_factory,
+    *,
+    n_steps: int = 2000,
+    lr: float = 0.02,
+) -> list[Array]:
+    """Fit a parametric basis so its forward circuit ≈ DCT_2D.
+
+    Parameters
+    ----------
+    basis_factory : callable
+        Zero-argument callable returning a basis instance. Must expose
+        ``.m, .n, .tensors, .code`` (any pdft basis class). The returned
+        tensor list has the same shapes as ``basis_factory().tensors``
+        and can be passed back as ``tensors=...`` for a DCT warm-start.
+    n_steps : int
+        Adam steps. 2000 is a generous default; convergence is typically
+        much faster when DCT lies in the parametric family.
+    lr : float
+        Adam learning rate.
+
+    The loss is the Frobenius² distance between the circuit's action on
+    a complete basis and ``DCT_{2^m} ⊗ DCT_{2^n}``. If the family does
+    not contain DCT, the loss plateaus at the closest reachable distance.
+    """
+    import time as _time
+
+    from ...manifolds import group_by_manifold, stack_tensors
+
+    basis = basis_factory()
+    m, n = basis.m, basis.n
+    code = basis.code
+    init_tensors = basis.tensors
+    D_row = _dct_matrix(2**m)
+    D_col = _dct_matrix(2**n)
+    target = jnp.kron(D_row, D_col).reshape(2**m, 2**n, 2**m, 2**n)
+
+    eye = jnp.eye(2 ** (m + n), dtype=jnp.complex128).reshape((2 ** (m + n),) + (2,) * (m + n))
+    target_flat = target.reshape(2**m, 2**n, 2 ** (m + n))
+
+    def loss_fn(tensors):
+        outs = jax.vmap(lambda x: code(*tensors, x))(eye)
+        outs_mat = outs.reshape(2 ** (m + n), 2**m, 2**n).transpose(1, 2, 0)
+        return jnp.real(jnp.sum(jnp.abs(outs_mat - target_flat) ** 2))
+
+    grad_fn = jax.value_and_grad(loss_fn)
+
+    tensors = list(init_tensors)
+    groups = group_by_manifold(tensors)
+    manifolds = list(groups.keys())
+    indices = [tuple(groups[mfd]) for mfd in manifolds]
+
+    # Adam state per group.
+    m_state, v_state = [], []
+    for idxs in indices:
+        pb = stack_tensors(tensors, list(idxs))
+        m_state.append(jnp.zeros_like(pb))
+        v_state.append(jnp.zeros(pb.shape, dtype=jnp.float64))
+
+    beta1, beta2, eps = 0.9, 0.999, 1e-8
+    t0 = _time.perf_counter()
+    for step in range(1, n_steps + 1):
+        loss_val, raw_grads = grad_fn(tensors)
+        grads = [jnp.conj(g) for g in raw_grads]
+
+        for k, (mfd, idxs) in enumerate(zip(manifolds, indices)):
+            pb = stack_tensors(tensors, list(idxs))
+            gb = stack_tensors(grads, list(idxs))
+            rg = mfd.project(pb, gb)
+            new_m = beta1 * m_state[k] + (1 - beta1) * rg
+            new_v = beta2 * v_state[k] + (1 - beta2) * jnp.real(jnp.conj(rg) * rg)
+            bc1 = 1.0 - beta1**step
+            bc2 = 1.0 - beta2**step
+            direction = (new_m / bc1) / (jnp.sqrt(new_v / bc2) + eps)
+            new_pb = mfd.retract(pb, -direction, lr)
+            new_m = mfd.transport(pb, new_pb, new_m)
+            m_state[k] = new_m
+            v_state[k] = new_v
+            for k2, idx in enumerate(idxs):
+                tensors[idx] = new_pb[..., k2]
+
+        if step % 200 == 0 or step == 1:
+            print(
+                f"  fit_to_dct step {step:>4d}: loss={float(loss_val):.4e} "
+                f"(elapsed {_time.perf_counter() - t0:.1f}s)",
+                flush=True,
+            )
+
+    return tensors
+
+
+__all__ = ["RichBasis", "fit_to_dct"]

pdft/bases/circuit/__init__.py

@@ -0,0 +1,15 @@
+"""Circuit-topology bases: QFT, EntangledQFT, TEBD, MERA."""
+
+from .qft import ft_mat, ift_mat, qft_code
+from .entangled_qft import entangled_qft_code
+from .mera import mera_code
+from .tebd import tebd_code
+
+__all__ = [
+    "entangled_qft_code",
+    "ft_mat",
+    "ift_mat",
+    "mera_code",
+    "qft_code",
+    "tebd_code",
+]