diff --git a/HISTORY.rst b/HISTORY.rst index 8b3c2187..7c5745b6 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,6 +2,22 @@ History ======= +Unreleased +---------- + +* New ``pyFDN.system`` module (draft, not yet exported from ``pyFDN``): the + ``pyfdn-system`` v1 JSON document, which holds either an unchanged + ``pyfdn-fdn-build`` or a node ``graph`` for systems a build cannot express. + Includes validation with path-named errors, ``build_to_graph`` / + ``graph_to_build`` and ``save_fdn_system`` / ``load_fdn_system``. + ``save_fdn_system`` stamps ``metadata.created`` and + ``metadata.pyfdn_version`` when they are missing. See + ``docs/design/system-format.md``. +* New ``pyFDN.auxiliary.graph_plot.plot_system_graph`` draws the signal flow of + a system document, a ``GraphNode`` tree or an ``FDNBuild`` without building + a FLAMO model. ``plot_flamo_graph`` now uses the same renderer; its output is + unchanged. + 0.5.0 (2026-09-28) ------------------ diff --git a/docs/design/system-format.md b/docs/design/system-format.md new file mode 100644 index 00000000..bafd8bc9 --- /dev/null +++ b/docs/design/system-format.md @@ -0,0 +1,553 @@ +# `pyfdn-system`: one JSON document for every pyFDN backend + +**Status:** proposal; PR 1 (schema, `pyFDN.system`) sketched; provenance +metadata added; results documents considered and rejected (§9) +**Date:** 2026-09-07, updated 2026-09-27 with review decisions +**Branch:** `docs/system-json-format` +**Supersedes:** nothing. `pyfdn-fdn-build` v2 stays valid and unchanged. + +--- + +## 1. The problem + +pyFDN can render the same feedback delay network three ways: + +| backend | entry point | domain | covers | +|---|---|---|---| +| `process_fdn` | `pyFDN.process_fdn` | NumPy, block-processed time domain | vanilla FDN with three filter hooks | +| `td` graph | `pyFDN.td.{Series,Parallel,Recursion}` + operators | NumPy, block-processed time domain | arbitrary topology, including time-varying and nonlinear | +| FLAMO | `pyFDN.dss_to_flamo`, `pyFDN.build_to_flamo` | torch, FFT domain | arbitrary topology, differentiable | + +`process_fdn` is not really a fourth thing: it drives `td.RecursionState` and +`td.MatrixFIR` under the hood, so it is the vanilla special case of the `td` +graph with a hand-written loop. + +Exactly one of those three is saveable today, and only for vanilla networks. +Everything that is not a vanilla FDN — a Kronecker or otherwise time-varying +feedback matrix, a shimmer FDN, an SDN, a reverberation-enhancement system, a +coupled-room model — exists only as live Python objects. It cannot be written +to disk, handed to a collaborator, cited by a paper, replayed a year later, or +compiled to FAUST. + +The ask is a settings format that carries **both** the baked parameters (what a +renderer needs) and the meta-parameters (what a *re*-generator or re-trainer +needs), and that is reusable across all three backends plus downstream +consumers. + +## 2. What exists today + +### 2.1 In pyFDN + +- **`src/pyFDN/build.py`** — `FDNBuild` and the `pyfdn-fdn-build` v2 JSON + format: `feedback_matrix`, `input_matrix`, `output_matrix`, + `direct_matrix`, `delays`, `sample_rate`, and the three optional SOS hooks + `post_delay`, `post_matrix`, `post_output`. Validated on both write and + read (`fdn_build_to_dict` deliberately round-trips through + `fdn_build_from_dict` before returning). +- **`src/pyFDN/preset.py`** — `FDNPreset` = `metadata` + `design` + `build`. + `design` is already the meta-parameter layer, and is validated against the + real generator vocabularies (`FEEDBACK_MATRIX_TYPES`, `IO_MATRIX_TYPES`, + `DELAY_DISTRIBUTIONS`, `EQ_DESIGNS`). +- **`src/pyFDN/auxiliary/flamo_graph.py`** — `flamo_model_to_nodes` walks a + FLAMO model into a node tree; `extract_build` bakes such a tree down to an + `FDNBuild`, refusing (rather than silently dropping) any hook that does not + bake to one SOS bank. +- **`src/pyFDN/auxiliary/flamo.py`** — `assemble_fdn_core` builds the graph in + the other direction, with the node names `extract_build` keys off. +- **`src/pyFDN/td/`** — no serialisation of any kind. + +So the *shape* of the answer already exists: baked parameters separated from +design information, with a document wrapper carrying metadata. It only covers +one topology. + +### 2.2 In adac + +[`adac`](https://github.com/cucuwritescode/adac) (local checkout: `rt-fdn`) +compiles a FLAMO model to FAUST through a JSON intermediate representation. +Relevant prior art: + +- `adac/codegen/flamo_to_json.py` is a fork of pyFDN's `flamo_model_to_nodes` + with parameter extraction added. adac does **not** depend on pyFDN — its + core is numpy-only — so the traversal is duplicated. +- Its IR has no `format` or `version` field. +- Node types are FLAMO class names (`parallelDelay`, `parallelSOSFilter`), + including a `SOSFilter` → `parallelSOSFilter` alias to absorb FLAMO version + drift. A library implementation detail is therefore part of the wire format. +- Two things it gets right and this proposal keeps: **effective vs. raw + parameters** (serialise `map(param)`, keep the trainable weights separately + under `flamo.param_raw`) and per-node constructor metadata for + reconstruction. +- `json_to_faust` already dispatches on types `flamo_to_json` never emits — + `variableDelay`, `fractionalDelay`, `onePole`, `allpassComb`, `dcBlocker`. + The FAUST side is *already* speaking a broader, backend-neutral vocabulary. + It has simply never been named or written down. + +## 3. Non-goals + +- Replacing `pyfdn-fdn-build`. It stays, byte-compatible, forever. +- A general dataflow IR. This describes the systems pyFDN builds, not + arbitrary DSP graphs. +- Model weights for training checkpoints. A document describes a *system*; + optimiser state belongs in a torch checkpoint. +- Rendered results (impulse responses, EDCs, RT curves). Any result is one + line of pyFDN code away from the system document, so it is not stored (§9). +- A binary format. Human-readable and diffable is a feature; a network with a + 100 000-tap FIR matrix is a legitimate reason to reach for `.npz` instead. + +## 4. Design + +### 4.1 Two tiers, both permanent + +| tier | used when | read by | +|---|---|---| +| `build` (exists) | the system is a vanilla FDN | `process_fdn`, `build_to_flamo`, `build_to_impz`, adac | +| `graph` (new) | anything else | `graph_to_td`, `graph_to_flamo`, adac | + +The flat `build` is what keeps the common case readable, diffable, and +citable — a reviewer can see the whole network in one screen. A graph IR must +not be allowed to eat that. `graph` is the escape hatch for everything the +flat form cannot express. + +A document carries exactly one of the two. `build_to_graph` lowers; `graph_to_build` +lifts and **refuses** a graph that is not vanilla, in the same spirit as +`extract_build` refusing an un-bakeable hook today. + +### 4.2 Document layout + +```json +{ + "format": "pyfdn-system", + "version": 1, + "metadata": {"name": "...", "description": "...", "tags": []}, + "sample_rate": 48000, + "render": {"nfft": 16384, "alias_decay_db": 0.0, "block_size": 64}, + "build": { ... }, + "graph": { ... } +} +``` + +- `format` / `version` are mandatory and checked on load. (adac's IR has + neither, which is the single cheapest thing to fix in this whole proposal.) +- `sample_rate` is top-level and singular. Nothing below it repeats a sample + rate. +- `metadata` is an open object. As in `FDNPreset`, `tags` is a list of + strings when present so catalogues can filter consistently. +- Two `metadata` fields record provenance, and `save_fdn_system` fills them + in when they are missing: + + | field | example | meaning | + |---|---|---| + | `created` | `"2026-09-27T14:03:00Z"` | when the document was first saved; ISO 8601, UTC | + | `pyfdn_version` | `"0.4.2"` | the pyFDN version that wrote it | + + Existing values are kept on re-save, so a document remembers when, and by + which pyFDN, it was first made. `pyfdn_version` also credits the library in + every file it writes. +- Exactly one of `build` / `graph`. + +### 4.3 One `type` namespace, owned by pyFDN + +Every node has a `type` drawn from a single flat namespace. Container types +carry children; leaf types carry `params`. + +**Containers** — `series`, `parallel`, `recursion`. + +**Leaves** — `identity`, `gain`, `matrix`, `matrix_fir`, `delay`, `sos`, +`kronecker_matrix`, `time_varying_kronecker_matrix`, `time_varying_matrix`, +`pitch_shift`, `dc_blocker`, … as the registry grows. A leaf's `params` are the +constructor arguments of the matching `td` operator, minus `fs` (which is the +document's `sample_rate`). + +This is a departure from adac's two-field split (`type` for structure, +`module_type` for leaves). One field is enough: a node has children or it has +params, never both, and a single namespace means a single registry. + +The names are pyFDN's, not FLAMO's. `sos`, not `parallelSOSFilter`. +`delay`, not `parallelDelay`. A FLAMO class rename must not be able to +invalidate documents on disk. + +The registry is **open** (`pyFDN.system.register_node_type`), not an enum: `td.TimeVaryingMatrix` has no FLAMO +*or* FAUST equivalent, and `td.GranularPitchShift` has neither either. Each +backend owns its own map from `type` to its construction, and a backend that +cannot build a type fails loudly, naming the type — never silently drops the +node. (adac's FAUST emitter currently emits a `//warning:` comment for an +unknown type, which produces a plausible-looking `.dsp` that is quietly the +wrong system. That is the failure mode to avoid.) + +### 4.4 `params` is baked, `design` is how it was made + +The `build` / `design` split that `FDNPreset` already has, pushed down one +level and applied per node: + +- **`params`** — the effective values a renderer needs, post-`map`. For a + FLAMO `Matrix(matrix_type="orthogonal")` that is the matrix exponential of + the stored skew-symmetric weights, not the weights. +- **`design`** — how the numbers were chosen, and the raw trainable + parametrisation they came from: `{"type": "orthogonal", "raw": [[...]]}`, + `{"type": "first_order_shelf", "rt": [1.8, 0.4]}`, + `{"type": "coprime", "range_ms": [20, 50], "seed": 42}`. + +A renderer ignores `design` entirely. Re-generation, re-training, and +round-tripping back into a trainable FLAMO model read it. This subsumes +adac's `flamo.param_raw` without pinning the format to FLAMO, and it reuses +the vocabularies `preset.py` already validates against. + +`design` is optional on every node. A node without it is fully renderable and +merely does not remember its provenance — which is exactly the status quo for +a bare `FDNBuild`. + +### 4.5 Units and array conventions + +One canonical form per quantity, chosen so that no backend's convention wins: + +| quantity | wire form | note | +|---|---|---| +| delays | samples, float permitted | FLAMO stores seconds internally; that is a FLAMO detail, converted in `graph_to_flamo`. adac currently emits both `samples` and `samples_fractional`; one field with a float is enough. | +| SOS | `(sections, 6, channels)`, `[b0,b1,b2,a0,a1,a2]` | pyFDN's existing layout. **Do not** normalise `a0` to 1 in the file — that is a FAUST `fi.tf2` concern and belongs in the emitter. | +| gains / matrices | nested lists, row-major, `(out, in)` | matches `FDNBuild` | +| FIR matrices | `(out, in, order)`, `z^-1` convention | matches `process_fdn`'s 3-D `A` | +| frequencies | Hz | | +| times | seconds | | +| levels | dB, sign as pyFDN uses it elsewhere | | + +All arrays are finite. `allow_nan=False` on write, as `save_fdn_build` +already does. + +**Cascaded matrices stay cascaded (decided in review).** A paraunitary matrix +built as a cascade of orthogonal stages and per-channel shifts is stored as that +cascade — a `series` of `matrix` and `delay` nodes — not as its expanded +`matrix_fir`. The expansion is always recoverable from the cascade, the reverse +is not, and a real-time or FAUST implementation wants the cascade. `design` on +the `series` node records how it was generated +(`{"type": "cascaded_paraunitary", "stages": 3, "seed": 1}`). `matrix_fir` +remains for FIR matrices that have no cascade, e.g. trained ones. To allow +this, `delay` accepts zero-length lines; only `graph_to_build` insists on +positive delays, because `FDNBuild` does. + +### 4.6 `render` is not the system + +`nfft` and `alias_decay_db` are FLAMO evaluation settings and are meaningless +in `td`. `block_size` is a `td` scheduling setting and is meaningless in +FLAMO. None of them changes *what system this is* — they change how one +backend approximates it. + +They therefore live in one top-level `render` block, and every field in it is +optional with a documented default. adac currently attaches `nfft` and +`alias_decay_db` to individual nodes, which makes a document unrenderable by +`td` without stripping fields. + +`alias_decay_db` in particular must not leak into `params`: `extract_build` +already takes care to return the *undamped* `A`/`B`/`C`, and a document +records the system, not the numerically-nudged version of it FLAMO evaluated. + +### 4.7 The loop-delay convention (decided) + +`td.Recursion(block_size=n)` inserts `n` samples of delay into its loop, and +today the caller compensates by hand — `examples/example_reverberation_enhancement.py` +shortens its delay lines with `td.Delay(delays - block_size)`. + +If a document stored those pre-shortened lengths, a `td` render and a FLAMO +render of the same file would be different systems. + +**The document stores the true delay lengths. `graph_to_td` performs the +block compensation.** An author never writes `delays - block_size` into a +file, and `td_to_graph` adds the compensation back when reading a live graph +whose `Recursion` has a non-zero `block_size`. This is the one place the +format takes a position on a backend's internals, and it is worth it: without +it the format silently means different things to different readers. + +### 4.8 Node names are part of the contract + +`extract_build` recognises a vanilla FDN by node name, and `assemble_fdn_core` +produces those names. They are normative in a `graph` document, not +decorative — they are what makes `graph_to_build` possible at all: + +``` +parallel(sum_output=true) +├── brA: series +│ ├── input_gain +│ ├── feedback_loop: recursion +│ │ ├── forward: series[ delay, post_delay ] (or a bare delay) +│ │ └── feedback: series[ mixing_matrix, post_matrix ] (or a bare matrix) +│ ├── output_gain +│ └── post_output +└── direct_gain +``` + +The feedback matrix is always named `mixing_matrix` in a document, including +when `post_matrix` is absent (where `assemble_fdn_core` leaves it unnamed on +`fB`). The direct path is always present, since `FDNBuild.D` always is. + +A hook given several modules keeps the existing `post_delay_0`, +`post_delay_1`, … convention from `hook_module`. + +`recursion` uses the field names `forward` / `feedback` in the document. +FLAMO's `fF` / `fB` are accepted as aliases on read for compatibility with +adac's existing IR, and never written. + +## 5. Backend mapping + +Each backend owns one table. Sketch, to be completed during implementation: + +| `type` | `td` | FLAMO | FAUST (adac) | +|---|---|---|---| +| `series` | `td.Series` | `system.Series` | `:` | +| `parallel` | `td.Parallel` | `system.Parallel` | `,` / `:>` | +| `recursion` | `td.Recursion` | `system.Recursion` | `~` | +| `gain` | `td.Gain` | `dsp.Gain` / `dsp.parallelGain` | `*(g)` | +| `matrix` | `td.Gain` | `dsp.Matrix` | matrix expansion | +| `matrix_fir` | `td.MatrixFIR` / `td.MatrixConvolver` | — | — | +| `delay` | `td.Delay` | `dsp.parallelDelay` | `@(n)` / `de.fdelay` | +| `sos` | `td.SOSBank` | `dsp.parallelSOSFilter` | `fi.tf2` | +| `time_varying_matrix` | `td.TimeVaryingMatrix` | — | — | +| `pitch_shift` | `td.PitchShift` | — | — | + +The empty cells are the point. They are not gaps to be filled; they are +systems that backend genuinely cannot render, and it must say so by name. + +## 6. Ownership + +**pyFDN owns the schema and every conversion in and out of it.** Nothing in +`pyFDN` core imports or depends on `adac`. + +- pyFDN: schema, validation, `flamo ↔ graph`, `td ↔ graph`, `build ↔ graph`, + and the packaged presets. +- adac: FAUST emission, stability certificate, JUCE export — reading + `pyfdn-system` documents. Since adac's numpy-only core is a selling point, + it vendors a small reader rather than taking a pyFDN dependency; the schema + is a versioned document format, so vendoring a reader is legitimate in a way + that forking the traversal was not. + +This is the direction the dependency already wants to run: adac's FAUST +dispatcher is the most complete leaf vocabulary that exists anywhere today, +and it is downstream of a traversal it copied from pyFDN. + +## 7. Implementation plan + +Four PRs, each independently useful and independently revertable. + +### PR 1 — the schema (this branch's follow-up) + +`src/pyFDN/system.py`, pure data, no backend imports: + +- `FDNSystemDoc` dataclass (or plain dicts — decide during review) +- `system_to_dict` / `system_from_dict`, `save_fdn_system` / `load_fdn_system` +- validation: `format`/`version`, exactly one of `build`/`graph`, node types + against the registry, container/leaf arity, finite arrays +- `build_to_graph` / `graph_to_build` +- `FDNPreset` accepts `build` **or** `graph`; every existing preset file on + disk stays valid and every existing test keeps passing + +Tests: round-trip each packaged preset through `build → graph → build`; +`graph_to_build` raises on a non-vanilla graph, naming what it could not bake. + +### PR 2 — the `td` backend + +`graph_to_td` / `td_to_graph`. This is the tier that can currently be saved in +no form at all, and it is where the Kronecker / time-varying feedback work +lands. Includes the §4.7 block compensation and its round-trip test. + +Tests: `graph_to_td(td_to_graph(g))` renders sample-identical output to `g` +for the FDN, shimmer, and RES graphs the examples already build. + +### PR 3 — the FLAMO backend + +`flamo_to_graph` / `graph_to_flamo`, generalising `extract_build` and +`build_to_flamo`. adac's `_serialise_leaf` is the reference for the extraction +logic (effective-vs-raw, channel counts, SOS handling), ported to the neutral +type names. `extract_build` becomes `graph_to_build(flamo_to_graph(model))` +and keeps its current signature and error messages. + +Tests: `flamo_to_graph → graph_to_flamo` reproduces the impulse response of +every model `example_train_fdn_to_rir` and `example_train_colorless_FDN` +build; a trained model round-trips with its `design.raw` weights intact and +still trains. + +### PR 4 — adac (separate repo) + +`graph_to_faust` reading the neutral types. adac's dispatcher already nearly +speaks them, so this is mostly a rename plus a reader. `flamo_to_json` stays +as a deprecated shim. + +## 8. Decisions from review, and what is still open + +1. **Channel counts — stored, and checked.** Every node is written with + `"channels": [in, out]`. They are still derived from the parameter shapes + (so they cannot be the source of truth), and a stored value that disagrees + is a load error naming the node's path, e.g. + `/root/brA/feedback_loop: stored channels [6, 6] disagree ...`. This gives + the validation and diagnosability asked for in review without a second way + to define a system. The system-level `[K, N, J]` triplet is not stored: `N` + is not a property of every graph (an SDN or a RES has no single loop size); + `K` and `J` are the root node's `channels`. +2. **Large arrays.** Still open. The `{"$ref": "sidecar.npz#A"}` hook is not + implemented; cascaded storage (§4.5) removes the most likely reason to need + it. +3. **`FDNSystem` name collision — rename agreed.** The sketch names the + document class `FDNSystemDoc`, which does not collide. Renaming the gallery's + `(A, B, C, D)` named tuple (e.g. to `DSSMatrices`, keeping `FDNSystem` as an + alias for one release) is left to its own small PR so it can be reviewed as + the breaking change it is. +4. **`render` — flat.** `nfft`, `alias_decay_db`, `block_size`; any other field + is a load error. Revisit if a name collides. +5. **`FDNPreset`.** Unchanged in the sketch. A `build` document takes an + optional top-level `design` in the preset vocabulary, so a preset converts to + a system document losslessly; `build_to_graph(build, design=preset.design)` + moves those records onto the named nodes. Whether `FDNPreset` should become a + thin alias of `FDNSystemDoc` is for PR 1 review. +6. **Provenance metadata — added.** `metadata.created` and + `metadata.pyfdn_version` (§4.2), stamped by `save_fdn_system`. + +7. **Diagrams — from the description, not a backend.** + `pyFDN.auxiliary.graph_plot.plot_system_graph` draws a system document, a + `GraphNode` tree or an `FDNBuild` without building a FLAMO model or a `td` + graph. `plot_flamo_graph` now uses the same renderer, so both draw the same + picture. + +## 9. Rendered results: considered, not stored + +A separate `pyfdn-results` document (impulse response, EDC, RT per band, echo +density), written next to each system file, was proposed in review for +comparing a trained network against a target, for sharing, and for teaching. + +**Decision: no results format.** A system document is complete: any result is +one line of pyFDN code away from it, e.g. `pyFDN.build_to_impz(doc.as_build(), +n)` today, or the backend-neutral `impulse_response` planned in #256. Storing +results as well would add a second file that can drift from the system it +claims to describe, and a second schema to version. + +What remains of the proposal: + +- **Provenance** (`metadata.created`, `metadata.pyfdn_version`, §4.2) is kept. + It records which pyFDN wrote a file, which is also what a reader needs to + reproduce a result exactly. +- **Viewing** works on the system itself (§8.7): the diagram needs no backend, + and the existing analysis plots take the rendered response. + +--- + +## Appendix A — worked example + +A 6-line FDN with a trained orthogonal feedback matrix, a designed +first-order-shelf absorption, and a time-varying matrix on the feedback path — +i.e. a system that has no representation on disk today. + +```json +{ + "format": "pyfdn-system", + "version": 1, + "metadata": { + "name": "tv-kronecker-N6", + "description": "Time-varying FDN, Kronecker feedback modulation.", + "tags": ["time-varying", "kronecker"] + }, + "sample_rate": 48000, + "render": {"nfft": 16384, "alias_decay_db": 0.0, "block_size": 64}, + "graph": { + "type": "parallel", + "name": "root", + "sum_output": true, + "children": [ + { + "type": "series", + "name": "brA", + "children": [ + { + "type": "gain", + "name": "input_gain", + "params": {"matrix": [[0.408], [0.408], [0.408], [0.408], [0.408], [0.408]]} + }, + { + "type": "recursion", + "name": "feedback_loop", + "forward": { + "type": "series", + "children": [ + { + "type": "delay", + "name": "delay", + "params": {"samples": [593, 743, 929, 1153, 1399, 1699]}, + "design": {"type": "coprime", "range_ms": [12, 36], "seed": 42} + }, + { + "type": "sos", + "name": "post_delay", + "params": {"sos": [[[0.98, "..."], "..."]]}, + "design": {"type": "first_order_shelf", "rt": [1.8, 0.4]} + } + ] + }, + "feedback": { + "type": "series", + "children": [ + { + "type": "matrix", + "name": "mixing_matrix", + "params": {"matrix": [["..."]]}, + "design": {"type": "orthogonal", "raw": [["..."]]} + }, + { + "type": "time_varying_matrix", + "name": "post_matrix", + "params": {"N": 6, "cycles_per_second": 1.2, "amplitude": 0.15, "spread": 0.3} + } + ] + } + }, + { + "type": "gain", + "name": "output_gain", + "params": {"matrix": [[0.408, 0.408, 0.408, 0.408, 0.408, 0.408]]} + } + ] + }, + { + "type": "gain", + "name": "direct_gain", + "params": {"matrix": [[0.0]]} + } + ] + } +} +``` + +Rendering this file: + +- `graph_to_td` builds it in full. +- `graph_to_flamo` raises, naming `time_varying_matrix`. +- `graph_to_build` raises, naming `time_varying_matrix` as the reason it does + not bake to an `FDNBuild`. +- adac's `graph_to_faust` raises, naming `time_varying_matrix`. + +Three of the four fail, all four say exactly why, and none of them silently +emits the wrong system. That is the whole design in one example. + +## Appendix B — the same network as a `build` + +Drop the `time_varying_matrix` node and the identical network is expressible +in the existing flat tier, which is what should be written whenever it is +possible: + +```json +{ + "format": "pyfdn-system", + "version": 1, + "metadata": {"name": "static-N6"}, + "sample_rate": 48000, + "build": { + "format": "pyfdn-fdn-build", + "version": 2, + "feedback_matrix": [["..."]], + "input_matrix": [["..."]], + "output_matrix": [["..."]], + "direct_matrix": [[0.0]], + "delays": [593, 743, 929, 1153, 1399, 1699], + "sample_rate": 48000.0, + "post_delay": [[["..."]]], + "post_matrix": null, + "post_output": null + } +} +``` + +Byte-for-byte the `build` object that `fdn_build_to_dict` writes today. diff --git a/src/pyFDN/__init__.py b/src/pyFDN/__init__.py index 286924da..7edac06b 100644 --- a/src/pyFDN/__init__.py +++ b/src/pyFDN/__init__.py @@ -520,11 +520,15 @@ mimo_rir_eigenvalues_per_frequency, model_response, param, - params, train_fdn, trainable_from_build, trainable_from_preset, ) + + # ``params`` is also the name of the submodule it lives in; importing it + # from ``.train`` lets mypy resolve it to the module under some import + # orders ("already defined" no-redef). + from .train.params import params from .translate.flamo_to_pr import ( FlamoDecompositionForPR, flamo_decompose_for_pr, diff --git a/src/pyFDN/auxiliary/flamo_graph.py b/src/pyFDN/auxiliary/flamo_graph.py index 7d8c807d..196de4b3 100644 --- a/src/pyFDN/auxiliary/flamo_graph.py +++ b/src/pyFDN/auxiliary/flamo_graph.py @@ -430,107 +430,6 @@ def _hook_sos(name: str) -> np.ndarray | None: ) -# --------------------------------------------------------------------------- -# Matplotlib renderer -# --------------------------------------------------------------------------- - -# Layout constants (abstract units; figure is scaled at the end) -_LEAF_H = 0.9 # leaf box height -_GAP_X = 0.6 # horizontal wire length between series elements -_GAP_Y = 0.5 # vertical gap between parallel branches / fF and fB -_NODE_R = 0.13 # radius of sum (+) and pickoff nodes -_PAD = 0.3 # padding inside recursion container -_STUB = 0.35 # split/merge wire stub length in Parallel - -_WIRE = "0.25" -_LOOP_EDGE = "#b05a4a" - -# Leaf box colors by module category: (facecolor, edgecolor) -_BLOCK_COLORS = { - "transform": ("#efefef", "#7a7a7a"), # FFT / iFFT - "gain": ("#dbe9f6", "#3d6d9e"), # gains - "matrix": ("#e8e1f2", "#7b5ea7"), # mixing / feedback matrices - "delay": ("#fdf1d6", "#c08a2d"), # delays - "filter": ("#e1efe3", "#4f8a5e"), # filters / EQs - "other": ("#f5f5f5", "#8c8c8c"), -} - - -def _block_category(mod_type: str) -> str: - t = mod_type.lower() - if "fft" in t: - return "transform" - if "delay" in t: - return "delay" - if any(k in t for k in ("filter", "biquad", "svf", "geq", "peq", "eq", "sos")): - return "filter" - if any(k in t for k in ("matrix", "scattering", "householder", "hadamard")): - return "matrix" - if "gain" in t: - return "gain" - return "other" - - -def _leaf_label(node: dict[str, Any]) -> str: - mod = node.get("module") - mod_type = type(mod).__name__ if mod is not None else "?" - name = node.get("name", "") - if name and name != mod_type and not name.isdigit(): - return f"{name}\n{mod_type}" - return mod_type - - -def _measure(node: dict[str, Any]) -> tuple[float, float, float]: - """ - Bottom-up size pass. Returns (w, h, ay) where ay is the distance from the - top of the bounding box to the signal-flow axis. Cached in node["_size"]. - """ - ntype = node.get("type", "Leaf") - - if ntype == "Shell": - parts = [] - if node.get("input_layer") is not None: - parts.append(node["input_layer"]) - parts.extend(node.get("children") or []) - if node.get("output_layer") is not None: - parts.append(node["output_layer"]) - node["_parts"] = parts - sizes = [_measure(ch) for ch in parts] - w = sum(s[0] for s in sizes) + _GAP_X * max(len(parts) - 1, 0) - ay = max((s[2] for s in sizes), default=_LEAF_H / 2) - below = max((s[1] - s[2] for s in sizes), default=_LEAF_H / 2) - size = (w, ay + below, ay) - - elif ntype == "Series": - sizes = [_measure(ch) for ch in node.get("children") or []] - w = sum(s[0] for s in sizes) + _GAP_X * max(len(sizes) - 1, 0) - ay = max((s[2] for s in sizes), default=_LEAF_H / 2) - below = max((s[1] - s[2] for s in sizes), default=_LEAF_H / 2) - size = (w, ay + below, ay) - - elif ntype == "Parallel": - sizes = [_measure(ch) for ch in node.get("children") or []] - w = max((s[0] for s in sizes), default=1.0) + 2 * (_STUB + _NODE_R * 2) - h = sum(s[1] for s in sizes) + _GAP_Y * max(len(sizes) - 1, 0) - size = (w, h, h / 2) - - elif ntype == "Recursion": - wf, hf, ayf = _measure(node["fF"]) if node.get("fF") else (1.0, _LEAF_H, 0.45) - wb, hb, _ = _measure(node["fB"]) if node.get("fB") else (1.0, _LEAF_H, 0.45) - inner_w = max(wf, wb) - w = _PAD + 2 * _NODE_R + _GAP_X + inner_w + _GAP_X + 2 * _NODE_R + _PAD - h = _PAD + hf + _GAP_Y + hb + _PAD - size = (w, h, _PAD + ayf) - - else: # Leaf - label = _leaf_label(node) - wmax = max(len(line) for line in label.split("\n")) - size = (max(1.1, 0.115 * wmax + 0.35), _LEAF_H, _LEAF_H / 2) - - node["_size"] = size - return size - - def plot_flamo_graph( model: Any, *, @@ -544,7 +443,9 @@ def plot_flamo_graph( Signal flows left to right; only the feedback path of a Recursion flows right to left, drawn below the forward path with a loop back to a sum - node at the forward path's input. + node at the forward path's input. To draw a saved ``pyfdn-system`` + document without building a FLAMO model, use + :func:`pyFDN.auxiliary.graph_plot.plot_system_graph`. Parameters ---------- @@ -563,219 +464,7 @@ def plot_flamo_graph( ------- ax : matplotlib Axes """ - import matplotlib.pyplot as plt - from matplotlib.patches import Circle, FancyArrowPatch, FancyBboxPatch + from pyFDN.auxiliary.graph_plot import plot_node_tree root = flamo_model_to_nodes(model, name=name, include_shell_io=True) - total_w, total_h, root_ay = _measure(root) - - if ax is None: - fig_w = (total_w + 2.0) * scale - fig_h = (total_h + 0.8) * scale - _, ax = plt.subplots(figsize=(fig_w, fig_h)) - - Point = tuple[float, float] - Ports = tuple[Point, Point] - - # y grows downward in layout; flip the axis at the end. - def route( - pts: list[Point], *, arrow: bool = True, color: str = _WIRE, ls: str = "-" - ) -> None: - xs = [p[0] for p in pts] - ys = [p[1] for p in pts] - if len(pts) > 2: - ax.plot(xs[:-1], ys[:-1], color=color, lw=1.2, ls=ls, zorder=1) - if arrow: - patch = FancyArrowPatch( - pts[-2], - pts[-1], - arrowstyle="-|>", - mutation_scale=11, - color=color, - lw=1.2, - linestyle=ls, - shrinkA=0, - shrinkB=0, - zorder=1, - ) - ax.add_patch(patch) - else: - ax.plot(xs[-2:], ys[-2:], color=color, lw=1.2, ls=ls, zorder=1) - - def sum_node(x: float, y: float, color: str = _WIRE) -> None: - ax.add_patch( - Circle( - (x, y), _NODE_R, facecolor="white", edgecolor=color, lw=1.2, zorder=3 - ) - ) - ax.text( - x, - y, - "+", - ha="center", - va="center_baseline", - fontsize=fontsize, - color=color, - zorder=4, - ) - - def pickoff(x: float, y: float, color: str = _WIRE) -> None: - ax.add_patch(Circle((x, y), 0.045, facecolor=color, edgecolor=color, zorder=3)) - - def xpos(x_left: float, w_total: float, off: float, w_el: float, d: int) -> float: - """Left edge of an element at offset `off` from the input side.""" - return x_left + off if d > 0 else x_left + w_total - off - w_el - - def draw(node: dict[str, Any], x: float, y_axis: float, d: int) -> Ports: - """Draw node with left edge x and flow axis y_axis, direction d (+1/-1). - Returns (p_in, p_out) port coordinates.""" - ntype = node.get("type", "Leaf") - w, h, ay = node["_size"] - y_top = y_axis - ay - - if ntype in ("Shell", "Series"): - children = ( - node["_parts"] if ntype == "Shell" else node.get("children") or [] - ) - if not children: - p = ((x, y_axis), (x + w, y_axis)) - return p if d > 0 else (p[1], p[0]) - off = 0.0 - ports = [] - prev_out = None - for ch in children: - cw, _, _ = ch["_size"] - cx = xpos(x, w, off, cw, d) - p_in, p_out = draw(ch, cx, y_axis, d) - if prev_out is not None: - route([prev_out, p_in]) - ports.append((p_in, p_out)) - prev_out = p_out - off += cw + _GAP_X - return ports[0][0], ports[-1][1] - - if ntype == "Parallel": - stub = _STUB + 2 * _NODE_R - x_in = x if d > 0 else x + w - x_merge = x + w - _NODE_R if d > 0 else x + _NODE_R - sum_node(x_merge, y_axis) - pickoff(x_in, y_axis) - cy = y_top - for ch in node.get("children") or []: - cw, chh, cay = ch["_size"] - cx = xpos(x, w, stub + (max(0.0, w - 2 * stub - cw)) / 2, cw, d) - p_in, p_out = draw(ch, cx, cy + cay, d) - route([(x_in, y_axis), (x_in, p_in[1]), p_in]) - route( - [ - p_out, - (x_merge, p_out[1]), - (x_merge, y_axis - _NODE_R * _sign(y_axis - p_out[1])), - ] - ) - cy += chh + _GAP_Y - p = ((x_in, y_axis), (x_merge + _NODE_R * d, y_axis)) - return p - - if ntype == "Recursion": - # dashed container - ax.add_patch( - FancyBboxPatch( - (x, y_top), - w, - h, - boxstyle="round,pad=0.06", - facecolor="none", - edgecolor=_LOOP_EDGE, - lw=1.0, - ls=(0, (4, 3)), - zorder=0, - ) - ) - nname = node.get("name", "") - if nname and not nname.isdigit(): - ax.text( - x + 0.12, - y_top + 0.06, - nname, - ha="left", - va="top", - fontsize=fontsize - 1.5, - color=_LOOP_EDGE, - style="italic", - zorder=3, - ) - fF, fB = node.get("fF"), node.get("fB") - if fF is None: # placeholder; matches the _measure() default size - fF = { - "type": "Leaf", - "name": "fF", - "module": None, - "_size": (1.0, _LEAF_H, _LEAF_H / 2), - } - wf, hf, _ = fF["_size"] - wb, hb, _ = fB["_size"] if fB else (0, 0, 0) - inner_w = max(wf, wb) - x_sum = xpos(x, w, _PAD + _NODE_R, 0, d) - x_pick = xpos(x, w, w - _PAD - _NODE_R, 0, d) - sum_node(x_sum, y_axis, color=_LOOP_EDGE) - pickoff(x_pick, y_axis) - fx = xpos(x, w, _PAD + 2 * _NODE_R + _GAP_X + (inner_w - wf) / 2, wf, d) - f_in, f_out = draw(fF, fx, y_axis, d) - route([(x_sum + _NODE_R * d, y_axis), f_in]) - route([f_out, (x_pick, y_axis)], arrow=False) - if fB is not None: - yb_axis = y_top + _PAD + hf + _GAP_Y + fB["_size"][2] - bx = xpos(x, w, _PAD + 2 * _NODE_R + _GAP_X + (inner_w - wb) / 2, wb, d) - b_in, b_out = draw(fB, bx, yb_axis, -d) - route([(x_pick, y_axis), (x_pick, b_in[1]), b_in], color=_LOOP_EDGE) - route( - [b_out, (x_sum, b_out[1]), (x_sum, y_axis + _NODE_R)], - color=_LOOP_EDGE, - ) - return (x_sum - _NODE_R * d, y_axis), (x_pick, y_axis) - - # Leaf - mod = node.get("module") - face, edge = _BLOCK_COLORS[_block_category(type(mod).__name__)] - ax.add_patch( - FancyBboxPatch( - (x, y_top), - w, - h, - boxstyle="round,pad=0.02", - facecolor=face, - edgecolor=edge, - lw=1.2, - zorder=2, - ) - ) - ax.text( - x + w / 2, - y_axis, - _leaf_label(node), - ha="center", - va="center", - fontsize=fontsize, - zorder=3, - linespacing=1.3, - ) - p = ((x, y_axis), (x + w, y_axis)) - return p if d > 0 else (p[1], p[0]) - - def _sign(v: float) -> float: - return 1.0 if v >= 0 else -1.0 - - x0 = 1.0 - p_in, p_out = draw(root, x0, root_ay, +1) - # model input / output stubs - route([(x0 - 0.8, p_in[1]), p_in]) - route([p_out, (p_out[0] + 0.8, p_out[1])]) - ax.text(x0 - 0.85, p_in[1], "in", ha="right", va="center", fontsize=fontsize) - ax.text(p_out[0] + 0.85, p_out[1], "out", ha="left", va="center", fontsize=fontsize) - - ax.set_xlim(x0 - 1.6, x0 + total_w + 1.6) - ax.set_ylim(total_h + 0.4, -0.4) # inverted: layout y grows downward - ax.set_aspect("equal") - ax.axis("off") - return ax + return plot_node_tree(root, ax=ax, scale=scale, fontsize=fontsize) diff --git a/src/pyFDN/auxiliary/graph_plot.py b/src/pyFDN/auxiliary/graph_plot.py new file mode 100644 index 00000000..82014641 --- /dev/null +++ b/src/pyFDN/auxiliary/graph_plot.py @@ -0,0 +1,457 @@ +"""Signal-flow diagrams of pyFDN systems, independent of any backend. + +:func:`plot_system_graph` draws a ``pyfdn-system`` document, a +:class:`~pyFDN.system.GraphNode` tree or an :class:`~pyFDN.FDNBuild` directly +from its description; no FLAMO model or ``td`` graph is built. +:func:`pyFDN.plot_flamo_graph` uses the same renderer for a live FLAMO model. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from pyFDN.build import FDNBuild + from pyFDN.system import FDNSystemDoc, GraphNode + + +# --------------------------------------------------------------------------- +# Matplotlib renderer +# --------------------------------------------------------------------------- + +# Layout constants (abstract units; figure is scaled at the end) +_LEAF_H = 0.9 # leaf box height +_GAP_X = 0.6 # horizontal wire length between series elements +_GAP_Y = 0.5 # vertical gap between parallel branches / fF and fB +_NODE_R = 0.13 # radius of sum (+) and pickoff nodes +_PAD = 0.3 # padding inside recursion container +_STUB = 0.35 # split/merge wire stub length in Parallel + +_WIRE = "0.25" +_LOOP_EDGE = "#b05a4a" + +# Leaf box colors by module category: (facecolor, edgecolor) +_BLOCK_COLORS = { + "transform": ("#efefef", "#7a7a7a"), # FFT / iFFT + "gain": ("#dbe9f6", "#3d6d9e"), # gains + "matrix": ("#e8e1f2", "#7b5ea7"), # mixing / feedback matrices + "delay": ("#fdf1d6", "#c08a2d"), # delays + "filter": ("#e1efe3", "#4f8a5e"), # filters / EQs + "other": ("#f5f5f5", "#8c8c8c"), +} + + +def _block_category(mod_type: str) -> str: + t = mod_type.lower() + if "fft" in t: + return "transform" + if "delay" in t: + return "delay" + if any(k in t for k in ("filter", "biquad", "svf", "geq", "peq", "eq", "sos")): + return "filter" + if any(k in t for k in ("matrix", "scattering", "householder", "hadamard")): + return "matrix" + if "gain" in t: + return "gain" + return "other" + + +def _leaf_type(node: dict[str, Any]) -> str: + """The leaf's type name: a system ``type`` or the FLAMO class name.""" + if node.get("kind"): + return str(node["kind"]) + mod = node.get("module") + return type(mod).__name__ if mod is not None else "?" + + +def _leaf_label(node: dict[str, Any]) -> str: + mod_type = _leaf_type(node) + name = node.get("name", "") + if name and name != mod_type and not name.isdigit(): + return f"{name}\n{mod_type}" + return mod_type + + +def _measure(node: dict[str, Any]) -> tuple[float, float, float]: + """ + Bottom-up size pass. Returns (w, h, ay) where ay is the distance from the + top of the bounding box to the signal-flow axis. Cached in node["_size"]. + """ + ntype = node.get("type", "Leaf") + + if ntype == "Shell": + parts = [] + if node.get("input_layer") is not None: + parts.append(node["input_layer"]) + parts.extend(node.get("children") or []) + if node.get("output_layer") is not None: + parts.append(node["output_layer"]) + node["_parts"] = parts + sizes = [_measure(ch) for ch in parts] + w = sum(s[0] for s in sizes) + _GAP_X * max(len(parts) - 1, 0) + ay = max((s[2] for s in sizes), default=_LEAF_H / 2) + below = max((s[1] - s[2] for s in sizes), default=_LEAF_H / 2) + size = (w, ay + below, ay) + + elif ntype == "Series": + sizes = [_measure(ch) for ch in node.get("children") or []] + w = sum(s[0] for s in sizes) + _GAP_X * max(len(sizes) - 1, 0) + ay = max((s[2] for s in sizes), default=_LEAF_H / 2) + below = max((s[1] - s[2] for s in sizes), default=_LEAF_H / 2) + size = (w, ay + below, ay) + + elif ntype == "Parallel": + sizes = [_measure(ch) for ch in node.get("children") or []] + w = max((s[0] for s in sizes), default=1.0) + 2 * (_STUB + _NODE_R * 2) + h = sum(s[1] for s in sizes) + _GAP_Y * max(len(sizes) - 1, 0) + size = (w, h, h / 2) + + elif ntype == "Recursion": + wf, hf, ayf = _measure(node["fF"]) if node.get("fF") else (1.0, _LEAF_H, 0.45) + wb, hb, _ = _measure(node["fB"]) if node.get("fB") else (1.0, _LEAF_H, 0.45) + inner_w = max(wf, wb) + w = _PAD + 2 * _NODE_R + _GAP_X + inner_w + _GAP_X + 2 * _NODE_R + _PAD + h = _PAD + hf + _GAP_Y + hb + _PAD + size = (w, h, _PAD + ayf) + + else: # Leaf + label = _leaf_label(node) + wmax = max(len(line) for line in label.split("\n")) + size = (max(1.1, 0.115 * wmax + 0.35), _LEAF_H, _LEAF_H / 2) + + node["_size"] = size + return size + + +def plot_node_tree( + root: dict[str, Any], + *, + ax: Any = None, + scale: float = 0.85, + fontsize: float = 9.0, +) -> Any: + """Draw a node tree as a signal-flow diagram with matplotlib. + + ``root`` is a tree of plain dicts as :func:`pyFDN.flamo_model_to_nodes` + returns: ``type`` is ``"Shell"``, ``"Series"``, ``"Parallel"``, + ``"Recursion"`` or ``"Leaf"``; containers carry ``children`` (a recursion + ``fF`` and ``fB``); a leaf is labelled by its ``name`` and by ``kind``, or + the class name of its ``module`` when there is no ``kind``. A ``Parallel`` + with ``sum_output`` false stacks its outputs instead of summing them. + + Signal flows left to right; only the feedback path of a Recursion flows + right to left, drawn below the forward path with a loop back to a sum + node at the forward path's input. + + Parameters + ---------- + root : dict + Root of the node tree. + ax : matplotlib Axes, optional + Draw into this axes; otherwise a new figure sized to the layout + is created. + scale : float + Inches per layout unit when creating a new figure. + fontsize : float + Base font size for leaf labels. + + Returns + ------- + ax : matplotlib Axes + """ + import matplotlib.pyplot as plt + from matplotlib.patches import Circle, FancyArrowPatch, FancyBboxPatch, Rectangle + + total_w, total_h, root_ay = _measure(root) + + if ax is None: + fig_w = (total_w + 2.0) * scale + fig_h = (total_h + 0.8) * scale + _, ax = plt.subplots(figsize=(fig_w, fig_h)) + + Point = tuple[float, float] + Ports = tuple[Point, Point] + + # y grows downward in layout; flip the axis at the end. + def route( + pts: list[Point], *, arrow: bool = True, color: str = _WIRE, ls: str = "-" + ) -> None: + xs = [p[0] for p in pts] + ys = [p[1] for p in pts] + if len(pts) > 2: + ax.plot(xs[:-1], ys[:-1], color=color, lw=1.2, ls=ls, zorder=1) + if arrow: + patch = FancyArrowPatch( + pts[-2], + pts[-1], + arrowstyle="-|>", + mutation_scale=11, + color=color, + lw=1.2, + linestyle=ls, + shrinkA=0, + shrinkB=0, + zorder=1, + ) + ax.add_patch(patch) + else: + ax.plot(xs[-2:], ys[-2:], color=color, lw=1.2, ls=ls, zorder=1) + + def sum_node(x: float, y: float, color: str = _WIRE) -> None: + ax.add_patch( + Circle( + (x, y), _NODE_R, facecolor="white", edgecolor=color, lw=1.2, zorder=3 + ) + ) + ax.text( + x, + y, + "+", + ha="center", + va="center_baseline", + fontsize=fontsize, + color=color, + zorder=4, + ) + + def stack_node(x: float, y: float, color: str = _WIRE) -> None: + """Where a non-summing Parallel stacks its branch outputs.""" + ax.add_patch( + Rectangle( + (x - 0.04, y - 2 * _NODE_R), + 0.08, + 4 * _NODE_R, + facecolor=color, + edgecolor=color, + zorder=3, + ) + ) + + def pickoff(x: float, y: float, color: str = _WIRE) -> None: + ax.add_patch(Circle((x, y), 0.045, facecolor=color, edgecolor=color, zorder=3)) + + def xpos(x_left: float, w_total: float, off: float, w_el: float, d: int) -> float: + """Left edge of an element at offset `off` from the input side.""" + return x_left + off if d > 0 else x_left + w_total - off - w_el + + def draw(node: dict[str, Any], x: float, y_axis: float, d: int) -> Ports: + """Draw node with left edge x and flow axis y_axis, direction d (+1/-1). + Returns (p_in, p_out) port coordinates.""" + ntype = node.get("type", "Leaf") + w, h, ay = node["_size"] + y_top = y_axis - ay + + if ntype in ("Shell", "Series"): + children = ( + node["_parts"] if ntype == "Shell" else node.get("children") or [] + ) + if not children: + p = ((x, y_axis), (x + w, y_axis)) + return p if d > 0 else (p[1], p[0]) + off = 0.0 + ports = [] + prev_out = None + for ch in children: + cw, _, _ = ch["_size"] + cx = xpos(x, w, off, cw, d) + p_in, p_out = draw(ch, cx, y_axis, d) + if prev_out is not None: + route([prev_out, p_in]) + ports.append((p_in, p_out)) + prev_out = p_out + off += cw + _GAP_X + return ports[0][0], ports[-1][1] + + if ntype == "Parallel": + stub = _STUB + 2 * _NODE_R + x_in = x if d > 0 else x + w + x_merge = x + w - _NODE_R if d > 0 else x + _NODE_R + if node.get("sum_output", True): + sum_node(x_merge, y_axis) + else: + stack_node(x_merge, y_axis) + pickoff(x_in, y_axis) + cy = y_top + for ch in node.get("children") or []: + cw, chh, cay = ch["_size"] + cx = xpos(x, w, stub + (max(0.0, w - 2 * stub - cw)) / 2, cw, d) + p_in, p_out = draw(ch, cx, cy + cay, d) + route([(x_in, y_axis), (x_in, p_in[1]), p_in]) + route( + [ + p_out, + (x_merge, p_out[1]), + (x_merge, y_axis - _NODE_R * _sign(y_axis - p_out[1])), + ] + ) + cy += chh + _GAP_Y + p = ((x_in, y_axis), (x_merge + _NODE_R * d, y_axis)) + return p + + if ntype == "Recursion": + # dashed container + ax.add_patch( + FancyBboxPatch( + (x, y_top), + w, + h, + boxstyle="round,pad=0.06", + facecolor="none", + edgecolor=_LOOP_EDGE, + lw=1.0, + ls=(0, (4, 3)), + zorder=0, + ) + ) + nname = node.get("name", "") + if nname and not nname.isdigit(): + ax.text( + x + 0.12, + y_top + 0.06, + nname, + ha="left", + va="top", + fontsize=fontsize - 1.5, + color=_LOOP_EDGE, + style="italic", + zorder=3, + ) + fF, fB = node.get("fF"), node.get("fB") + if fF is None: # placeholder; matches the _measure() default size + fF = { + "type": "Leaf", + "name": "fF", + "module": None, + "_size": (1.0, _LEAF_H, _LEAF_H / 2), + } + wf, hf, _ = fF["_size"] + wb, hb, _ = fB["_size"] if fB else (0, 0, 0) + inner_w = max(wf, wb) + x_sum = xpos(x, w, _PAD + _NODE_R, 0, d) + x_pick = xpos(x, w, w - _PAD - _NODE_R, 0, d) + sum_node(x_sum, y_axis, color=_LOOP_EDGE) + pickoff(x_pick, y_axis) + fx = xpos(x, w, _PAD + 2 * _NODE_R + _GAP_X + (inner_w - wf) / 2, wf, d) + f_in, f_out = draw(fF, fx, y_axis, d) + route([(x_sum + _NODE_R * d, y_axis), f_in]) + route([f_out, (x_pick, y_axis)], arrow=False) + if fB is not None: + yb_axis = y_top + _PAD + hf + _GAP_Y + fB["_size"][2] + bx = xpos(x, w, _PAD + 2 * _NODE_R + _GAP_X + (inner_w - wb) / 2, wb, d) + b_in, b_out = draw(fB, bx, yb_axis, -d) + route([(x_pick, y_axis), (x_pick, b_in[1]), b_in], color=_LOOP_EDGE) + route( + [b_out, (x_sum, b_out[1]), (x_sum, y_axis + _NODE_R)], + color=_LOOP_EDGE, + ) + return (x_sum - _NODE_R * d, y_axis), (x_pick, y_axis) + + # Leaf + face, edge = _BLOCK_COLORS[_block_category(_leaf_type(node))] + ax.add_patch( + FancyBboxPatch( + (x, y_top), + w, + h, + boxstyle="round,pad=0.02", + facecolor=face, + edgecolor=edge, + lw=1.2, + zorder=2, + ) + ) + ax.text( + x + w / 2, + y_axis, + _leaf_label(node), + ha="center", + va="center", + fontsize=fontsize, + zorder=3, + linespacing=1.3, + ) + p = ((x, y_axis), (x + w, y_axis)) + return p if d > 0 else (p[1], p[0]) + + def _sign(v: float) -> float: + return 1.0 if v >= 0 else -1.0 + + x0 = 1.0 + p_in, p_out = draw(root, x0, root_ay, +1) + # model input / output stubs + route([(x0 - 0.8, p_in[1]), p_in]) + route([p_out, (p_out[0] + 0.8, p_out[1])]) + ax.text(x0 - 0.85, p_in[1], "in", ha="right", va="center", fontsize=fontsize) + ax.text(p_out[0] + 0.85, p_out[1], "out", ha="left", va="center", fontsize=fontsize) + + ax.set_xlim(x0 - 1.6, x0 + total_w + 1.6) + ax.set_ylim(total_h + 0.4, -0.4) # inverted: layout y grows downward + ax.set_aspect("equal") + ax.axis("off") + return ax + + +_NODE_KIND = {"series": "Series", "parallel": "Parallel", "recursion": "Recursion"} + + +def graph_to_plot_nodes(graph: GraphNode, name: str | None = None) -> dict[str, Any]: + """Convert a :class:`~pyFDN.system.GraphNode` tree to the dicts + :func:`plot_node_tree` draws.""" + node: dict[str, Any] = { + "type": _NODE_KIND.get(graph.type, "Leaf"), + "name": graph.name or name or "", + "children": [], + } + if graph.type == "recursion": + forward, feedback = graph.children + node["fF"] = graph_to_plot_nodes(forward, "forward") + node["fB"] = graph_to_plot_nodes(feedback, "feedback") + elif graph.type in _NODE_KIND: + node["children"] = [graph_to_plot_nodes(child) for child in graph.children] + if graph.type == "parallel": + node["sum_output"] = graph.options.get("sum_output", True) + else: + node["kind"] = graph.type + return node + + +def plot_system_graph( + system: FDNSystemDoc | GraphNode | FDNBuild, + *, + ax: Any = None, + scale: float = 0.85, + fontsize: float = 9.0, +) -> Any: + """Draw the signal flow of a system, without choosing a backend. + + Parameters + ---------- + system : FDNSystemDoc, GraphNode or FDNBuild + What to draw. A ``build`` document and a bare build are drawn as the + FDN graph :func:`pyFDN.system.build_to_graph` produces. + ax : matplotlib Axes, optional + Draw into this axes; otherwise a new figure sized to the layout + is created. + scale : float + Inches per layout unit when creating a new figure. + fontsize : float + Base font size for leaf labels. + + Returns + ------- + ax : matplotlib Axes + """ + from pyFDN.build import FDNBuild + from pyFDN.system import FDNSystemDoc, build_to_graph + + if isinstance(system, FDNSystemDoc): + graph = system.as_graph() + elif isinstance(system, FDNBuild): + graph = build_to_graph(system) + else: + graph = system + root = graph_to_plot_nodes(graph) + return plot_node_tree(root, ax=ax, scale=scale, fontsize=fontsize) + + +__all__ = ["graph_to_plot_nodes", "plot_node_tree", "plot_system_graph"] diff --git a/src/pyFDN/system.py b/src/pyFDN/system.py new file mode 100644 index 00000000..79768870 --- /dev/null +++ b/src/pyFDN/system.py @@ -0,0 +1,742 @@ +"""``pyfdn-system``: one JSON document for every pyFDN backend. + +A system document describes a feedback delay network well enough for any +backend to render it, and well enough to re-generate or re-train it. It has two +tiers, and a document carries exactly one of them: + +* ``build`` -- the existing ``pyfdn-fdn-build`` v2 object, unchanged, for a + vanilla FDN. It is what should be written whenever it is possible. +* ``graph`` -- a tree of :class:`GraphNode` for everything the flat form cannot + express: time-varying mixing, nested structures, cascaded FIR matrices. + +:func:`build_to_graph` lowers a build to a graph; :func:`graph_to_build` lifts a +graph back and refuses, by name, any node that does not bake into an +:class:`~pyFDN.FDNBuild`. + +This module is pure data: it imports no backend. Rendering a graph is the job +of the backend converters (``graph_to_td``, ``graph_to_flamo``), which each +raise on a node type they cannot build rather than drop it. + +See ``docs/design/system-format.md`` for the rationale. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field, replace +from datetime import datetime, timezone +from os import PathLike +from pathlib import Path +from typing import Any, Literal + +import numpy as np + +from .build import FDNBuild, fdn_build_from_dict, fdn_build_to_dict + +SYSTEM_FORMAT = "pyfdn-system" +SYSTEM_VERSION = 1 + +ParamKind = Literal["array", "int", "float", "bool", "str", "json"] +"""How a ``params`` entry is stored: an ``array`` is a finite float array; a +``json`` value is copied as is (e.g. a kernel type that is a string or a list). +""" + +CONTAINER_TYPES = ("series", "parallel", "recursion") + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class NodeType: + """How to read, validate and count the channels of one leaf ``type``. + + ``params`` maps each parameter to its kind; ``required`` lists the ones a + node must carry. ``channels`` returns ``(in_channels, out_channels)`` from + the converted parameters and raises ``ValueError`` on a malformed shape. + """ + + params: Mapping[str, ParamKind] + required: tuple[str, ...] + channels: Callable[[Mapping[str, Any]], tuple[int, int]] + + +def _square(n: int) -> tuple[int, int]: + return n, n + + +def _gain_channels(p: Mapping[str, Any]) -> tuple[int, int]: + matrix = p["matrix"] + if matrix.ndim != 2: + raise ValueError("matrix must be 2-D, shape (out, in)") + return matrix.shape[1], matrix.shape[0] + + +def _matrix_channels(p: Mapping[str, Any]) -> tuple[int, int]: + in_ch, out_ch = _gain_channels(p) + if in_ch != out_ch: + raise ValueError("matrix must be square") + return in_ch, out_ch + + +def _matrix_fir_channels(p: Mapping[str, Any]) -> tuple[int, int]: + coeffs = p["coeffs"] + if coeffs.ndim != 3: + raise ValueError("coeffs must be 3-D, shape (out, in, taps)") + return coeffs.shape[1], coeffs.shape[0] + + +def _delay_channels(p: Mapping[str, Any]) -> tuple[int, int]: + samples = p["samples"] + if samples.ndim != 1 or np.any(samples < 0): + raise ValueError("samples must be a 1-D array of non-negative delays") + return _square(samples.size) + + +def _sos_channels(p: Mapping[str, Any]) -> tuple[int, int]: + sos = p["sos"] + if sos.ndim != 3 or sos.shape[1] != 6: + raise ValueError("sos must have shape (sections, 6, channels)") + return _square(sos.shape[2]) + + +def _kronecker_channels(p: Mapping[str, Any]) -> tuple[int, int]: + return _square(1 << np.asarray(p["angles"]).size) + + +def _channels_param(p: Mapping[str, Any]) -> tuple[int, int]: + if p["channels"] <= 0: + raise ValueError("channels must be positive") + return _square(p["channels"]) + + +def _n_param(p: Mapping[str, Any]) -> tuple[int, int]: + if p["N"] <= 0 or p["N"] % 2: + raise ValueError("N must be a positive even integer") + return _square(p["N"]) + + +NODE_TYPES: dict[str, NodeType] = { + "identity": NodeType({"channels": "int"}, ("channels",), _channels_param), + "gain": NodeType({"matrix": "array"}, ("matrix",), _gain_channels), + "matrix": NodeType({"matrix": "array"}, ("matrix",), _matrix_channels), + "matrix_fir": NodeType({"coeffs": "array"}, ("coeffs",), _matrix_fir_channels), + "delay": NodeType({"samples": "array"}, ("samples",), _delay_channels), + "sos": NodeType({"sos": "array"}, ("sos",), _sos_channels), + "kronecker_matrix": NodeType( + {"angles": "array", "kernel_type": "json"}, ("angles",), _kronecker_channels + ), + "time_varying_kronecker_matrix": NodeType( + { + "angles": "array", + "rate": "array", + "depth": "array", + "phase": "array", + "kernel_type": "json", + "waveform": "str", + }, + ("angles",), + _kronecker_channels, + ), + "time_varying_matrix": NodeType( + { + "N": "int", + "cycles_per_second": "float", + "amplitude": "float", + "spread": "float", + }, + ("N", "cycles_per_second", "amplitude", "spread"), + _n_param, + ), + "dc_blocker": NodeType( + { + "channels": "int", + "R": "float", + "correct_loss": "bool", + "env_tau_s": "float", + "gain_tau_s": "float", + "max_gain": "float", + }, + ("channels",), + _channels_param, + ), + "pitch_shift": NodeType( + { + "channels": "int", + "max_delay_samps": "int", + "window_size": "int", + "transpose_cents": "float", + "active_channels": "array", + "min_delay_samps": "int", + }, + ("channels", "max_delay_samps", "window_size", "transpose_cents"), + _channels_param, + ), +} +"""The leaf vocabulary. Names are pyFDN's own, never a backend's class names.""" + + +def register_node_type(name: str, node_type: NodeType) -> None: + """Add a leaf type to the registry. + + The registry is open so that a new operator can be saved before every + backend can render it. Container names and existing types cannot be + replaced: a document on disk must keep meaning what it meant. + """ + if name in CONTAINER_TYPES or name in NODE_TYPES: + raise ValueError(f"node type {name!r} is already defined") + NODE_TYPES[name] = node_type + + +# --------------------------------------------------------------------------- +# Graph nodes +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class GraphNode: + """One node of a ``graph`` document. + + A container (``series``, ``parallel``, ``recursion``) has ``children`` and + no ``params``; a leaf has ``params`` and no ``children``. A ``recursion`` + has exactly two children, ``(forward, feedback)``. + + ``params`` holds the effective values a renderer needs. ``design`` holds how + they were made (generator type, seed, raw trainable weights, target RT); + renderers ignore it. Any node may carry ``design``, including a container: + a cascaded paraunitary matrix is a ``series`` of ``matrix`` and ``delay`` + stages whose ``design`` records the cascade. + + ``options`` holds container settings: ``sum_output`` for ``parallel``, + ``delay_position`` for ``recursion``. + """ + + type: str + name: str | None = None + params: dict[str, Any] = field(default_factory=dict) + design: dict[str, Any] = field(default_factory=dict) + children: tuple[GraphNode, ...] = () + options: dict[str, Any] = field(default_factory=dict) + + @property + def channels(self) -> tuple[int, int]: + """``(in_channels, out_channels)``, derived from the parameters.""" + return node_channels(self) + + def walk(self, path: str = "") -> Iterator[tuple[str, GraphNode]]: + """Yield ``(path, node)`` for this node and every descendant.""" + here = f"{path}/{self.name or self.type}" + yield here, self + for child in self.children: + yield from child.walk(here) + + def find(self, name: str) -> GraphNode | None: + """Return the first node named ``name``, depth first.""" + return next((node for _, node in self.walk() if node.name == name), None) + + +def node_channels(node: GraphNode) -> tuple[int, int]: + """``(in_channels, out_channels)`` of a node, checking that children fit.""" + if node.type == "series": + if not node.children: + raise ValueError("series needs at least one child") + counts = [node_channels(child) for child in node.children] + for i, ((_, out_a), (in_b, _)) in enumerate( + zip(counts, counts[1:], strict=False) + ): + if out_a != in_b: + raise ValueError( + f"series child {i} outputs {out_a} channels but child " + f"{i + 1} takes {in_b}" + ) + return counts[0][0], counts[-1][1] + if node.type == "parallel": + if not node.children: + raise ValueError("parallel needs at least one child") + counts = [node_channels(child) for child in node.children] + if len({c[0] for c in counts}) != 1: + raise ValueError("parallel children must take the same input count") + if node.options.get("sum_output", True): + if len({c[1] for c in counts}) != 1: + raise ValueError("summed parallel children must match output counts") + return counts[0] + return counts[0][0], sum(c[1] for c in counts) + if node.type == "recursion": + if len(node.children) != 2: + raise ValueError("recursion needs exactly (forward, feedback)") + (f_in, f_out), (b_in, b_out) = map(node_channels, node.children) + if f_out != b_in or b_out != f_in: + raise ValueError( + f"recursion forward is {f_in}->{f_out} but feedback is {b_in}->{b_out}" + ) + return f_in, f_out + spec = NODE_TYPES.get(node.type) + if spec is None: + raise ValueError(f"unknown node type {node.type!r}") + return spec.channels(node.params) + + +# --------------------------------------------------------------------------- +# JSON conversion +# --------------------------------------------------------------------------- + + +def _json_value(value: Any) -> Any: + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Mapping): + return {str(k): _json_value(v) for k, v in value.items()} + if isinstance(value, tuple | list): + return [_json_value(v) for v in value] + return value + + +def _read_param(value: Any, kind: ParamKind, where: str) -> Any: + if kind == "array": + array = np.asarray(value, dtype=float) + if not np.isfinite(array).all(): + raise ValueError(f"{where} must contain only finite values") + return array + if kind == "int": + if isinstance(value, bool) or int(value) != value: + raise ValueError(f"{where} must be an integer") + return int(value) + if kind == "float": + number = float(value) + if not np.isfinite(number): + raise ValueError(f"{where} must be finite") + return number + if kind == "bool": + if not isinstance(value, bool): + raise ValueError(f"{where} must be true or false") + return value + if kind == "str": + if not isinstance(value, str): + raise ValueError(f"{where} must be a string") + return value + return value + + +def graph_to_dict(node: GraphNode) -> dict[str, Any]: + """Convert a :class:`GraphNode` tree to JSON-compatible dictionaries. + + Every node is written with its ``channels`` as ``[in, out]``. They are + derivable, and are checked on read, so a hand-edited document that no + longer fits together is reported at the node that broke it. + """ + in_ch, out_ch = node_channels(node) + data: dict[str, Any] = {"type": node.type} + if node.name is not None: + data["name"] = node.name + data["channels"] = [in_ch, out_ch] + data.update(_json_value(node.options)) + if node.type == "recursion": + forward, feedback = node.children + data["forward"] = graph_to_dict(forward) + data["feedback"] = graph_to_dict(feedback) + elif node.type in CONTAINER_TYPES: + data["children"] = [graph_to_dict(child) for child in node.children] + else: + data["params"] = _json_value(node.params) + if node.design: + data["design"] = _json_value(node.design) + return data + + +_OPTIONS = {"parallel": ("sum_output",), "recursion": ("delay_position",)} +_NODE_KEYS = {"type", "name", "channels", "params", "design", "children"} + + +def graph_from_dict(data: Mapping[str, Any], *, path: str = "") -> GraphNode: + """Construct and validate a :class:`GraphNode` tree from parsed JSON. + + Errors name the path of the offending node, e.g. + ``/root/brA/feedback_loop/forward``. + """ + if not isinstance(data, Mapping): + raise ValueError(f"{path or '/'}: node must be an object") + node_type = data.get("type") + name = data.get("name") + where = f"{path}/{name or node_type}" + if not isinstance(node_type, str): + raise ValueError(f"{where}: node needs a string 'type'") + if name is not None and not isinstance(name, str): + raise ValueError(f"{where}: 'name' must be a string") + design = data.get("design", {}) + if not isinstance(design, Mapping): + raise ValueError(f"{where}: 'design' must be an object") + + options = {key: data[key] for key in _OPTIONS.get(node_type, ()) if key in data} + allowed = _NODE_KEYS | set(_OPTIONS.get(node_type, ())) + if node_type == "recursion": + # fF / fB are FLAMO's names, accepted from adac's IR but never written. + forward = data.get("forward", data.get("fF")) + feedback = data.get("feedback", data.get("fB")) + if forward is None or feedback is None: + raise ValueError(f"{where}: recursion needs 'forward' and 'feedback'") + children: tuple[GraphNode, ...] = ( + graph_from_dict(forward, path=where), + graph_from_dict(feedback, path=where), + ) + params: dict[str, Any] = {} + allowed |= {"forward", "feedback", "fF", "fB"} + elif node_type in CONTAINER_TYPES: + raw_children = data.get("children") + if not isinstance(raw_children, list): + raise ValueError(f"{where}: {node_type} needs a 'children' list") + children = tuple(graph_from_dict(child, path=where) for child in raw_children) + params = {} + else: + spec = NODE_TYPES.get(node_type) + if spec is None: + raise ValueError(f"{where}: unknown node type {node_type!r}") + raw_params = data.get("params", {}) + if not isinstance(raw_params, Mapping): + raise ValueError(f"{where}: 'params' must be an object") + unknown = set(raw_params) - set(spec.params) + missing = set(spec.required) - set(raw_params) + if unknown or missing: + detail = [f"unknown {sorted(unknown)}"] if unknown else [] + detail += [f"missing {sorted(missing)}"] if missing else [] + raise ValueError(f"{where}: {node_type} params: {'; '.join(detail)}") + params = { + key: _read_param(value, spec.params[key], f"{where}.params.{key}") + for key, value in raw_params.items() + } + children = () + + extra = set(data) - allowed + if extra: + raise ValueError(f"{where}: unexpected fields {sorted(extra)}") + + node = GraphNode(node_type, name, params, dict(design), children, options) + try: + channels = node_channels(node) + except ValueError as err: + raise ValueError(f"{where}: {err}") from None + stored = data.get("channels") + if stored is not None and list(stored) != list(channels): + raise ValueError( + f"{where}: stored channels {list(stored)} disagree with the " + f"parameters, which give {list(channels)}" + ) + return node + + +# --------------------------------------------------------------------------- +# build <-> graph +# --------------------------------------------------------------------------- + + +_DESIGN_NODE = { + "feedback_matrix": "mixing_matrix", + "input_matrix": "input_gain", + "output_matrix": "output_gain", + "direct_matrix": "direct_gain", + "delays": "delay", +} + + +def build_to_graph( + build: FDNBuild, *, design: Mapping[str, Mapping[str, Any]] | None = None +) -> GraphNode: + """Lower an :class:`~pyFDN.FDNBuild` to the canonical FDN graph. + + The node names are the ones :func:`pyFDN.auxiliary.flamo.assemble_fdn_core` + produces, and they are normative: :func:`graph_to_build` recognises a + vanilla FDN by them. ``design`` takes the per-component dictionaries of an + :class:`~pyFDN.FDNPreset` and attaches each to its node. + """ + # Preset design components are keyed by FDNBuild field, nodes by name. + by_node = { + _DESIGN_NODE.get(key, key): value for key, value in (design or {}).items() + } + + def leaf(node_type: str, name: str, **params: Any) -> GraphNode: + return GraphNode(node_type, name, params, dict(by_node.get(name, {}))) + + def with_hook(first: GraphNode, hook: np.ndarray | None, name: str) -> GraphNode: + if hook is None: + return first + return GraphNode( + "series", children=(first, leaf("sos", name, sos=np.asarray(hook))) + ) + + forward = with_hook( + leaf("delay", "delay", samples=np.asarray(build.delays, dtype=float)), + build.post_delay, + "post_delay", + ) + feedback = with_hook( + leaf("matrix", "mixing_matrix", matrix=np.asarray(build.A)), + build.post_matrix, + "post_matrix", + ) + wet: list[GraphNode] = [ + leaf("gain", "input_gain", matrix=np.asarray(build.B)), + GraphNode("recursion", "feedback_loop", children=(forward, feedback)), + leaf("gain", "output_gain", matrix=np.asarray(build.C)), + ] + if build.post_output is not None: + wet.append(leaf("sos", "post_output", sos=np.asarray(build.post_output))) + return GraphNode( + "parallel", + "root", + children=( + GraphNode("series", "brA", children=tuple(wet)), + leaf("gain", "direct_gain", matrix=np.asarray(build.D)), + ), + options={"sum_output": True}, + ) + + +def _single_leaf(node: GraphNode, want: str, role: str) -> np.ndarray: + if node.type != want: + raise ValueError( + f"{role} is a {node.type!r} node; an FDNBuild needs a single {want!r}" + ) + return next(iter(node.params.values())) + + +def _split_hook( + path: GraphNode, first_name: str, first_type: str, hook: str +) -> tuple[np.ndarray, np.ndarray | None]: + if path.type != "series": + return _single_leaf(path, first_type, first_name), None + by_name = {child.name: child for child in path.children} + extra = [ + c.name or c.type for c in path.children if c.name not in (first_name, hook) + ] + if extra or first_name not in by_name or len(path.children) > 2: + raise ValueError( + f"{first_name} path has nodes an FDNBuild cannot hold: " + f"{', '.join(map(str, extra)) or 'unexpected layout'}" + ) + first = _single_leaf(by_name[first_name], first_type, first_name) + return first, ( + _single_leaf(by_name[hook], "sos", hook) if hook in by_name else None + ) + + +def graph_to_build(graph: GraphNode, fs: float) -> FDNBuild: + """Lift a graph back to an :class:`~pyFDN.FDNBuild`, or say why it cannot. + + Accepts exactly the layout :func:`build_to_graph` writes. Anything else -- + a time-varying mixing stage, a nested structure in a hook, an FIR matrix -- + raises ``ValueError`` naming the node that does not bake. Nothing is + silently dropped. + """ + if graph.type != "parallel" or len(graph.children) != 2: + raise ValueError("an FDNBuild graph is parallel(brA, direct_gain)") + wet, direct = graph.children + D = _single_leaf(direct, "gain", "direct_gain") + if wet.type != "series": + raise ValueError("brA must be a series") + names = [child.name for child in wet.children] + if names[:3] != ["input_gain", "feedback_loop", "output_gain"] or names[3:] not in ( + [], + ["post_output"], + ): + raise ValueError( + "brA must be input_gain, feedback_loop, output_gain[, post_output]; " + f"got {names}" + ) + B = _single_leaf(wet.children[0], "gain", "input_gain") + loop = wet.children[1] + if loop.type != "recursion": + raise ValueError("feedback_loop must be a recursion") + C = _single_leaf(wet.children[2], "gain", "output_gain") + post_output = ( + _single_leaf(wet.children[3], "sos", "post_output") + if len(wet.children) == 4 + else None + ) + forward, feedback = loop.children + delays, post_delay = _split_hook(forward, "delay", "delay", "post_delay") + A, post_matrix = _split_hook(feedback, "mixing_matrix", "matrix", "post_matrix") + if not np.array_equal(delays, np.rint(delays)): + raise ValueError("an FDNBuild needs integer delays") + build = FDNBuild( + A, B, C, D, delays.astype(np.int64), float(fs), + post_delay=post_delay, post_matrix=post_matrix, post_output=post_output, + ) # fmt: skip + fdn_build_from_dict(fdn_build_to_dict(build)) # shape checks, one place + return build + + +# --------------------------------------------------------------------------- +# The document +# --------------------------------------------------------------------------- + + +@dataclass +class FDNSystemDoc: + """A ``pyfdn-system`` document: metadata plus exactly one of build or graph. + + ``render`` holds backend evaluation settings -- ``nfft`` and + ``alias_decay_db`` for FLAMO, ``block_size`` for ``td`` -- which change how a + backend approximates the system, never what the system is. + + ``design`` is the per-component design record of a ``build`` document, in + the vocabulary :class:`~pyFDN.FDNPreset` validates. A ``graph`` document + keeps its design on the nodes instead. + """ + + sample_rate: float + build: FDNBuild | None = None + graph: GraphNode | None = None + metadata: dict[str, Any] = field(default_factory=dict) + render: dict[str, Any] = field(default_factory=dict) + design: dict[str, dict[str, Any]] = field(default_factory=dict) + + def __post_init__(self) -> None: + if (self.build is None) == (self.graph is None): + raise ValueError("a system document carries exactly one of build, graph") + if self.graph is not None and self.design: + raise ValueError("a graph document keeps its design on the nodes") + if self.build is not None and float(self.build.fs) != float(self.sample_rate): + raise ValueError("build.sample_rate disagrees with the document's") + + def as_graph(self) -> GraphNode: + """The system as a graph, lowering a build when necessary.""" + if self.graph is not None: + return self.graph + assert self.build is not None + return build_to_graph(self.build, design=self.design) + + def as_build(self) -> FDNBuild: + """The system as a build; raises if the graph is not a vanilla FDN.""" + if self.build is not None: + return self.build + assert self.graph is not None + return graph_to_build(self.graph, self.sample_rate) + + +RENDER_FIELDS: dict[str, ParamKind] = { + "nfft": "int", + "alias_decay_db": "float", + "block_size": "int", +} + + +def system_to_dict(doc: FDNSystemDoc) -> dict[str, Any]: + """Convert an :class:`FDNSystemDoc` to its JSON-compatible form.""" + data: dict[str, Any] = {"format": SYSTEM_FORMAT, "version": SYSTEM_VERSION} + if doc.metadata: + data["metadata"] = _json_value(doc.metadata) + data["sample_rate"] = float(doc.sample_rate) + if doc.render: + data["render"] = _json_value(doc.render) + if doc.build is not None: + data["build"] = fdn_build_to_dict(doc.build) + if doc.design: + data["design"] = _json_value(doc.design) + else: + assert doc.graph is not None + data["graph"] = graph_to_dict(doc.graph) + system_from_dict(data) + return data + + +def system_from_dict(data: Mapping[str, Any]) -> FDNSystemDoc: + """Construct and validate an :class:`FDNSystemDoc` from parsed JSON.""" + if data.get("format") != SYSTEM_FORMAT: + raise ValueError(f"Expected format '{SYSTEM_FORMAT}'") + if data.get("version") != SYSTEM_VERSION: + raise ValueError(f"Unsupported pyfdn-system version: {data.get('version')!r}") + extra = set(data) - { + "format", "version", "metadata", "sample_rate", "render", "build", "graph", + "design", + } # fmt: skip + if extra: + raise ValueError(f"unexpected top-level fields {sorted(extra)}") + sample_rate = _read_param(data.get("sample_rate"), "float", "sample_rate") + if sample_rate <= 0: + raise ValueError("sample_rate must be positive") + metadata = data.get("metadata", {}) + if not isinstance(metadata, Mapping): + raise ValueError("metadata must be an object") + tags = metadata.get("tags") + if tags is not None and not ( + isinstance(tags, list) and all(isinstance(t, str) for t in tags) + ): + raise ValueError("metadata.tags must be a list of strings") + render_in = data.get("render", {}) + if not isinstance(render_in, Mapping) or set(render_in) - set(RENDER_FIELDS): + raise ValueError(f"render may only hold {', '.join(RENDER_FIELDS)}") + render = { + key: _read_param(value, RENDER_FIELDS[key], f"render.{key}") + for key, value in render_in.items() + } + build = graph = None + if "build" in data: + build = fdn_build_from_dict(data["build"]) + if "graph" in data: + graph = graph_from_dict(data["graph"]) + design = data.get("design", {}) + if design: + from .preset import _design_from_dict + + design = _design_from_dict(design) + return FDNSystemDoc( + sample_rate=sample_rate, + build=build, + graph=graph, + metadata=dict(metadata), + render=render, + design=dict(design), + ) + + +def save_fdn_system(path: str | PathLike[str], doc: FDNSystemDoc) -> None: + """Write a system document as indented, human-readable JSON. + + ``metadata.created`` and ``metadata.pyfdn_version`` are added unless the + document already carries them. + """ + from . import __version__ + + created = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + metadata = {"created": created, "pyfdn_version": __version__, **doc.metadata} + Path(path).write_text( + json.dumps( + system_to_dict(replace(doc, metadata=metadata)), indent=2, allow_nan=False + ) + + "\n", + encoding="utf-8", + ) + + +def load_fdn_system(path: str | PathLike[str]) -> FDNSystemDoc: + """Load a system document from a JSON file.""" + data = json.loads(Path(path).read_text(encoding="utf-8")) + if not isinstance(data, dict): + raise ValueError("pyfdn-system JSON must contain an object") + return system_from_dict(data) + + +__all__ = [ + "CONTAINER_TYPES", + "FDNSystemDoc", + "GraphNode", + "NODE_TYPES", + "NodeType", + "RENDER_FIELDS", + "SYSTEM_FORMAT", + "SYSTEM_VERSION", + "build_to_graph", + "graph_from_dict", + "graph_to_build", + "graph_to_dict", + "load_fdn_system", + "node_channels", + "register_node_type", + "save_fdn_system", + "system_from_dict", + "system_to_dict", +] diff --git a/tests/test_graph_plot.py b/tests/test_graph_plot.py new file mode 100644 index 00000000..70131f31 --- /dev/null +++ b/tests/test_graph_plot.py @@ -0,0 +1,92 @@ +"""Tests for backend-independent signal-flow diagrams (``graph_plot``).""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 +import pytest # noqa: E402 + +from pyFDN.auxiliary.graph_plot import ( # noqa: E402 + graph_to_plot_nodes, + plot_system_graph, +) +from pyFDN.build import FDNBuild # noqa: E402 +from pyFDN.system import FDNSystemDoc, GraphNode, build_to_graph # noqa: E402 + + +def _build() -> FDNBuild: + N = 4 + sos = np.zeros((1, 6, N)) + sos[0, 0] = 0.9 + sos[0, 3] = 1.0 + return FDNBuild( + A=np.eye(N), + B=np.ones((N, 1)), + C=np.ones((1, N)), + D=np.zeros((1, 1)), + delays=np.array([101, 131, 173, 211]), + fs=48000.0, + post_delay=sos, + ) + + +def _texts(ax) -> set[str]: + return {t.get_text() for t in ax.texts} + + +@pytest.mark.parametrize("as_doc", [False, True]) +def test_plot_system_graph_draws_a_build_without_a_backend(as_doc): + build = _build() + system = FDNSystemDoc(sample_rate=build.fs, build=build) if as_doc else build + ax = plot_system_graph(system) + labels = _texts(ax) + assert {"input_gain\ngain", "delay", "post_delay\nsos"} <= labels + assert "mixing_matrix\nmatrix" in labels + assert "feedback_loop" in labels + plt.close(ax.figure) + + +def test_plot_system_graph_draws_types_no_backend_renders(): + tv = GraphNode( + "time_varying_matrix", + "post_matrix", + {"N": 4, "cycles_per_second": 1.2, "amplitude": 0.1, "spread": 0.2}, + ) + ax = plot_system_graph(GraphNode("series", children=(tv,))) + assert "post_matrix\ntime_varying_matrix" in _texts(ax) + plt.close(ax.figure) + + +def test_graph_to_plot_nodes_keeps_structure(): + root = graph_to_plot_nodes(build_to_graph(_build())) + assert root["type"] == "Parallel" and root["sum_output"] is True + wet, direct = root["children"] + assert direct["kind"] == "gain" + loop = wet["children"][1] + assert loop["type"] == "Recursion" + assert [c["name"] for c in loop["fF"]["children"]] == ["delay", "post_delay"] + assert loop["fB"]["kind"] == "matrix" + + +def test_non_summing_parallel_is_drawn_as_a_stack(): + branch = GraphNode("gain", "g", {"matrix": np.ones((1, 1))}) + summed = plot_system_graph(GraphNode("parallel", children=(branch, branch))) + stacked = plot_system_graph( + GraphNode("parallel", children=(branch, branch), options={"sum_output": False}) + ) + assert "+" in _texts(summed) + assert "+" not in _texts(stacked) + plt.close("all") + + +def test_plot_flamo_graph_uses_the_same_renderer(): + pytest.importorskip("flamo") + from pyFDN import build_to_flamo, plot_flamo_graph + + ax = plot_flamo_graph(build_to_flamo(_build(), nfft=256)) + assert "feedback_loop" in _texts(ax) + plt.close(ax.figure) diff --git a/tests/test_system.py b/tests/test_system.py new file mode 100644 index 00000000..a3718b93 --- /dev/null +++ b/tests/test_system.py @@ -0,0 +1,220 @@ +"""Tests for the ``pyfdn-system`` document format (``pyFDN.system``).""" + +from __future__ import annotations + +import copy +import json + +import numpy as np +import pytest + +import pyFDN +from pyFDN.build import FDNBuild, fdn_build_to_dict +from pyFDN.system import ( + FDNSystemDoc, + GraphNode, + build_to_graph, + graph_from_dict, + graph_to_build, + graph_to_dict, + load_fdn_system, + save_fdn_system, + system_from_dict, + system_to_dict, +) + + +def _build(hooks: bool = True) -> FDNBuild: + rng = np.random.default_rng(0) + N = 4 + sos = np.zeros((1, 6, N)) + sos[0, 0] = 0.9 + sos[0, 3] = 1.0 + return FDNBuild( + A=np.linalg.qr(rng.standard_normal((N, N)))[0], + B=np.ones((N, 1)), + C=np.ones((2, N)), + D=np.zeros((2, 1)), + delays=np.array([101, 131, 173, 211]), + fs=48000.0, + post_delay=sos if hooks else None, + post_matrix=sos if hooks else None, + post_output=sos[:, :, :2] if hooks else None, + ) + + +def _assert_builds_equal(a: FDNBuild, b: FDNBuild) -> None: + assert fdn_build_to_dict(a) == fdn_build_to_dict(b) + + +@pytest.mark.parametrize("hooks", [False, True]) +def test_build_graph_build_round_trip(hooks): + build = _build(hooks) + graph = build_to_graph(build) + assert graph.channels == (1, 2) + _assert_builds_equal(graph_to_build(graph, build.fs), build) + + +def test_graph_json_round_trip_is_stable(): + data = graph_to_dict(build_to_graph(_build())) + again = graph_to_dict(graph_from_dict(json.loads(json.dumps(data)))) + assert again == data + assert data["name"] == "root" + assert data["children"][0]["children"][1]["forward"]["children"][1]["name"] == ( + "post_delay" + ) + + +@pytest.mark.parametrize("name", pyFDN.available_fdn_presets()) +def test_packaged_presets_round_trip_through_graph(name): + preset = pyFDN.get_fdn_preset(name) + graph = build_to_graph(preset.build, design=preset.design) + lifted = graph_to_build(graph_from_dict(graph_to_dict(graph)), preset.build.fs) + _assert_builds_equal(lifted, preset.build) + + +def test_preset_design_lands_on_nodes(): + graph = build_to_graph(_build(), design={"feedback_matrix": {"type": "orthogonal"}}) + node = graph.find("mixing_matrix") + assert node is not None and node.design == {"type": "orthogonal"} + + +def test_graph_to_build_refuses_time_varying_mixing_by_name(): + graph = build_to_graph(_build(hooks=False)) + wet, direct = graph.children + loop = wet.children[1] + tv = GraphNode( + "time_varying_matrix", + "post_matrix", + {"N": 4, "cycles_per_second": 1.2, "amplitude": 0.1, "spread": 0.2}, + ) + feedback = GraphNode("series", children=(loop.children[1], tv)) + loop = GraphNode( + "recursion", "feedback_loop", children=(loop.children[0], feedback) + ) + wet = GraphNode("series", "brA", children=(wet.children[0], loop, wet.children[2])) + graph = GraphNode("parallel", "root", children=(wet, direct), options=graph.options) + + assert graph.channels == (1, 2) + with pytest.raises(ValueError, match="time_varying_matrix"): + graph_to_build(graph, 48000.0) + + +def test_cascaded_matrix_is_a_series_of_stages(): + # A two-stage paraunitary cascade, stored as cascaded rather than as FIR. + rng = np.random.default_rng(1) + stage = lambda: GraphNode( # noqa: E731 + "matrix", params={"matrix": np.linalg.qr(rng.standard_normal((4, 4)))[0]} + ) + shift = GraphNode("delay", params={"samples": np.array([0.0, 1, 2, 3])}) + cascade = GraphNode( + "series", + "mixing_matrix", + design={"type": "cascaded_paraunitary", "stages": 2, "seed": 1}, + children=(stage(), shift, stage()), + ) + data = graph_to_dict(cascade) + assert data["design"]["type"] == "cascaded_paraunitary" + assert graph_from_dict(data).channels == (4, 4) + + +def test_stored_channels_are_checked(): + data = graph_to_dict(build_to_graph(_build())) + data["children"][1]["channels"] = [1, 3] + with pytest.raises(ValueError, match="/root/direct_gain.*disagree"): + graph_from_dict(data) + + +def test_mismatched_series_names_the_node(): + data = graph_to_dict(build_to_graph(_build())) + data["children"][0]["children"][0]["params"]["matrix"] = [[1.0]] * 3 + with pytest.raises(ValueError, match="/root/brA"): + graph_from_dict(data) + + +def test_unknown_type_is_refused_by_name(): + with pytest.raises(ValueError, match="parallelSOSFilter"): + graph_from_dict({"type": "parallelSOSFilter", "params": {}}) + + +def test_flamo_recursion_aliases_are_read_but_not_written(): + data = graph_to_dict(build_to_graph(_build(hooks=False))) + loop = data["children"][0]["children"][1] + loop["fF"] = loop.pop("forward") + loop["fB"] = loop.pop("feedback") + written = graph_to_dict(graph_from_dict(data)) + assert "forward" in written["children"][0]["children"][1] + assert "fF" not in json.dumps(written) + + +def test_document_build_tier_embeds_build_unchanged(tmp_path): + build = _build() + doc = FDNSystemDoc( + sample_rate=build.fs, + build=build, + metadata={"name": "static-N4", "tags": ["test"]}, + render={"nfft": 16384, "block_size": 64}, + design={"feedback_matrix": {"type": "orthogonal"}}, + ) + data = system_to_dict(doc) + assert data["build"] == fdn_build_to_dict(build) + + path = tmp_path / "static.json" + save_fdn_system(path, doc) + loaded = load_fdn_system(path) + _assert_builds_equal(loaded.as_build(), build) + assert loaded.render == {"nfft": 16384, "block_size": 64} + assert loaded.as_graph().find("mixing_matrix").design == {"type": "orthogonal"} + + +def test_document_graph_tier_round_trips(): + build = _build() + doc = FDNSystemDoc(sample_rate=build.fs, graph=build_to_graph(build)) + loaded = system_from_dict(json.loads(json.dumps(system_to_dict(doc)))) + _assert_builds_equal(loaded.as_build(), build) + + +@pytest.mark.parametrize( + "mutate, match", + [ + (lambda d: d.update(format="adac"), "format"), + (lambda d: d.update(version=2), "version"), + (lambda d: d.update(graph={"type": "identity", "params": {"channels": 1}}), + "exactly one"), + (lambda d: d.update(render={"nfft": 1024, "hop": 2}), "render"), + (lambda d: d.update(metadata={"tags": "reverb"}), "tags"), + ], +) # fmt: skip +def test_document_validation(mutate, match): + build = _build(hooks=False) + data = system_to_dict(FDNSystemDoc(sample_rate=build.fs, build=build)) + data = copy.deepcopy(data) + mutate(data) + with pytest.raises(ValueError, match=match): + system_from_dict(data) + + +def test_save_stamps_created_and_pyfdn_version(tmp_path): + build = _build(hooks=False) + doc = FDNSystemDoc(sample_rate=build.fs, build=build, metadata={"name": "N4"}) + path = tmp_path / "n4.json" + save_fdn_system(path, doc) + + metadata = load_fdn_system(path).metadata + assert metadata["name"] == "N4" + assert metadata["pyfdn_version"] == pyFDN.__version__ + assert metadata["created"].endswith("Z") + assert doc.metadata == {"name": "N4"} + + # Re-saving keeps when, and by which pyFDN, the document was first made. + stamped = {"created": "2026-09-27T14:03:00Z", "pyfdn_version": "0.4.1"} + save_fdn_system( + path, FDNSystemDoc(sample_rate=build.fs, build=build, metadata=stamped) + ) + assert load_fdn_system(path).metadata == stamped + + +def test_non_finite_params_are_refused(): + node = {"type": "gain", "params": {"matrix": [[float("nan")]]}} + with pytest.raises(ValueError, match="finite"): + graph_from_dict(node)