learn_mode_select_interior.py

"""Recover an *interior* Faust ``nentry`` menu entry by annealing the temperature.

``examples/learnable_structure/learn_mode_select.py`` recovers a menu's *endpoint* entry, which is
well-posed: a Gumbel-softmax sample is a convex mixture of grid values, and an
extreme value is reachable only by putting all mass on that one entry. An
**interior** target is ambiguous — 12 dB is also "half of 6 dB, half of 18 dB",
so an audio loss cannot tell "select entry 2" from "blend entries 1 and 3", and
at a fixed temperature ``argmax`` need not land on the intended entry.

**Annealing the temperature resolves it.** The soft value is
``dot(softmax((logits + gumbel) / tau), grid)``:

* High ``tau`` early — the softmax is diffuse and the value is genuinely a
  mixture, so gradients move the logits into the right *region* of the grid.
* Low ``tau`` late — each sample collapses toward a single grid entry (the
  Gumbel-perturbed argmax). A 6/18 split now samples ``~6`` or ``~18`` on any
  given step, both far from the 12 dB target, so the loss *punishes* the
  mixture and forces the logits to commit to entry 2, whose samples sit on 12.

This example drives ``tau`` on a geometric schedule from ``tau_hi`` to
``tau_lo`` (the menu's ``[tau:learnable]`` value is overwritten each step, so
the schedule wins over any gradient on ``tau``) and shows that annealing lands
on the interior entry while a fixed high temperature does not. Reuses
``src/faustax/dsp/mode_drive.dsp``; grid ``[0, 6, 12, 18]`` dB.

Run:
    uv run python examples/learnable_structure/learn_mode_select_interior.py
"""

from typing import Tuple

import jax
import jax.numpy as jnp
from flax import nnx

from faustax.modules import ModeDrive

SAMPLE_RATE = 44100
NUM_SAMPLES = 2048
GRID = jnp.array([0.0, 6.0, 12.0, 18.0])  # mode_drive.dsp nentry grid, in dB
INTERIOR_INDEX = 2  # 12 dB: a mixture of entries 1 and 3 reproduces it


def _make_target() -> jax.Array:
    """Render the 12 dB preset with a deterministic (hard-argmax) teacher."""
    x = _input()
    teacher = ModeDrive(sample_rate=SAMPLE_RATE, rngs=nnx.Rngs(0), deterministic=True)
    return teacher(x, params={"drive": float(GRID[INTERIOR_INDEX])})


def _input() -> jax.Array:
    """Fixed low-amplitude input; shared by target render and training."""
    return 0.05 * jax.random.normal(jax.random.key(0), (1, NUM_SAMPLES))


def train(
    anneal: bool,
    seed: int = 7,
    num_steps: int = 300,
    lr: float = 8.0,
    tau_hi: float = 3.0,
    tau_lo: float = 0.25,
) -> Tuple[int, float]:
    """Train the menu logits toward the 12 dB target.

    Args:
        anneal: If True, drive ``tau`` geometrically from ``tau_hi`` to
            ``tau_lo``; if False, hold it at ``tau_hi`` throughout.
        seed: Seed for the module's ``nentry`` Gumbel stream.
        num_steps: Number of gradient-descent steps.
        lr: Learning rate for the plain SGD updates.
        tau_hi: Starting (and, when not annealing, fixed) temperature.
        tau_lo: Final temperature when annealing.

    Returns:
        Tuple ``(recovered_index, eval_loss)`` where ``eval_loss`` is the
        hard-argmax (eval-mode) loss of the committed selection.
    """
    x, target = _input(), _make_target()
    model = ModeDrive(sample_rate=SAMPLE_RATE, rngs=nnx.Rngs(0, nentry=seed))
    zone = "fEntry0"
    setattr(model, f"{zone}_logits", nnx.Param(jnp.zeros(GRID.shape[0]).at[0].set(2.0)))

    @nnx.jit
    def step(m: ModeDrive, tau: jax.Array) -> None:
        def loss_fn(mm: ModeDrive) -> jax.Array:
            return jnp.mean((mm(x) - target) ** 2)

        _, grads = nnx.value_and_grad(loss_fn)(m)
        params = nnx.state(m, nnx.Param)
        nnx.update(m, jax.tree.map(lambda p, g: p - lr * g, params, grads))
        # The schedule overrides tau, discarding the gradient step on it.
        setattr(m, f"{zone}_tau", nnx.Param(tau))

    for i in range(num_steps):
        tau = tau_hi * (tau_lo / tau_hi) ** (i / (num_steps - 1)) if anneal else tau_hi
        step(model, jnp.array(tau, dtype=jnp.float32))

    recovered_index = int(jnp.argmax(getattr(model, f"{zone}_logits")[...]))
    model.deterministic = True  # eval mode: hard argmax, no Gumbel noise
    eval_loss = float(jnp.mean((model(x) - target) ** 2))
    return recovered_index, eval_loss


def main() -> Tuple[int, int, float]:
    """Contrast annealed vs fixed temperature on the interior target.

    Returns:
        Tuple ``(annealed_index, true_index, annealed_eval_loss)``.
    """
    annealed_index, annealed_loss = train(anneal=True)
    fixed_index, fixed_loss = train(anneal=False)
    print(
        f"interior target: entry {INTERIOR_INDEX} ({float(GRID[INTERIOR_INDEX]):.0f} dB)\n"
        f"  annealed tau : entry {annealed_index} "
        f"({float(GRID[annealed_index]):.0f} dB)  eval_loss {annealed_loss:.6f}\n"
        f"  fixed tau    : entry {fixed_index} "
        f"({float(GRID[fixed_index]):.0f} dB)  eval_loss {fixed_loss:.6f}"
    )
    return annealed_index, INTERIOR_INDEX, annealed_loss


if __name__ == "__main__":
    main()