src/hamiltonzero/compiled/model.py

# Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import Any

import equinox as eqx
import jax
import jax.numpy as jnp
from jaxtyping import Array

from hamiltonzero.model import (
    normalize_leaf_carriers,
    quadrilinear_merge,
    tagged_dense_no_bias,
)
from hamiltonzero.optim import register_scale_and_shift

from .execute import (
    _compiled_quadrilinear_merge,
    _factorized_apply,
    _gate_reference,
    _single,
)
from .types import EMPTY, MERGE, OPCODE_DTYPE, CompiledTree, SharedKernel, level_widths


def _scale_tag(y_flat, x_flat, scale_param, *, tag_id: str):
    return register_scale_and_shift(
        y_flat,
        x_flat,
        scale=scale_param,
        tag_id=tag_id,
    )


class CompiledFinetuneWaveFunction(eqx.Module):
    kernel: SharedKernel
    leaf_h: Array
    merge_h: Array
    readout_h: Array
    perm: Array
    inv_perm: Array
    leaf_real: Array
    opcodes: Array
    n_sites: int = eqx.field(static=True)
    r_leaf: int = eqx.field(static=True)
    r_merge: int = eqx.field(static=True)

    def leaf_h_rows(self) -> Array:
        return self.leaf_h.reshape(self.n_sites, self.r_leaf)

    def merge_h_level(self, level: int, width: int) -> Array:
        return self.merge_h[level, : width * self.r_merge].reshape(width, self.r_merge)

    @classmethod
    def from_compiled(cls, kernel: SharedKernel, tree: CompiledTree):
        leaf_h = _single(tree.leaf_h, "leaf conditioner")
        readout_h = _single(tree.readout_h, "readout conditioner")
        n_sites = int(tree.perm.shape[-1])
        widths = level_widths(n_sites)
        w_max = widths[0]
        r_leaf = int(leaf_h.shape[-1])
        r_merge = int(tree.merge_h[0].shape[-1])
        merge_rows = []
        opcode_rows = []
        for width, h_l, ops_l in zip(widths, tree.merge_h, tree.opcodes, strict=True):
            if h_l.shape[-2] != width or ops_l.shape[-1] != width:
                raise ValueError(
                    f"level width mismatch: expected {width}, got "
                    f"{h_l.shape[-2]}/{ops_l.shape[-1]}"
                )
            pad = w_max - width
            merge_rows.append(
                jnp.pad(h_l, ((0, pad), (0, 0))).reshape(-1) if pad else h_l.reshape(-1)
            )
            opcode_rows.append(
                jnp.pad(ops_l, (0, pad), constant_values=EMPTY) if pad else ops_l
            )
        return cls(
            kernel=kernel,
            leaf_h=leaf_h.reshape(-1),
            merge_h=jnp.stack(merge_rows),
            readout_h=readout_h,
            perm=tree.perm,
            inv_perm=tree.inv_perm,
            leaf_real=tree.leaf_real,
            opcodes=jnp.stack(opcode_rows).astype(OPCODE_DTYPE),
            n_sites=n_sites,
            r_leaf=r_leaf,
            r_merge=r_merge,
        )

    def as_compiled_tree(self) -> CompiledTree:
        widths = level_widths(self.n_sites)
        return CompiledTree(
            perm=self.perm,
            inv_perm=self.inv_perm,
            leaf_real=self.leaf_real,
            leaf_h=(self.leaf_h_rows(),),
            leaf_combiner_h=(),
            merge_h=tuple(
                self.merge_h_level(i, width) for i, width in enumerate(widths)
            ),
            opcodes=tuple(self.opcodes[i, :width] for i, width in enumerate(widths)),
            readout_h=(self.readout_h,),
            readout_combiner_h=(),
        )

    def route_q(self, q: Array) -> Array:
        return jnp.take(q, self.perm, axis=-2)

    def param_counts(self) -> dict:
        kernel_leaves = {
            "q_to_odd": self.kernel.q_to_odd.weight.size,
            "leaf_V": self.kernel.leaf_factors[0].V.size,
            "leaf_U": self.kernel.leaf_factors[0].U.size,
            "merge_T": self.kernel.merge_T.size,
            "merge_V": self.kernel.merge_factors[0].V.size,
            "merge_U": self.kernel.merge_factors[0].U.size,
            "readout_V": self.kernel.readout_factors[0].V.size,
            "readout_U": self.kernel.readout_factors[0].U.size,
        }
        tree_leaves = {
            "leaf_h": self.leaf_h.size,
            "merge_h": self.merge_h.size,
            "readout_h": self.readout_h.size,
        }
        return {
            "kernel": kernel_leaves,
            "tree": tree_leaves,
            "kernel_total": sum(kernel_leaves.values()),
            "tree_total": sum(tree_leaves.values()),
            "total": sum(kernel_leaves.values()) + sum(tree_leaves.values()),
        }

    def __call__(self, q: Array, ctx: Any = None, t: Any = 0.0):
        del ctx, t
        return self._forward_plain(q)

    def _forward_plain(self, q: Array):
        kernel = self.kernel
        odd_dtype = jnp.float32
        q_c = q if q.dtype == odd_dtype else q.astype(odd_dtype)
        weight = kernel.q_to_odd.weight
        weight = weight if weight.dtype == odd_dtype else weight.astype(odd_dtype)
        z = q_c @ weight
        u_raw = _factorized_apply(kernel.leaf_factors[0], self.leaf_h_rows(), z)
        u, log_rms = normalize_leaf_carriers(u_raw)
        s = jnp.zeros(u.shape[:-1], dtype=odd_dtype) + log_rms.astype(odd_dtype)
        widths = level_widths(self.n_sites)
        for level, width in enumerate(widths):
            h_level = self.merge_h_level(level, width)
            opcodes = self.opcodes[level, :width]
            u_left, u_right = u[..., 0::2, :], u[..., 1::2, :]
            s_left, s_right = s[..., 0::2], s[..., 1::2]
            raw = _compiled_quadrilinear_merge(kernel.merge_T, u_left, u_right)
            out = raw + _factorized_apply(kernel.merge_factors[0], h_level, raw)
            scale = jnp.sqrt(jnp.mean(out * out, axis=-1) + kernel.merge_eps)
            candidate_u = out / scale[..., None]
            candidate_s = s_left + s_right + jnp.log(scale)
            u = _gate_reference(
                candidate_u, u_left, u_right, opcodes, feature_axis=True
            )
            s = _gate_reference(
                candidate_s, s_left, s_right, opcodes, feature_axis=False
            )
        return self._readout(
            _factorized_apply(kernel.readout_factors[0], self.readout_h, u[..., 0, :]),
            s[..., 0],
        )

    def _readout(self, psi: Array, s_root: Array):
        psi_re, psi_im = psi[..., 0], psi[..., 1]
        log_abs = 0.5 * jnp.log(psi_re * psi_re + psi_im * psi_im) + s_root
        return log_abs, jnp.arctan2(psi_im, psi_re)

    def call_tagged(self, q: Array, ctx: Any = None, t: Any = 0.0):
        del ctx, t
        if q.ndim != 2:
            raise ValueError(
                f"call_tagged is per-walker: expected q [P, 4], got {q.shape}"
            )
        kernel = self.kernel
        odd_dtype = jnp.float32
        q_c = q if q.dtype == odd_dtype else q.astype(odd_dtype)
        real = self.leaf_real
        z = tagged_dense_no_bias(
            kernel.q_to_odd.weight,
            q_c,
            tag_id="compiled.q_to_odd",
            pathway="odd",
            kfac_structural_mask=real,
            kfac_scan_shared=False,
            kfac_repeat_ndim=1,
        )
        leaf = kernel.leaf_factors[0]
        vz = tagged_dense_no_bias(
            leaf.V,
            z,
            tag_id="compiled.leaf.V",
            pathway="odd",
            kfac_structural_mask=real,
            kfac_scan_shared=False,
            kfac_repeat_ndim=1,
        )
        vz_flat = vz.reshape(-1)
        mixed = _scale_tag(
            vz_flat * self.leaf_h,
            vz_flat,
            self.leaf_h,
            tag_id="compiled.leaf.h",
        ).reshape(vz.shape)
        u_raw = tagged_dense_no_bias(
            leaf.U,
            mixed,
            tag_id="compiled.leaf.U",
            pathway="odd",
            kfac_structural_mask=real,
            kfac_scan_shared=False,
            kfac_repeat_ndim=1,
        )
        u, log_rms = normalize_leaf_carriers(u_raw)
        s = jnp.zeros(u.shape[:-1], dtype=odd_dtype) + log_rms.astype(odd_dtype)
        merge = kernel.merge_factors[0]
        merge_T = kernel.merge_T
        merge_eps = kernel.merge_eps

        def level_body(carry, xs):
            u_buffer, s_buffer = carry
            h_level, opcodes = xs
            u_left, u_right = u_buffer[0::2], u_buffer[1::2]
            s_left, s_right = s_buffer[0::2], s_buffer[1::2]
            merge_rows = opcodes == MERGE
            raw = quadrilinear_merge(
                merge_T,
                u_left,
                u_right,
                tag_id="compiled.merge.T",
                pathway="odd",
                kfac_structural_mask=merge_rows,
                kfac_scan_shared=True,
                kfac_repeat_ndim=1,
            )
            vx = tagged_dense_no_bias(
                merge.V,
                raw,
                tag_id="compiled.merge.V",
                pathway="odd",
                kfac_structural_mask=merge_rows,
                kfac_scan_shared=True,
                kfac_repeat_ndim=1,
            )
            vx_flat = vx.reshape(-1)
            mixed_level = _scale_tag(
                vx_flat * h_level,
                vx_flat,
                h_level,
                tag_id="compiled.merge.h",
            ).reshape(vx.shape)
            correction = tagged_dense_no_bias(
                merge.U,
                mixed_level,
                tag_id="compiled.merge.U",
                pathway="odd",
                kfac_structural_mask=merge_rows,
                kfac_scan_shared=True,
                kfac_repeat_ndim=1,
            )
            out = raw + correction
            scale = jnp.sqrt(jnp.mean(out * out, axis=-1) + merge_eps)
            candidate_u = out / scale[..., None]
            candidate_s = s_left + s_right + jnp.log(scale)
            u_next = _gate_reference(
                candidate_u, u_left, u_right, opcodes, feature_axis=True
            )
            s_next = _gate_reference(
                candidate_s, s_left, s_right, opcodes, feature_axis=False
            )
            return (
                jnp.concatenate([u_next, jnp.zeros_like(u_next)], axis=0),
                jnp.concatenate([s_next, jnp.zeros_like(s_next)], axis=0),
            ), None

        (u_buffer, s_buffer), _ = jax.lax.scan(
            level_body,
            (u, s),
            (self.merge_h, self.opcodes),
        )
        readout = kernel.readout_factors[0]
        vr = tagged_dense_no_bias(
            readout.V,
            u_buffer[0],
            tag_id="compiled.readout.V",
            pathway="odd",
        )
        mixed_readout = _scale_tag(
            vr * self.readout_h,
            vr,
            self.readout_h,
            tag_id="compiled.readout.h",
        )
        psi = tagged_dense_no_bias(
            readout.U,
            mixed_readout,
            tag_id="compiled.readout.U",
            pathway="odd",
        )
        return self._readout(psi, s_buffer[0])


