pdft 0.2.1__tar.gz

pdft-0.2.1/.claude/CLAUDE.md

@@ -0,0 +1,184 @@
+# CLAUDE.md
+
+Project-specific guidance for Claude Code working in this repo. Read this before making non-trivial changes — several conventions are *not* discoverable from the code without painful debugging, and at least one is documented here only because the original Python port spent hours rediscovering it.
+
+## What this is
+
+`pdft` is a faithful Python port of [ParametricDFT.jl](https://github.com/nzy1997/ParametricDFT.jl) using JAX. Goal: a Python user can reproduce Julia results bit-for-bit (or within documented tolerances) on the same input. **Behavior parity with Julia is the primary correctness criterion** — never sacrifice it for Pythonic-ness.
+
+- Design spec: `docs/superpowers/specs/2026-04-24-pdft-migration-design.md`
+- Roadmap: GitHub issue #1
+- Reliability hardening backlog: GitHub issue #2
+- Upstream is pinned to `nzy1997/ParametricDFT.jl@a201a27e47df2f0f3ab460f83d49b6e5f5d1e9ef`. The pin lives in two places that must stay in sync: `reference/julia/generate_goldens.jl` (`UPSTREAM_SHA`) and `src/pdft/__init__.py` (`__upstream_ref__`).
+
+## Critical conventions (do not break)
+
+These are the non-obvious invariants that took the original port multiple iterations to find. Violating any of them silently breaks parity with Julia.
+
+### 1. Wirtinger gradient conjugation
+
+**JAX returns `∂f/∂z̄` while Julia's Zygote returns `∂f/∂z`** for real-valued functions of complex inputs. They are complex conjugates. `optimize()` in `src/pdft/optimizers.py` conjugates `raw_grads` immediately after `grad_fn(...)`:
+
+```python
+raw_grads = grad_fn(state.current_tensors)
+raw_grads = [jnp.conj(g) for g in raw_grads]   # MUST stay
+```
+
+Without this line, GD trajectories drift ~10% over 50 steps and Adam outright diverges. Any new optimizer added to this module must apply the same conjugation, or define a custom JAX `vjp` that matches Julia's convention.
+
+### 2. Yao little-endian qubit ordering
+
+`pic.reshape((2,)*(m+n))` gives axes in big-endian bit order, but Julia's Yao treats qubit 1 as the LSB. `_circuit.build_circuit_einsum` reverses qubit-to-axis mapping within each block:
+
+```python
+row_pic = [input_labels[q - 1] for q in range(m, 0, -1)]   # axis 0 ↔ qubit m
+col_pic = [input_labels[q - 1] for q in range(m + n, m, -1)]
+```
+
+Same reversal applies to `out_labels`. Don't undo this without first making `ft_mat` parity tests fail in a controlled way.
+
+### 3. Compact 2×2 CP tensor representation
+
+A controlled-phase gate is `diag(1, 1, 1, exp(iφ))` as a 4×4 matrix, but Yao's `yao2einsum` emits the **compact 2×2 form**:
+
+```python
+controlled_phase_diag(φ) = [[1, 1], [1, exp(iφ)]]
+```
+
+This shares the wire labels of control and target qubits and does NOT introduce new labels (the gate is diagonal so each wire's value is unchanged; only a multiplicative phase is injected). All circuit modules use this representation; never substitute the 4×4 form even though it would be mathematically equivalent — Julia goldens won't match.
+
+### 4. Hadamard-first tensor sort
+
+Julia's `qft_code` does `perm_vec = sortperm(tn.tensors, by=x -> !(x ≈ mat(H)))` — Hadamards come first, CPs after. Python's `_circuit.build_circuit_einsum` applies the same sort. The einsum is rebuilt with the permuted tensor list AND permuted subscripts so it stays valid. **JSON serialization assumes this order**; cross-language interop breaks if you change it.
+
+### 5. JAX x64 mode at import time
+
+`src/pdft/__init__.py` calls `jax.config.update("jax_enable_x64", True)` before any other import. `tests/conftest.py` does the same. Without x64, JAX uses complex64 / float32 and parity tolerances become unreachable. Any new test file that uses `jax.numpy` directly without importing `pdft` first must include the same call (or import pdft to force it).
+
+### 6. Column-major iteration for cross-language hashes
+
+`io_json.basis_hash` and `compression` use **Julia 1-based column-major** flat indices. NumPy is row-major by default. Use `array.flatten(order="F")` or `np.unravel_index(..., order="F")`. Never use plain `array.flatten()` for serialization.
+
+### 7. Julia-compatible float formatting
+
+Python's `repr(5e-7)` gives `"5e-07"`; Julia's `string(5e-7)` gives `"5.0e-7"`. `io_json._format_float_julia_like` post-processes `repr()` to match Julia. Use it (not `repr` or `str`) for any value that participates in a cross-language hash or JSON byte-comparison.
+
+### 8. L1-cusp horizon for trajectory parity
+
+GD on L1 loss is bit-exact for ~50 steps and stays within `atol=1e-3` over 200 steps. Beyond that, FP accumulation eventually pushes Python and Julia across the same loss cusp at slightly different alphas in Armijo line search. Both still converge to the same loss basin. Don't tighten `tests/test_parity_long_run.py` to `atol=1e-10` over 200 steps — this is mathematically expected on non-smooth losses, not a bug. Smooth losses (MSE without truncation) don't have this problem.
+
+### 9. Phase extractors classify by tensor shape, not gate-list metadata
+
+`get_*_gate_indices` walks the tensor list and tags any 2×2 tensor whose four entries have unit-modulus magnitudes (within `atol=0.15`) as a CP gate, then returns the LAST `n_gates` such positions. After training, individual entries can drift slightly off the unit circle; the moderate tolerance accommodates that. If you tighten the tolerance, do so in tandem with a regression test on a trained basis.
+
+## Repo layout
+
+```
+src/pdft/
+├── manifolds.py            Mathematical core (UnitaryManifold, PhaseManifold, batched ops)
+├── loss.py                 L1Norm, MSELoss, topk_truncate, loss_function (public)
+├── profiling.py            Cross-cutting profiling helpers
+├── bases/
+│   ├── base.py             AbstractSparseBasis + bases_allclose + 4 basis dataclasses
+│   ├── circuit/            QFT, EntangledQFT, TEBD, MERA (full circuit topologies)
+│   └── block/              Blocked, Rich, RealRich (parameter-efficient)
+├── circuit/                Einsum builder (builder.py) + JIT closure cache (cache.py)
+├── optimizers/             core, gd (RiemannianGD + Armijo), adam (RiemannianAdam), loop
+├── training/               schedules, single (train_basis), batched, adam_step, eval_loop
+├── io/                     serialize (JSON), compression
+└── viz/                    loss (matplotlib loss plots), circuit (schematic)
+
+reference/julia/            Julia harness — needed only to regenerate goldens
+reference/goldens/          Committed .npz + .json files (<200 KB total)
+examples/                   3 runnable demos, each <10s
+tests/                      pytest; mirrors src/pdft/ layout (tests/bases/, tests/optimizers/, ...)
+```
+
+Benchmarks live in a separate repo: https://github.com/zazabap/pdft-benchmarks (split from this repo at pdft v0.2.0; pinned via its `pyproject.toml`).
+
+`DCTBasis` was removed during the modular-src refactor (PR #11). Use a parametric basis (`RichBasis` / `RealRichBasis`) for similar use cases.
+
+`pdft/__init__.py` re-exports a small set of most-used names (basis classes,
+trainer, optimizer, loss). Less-public names live in their subpackages —
+e.g. `pdft.io.save_basis`, `pdft.manifolds.UnitaryManifold`,
+`pdft.profiling.profile_training`, `pdft.bases.circuit.qft_code`.
+
+`AbstractRiemannianOptimizer` is `RiemannianGD | RiemannianAdam` — a structural union, not a Protocol. Adding a third optimizer means extending that union *and* adding an `isinstance` branch in `optimize()`.
+
+`train_basis` is generic over basis type via JAX pytree flatten/unflatten. The convention: each registered basis pytree must have leaves ordered as `tuple(tensors) + tuple(inv_tensors)`, with everything else (m, n, code, inv_code, counts) in aux data. New basis types must follow this convention.
+
+## Dev workflow
+
+```bash
+# Install (Python 3.11+; conda env at /opt/conda/envs/pdft on the dev box)
+pip install -e ".[dev]"
+
+# Run tests
+pytest                                    # full suite
+pytest --cov=pdft --cov-fail-under=90     # CI gate
+pytest tests/test_parity_*.py             # parity-only
+
+# Lint (CI fails if this is dirty — check before pushing)
+ruff check src tests
+ruff format src tests
+
+# Run examples
+python examples/basis_demo.py
+python examples/optimizer_benchmark.py
+python examples/mera_demo.py
+
+# Regenerate Julia goldens (requires Julia 1.10+)
+make goldens
+```
+
+After regenerating goldens: also update `__upstream_ref__` in `src/pdft/__init__.py` to match `manifest.json`'s `upstream_sha`.
+
+## Tests and parity
+
+Three layers (per spec section 7):
+
+1. **Parity tests** (`tests/test_parity_*.py`) — load committed `.npz` / `.json` goldens from `reference/goldens/` and assert Python matches Julia. These are the load-bearing correctness tests.
+2. **Property tests** (`tests/test_<module>.py`) — math-invariant checks (unitarity preserved, round-trip identity, monotone descent, …). Don't depend on Julia.
+3. **Smoke / integration** (`tests/test_smoke.py`, `tests/test_training_integration.py`).
+
+Coverage gate is `--cov-fail-under=90`. Don't add tests that reduce per-module coverage below the line; if a new module legitimately needs more code, also add the property tests for it.
+
+`tests/conftest.py` enables x64. Don't import jax in a test before pdft unless the test also calls `jax.config.update("jax_enable_x64", True)`.
+
+## When to break parity
+
+The default answer is "never". The only acceptable cases:
+
+- **Float-precision noise**: test tolerances are documented per case; if you cross a tolerance boundary, investigate before relaxing.
+- **Tie-breaking divergences** (e.g. argpartition vs partialsortperm): document and assert downstream behavior matches instead.
+
+If you find another mismatch:
+1. Diagnose with `reference/julia/` (write a small Julia script to dump intermediate values).
+2. Compare to Python at the same intermediate point.
+3. **The bug is in Python**, almost always.
+4. If the bug is in upstream Julia, file an issue there and document the divergence in this file.
+
+## CI gotchas
+
+- Ruff is strict (`F401`, `F841`, `E741`, `E731` all errors). Auto-`--fix` handles most.
+- Coverage gate runs only on the actual test suite, not examples; don't rely on examples for coverage.
+- The `verify-upstream-pin.yml` workflow runs only on PRs touching `reference/`. It uses the GitHub API to confirm the pinned sha exists in upstream — don't push a sha that's only in a fork.
+- The matrix is 3.11 / 3.12 / 3.13. JAX dropped 3.10 in 0.10.0; do not lower the floor.
+
+## What NOT to do
+
+- **Don't add `optax`** as a dependency. The optimizer logic is hand-rolled to match Julia's exact moment-update math. `optax`'s defaults and FP order will diverge from Julia.
+- **Don't switch from `jnp.einsum_path("greedy")` to `"optimal"`.** "Optimal" is exponential in tensor count and hangs on the 3×3 QFT (12 tensors).
+- **Don't add explicit JIT to `train_basis`.** It calls a basis-typed loss closure with Python-list pytrees; JIT decisions are best left to inner functions where the static-vs-leaf split is clearer.
+- **Don't introduce backwards-compat shims** for the JSON schema. We're at v0.1.0; if the schema changes, bump the version and regenerate goldens.
+- **Don't add ML scaffolding** (no DataLoader, no Trainer-like classes, no Lightning). Upstream is one-target-image-at-a-time and we mirror that. Batched training is open work in #2.
+- **Don't run examples in CI.** They write to `out/` (gitignored) and are not coverage-relevant.
+
+## When making changes
+
+1. Make the change.
+2. Run `ruff check src tests` and `ruff format src tests`.
+3. Run `pytest --cov=pdft --cov-fail-under=90`.
+4. If you touched a circuit, basis, optimizer, or io_json module: also run the relevant `tests/test_parity_*.py` to confirm Julia-parity hasn't drifted.
+5. If you changed the upstream pin: update both `reference/julia/generate_goldens.jl` (`UPSTREAM_SHA`) and `src/pdft/__init__.py` (`__upstream_ref__`), regenerate goldens, and verify all parity tests still pass.
+6. Push and watch the GitHub Actions matrix — three Python versions must all be green before merging.

pdft-0.2.1/.github/workflows/ci.yml

@@ -0,0 +1,26 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+      - name: Install
+        run: pip install -e ".[dev]"
+      - name: Ruff lint
+        run: ruff check src tests
+      - name: Pytest
+        run: pytest --cov=pdft --cov-report=term --cov-fail-under=90

pdft-0.2.1/.github/workflows/publish.yml

@@ -0,0 +1,63 @@
+name: publish
+
+# Builds and publishes pdft to PyPI on every v*.*.* tag push.
+#
+# Uses PyPI Trusted Publishers (OIDC) — NO API token stored in GitHub
+# secrets. To activate: on https://pypi.org/manage/project/pdft/settings/publishing/
+# add a Trusted Publisher with:
+#   PyPI Project Name: pdft
+#   Owner:             zazabap
+#   Repository name:   pdft
+#   Workflow name:     publish.yml
+#   Environment name:  pypi
+#
+# After the first successful publish, future releases are:
+#   git tag v0.2.1 && git push --tags
+#
+# Docs: https://docs.pypi.org/trusted-publishers/
+
+on:
+  push:
+    tags:
+      - "v*.*.*"
+
+jobs:
+  build:
+    name: Build sdist + wheel
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install build backend
+        run: python -m pip install --upgrade pip build
+      - name: Verify pyproject version matches the git tag
+        run: |
+          tag="${GITHUB_REF#refs/tags/v}"
+          file_ver=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml','rb'))['project']['version'])")
+          if [ "$tag" != "$file_ver" ]; then
+            echo "::error::Tag v$tag does not match pyproject.toml version $file_ver"
+            exit 1
+          fi
+      - name: Build
+        run: python -m build
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    name: Publish to PyPI
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write   # mandatory for Trusted Publisher OIDC
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

pdft-0.2.1/.github/workflows/verify-upstream-pin.yml

@@ -0,0 +1,28 @@
+name: Verify upstream pin
+
+on:
+  pull_request:
+    paths:
+      - "reference/**"
+      - ".github/workflows/verify-upstream-pin.yml"
+
+jobs:
+  verify:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Extract pinned sha
+        id: pin
+        run: |
+          sha=$(grep -oP 'UPSTREAM_SHA = "\K[0-9a-f]+' reference/julia/generate_goldens.jl | head -n1)
+          echo "sha=$sha" >> "$GITHUB_OUTPUT"
+      - name: Assert sha exists upstream
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          sha="${{ steps.pin.outputs.sha }}"
+          if [ -z "$sha" ]; then
+            echo "Could not extract UPSTREAM_SHA from generate_goldens.jl"
+            exit 1
+          fi
+          gh api "repos/nzy1997/ParametricDFT.jl/commits/$sha" >/dev/null

pdft-0.2.1/.gitignore

@@ -0,0 +1,59 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+*.egg
+MANIFEST
+
+# Virtual environments
+.venv/
+venv/
+env/
+ENV/
+
+# Tooling caches
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.tox/
+.coverage
+.coverage.*
+htmlcov/
+coverage.xml
+*.cover
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Editors / OS
+.idea/
+.vscode/
+.DS_Store
+*.swp
+
+# Claude Code — track CLAUDE.md, ignore everything else
+.claude/*
+!.claude/CLAUDE.md
+
+# Example outputs
+out/
+/out/
+
+# Benchmark outputs
+benchmarks/results/

pdft-0.2.1/.pre-commit-config.yaml

@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.5.0
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format

pdft-0.2.1/LICENSE

@@ -0,0 +1,25 @@
+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.

pdft-0.2.1/Makefile

@@ -0,0 +1,8 @@
+.PHONY: goldens test
+
+test:
+	pytest
+
+goldens:
+	cd reference/julia && julia --project=. -e 'using Pkg; Pkg.instantiate()'
+	julia --project=reference/julia reference/julia/generate_goldens.jl

pdft-0.2.1/PKG-INFO

@@ -0,0 +1,100 @@
+Metadata-Version: 2.4
+Name: pdft
+Version: 0.2.1
+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.5; 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
+
+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.
+
+> Status: early scaffold. The Julia package is the reference implementation;
+> this repository will grow the Python equivalent incrementally.
+
+## Installation
+
+Once published on PyPI:
+
+```bash
+pip install pdft
+```
+
+From source:
+
+```bash
+git clone https://github.com/zazabap/pdft.git
+cd pdft
+pip install -e ".[dev]"
+```
+
+## Quick start
+
+_(coming soon — mirrors the `make example` demo in the upstream Julia
+package)_
+
+## Background
+
+See the upstream notes for the theory:
+- [`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)
+
+## License
+
+MIT. See [LICENSE](LICENSE). This project is a derivative port of
+ParametricDFT.jl (Copyright © 2025 nzy1997, MIT).

pdft-0.2.1/README.md

@@ -0,0 +1,41 @@
+# pdft
+
+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.
+
+> Status: early scaffold. The Julia package is the reference implementation;
+> this repository will grow the Python equivalent incrementally.
+
+## Installation
+
+Once published on PyPI:
+
+```bash
+pip install pdft
+```
+
+From source:
+
+```bash
+git clone https://github.com/zazabap/pdft.git
+cd pdft
+pip install -e ".[dev]"
+```
+
+## Quick start
+
+_(coming soon — mirrors the `make example` demo in the upstream Julia
+package)_
+
+## Background
+
+See the upstream notes for the theory:
+- [`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)
+
+## License
+
+MIT. See [LICENSE](LICENSE). This project is a derivative port of
+ParametricDFT.jl (Copyright © 2025 nzy1997, MIT).

pdft-0.2.1/examples/basis_demo.py

@@ -0,0 +1,49 @@
+"""basis_demo.py — train a QFTBasis on a random 4x4 image and plot the loss.
+
+Run: python examples/basis_demo.py
+Requires: pdft + pdft[plot] extra.
+"""
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+
+import pdft
+from pdft.viz import TrainingHistory, plot_training_loss
+
+
+def main(out_dir: str | Path = "out") -> None:
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    target = jax.random.normal(jax.random.PRNGKey(7), (4, 4)).astype(jnp.complex128)
+    basis = pdft.QFTBasis(m=2, n=2)
+
+    print(f"Training QFTBasis on a 4x4 target image for 50 steps with RiemannianGD...")
+    t0 = time.perf_counter()
+    result = pdft.train_basis(
+        basis,
+        target=target,
+        loss=pdft.L1Norm(),
+        optimizer=pdft.RiemannianGD(lr=0.01),
+        steps=50,
+        seed=0,
+    )
+    elapsed = time.perf_counter() - t0
+    print(f"done in {elapsed:.2f}s")
+    print(f"  initial loss: {result.loss_history[0]:.4f}")
+    print(f"  final   loss: {result.loss_history[-1]:.4f}")
+    print(f"  reduction   : {result.loss_history[0] - result.loss_history[-1]:.4f}")
+
+    hist = TrainingHistory(losses=result.loss_history, label="RiemannianGD")
+    fig_path = out / "basis_demo_loss.png"
+    plot_training_loss(hist, output_path=fig_path, smooth_alpha=0.2,
+                       title="QFTBasis training loss (L1, RiemannianGD, lr=0.01)")
+    print(f"  wrote plot:  {fig_path}")
+
+
+if __name__ == "__main__":
+    main()

pdft-0.2.1/examples/mera_demo.py

@@ -0,0 +1,51 @@
+"""mera_demo.py — compare QFTBasis, EntangledQFTBasis, TEBDBasis, MERABasis on a single target.
+
+Run: python examples/mera_demo.py
+"""
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+
+import pdft
+from pdft.viz import TrainingHistory, plot_training_comparison
+
+
+def main(out_dir: str | Path = "out") -> None:
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    target = jax.random.normal(jax.random.PRNGKey(11), (4, 4)).astype(jnp.complex128)
+
+    histories: list[TrainingHistory] = []
+    for label, basis in [
+        ("QFTBasis", pdft.QFTBasis(m=2, n=2)),
+        ("EntangledQFTBasis", pdft.EntangledQFTBasis(m=2, n=2)),
+        ("TEBDBasis", pdft.TEBDBasis(m=2, n=2)),
+        ("MERABasis", pdft.MERABasis(m=2, n=2)),
+    ]:
+        print(f"Training {label}...")
+        t0 = time.perf_counter()
+        result = pdft.train_basis(
+            basis,
+            target=target,
+            loss=pdft.L1Norm(),
+            optimizer=pdft.RiemannianAdam(lr=0.01),
+            steps=40,
+            seed=0,
+        )
+        elapsed = time.perf_counter() - t0
+        print(f"  {elapsed:.2f}s, initial {result.loss_history[0]:.4f}, final {result.loss_history[-1]:.4f}")
+        histories.append(TrainingHistory(losses=result.loss_history, label=label))
+
+    fig_path = out / "mera_demo.png"
+    plot_training_comparison(histories, output_path=fig_path,
+                             title="Four bases on the same 4x4 target (RiemannianAdam, L1 loss)")
+    print(f"wrote: {fig_path}")
+
+
+if __name__ == "__main__":
+    main()