def _expand_stage(V, U, h_2d, new_rank: int, key):
    old_rank = V.shape[-1]
    if new_rank < old_rank:
        raise ValueError(f"cannot shrink rank {old_rank} -> {new_rank}")
    if new_rank == old_rank:
        return V, U, h_2d
    extra = new_rank - old_rank
    key_u, key_h = jax.random.split(key)
    v_new = jnp.zeros((*V.shape[:-1], extra), dtype=V.dtype)
    u_new = jnp.std(U) * jax.random.normal(key_u, (extra, *U.shape[1:]), dtype=U.dtype)
    h_new = jnp.std(h_2d) * jax.random.normal(
        key_h, (*h_2d.shape[:-1], extra), dtype=h_2d.dtype
    )
    return (
        jnp.concatenate([V, v_new], axis=-1),
        jnp.concatenate([U, u_new], axis=0),
        jnp.concatenate([h_2d, h_new], axis=-1),
    )


def expand_rank(
    model: CompiledFinetuneWaveFunction,
    *,
    leaf_rank: int,
    merge_rank: int,
    key,
) -> CompiledFinetuneWaveFunction:
    key_leaf, key_merge = jax.random.split(jnp.asarray(key), 3)[:2]
    kernel = model.kernel
    leaf = kernel.leaf_factors[0]
    merge = kernel.merge_factors[0]
    n_sites = model.n_sites
    max_width = n_sites // 2
    n_levels = model.merge_h.shape[0]
    leaf_h = model.leaf_h.reshape(n_sites, model.r_leaf)
    merge_h = model.merge_h.reshape(n_levels, max_width, model.r_merge)
    readout_h = model.readout_h
    r_leaf, r_merge = model.r_leaf, model.r_merge
    V, U, leaf_h = _expand_stage(leaf.V, leaf.U, leaf_h, int(leaf_rank), key_leaf)
    leaf = eqx.tree_at(lambda factor: (factor.V, factor.U), leaf, (V, U))
    r_leaf = int(leaf_rank)
    V, U, merge_h = _expand_stage(merge.V, merge.U, merge_h, int(merge_rank), key_merge)
    merge = eqx.tree_at(lambda factor: (factor.V, factor.U), merge, (V, U))
    r_merge = int(merge_rank)
    kernel = eqx.tree_at(
        lambda value: (
            value.leaf_factors,
            value.merge_factors,
        ),
        kernel,
        ((leaf,), (merge,)),
    )
    return CompiledFinetuneWaveFunction(
        kernel=kernel,
        leaf_h=leaf_h.reshape(-1),
        merge_h=merge_h.reshape(n_levels, -1),
        readout_h=readout_h,
        perm=model.perm,
        inv_perm=model.inv_perm,
        leaf_real=model.leaf_real,
        opcodes=model.opcodes,
        n_sites=n_sites,
        r_leaf=r_leaf,
        r_merge=r_merge,
    )


def compile_finetune_model(
    eager_model,
    ctx_row,
    *,
    leaf_rank: int,
    merge_rank: int,
    physical_perm,
    key,
) -> CompiledFinetuneWaveFunction:
    from .tree import compile_physical_tree_reference
    from .trunk import bind_shared_kernel, compile_shared_trunk

    shared_trunk = compile_shared_trunk(eager_model, ctx_row)
    n_sites = int(shared_trunk.real_mask.shape[-1])
    identity = jnp.arange(n_sites, dtype=jnp.int32)
    tree = compile_physical_tree_reference(eager_model, shared_trunk, identity)
    model = CompiledFinetuneWaveFunction.from_compiled(
        bind_shared_kernel(eager_model), tree
    )
    model = expand_rank(
        model,
        leaf_rank=leaf_rank,
        merge_rank=merge_rank,
        key=key,
    )
    physical_perm = jnp.asarray(physical_perm, dtype=jnp.int32)
    if physical_perm.shape != (n_sites,):
        raise ValueError(
            f"physical_perm must have shape {(n_sites,)}, got {physical_perm.shape}"
        )
    model = eqx.tree_at(
        lambda value: (value.perm, value.inv_perm),
        model,
        (
            physical_perm,
            jnp.argsort(physical_perm).astype(jnp.int32),
        ),
    )
    return model


def build_finetune_template_model(
    eager_model,
    n_sites: int,
    *,
    leaf_rank: int,
    merge_rank: int,
) -> CompiledFinetuneWaveFunction:
    from .trunk import bind_shared_kernel

    kernel = bind_shared_kernel(eager_model)
    widths = level_widths(int(n_sites))
    r_leaf = int(kernel.leaf_factors[0].V.shape[-1])
    r_merge = int(kernel.merge_factors[0].V.shape[-1])
    r_readout = int(kernel.readout_factors[0].V.shape[-1])
    tree = CompiledTree(
        perm=jnp.arange(n_sites, dtype=jnp.int32),
        inv_perm=jnp.arange(n_sites, dtype=jnp.int32),
        leaf_real=jnp.ones((n_sites,), dtype=jnp.bool_),
        leaf_h=(jnp.ones((n_sites, r_leaf), dtype=jnp.float32),),
        leaf_combiner_h=(),
        merge_h=tuple(
            jnp.ones((width, r_merge), dtype=jnp.float32) for width in widths
        ),
        opcodes=tuple(
            jnp.full((width,), MERGE, dtype=OPCODE_DTYPE) for width in widths
        ),
        readout_h=(jnp.ones((r_readout,), dtype=jnp.float32),),
        readout_combiner_h=(),
    )
    model = CompiledFinetuneWaveFunction.from_compiled(kernel, tree)
    return expand_rank(
        model,
        leaf_rank=leaf_rank,
        merge_rank=merge_rank,
        key=jax.random.PRNGKey(0),
    )